Training

From the training script to the screen

Everything a run shows in LaunchHelm comes from its training script. Each numbered part of the script below creates something you see in the app: the strip, a control, a stream, a chart, the model's tensors. This guide goes through the script one part at a time and says where each part shows up.

The script

script_to_screen.py trains a small model on synthetic data drawn from four sources, at a tenth of a second a step, so you have time to watch the run and change it. It needs the torch and system extras and the four LAUNCHHELM_* variables from Getting started:

Shell

pip install "launchhelm[torch,system]"
python script_to_screen.py

PythonDownload

import time

import torch

import launchhelm.torch as lht
from launchhelm.models import Capability, CommandContext, opaque_schema

STEPS = 600
SOURCES = ["web", "code", "math", "books"]
mixture = {"web": 0.4, "code": 0.3, "math": 0.2, "books": 0.1}

torch.manual_seed(0)
# Each source is a regression target of its own.
targets = {name: torch.randn(16, 1) for name in SOURCES}
model = torch.nn.Sequential(
    torch.nn.Linear(16, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 1),
)
layers = {
    index: module for index, module in enumerate(model) if isinstance(module, torch.nn.Linear)
}
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=STEPS)


def draw_batch(size: int = 64) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
    """A batch drawn by the current mixture, and the share each source got."""
    weights = torch.tensor([mixture[name] for name in SOURCES])
    picks = torch.multinomial(weights, size, replacement=True)
    x = torch.randn(size, 16)
    y = torch.empty(size, 1)
    for index, name in enumerate(SOURCES):
        rows = picks == index
        y[rows] = x[rows] @ targets[name]
    shares = {name: float((picks == index).float().mean()) for index, name in enumerate(SOURCES)}
    return x, y, shares


def run_benchmarks() -> dict[str, float]:
    """Held-out loss on each source. The eval control runs it on request."""
    scores = {}
    with torch.no_grad():
        for name in SOURCES:
            x = torch.randn(256, 16)
            loss = torch.nn.functional.mse_loss(model(x), x @ targets[name])
            scores[f"loss/{name}"] = float(loss)
    return scores


# 1. The run.
tr = lht.init(
    "quickstart",
    name="script-to-screen",
    config={"lr": 3e-3, "weight_decay": 0.1, "batch_size": 64, "steps": STEPS},
    system_metrics=True,
)

# 2. Controls and their streams.
tr.watch(
    model,
    optimizer,
    scheduler=scheduler,
    controls=["lr", "weight_decay", "grad_clip", "checkpoint"],
    eval_fn=run_benchmarks,
    captures=True,  # 3. the model's tensors, and a control to capture them
)
tr.plan(step=STEPS)


# 4. Your own control: a typed command with a mixture of four weights.
def set_mixture(value: dict, context: CommandContext) -> dict:
    """Draw the next batches by these weights, scaled to sum to 1."""
    requested = value["weights"]["weights"]
    total = sum(float(share) for share in requested.values())
    if total > 0:
        mixture.update({name: float(share) / total for name, share in requested.items()})
    return {name: round(share, 4) for name, share in mixture.items()}


tr.run.on(
    "after_step",
    Capability(
        name="set_mixture",
        provider="script",
        version="1.0.0",
        input_schema={
            "revision": "1",
            "root": {
                "kind": "record",
                "renderer": "form",
                "fields": [
                    {
                        "name": "weights",
                        "type": {
                            "kind": "mixture",
                            "components": SOURCES,
                            "sampled": "data/mix/{component}",
                        },
                    },
                ],
            },
        },
        output_schema=opaque_schema(),
        safe_points=("after_step",),
        starting={"weights": {"weights": {name: str(share) for name, share in mixture.items()}}},
    ),
    set_mixture,
)

for step in range(STEPS):
    x, y, shares = draw_batch()
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    tr.clip_grad_norm_()  # uses the grad_clip control
    health = step % 10 == 0
    if health:
        before = {index: layer.weight.detach().clone() for index, layer in layers.items()}
    optimizer.step()
    scheduler.step()

    # 5. Many streams under one name pattern, every ten steps.
    if health:
        ratios = {
            f"health/layer{index}/update_ratio": float(
                (layer.weight.detach() - before[index]).norm() / before[index].norm()
            )
            for index, layer in layers.items()
        }
        drawn = {f"data/mix/{name}": share for name, share in shares.items()}
        tr.log(ratios | drawn, step=step)

    # 6. The core metrics, and the safe points.
    tr.step(step=step, loss=loss.item(), samples=len(x))
    tr.run.sync("after_step", step=step)
    time.sleep(0.1)

