Distributed training

Log rank-aware metrics for reducers, coverage, heatmaps, and outlier inspection.

Open .md

Log per-rank metrics from distributed jobs so the Distributed dashboard tab can show reducers, rank coverage, heatmaps, and outliers instead of one flattened scalar.

Log per-rank metrics

Each call records one rank's metric batch for one global step.

python
run.log_rank_metrics(
    {"train/loss": 0.12, "optimizer/grad_norm": 3.4},
    step=100,
    rank=0,
    world_size=8,
    local_rank=0,
    weight=1024,
)
ArgumentMeaning
metricsScalar metric dictionary for one rank.
stepGlobal step shared by all ranks.
rankZero-based global rank.
world_sizeExpected number of ranks, capped at 512.
local_rankOptional zero-based rank within a node.
weightOptional sample count or work weight for weighted reducers. Defaults to 1.0.

Inspect ranks in the Distributed tab

The Distributed tab reads a run-scoped summary and renders:

  • Mean and weighted mean.
  • Min, max, range, and standard deviation.
  • p05, p50, and p95 percentiles.
  • Rank coverage per step.
  • Heatmap cells.
  • Outlier rows by z-score and delta from mean.

Use these views to answer questions like:

  • Did one rank stop reporting?
  • Is a rank consistently slower or noisier?
  • Did a data shard or process produce outlier losses?
  • Are weighted reducers different from unweighted reducers?

Use spool mode for rank metrics

Process-isolated spool mode writes rank metric calls as JSONL-style request events and replays them with the event ID as Idempotency-Key.

python
import instantml as im

run = im.init(
    project="distributed-job",
    upload_mode="spool",
    spool_dir=".instantml/spool",
)

run.log_rank_metrics(
    {"train/loss": loss},
    step=step,
    rank=rank,
    world_size=world_size,
    weight=batch_size,
)

This keeps post-init HTTP calls out of the training process while preserving retry safety for compatible servers.

Choose which keys to log per rank

Log rank-aware metrics only for the keys that help debug distributed behavior. For example, log ordinary train/loss and eval/accuracy from rank 0 or a trainer process, then log rank-aware train/loss, optimizer/grad_norm, and system/gpu_memory_mb from all ranks.

That keeps scalar comparison fast while still giving the Distributed tab enough information to detect coverage and outlier problems.

Next steps