> ## Documentation Index
> Fetch the complete documentation index at: https://docs.warpscale.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Report training progress

> Report run intent and step, epoch, and checkpoint progress from your training script.

The probe and the CLI capture GPU and system telemetry. The SDK adds what only your code knows: where the run is, and what it is training. Without it a stall shows up in the hardware but cannot be lined up against a step.

## Before you start

* Training already reporting — see [Setup](/training/wrap)
* Python 3.9 or newer

## 1. Install the SDK

Install it into the same environment as your training script.

```bash theme={null}
pip install warpscale
```

## 2. Add the calls

Call `init` once after the process group is up, then mark progress as the run proceeds.

```python train.py theme={null}
import warpscale

# After dist.init_process_group(), before your DataLoader.
warpscale.init(
    total_steps=40_000,
    total_epochs=3,
    framework="fsdp",
    precision="bf16",
    model=model,
)

for epoch in range(epochs):
    for batch in loader:
        ...                          # one optimizer step
        warpscale.step()
    warpscale.epoch()

    warpscale.checkpoint_started()
    save_checkpoint()
    warpscale.checkpoint_persisted()
```

<Warning>
  Call `warpscale.init()` **after** `dist.init_process_group()`. The emitter is elected inside `init`, so calling it earlier leaves every rank a non-emitter and the SDK silently inactive for the whole run.
</Warning>

Passing `model=` derives `param_count` from the module. That, with `flops_per_token`, is what the MFU panels are computed from — leave both out and those panels stay empty.

### A complete example

The snippet above is the shape. Below is a script that runs as-is: a small GPT trained with DDP on two GPUs, on synthetic tokens, so there is no dataset to fetch first.