tr.finish()
print("run finished:", tr.run.run_id)

Open the project in LaunchHelm and choose the run while it trains.

1. The run

PythonDownload

tr = lht.init(
    "quickstart",
    name="script-to-screen",
    config={"lr": 3e-3, "weight_decay": 0.1, "batch_size": 64, "steps": STEPS},
    system_metrics=True,
)

lht.init creates the run in the project named by its first argument and joins it as one process, rank-0, with the role trainer.

The strip, the one line of live state above the run's board: the run is live while this process checks in, its running time, its processes and machines, and the Controls button.
The runs panel and Details. name is the run's name everywhere. The config values are listed under Details in the run's dock, and are settings you can search, filter on and add as columns in the runs panel when you compare this run with others.
System streams. system_metrics=True starts the host collector on this process: system/cpu_percent, system/memory_used_bytes, disk and network rates and, where NVML is present, system/gpu<i>/utilization_percent, memory_used_bytes, temperature_c and power_w for each GPU, every 15 seconds. They land on the board with every other metric. No infrastructure page is preset: to keep them apart, add a page with "+" beside the page tabs, put the system/ charts on it, and draw the GPUs as heatmap rows the way part 5 draws layers.
Sources. "Sources…" in the bar's ⚙ menu lists everything the run's processes declared or wrote: its Streams with who writes each one, its Controls and its Processes.

2. Controls and their streams

PythonDownload

tr.watch(
    model,
    optimizer,
    scheduler=scheduler,
    controls=["lr", "weight_decay", "grad_clip", "checkpoint"],
    eval_fn=run_benchmarks,
    captures=True,  # 3. the model's tensors, and a control to capture them
)
tr.plan(step=STEPS)

tr.watch declares the controls named in controls, and eval because an eval_fn is given. Those five are what the strip's Controls and the Controls section of the dock list for this run, each card with its current value or the value the run started with. The script leaves out pause, resume, dropout and rollback, so the app offers none of them; leave controls out and you get all eight. Control and tunables has each control's input.

lr. Its card shows the rate set and the rate the run actually used ("Ran at …"), read from optim/lr/group0. Because the scheduler is watched, a change scales the schedule's base rates, so the schedule carries it on instead of overwriting it at the next step.
The plan ahead. tr.plan(step=STEPS) gives the strip its progress ("step N of 600") and time left, and publishes the rate the scheduler will set from now to the end as optim/lr/planned. The lr chart draws it dashed past now, named "planned ahead". While the run is live, the chart's header offers "Change lr": the plan becomes editable on the chart and the change bar at the bottom collects the draft, the approver and the reason. Without tr.plan there is no plan to draw or edit.
weight_decay and grad_clip. Cards with their values; optim/weight_decay/group0 shows what the run used. grad_clip sets the limit tr.clip_grad_norm_() applies in the loop.
checkpoint. Saves and uploads a checkpoint when you send it. Each saved checkpoint is a ▮ mark on every chart's x axis, the strip's last checkpoint and its age, and a line under Activity.
eval. Runs run_benchmarks when you send it and logs what it returns under eval/: eval/loss/web, eval/loss/code and so on. Those streams exist only once you have run it.

3. Inspectable tensors

captures=True (the default) is the third part of the same call. tr.watch publishes the model's parameters as the run's catalog of tensors and declares the capture control. There is no separate call for it.

Model, under Details in the run's dock: each tensor's path, shape and type, with the topology it belongs to. "Capture tensors…" opens the New capture view: the model on the left, when and what to take in the middle, its cost and whether it needs approval on the right.
The capture itself. Its stages come from the loop: after_backward and after_clip can be taken because the loop calls tr.clip_grad_norm_(), before_step and after_step because the optimizer is watched. The PyTorch guide lists what each stage offers and each summary.

4. Your own control

PythonDownload

def set_mixture(value: dict, context: CommandContext) -> dict:
    """Draw the next batches by these weights, scaled to sum to 1."""
    requested = value["weights"]["weights"]
    total = sum(float(share) for share in requested.values())
    if total > 0:
        mixture.update({name: float(share) / total for name, share in requested.items()})
    return {name: round(share, 4) for name, share in mixture.items()}

