benchmark-v1

[abhinavsriva]

What changed

Added a simple end-to-end DDP MLP benchmark under benchmarks/ for measuring
TraceML overhead using normal shell wall time.

The benchmark workload is intentionally close to examples/ddp_minimal.py, but
is parameterized so we can tune compute cost:

• --duration-sec
• --batch-size
• --input-dim
• --hidden-dim
• --hidden-layers
• --num-samples

It runs as plain PyTorch under torchrun and only enables TraceML when launched
through traceml run.

Also added the missing traceml_ai.diagnostics.step_time.names module required
by existing step-time diagnostics/reporting imports.

Why

This gives us a practical first benchmark that is easy to explain:

1. Run normal DDP MLP training.
2. Run the same training with traceml run.
3. Compare end-to-end shell real time.

This avoids custom per-step timing in the first benchmark while still allowing
batch size and model size to be increased when we need slower, more
compute-heavy steps.

The diagnostics names module is needed because the TraceML aggregator can fail
at startup when reporting imports traceml_ai.diagnostics.step_time.names and
the file is missing from a clean checkout.

How I tested

• Tests added or updated
• Existing tests run
• Manual smoke test
• Not run; reason:

Compile check:

python -m py_compile benchmarks/workloads/ddp_mlp_e2e.py

Native tiny smoke:

PYTHONPATH=src python -m torch.distributed.run \
  --nnodes=1 \
  --nproc_per_node=1 \
  --node_rank=0 \
  --master_addr=127.0.0.1 \
  --master_port=29610 \
  benchmarks/workloads/ddp_mlp_e2e.py \
  --duration-sec 2 \
  --batch-size 2 \
  --input-dim 8 \
  --hidden-dim 8 \
  --hidden-layers 1 \
  --num-classes 4 \
  --num-samples 32 \
  --log-every 0

TraceML tiny smoke:

PYTHONPATH=src python -c 'from traceml_ai.launcher.cli import main; main()' run \
  benchmarks/workloads/ddp_mlp_e2e.py \
  --mode=summary \
  --logs-dir /tmp/traceml-e2e-smoke \
  --run-name e2e_smoke \
  --nproc-per-node=1 \
  --master-port=29611 \
  --aggregator-port=29811 \
  --args \
  --duration-sec 2 \
  --batch-size 2 \
  --input-dim 8 \
  --hidden-dim 8 \
  --hidden-layers 1 \
  --num-classes 4 \
  --num-samples 32 \
  --log-every 0

Aggregator import check:

PYTHONPATH=src python -c "from traceml_ai.aggregator.trace_aggregator import TraceMLAggregator; print(TraceMLAggregator.__name__)"

Step-time diagnostics/reporting tests:

PYTHONPATH=src python -m pytest tests/diagnostics/test_step_time.py tests/reporting/summary/test_step_time_card.py

Result:

16 passed

Runtime impact

• No training-path impact
• May affect tracing, runtime, telemetry, or distributed behavior
• Not sure

Notes:

The benchmark workload only affects users who explicitly run it. It is not
imported by TraceML runtime code.

The step_time.names module is used by existing diagnostics/reporting imports.
Including it fixes aggregator startup from clean checkouts; it does not change
training-loop instrumentation by itself.

Docs

• Updated
• Not needed

Extra context

Example intended benchmark usage:

time torchrun --nproc_per_node=1 benchmarks/workloads/ddp_mlp_e2e.py \
  --duration-sec 600 \
  --batch-size 256 \
  --hidden-dim 4096

time traceml run benchmarks/workloads/ddp_mlp_e2e.py \
  --mode=summary \
  --run-name=e2e_1gpu_traceml_r01 \
  --nproc-per-node=1 \
  --args \
  --duration-sec 600 \
  --batch-size 256 \
  --hidden-dim 4096

For publishable numbers, run paired repeats and alternate order:

repeat 1: native, TraceML
repeat 2: TraceML, native
repeat 3: native, TraceML
repeat 4: TraceML, native
repeat 5: native, TraceML

[Pendu]

