Skip to content

Debug slow DDP training and rank stragglers

Use this guide when PyTorch DDP training is slow, uneven, or controlled by one culprit rank.

TraceML looks for visible rank skew in distributed runs, identifies a culprit rank and a victim/reference rank, then attributes material culprit excess to input, H2D, DDP forward compute, or sync/unattributed work.

For whole-run triage, start with Find why PyTorch training is slow. This page stays focused on DDP rank skew.

Before you run

Standard Hugging Face Trainer, Lightning Trainer.fit(), and RF-DETR model.train() scripts are instrumented automatically when launched with traceml run. For a custom PyTorch loop, call traceml.init() once and mark each training step with traceml.trace_step(model). See the quickstart for both setup paths.

Run DDP with TraceML

Your training script or framework configuration must already support DDP. TraceML launches the requested workers; it does not convert a single-process script into distributed training.

For single-node DDP:

traceml run train.py --nproc-per-node=4

For multi-node DDP, use the same --run-name, --nnodes, --nproc-per-node, --master-addr, --master-port, and script args on every node. See Distributed Training.

On Slurm, use the Slurm guide instead of writing the node flags by hand.

Reproduce rank stragglers

The repo includes a small DDP demo with a compute-heavy baseline and two controlled rank-local delays. The input scenario exercises rank-straggler attribution. The historical compute-straggler scenario adds optimizer-side work, so use it to inspect phase timing rather than to require a specific straggler verdict.

Balanced baseline:

traceml run examples/distributed/ddp_rank_straggler_demo.py --mode=summary --nproc-per-node=2 --run-name ddp_balanced --args --scenario balanced

Input straggler:

traceml run examples/distributed/ddp_rank_straggler_demo.py --mode=summary --nproc-per-node=2 --run-name ddp_input_straggler --args --scenario input-straggler --straggler-rank 0 --input-sleep-ms 200

Optimizer-side delay (compute-straggler CLI name):

traceml run examples/distributed/ddp_rank_straggler_demo.py --mode=summary --nproc-per-node=2 --run-name ddp_compute_straggler --args --scenario compute-straggler --straggler-rank 0 --compute-extra-matmuls 8

For two nodes with one GPU each, run the same script on both nodes with --nnodes=2 --nproc-per-node=1, the same --master-addr, --master-port, and --run-name, and a different --node-rank on each node.

Example node 0:

traceml run examples/distributed/ddp_rank_straggler_demo.py --mode=summary --nnodes=2 --nproc-per-node=1 --node-rank=0 --master-addr <NODE_0_PRIVATE_IP> --master-port 29546 --run-name ddp_compute_straggler --args --scenario compute-straggler --straggler-rank 0 --compute-extra-matmuls 8

Example node 1:

traceml run examples/distributed/ddp_rank_straggler_demo.py --mode=summary --nnodes=2 --nproc-per-node=1 --node-rank=1 --master-addr <NODE_0_PRIVATE_IP> --master-port 29546 --run-name ddp_compute_straggler --args --scenario compute-straggler --straggler-rank 0 --compute-extra-matmuls 8

What the demo should show

Scenario Expected observation Interpretation
Balanced Similar phase times across ranks Control run with no injected rank-local delay
Input straggler Higher Input Wait on the selected rank Can produce INPUT STRAGGLER when the delay is material
Optimizer-side delay Higher Optimizer Step on the selected rank Demonstrates phase timing; it is not a COMPUTE STRAGGLER acceptance test

Read the balanced run as the control: both ranks have similar Step Time, Input Wait is small, and the run is mostly compute.

Read the input-straggler run as a rank-local input issue: rank 0 reaches compute late because its input path is slower.

Read the optimizer-delay run as a phase-timing example: the selected rank spends more time in optimizer work than its peers. The current DDP COMPUTE STRAGGLER attribution requires excess forward time, so this scenario does not guarantee that verdict.

What to look for

Start with the Step Time diagnosis.

Diagnosis Meaning
INPUT STRAGGLER one rank spends much longer waiting for input than peer ranks
COMPUTE STRAGGLER in DDP, the culprit rank has materially more forward-time burden than the victim rank
H2D STRAGGLER the culprit rank has meaningfully more host-to-device transfer burden than the victim rank
STRAGGLER visible rank skew exists, but input wait, H2D, and DDP forward do not explain it
INPUT-BOUND input work is broad, not just one bad rank
RESIDUAL-HEAVY meaningful residual time is not attributed to input wait, H2D, forward, backward, or optimizer work

For rank stragglers, TraceML first looks for visible wait cost: in DDP this is the gap between the upper-median backward time and the rank with the smallest backward time. The smallest-backward rank is the likely culprit because it arrived late and waited least. TraceML then compares that culprit with the victim rank to see whether input wait, H2D, or DDP forward time explains the cost.

Then inspect:

  • culprit rank
  • victim/reference rank
  • visible rank skew
  • Input Wait
  • Forward
  • Backward
  • Optimizer Step
  • Residual

How to triage the culprit rank

If the diagnosis is INPUT STRAGGLER, inspect input loading on the culprit rank:

  • uneven shards or batch construction
  • rank-local preprocessing jitter
  • slow storage path on one host
  • custom collation or tokenization on one rank

If the diagnosis is COMPUTE STRAGGLER, inspect DDP forward work on the culprit rank:

  • uneven input shapes
  • rank-local branching
  • extra forward work
  • framework hooks or callbacks that run on one rank

If the diagnosis is H2D STRAGGLER, inspect host-to-device transfer on the culprit rank: batch shapes, transfer placement, pinned memory, and transfer jitter.

If the diagnosis is STRAGGLER, inspect synchronization, collectives, and unattributed work around the culprit rank. Explicit collective timing is not available yet, so this is the honest fallback when input, H2D, and DDP forward do not explain the visible wait cost.

If the diagnosis is RESIDUAL-HEAVY, remember that TraceML reports residual time as a derived bucket. It is not direct NCCL or all-reduce timing. Inspect logging, checkpointing, validation, CPU stalls, framework orchestration, and unobserved transfer paths before assuming the cause.

FSDP note

TraceML supports single-node multi-GPU DDP and FSDP in the current public docs. Multi-node DDP summary reports are supported. Multi-node FSDP uses the same distributed launch path as DDP, but should be validated in your environment.

TraceML uses backward as the visible rank-skew phase for DDP/default strategy. For FSDP, it uses forward + backward because sharding communication can appear in both phases. FSDP rank stragglers can still be attributed to input wait or H2D, but TraceML does not emit COMPUTE STRAGGLER from the FSDP rank-skew rule without explicit all-gather or reduce-scatter timing.

Compare a fix

After changing data loading, batching, model logic, or host placement, compare the old and new summaries:

traceml compare old_run/final_summary.json new_run/final_summary.json

Look for changes in common-clock Step Time, visible rank skew, Input Wait, compute time, residual time, and diagnosis.