tr.run.on("after_step", Capability(...), set_mixture) declares a command of your own, set_mixture, at a safe point named after_step. Its input schema has one field, weights, of the mixture type over the four sources. A mixture value is {"weights": {"web": "0.4", ...}}, each weight a decimal string, so the handler reads value["weights"]["weights"].

The card and the editor. set_mixture appears with the built-in controls. Its form is the mixture editor, with three bars: Now (what the last change set, or starting until anything has), Draft (the weights you are setting; drag a divider between two sources) and Sampled.
Sampled reads the streams named by "sampled": "data/mix/{component}": data/mix/web, data/mix/code and so on, which part 5 logs. The editor and the card show the mix the run actually drew beside the one set ("60% · was 30% · ran 28%").
Starting value. starting= is the input as the run began, so the card has a value before anyone changes it.

A command with plain annotations needs no Capability: see run.command in Control and tunables, where int, float, str and bool parameters become a form. The mixture editor needs the mixture type, which only an explicit input schema declares.

set_mixture runs only where the loop calls tr.run.sync("after_step", ...), in the process that declared it. In a distributed job, the trainer's own controls and tunables change every rank at the same step; a command you declare with tr.run.on does not.

5. Many streams under one name pattern

PythonDownload

    if health:
        ratios = {
            f"health/layer{index}/update_ratio": float(
                (layer.weight.detach() - before[index]).norm() / before[index].norm()
            )
            for index, layer in layers.items()
        }
        drawn = {f"data/mix/{name}": share for name, share in shares.items()}
        tr.log(ratios | drawn, step=step)

Every ten steps the loop logs seven streams: each linear layer's update ratio (how far the step moved its weights, relative to their size) as health/layer0/update_ratio, health/layer2/update_ratio and health/layer4/update_ratio, and the share of the batch each source supplied as data/mix/<source>. tr.log writes from rank 0 only, unless lht.init(..., log_ranks="all").

Charts. The run's board is made from the streams it logs, one chart each: in one section up to twelve streams, then a section per name prefix (health, data, optim, …), and past 48 streams a page for each large prefix group, titled by the prefix.
Heatmap rows. Open one of the update-ratio charts with ✎. In the chart editor's Variables, choose "+ Variable" and click the 0 in health/layer0/update_ratio. That part becomes the variable layer, with the values 0, 2 and 4 found in the logged names, and "Heatmap rows" draws one row per layer. A real model with dozens of layers is one chart this way.
From a row to the model. On a run that offers captures, each row is a button ("Inspect layer 2 in the model") that opens New capture with the model searched for that layer. The script names layers by their index in nn.Sequential, so 2 finds 2.weight and 2.bias. Name per-layer streams after your modules' own paths for the same reason.
Sources shows each stream's spacing (every ten steps) and its last write; the Metrics section of the dock lists the health/layer… names as one family, with the layer number as a facet.

6. The core metrics, and the safe points

PythonDownload

    tr.step(step=step, loss=loss.item(), samples=len(x))
    tr.run.sync("after_step", step=step)

tr.step logs the run's core streams: train/loss, optim/grad_norm (measured by tr.clip_grad_norm_()), optim/lr/group0, optim/weight_decay/group0, perf/step_time_s and perf/samples_per_s (pass tokens= for perf/tokens_per_s), plus cuda/max_memory_allocated_bytes on a GPU. step= is the x axis of every chart and the position the strip counts against the plan.

It is also the trainer's safe point: it applies at most one waiting control per call, the same step on every rank. tr.run.sync("after_step", ...) is the script's own safe point, for set_mixture. A command changes training only at one of these two lines; a capture taken at a stage inside the step reads tensors and changes nothing.

A change, from request to effect. A change you send waits until the run reaches its safe point, is applied there, and reports the step it applied at. The change shows its track under Controls, a ◆ mark on every chart at that step, and a line under Activity.
Rules watch these streams: "+ Rule" in the Watch section of the dock. The script declares none; see Rules and notifications.

Nothing is preset

Nothing on the run's screens is preset by the app. No control, metric or page exists unless the script declared or logged it. The controls are the five tr.watch declared, its capture control and set_mixture. The charts are the streams the loop logged, and the board is laid out from their names. The mixture editor is the type set_mixture declared. Every page, heatmap and rule beyond that is one you made from those streams. Remove a line from the script and what it created is gone from the next run.

Next: PyTorch for the trainer in full, Control and tunables for every control's input, Checkpoints and resume for what checkpoint saves.