PyTorch training
launchhelm.torch (import it as lht) wraps an ordinary PyTorch loop. Install the extra:
pip install "launchhelm[torch]"
The loop
torch.manual_seed(0)
model = torch.nn.Sequential(
torch.nn.Linear(8, 16), torch.nn.ReLU(), torch.nn.Linear(16, 1)
)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x = torch.randn(64, 8)
y = x.sum(dim=1, keepdim=True)
tr = lht.init("quickstart", name="torch-quickstart", config={"lr": 0.01})
tr.watch(model, optimizer)
dropout_p = tr.tunable("dropout_p", 0.0, minimum=0.0, maximum=0.9)
for step in range(20):
for module in model.modules():
if isinstance(module, torch.nn.Dropout):
module.p = dropout_p.value
optimizer.zero_grad()
loss = torch.nn.functional.mse_loss(model(x), y)
loss.backward()
tr.clip_grad_norm_()
optimizer.step()
tr.step(step=step, loss=loss.item())
tr.finish()
print("run finished:", tr.run.run_id)python-torch_quickstart.py · verbatim from sdks/python/examples/torch_quickstart.py
The three calls that matter:
tr.watch(model, optimizer)— attaches the built-in controls (lr,weight_decay,grad_clip,dropout,checkpoint,rollback,pause,resume) and enablesoptim/lr/group0,optim/grad_normandperf/*streams. Passscheduler=too when you use one — thelrcontrol then scales the schedule's base lrs so closed-form schedulers continue proportionally.tr.step(step=..., loss=...)— call once per optimizer step. It logstrain/loss,perf/step_time_s, throughput, then services control.tr.tunable(name, default, minimum=..., maximum=...)— a dashboard-editable value you read each step. Alsochoices=for string enumerations.
tr.plan(step=total_steps) declares the run's planned last step. Besides the progress and time-left estimate on the run page, it lets the trainer publish the learning rate the scheduler will set from now to that step: it copies the scheduler onto a stand-in optimizer (the real one, its weights and the real scheduler are untouched), steps the copy forward and logs the result as one point on optim/lr/planned. The lr control declares that stream, so the run page draws the real schedule ahead and what a change does to it. A new plan is logged after every applied lr change and after load_checkpoint. There is no plan for ReduceLROnPlateau (its future depends on metrics), and none when the scheduler is not stepped once per training step (the trainer warns once). Sampling takes about 0.2 s for a 32,000-step horizon; horizons past 200,000 steps are cut there and marked incomplete.
tr.clip_grad_norm_() wraps torch.nn.utils.clip_grad_norm_ on the watched parameters using the current grad_clip control value — the UI can tighten or lift the cap mid-run.
System metrics
lht.init(..., system_metrics=True) starts the host collector on the run (launchhelm.system for non-torch processes — see Examples): CPU, RSS, memory, disk and network rates, and per-GPU utilisation, memory, temperature and power when NVML is present. Without psutil/NVML it logs one status line and the run is unaffected.
What gets logged
train/loss, optim/lr/group<i>, optim/weight_decay/group<i>, optim/grad_norm, optim/loss_scale (with a watched GradScaler), perf/step_time_s, perf/samples_per_s, perf/tokens_per_s (pass samples=/tokens= to step), cuda/max_memory_allocated_bytes on CUDA.
The model's tensors
tr.watch also publishes the model's parameters as the run's catalog of inspectable tensors: each parameter's path, shape, dtype, element count, whether it trains, its optimizer param group, and the path it is tied to when one tensor appears under two names. They appear under Model on the run's Details tab. The catalog has a topology revision (a hash of the paths, shapes and dtypes, readable as tr.topology_revision); a model that changes publishes a new revision, and tensors it no longer has stay readable for captures already taken. A failed publish warns and training continues.
Captures
A capture summarizes some of those tensors at one point in a step, on the trainer, when it is taken. Only the summaries leave the process: a small JSON manifest uploaded as a capture artifact. tr.watch registers the capture capability on the coordinator (pass captures=False to leave it out).
You ask for one from the run's Model view or POST /captures, naming the tensors, a stage, what to take there, and the summaries:
| Stage | Where in the step | What it offers |
|---|---|---|
after_backward | entering tr.clip_grad_norm_, before clipping | gradients |
after_clip | leaving tr.clip_grad_norm_ | gradients |
before_step | just before optimizer.step() | weights, gradients, optimizer_state |
after_step | just after optimizer.step() | weights, gradients, optimizer_state, updates |
optimizer_state lists each per-element state tensor (Adam's exp_avg and exp_avg_sq); updates is the change the step made to the weights. Under DDP the gradients are the reduced ones every rank shares.
| Summary | What it holds |
|---|---|
stats | count, min, max, mean, standard deviation, L2 norm, zero and non-finite fractions |
histogram | 128 bins (bins up to 512) between the finite min and max; NaN and infinities counted apart |
top_rows | the 32 rows (k up to 256) with the largest L2 norm, with their values when a row has at most 512 |
block_max | the tensor as a grid of at most 128 by 128 blocks (size up to 256), each the largest absolute value in it |
singular_values | the 128 largest (k up to 1024), with the share of the matrix's energy they hold |
full | the tensor itself, in its own dtype, as tiles of at most 256 by 256 values |
The summaries are exact (computed in float64; block_max compares in float32, exact for tensors of 32 bits or fewer) and say how they were computed. singular_values are exact (float64) when the matrix's smaller side is at most 1024 and it holds at most 4M values; larger matrices get a randomized estimate (10 extra sketch columns, 2 power iterations), labelled estimated with those controls. For block_max, singular_values and full a tensor is read as rows and columns: its first dimension by the rest flattened, a vector as one row.
full copies weights out of the run, so asking for it, and reading it, needs the capture.raw permission. One capture copies at most 1 GiB. The tiles go to a second file, the capture's data, whose upload starts during the pause and finishes in the background after training resumes; the manifest lists each tile's byte offset, length and SHA-256, so a viewer reads only the tiles it shows. The capture is ready once both files are verified; if the data upload fails it keeps its summaries and is marked partial. tr.finish waits for an unfinished upload within its timeout. A capture is taken at the next point of its stage: one for after_step weights runs at the tr.step that receives it; the others wait for that stage in the next step, and other controls wait with them. after_backward and after_clip need the loop to call tr.clip_grad_norm_ (with no limit set it measures the norm without clipping); without it the capture is refused when the optimizer steps. With a GradScaler, gradients before scaler.unscale_ are scaled; the manifest records the loss scale. Training pauses while the summaries are computed and the manifest uploads, usually milliseconds per tensor plus one upload. A tensor with nothing to take (no gradient yet, no optimizer state before the first step) is listed as an omission and the capture is marked partial. A capture whose tensors were chosen from an older topology revision is refused.
A capture that kept full tensors can be summarized again later without asking the run: "Recompute…" in the viewer, or POST /captures/{id}/derive with the summaries to compute. The platform reads the stored tiles, checks their digests, and computes the summaries with the same definitions and parameters, as a new capture that names the one it came from. Training does not pause for it. It needs the same capture.raw permission as reading the tensors.
Next: distributed training for DDP, control for the full control surface.