<Accordion title="train.py">
  ```python train.py theme={null}
  """A small GPT-style DDP run, instrumented with the Warpscale Python SDK.

  Synthetic tokens only — no dataset download — so the run is reproducible on any
  node with two GPUs.
  """

  import argparse
  import os
  import time

  import torch
  import torch.distributed as dist
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader, Dataset, DistributedSampler

  import warpscale


  class SyntheticTokens(Dataset):
      """Random token ids built on the CPU worker so every batch costs a real
      host-to-device copy, which is what the data-transfer features measure."""

      def __init__(self, n_samples, seq_len, vocab_size, seed=0):
          self.n_samples = n_samples
          self.seq_len = seq_len
          self.vocab_size = vocab_size
          self.seed = seed

      def __len__(self):
          return self.n_samples

      def __getitem__(self, idx):
          g = torch.Generator().manual_seed(self.seed + idx)
          ids = torch.randint(0, self.vocab_size, (self.seq_len + 1,), generator=g)
          return ids[:-1], ids[1:]


  class Block(nn.Module):
      def __init__(self, d_model, n_head):
          super().__init__()
          self.n_head = n_head
          self.ln1 = nn.LayerNorm(d_model)
          self.qkv = nn.Linear(d_model, 3 * d_model)
          self.proj = nn.Linear(d_model, d_model)
          self.ln2 = nn.LayerNorm(d_model)
          self.mlp = nn.Sequential(
              nn.Linear(d_model, 4 * d_model),
              nn.GELU(),
              nn.Linear(4 * d_model, d_model),
          )

      def forward(self, x):
          b, t, c = x.shape
          q, k, v = self.qkv(self.ln1(x)).split(c, dim=2)
          q, k, v = (
              z.view(b, t, self.n_head, c // self.n_head).transpose(1, 2) for z in (q, k, v)
          )
          a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
          x = x + self.proj(a.transpose(1, 2).reshape(b, t, c))
          return x + self.mlp(self.ln2(x))


  class TinyGPT(nn.Module):
      def __init__(self, vocab_size, seq_len, d_model, n_layer, n_head):
          super().__init__()
          self.tok = nn.Embedding(vocab_size, d_model)
          self.pos = nn.Embedding(seq_len, d_model)
          self.blocks = nn.ModuleList(Block(d_model, n_head) for _ in range(n_layer))
          self.ln_f = nn.LayerNorm(d_model)
          self.head = nn.Linear(d_model, vocab_size, bias=False)

      def forward(self, idx):
          pos = torch.arange(idx.shape[1], device=idx.device)
          x = self.tok(idx) + self.pos(pos)
          for block in self.blocks:
              x = block(x)
          return self.head(self.ln_f(x))


  def parse_args():
      p = argparse.ArgumentParser()
      p.add_argument("--epochs", type=int, default=3)
      p.add_argument("--steps-per-epoch", type=int, default=800)
      p.add_argument("--batch-size", type=int, default=8, help="sequences per rank")
      p.add_argument("--seq-len", type=int, default=1024)
      p.add_argument("--vocab-size", type=int, default=32000)
      p.add_argument("--d-model", type=int, default=768)
      p.add_argument("--n-layer", type=int, default=12)
      p.add_argument("--n-head", type=int, default=12)
      p.add_argument("--num-workers", type=int, default=4)
      p.add_argument("--log-every", type=int, default=50)
      p.add_argument("--ckpt-dir", default="./ckpt")
      return p.parse_args()


  def save_checkpoint(model, optimizer, ckpt_dir, epoch):
      """Write one full checkpoint, overwriting the last so repeated epochs keep
      exercising checkpoint I/O without growing on disk."""
      os.makedirs(ckpt_dir, exist_ok=True)
      path = os.path.join(ckpt_dir, "latest.pt")
      tmp = path + ".tmp"
      torch.save(
          {
              "epoch": epoch,
              "model": model.module.state_dict(),
              "optimizer": optimizer.state_dict(),
          },
          tmp,
      )
      os.replace(tmp, path)
      return path


  def main():
      args = parse_args()

      local_rank = int(os.environ["LOCAL_RANK"])
      torch.cuda.set_device(local_rank)
      device = torch.device("cuda", local_rank)
      dist.init_process_group(backend="nccl", device_id=device)
      rank = dist.get_rank()
      world_size = dist.get_world_size()
      torch.manual_seed(1234 + rank)

      model = TinyGPT(
          args.vocab_size, args.seq_len, args.d_model, args.n_layer, args.n_head
      ).to(device)
      param_count = sum(p.numel() for p in model.parameters())
      ddp = DistributedDataParallel(model, device_ids=[local_rank])
      optimizer = torch.optim.AdamW(ddp.parameters(), lr=3e-4, fused=True)

      # Elects rank 0 as the sole emitter, so it must follow init_process_group and
      # precede the DataLoader whose workers fork.
      warpscale.init(
          total_steps=args.steps_per_epoch * args.epochs,
          total_epochs=args.epochs,
          global_batch_size=args.batch_size * world_size,
          precision="bf16",
          framework="ddp",
          param_count=param_count,
          flops_per_token=6 * param_count,
      )

      dataset = SyntheticTokens(
          args.steps_per_epoch * args.batch_size * world_size, args.seq_len, args.vocab_size
      )
      sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, drop_last=True)
      loader = DataLoader(
          dataset,
          batch_size=args.batch_size,
          sampler=sampler,
          num_workers=args.num_workers,
          pin_memory=True,
          drop_last=True,
          persistent_workers=args.num_workers > 0,
      )

      if rank == 0:
          print(
              f"[train] world_size={world_size} params={param_count/1e6:.1f}M "
              f"global_batch={args.batch_size * world_size} seq_len={args.seq_len} "
              f"run_id={os.environ.get('WS_RUN_ID', '<unset>')}",
              flush=True,
          )

      started = time.perf_counter()
      for epoch in range(args.epochs):
          sampler.set_epoch(epoch)
          epoch_started = time.perf_counter()
          for step, (x, y) in enumerate(loader):
              x = x.to(device, non_blocking=True)
              y = y.to(device, non_blocking=True)
              with torch.autocast("cuda", dtype=torch.bfloat16):
                  loss = F.cross_entropy(ddp(x).flatten(0, 1), y.flatten())
              loss.backward()
              optimizer.step()
              optimizer.zero_grad(set_to_none=True)
              warpscale.step()

              if rank == 0 and (step + 1) % args.log_every == 0:
                  elapsed = time.perf_counter() - epoch_started
                  print(
                      f"[train] epoch {epoch} step {step + 1}/{args.steps_per_epoch} "
                      f"loss {loss.item():.3f} {(step + 1) / elapsed:.2f} steps/s",
                      flush=True,
                  )

          warpscale.epoch()

          warpscale.checkpoint_started()
          if rank == 0:
              path = save_checkpoint(ddp, optimizer, args.ckpt_dir, epoch)
              print(f"[train] epoch {epoch} checkpoint -> {path}", flush=True)
          dist.barrier()
          warpscale.checkpoint_persisted()

      if rank == 0:
          print(f"[train] done in {time.perf_counter() - started:.1f}s", flush=True)
      dist.destroy_process_group()


  if __name__ == "__main__":
      main()
  ```

  Launch it the same way as any other run:

  ```bash theme={null}
  warpscale run -- torchrun --standalone --nproc_per_node=2 train.py
  ```
</Accordion>

## 3. Verify

The SDK is inert unless it is running under the CLI — it looks for the `WS_RUN_ID` that `warpscale run` sets.

```bash theme={null}
warpscale run -- python train.py
```

Two checks:

* Run your script **without** the CLI. The SDK logs `WS_RUN_ID not set — not running under 'warpscale run'; warpscale inactive` at INFO and does nothing else. That is the expected inactive path.
* Run it **with** the CLI. Step and epoch progress appear on the run at `https://<your-organization-host>/runs`.

## Troubleshooting

<AccordionGroup>
  <Accordion title="No progress data, but the run appears">
    The SDK is inactive. Confirm the run is launched through `warpscale run --`, and that `warpscale.init()` is called after `dist.init_process_group()`.
  </Accordion>

  <Accordion title="Progress stalls partway through the run">
    Progress is emitted by rank 0 only. If rank 0 exits or is restarted mid-run, reporting stops with it.
  </Accordion>

  <Accordion title="The MFU panels are empty">
    MFU needs `param_count` and `flops_per_token`. Pass `model=` to derive the first, and `flops_per_token=` for the second.
  </Accordion>
</AccordionGroup>