Ran the overhead matrix on the AWS g4dn pair (2× T4), per the README method:
5 paired repeats × 2 modes × {single-GPU, 2-node DDP}, 600 s/trial, alternating
order, identical args. 20/20 trials, 0 failures. Results added to
benchmarks/README.md + a full write-up with plots under
benchmarks/analysis/2026-06-11_pr153_ddp_mlp_g4dn/report.md.

• Single GPU (T4): +0.90% wall, +1.02% throughput
• 2-node DDP (T4, TCP): +0.64% wall, ~0% throughput — the step is
  network-bound (~92% backward/all-reduce), so TraceML's per-step cost is
  masked. Base aggregator/display layer is within the <2% budget → looks safe
  to merge.

Two methodology notes for before-merge (also in the report §7):

1. The fixed---duration-sec design makes shell wall-time blind to
per-step overhead. Both modes train for the same 600 s, so they finish at
~the same wall time by construction — native does ~14,300 steps, TraceML
~14,180. The wall difference is therefore almost entirely fixed startup
(~5.5 s aggregator boot), not per-step cost; the per-step cost shows up as
fewer steps completed (lower throughput). Concretely: a workload where
TraceML made every step 30% slower would still report ~0% wall overhead (both
ran 600 s) while throughput overhead would be +30%. That's why I report both,
and it's why single-GPU shows wall +0.90% (mostly startup) but throughput
+1.02% (the real per-step cost).

This is also what the DDP row's split means, so it isn't misread: its
+0.64% wall is the ~4 s fixed startup (aggregator boot + 2-rank telemetry
setup + report writing) spread over the 600 s run — a one-time cost, not a
per-step tax — while its ~0% throughput is the true per-step cost
(negligible because the DDP step is mostly NCCL network wait that hides
TraceML's ~1 ms of sampling). Both numbers are correct; they measure different
things (one-time startup vs per-step), which is why they diverge. The
startup-dominated wall figure also shrinks with run length: ~0.64% at 600 s →
~0.06% at 6000 s, i.e. rounding error on a real multi-hour job.

A fixed---steps N mode fixes this: both
modes do identical work, so wall = N × per-step-time and wall time directly
captures per-step overhead (and it also removes note 2).

2. Rank-divergent stop can hang multi-node runs.

Why DDP needs lockstep: in DDP the two ranks must execute the same number of
steps. Every backward pass does an all-reduce to average gradients across
ranks, and an all-reduce is a collective — it only completes when every
rank calls it. If one rank calls all-reduce and the other has already exited,
the calling rank blocks until the NCCL watchdog aborts it (default ~10 min);
on a 2-node setup the surviving rank cannot recover because its peer is gone.

The bug: the workload stops based on perf_counter() - start >= duration,
checked independently on each rank, each with its own clock and its own
start time. Nothing synchronizes that decision. At the 600 s boundary the two
ranks can land on different sides of the threshold within the same step:

• Rank 0 reads 600.01 s → breaks out of the loop, exits, destroys its process group.
• Rank 1 reads 599.99 s → runs one more step → calls all-reduce → its partner
  is gone → blocks until the watchdog kills the job ~10 min later.

Why it didn't bite us: it's a race between the inter-node clock skew and
the step period (~423 ms here). The skew has to fall in the narrow window
where one rank crosses 600 s and the other doesn't on the same step, so it
fires only occasionally — 0 times in these 20 runs. But it is a real latent
hang, and the consequence is asymmetric: when it does fire in TraceML mode,
the run looks like TraceML hung the training, when the cause is the workload's
unsynchronized stop. That's a reputational landmine for an overhead benchmark.

Two fixes (either suffices):

• Collective stop flag — each rank computes its local "should stop", then
  dist.all_reduce(stop, op=ReduceOp.MAX); all ranks then stop on the same
  step and exit together, so there's never an unmatched all-reduce. ~3 lines,
  and it's mode-symmetric (added to native and TraceML alike, so it doesn't
  bias the comparison).
• --steps N — both ranks run exactly N steps, so the stop is
  deterministic and identical across ranks; no clock, no skew, no race. (This
  is the same change that fixes note 1, which is why I'd lean toward it.)