[Benchmark] Add all-reduce rank-skew calibration

[samnordmann]

Purpose

All-reduce duration includes time waiting for late ranks. Add an opt-in calibration mode to the existing communicator benchmark: inject a GPU delay on one rank and report per-rank timings for exact BF16 shapes in eager or CUDA-graph mode.

Trial summaries expose maximum residence, injected-rank residence, measured delay, and their residence gap. The injected rank is not assumed to arrive last, especially at zero requested delay; use measured delay for skew-sensitivity analysis.

Implements the calibration step of #52777; no duplicate implementation was found in open-PR searches. Benchmark and tests only; default output and serving behavior are unchanged.

Test Plan

.venv/bin/python -m pytest tests/benchmarks/test_device_communicator_benchmark.py -q
.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=2 \
  benchmarks/kernels/benchmark_device_communicators.py \
  --calibration-backends pynccl --calibration-execution-mode cuda_graph \
  --calibration-shapes 8x8192 8627x8192 16384x8192 \
  --skew-delays-us 0 50 250 1000 --num-warmup 10 --num-trials 50 \
  --output-json skew.json

Repeat with 4 ranks, eager mode, and alternate injected ranks.

Test Result

• 22 focused tests and all local pre-commit checks pass, including Ruff and mypy.
• New summaries replayed successfully over all 10,560 retained GPU trials; existing fields are unchanged. The GPU timing loop is unchanged by the review fixes.
• Fresh rebased-benchmark validation: 720 GPU trials across TP2/TP4, eager/graph and six backends, plus both legacy CLI runs, passed on GB200. Used the same pinned vLLM 44fe2a392 runtime as the original matrix, not a full rebuild of main.
• Original GB200 calibration: three independent runs/configuration, PyTorch 2.13.0+cu130, PyNCCL's NCCL 2.30.7, six backends. The legacy CLI matched upstream's JSON structure and backend coverage at TP2/TP4.

For BF16 8627×8192 per rank, original median maximum-rank GPU residence:

TP / mode	No injected delay	Requested 1 ms delay	Added residence, range across 3 runs
TP2 / graph	0.318 ms	1.315 ms	0.997–0.999 ms
TP2 / eager	0.352 ms	1.309 ms	0.957–0.961 ms
TP4 / graph	0.388 ms	1.385 ms	0.990–1.011 ms
TP4 / eager	0.434 ms	1.396 ms	0.933–0.964 ms

Actual injected delays were ≈1.01 ms. TP4 used eight GPU-local CPU cores/rank; unrestricted runs exposed millisecond-scale host enqueue skew. GPU residence includes launch gaps and peer waiting, not just transport; the residence gap is diagnostic, not an isolated wait-time measurement.

AI assistance: OpenAI Codex. Model evaluation is not applicable to this benchmark-only change.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @samnordmann.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[calvin-fei]

Thanks for picking up the calibration step of #52777. The approach matches what the RFC intended. There is
a CPU barrier before each trial but no extra collective inside the timed region. The delay is injected
on the GPU with the measured delay recorded next to the requested one. Rank records are merged
afterwards without comparing CUDA clocks across ranks.

I ran it on 2× H100 80GB (NVLink, TP2): your test-plan command, the same in eager mode, all six
all-reduce backends in graph mode, and --skew-rank 1. test_device_communicator_benchmark.py passes
(20), and the default benchmark mode still works. Your result reproduces:

• pynccl, 8627×8192, with a requested 1 ms delay adds 0.989 ms of max-rank residence in graph mode
  and 0.981 ms in eager.
• The other five backends, including ca_1stage/ca_2stage, add 0.965–0.978 ms.

Suggestions:

1. Expose the last-arriving rank's residence in the trial summary. max_gpu_collective_ms is the
   waiting rank's residence. On the delayed rank the residence stays flat as the delay grows, so it is
   the zero-wait cost (pynccl graph, 8627×8192, medians):

   injected delay	max-rank	delayed (last-arriving) rank	measured delay
   0	0.492 ms	0.491 ms	0.003 ms
   250 µs	0.740 ms	0.488 ms	0.259 ms
   1 ms	1.481 ms	0.488 ms	1.009 ms

   ca_1stage behaves the same way: 0.528 → 0.513 ms on the late rank, while max rises to 1.498 ms.
   So max − last-arriver ≈ the waiting time, which is the split #52777 is after. It also matches
   @wickist's observation there that the waiting rank's spin inside custom AR gets counted as
   communication. The per-rank values are already in ranks, so this is only a summary/print change.

2. Report the measured delay. _delay_us_to_cycles uses the rated clock_rate, and on H100 the
   requested and measured delays drift apart most at small values:

   • 1000 µs → 1.009 ms (+1%)
   • 250 µs → 0.259 ms (+4%)
   • 50 µs → 0.059 ms (+18%)

   Since you already record measured_gpu_delay_ms, any fitted skew-sensitivity curve should use it
   rather than the requested value.

3. Zero-delay baseline in eager mode. At delay 0, host enqueue spread is about 22–25 µs in eager
   vs 3–15 µs in graph mode (unpinned CPUs). So the "delayed" rank sometimes arrives first at delay 0,
   and its eager baseline is about 0.03 ms higher than at larger delays. It may be worth noting that the
   delay-0 row in eager mode already includes some natural skew.

4. Schema alignment (optional). The field names already follow the RFC's record. A
   collective_ordinal/signature key would make the later trace-join phase easier.

5. Rebase. The conflict with main is from #52136 (pydocstyle in ruff) and is mechanical: keep this
   PR's side in both hunks. After that, ruff 0.14.0 reports two of the new rules on this file:

   • D413 on the module docstring's "Example" section;
   • D202 on the blank line after the benchmark_allreduce docstring.

   ruff check --fix handles both.

Not tested: more than 2 GPUs, multi-node, and GB200.

[samnordmann]

Thanks for reproducing this on H100. Addressed in fbd7dd7:

• Rebased on main and fixed both docstring lint issues.
• Added injected-rank residence, measured delay, and max-minus-injected residence to trial summaries and printed medians.
• Explicitly note that the injected rank need not arrive last, zero-delay baselines include natural skew, and sensitivity analysis should use measured delay. The residence gap is diagnostic, not isolated wait time.
• Deferred the optional trace-join key to that follow-up to keep this scoped.

Validation: 22 tests and local hooks pass; 720 fresh GB200 trials cover TP2/TP4, eager/graph and six backends, plus both legacy CLI runs. The new summaries also pass replay over all 10,560 retained trials. GPU timing logic is unchanged.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.