[alpha0422]

perf: remove 104 host synchronizations per training iteration

Every training iteration blocked the host far more often than it needed to. Nearly all of
those blocks were diagnostic — values produced on device, pulled to the host immediately,
and then not consumed until the next log line. This removes 104 host synchronizations per
iteration without changing any kernel, any collective, any numerics, or any optimizer math.

7 files · +433 / −100 · 5 commits

---

The problem

Measured on rank 0 with a CUPTI-based tool counting driver-domain cu*Synchronize calls
over a 6-iteration region of interest:

	before	after
host syncs / iteration	142	38
distinct sync sites	18	4
dynamic-shape sites (out of scope, see below)	35	35
everything else	107	3

The 35 dynamic-shape syncs come from an MoE dispatcher outside this tree; they are counted
here only so the totals reconcile. The removable population is 107, and 104 of them go.

Where they came from, per iteration:

count	site	what it was
74	rerun_state_machine.py validate_result	NaN/inf check reading a device predicate into Python control flow — one .item() per gradient bucket, on the autograd thread inside the grad-ready hook
8	common_utils.py logical_and_across_model_parallel_group	update_successful — pageable H2D construction plus a host read
7	common_utils.py reduce_max_stat_across_model_parallel_group	grad norm, num zeros, learning rate
5	training.py	torch.tensor([0.0], device='cuda') seeding the loss accumulator
4	common_utils.py calc_params_l2_norm	reading its result
4	timers.py + barrier()	the interval timer: host clock behind torch.cuda.synchronize() and a world barrier, in both start() and stop()
2	rerun_state_machine.py _reduce_any	the rerun decision's own construction and read
2	training.py checkpoint_and_decide_exit	--exit-duration-in-mins, tested every iteration
1	clip_grads.py count_zeros_fp32	.item() on the zero count
1	training.py	get_loss_scale().item()

Three recurring root causes:

1. A device value pulled into Python control flow where nothing consumed it until much
   later (the NaN checks, update_successful).
2. torch.tensor([x], device='cuda') for a host scalar — a pageable host-to-device
   copy, which blocks. Measured 20/20 blocking over 20 repetitions behind pending device
   work; torch.full / zeros / fill_ are 0/20, because the value rides in the kernel
   argument buffer with no copy at all.
3. Diagnostics on the critical path — timing, logging and the exit check reading values
   every iteration that are only ever printed.

---

How they were removed

commit	change	syncs/iter
1	Bank validate_result verdicts on device; merge them into the all-reduce should_run_forward_backward already performs at the loop exit. Detection timing is unchanged — still before optimizer.step().	74
2	Keep the logged stats on device and read the whole bundle — grad norm, num zeros, params norm, loss scale and the per-key loss averages — in one device-to-host copy at the log point.	17
3	Build device scalars with torch.full / zeros / ones instead of torch.tensor([x], device='cuda'); defer the learning rate into the same batched read.	7
4	Time with CUDA events instead of host clock plus torch.cuda.synchronize(), and drop the per-iteration barrier(). Reporting becomes lagged but lossless — every interval reported exactly once, in order.	4
5	Fold the --exit-duration-in-mins vote into the same batched read.	2
	total	104

Three design points worth a reviewer's attention:

• Batching, not just deferring. Deferring each read individually removed the six target
  sites but netted only 10 of 17 syncs: PyTorch elides the sync when the stream is already
  idle, so the old early reads had been hiding the later ones. One batched
  torch.cat(...).tolist() is what actually collapsed them.
• The timer's first report blocks once, deliberately. It has no earlier value to stand
  in for it and the caller divides by it for throughput; without that the run dies with
  ZeroDivisionError on the first log line.
• non_blocking is gated on is_pinned(), not forced. A pageable source ignores the
  flag and stalls anyway while emitting no synchronize call — forcing it would move the
  stall and hide it from the trace.

The 3 remaining removable syncs are load-bearing: the rerun decision (_reduce_any),
update_successful (logical_and), and the single batched logging read.

---

Results

Measured on an internal large-scale MoE pretraining workload; the configuration is not
disclosed here, so treat the percentage as indicative rather than reproducible. Both arms
ran in a single allocation as adjacent phases, because rack placement moves step time by
more than the effect being measured — an early attempt across separate allocations inverted
the sign.

At 256 GPUs, over iterations 40→160:

	ms/iteration
baseline	2646.96
this branch	2595.99
	−1.93%

Measured as wall-clock between log-line timestamps with identical machinery in both arms —
the arms' internal timers differ by construction after commit 4 and are not comparable
across arms. Per-20-iteration windows: five of six branch windows sit below every baseline
window (Mann-Whitney 6v6, one inversion, p ≈ 0.004). This was measured at --log-interval 20, where the baseline pays its logging syncs only one iteration in twenty — the setting
that gives this change the least to remove.

Nsight attribution on a smaller comm-bound configuration shows the gain landing entirely in
launch starvation: launch-starved time −108.3 ms/iteration against a step-time delta of
−105.8 ms, while blocking (98% communication) and dependency-stalled time are unchanged
(−8.0 ms and +0.1 ms). The two deltas agree to within 3 ms.

Expect the benefit to scale with how much exposed host time a configuration has. A
comm-bound run with a long log interval recovers the least; a run with --log-interval 1,
or one whose device work does not already cover the host, recovers more.

---

Correctness

• Iteration-1 lm loss identical between arms at 256 GPUs, with every resolved argument
  identical across the two arms.
• Iterations 1–2: lm loss and params norm bit-identical across all six runs of both
  arms at a smaller scale — spread exactly 0 within each arm and between arms. This is
  the sharp test that the computation is unchanged.
• grad norm and num zeros vary in the last digits — but they vary that much within
  each arm too, and the between-arm mean difference is smaller than the within-arm spread
  at every field. Under 64-way data parallelism NCCL reduction order varies run to run; this
  is a property of the reference, not of the change.
• Learning-rate sequence bit-identical over 200 iterations, all six runs.
• Zero NaN, zero skipped iterations in every run of both arms at both scales — relevant
  because commit 1 changed when NaN verdicts are read.
• Unit tests: 66 checks for the logging/exit changes, 21 for the CUDA-event timer.
  Multi-GPU model-parallel smoke: PP=2, and TP=2 with --sequence-parallel.

Limits of the convergence evidence, stated plainly. Beyond iteration 2 the loss
trajectories cannot be compared: the reference workload disagrees with itself by ~13% mean
and ~37% max relative loss deviation over iterations 50–200, and a held-out reference run
falls outside the envelope of two other reference runs 39–87% of the time. A distributional
test at n=3 has no power there. A bit-exact deterministic convergence test was also
unavailable — --deterministic-mode aborts on this hardware with SM100 backward with head_dim=256 does not support deterministic mode. The evidence is therefore the
bit-identical early iterations and the paired-run analysis above.

---

Reviewer notes — API changes

• count_zeros_fp32 now returns a device tensor rather than a Python float; its return
  annotation changed with it, and MegatronOptimizer.count_zeros() follows.
  LayerWiseOptimizer.count_zeros() forwards the tensor unchanged;
  MimoOptimizer.count_zeros() already reduces to a Python int and is unaffected — the
  batched read accepts host values transparently.
• reduce_max_stat_across_model_parallel_group gained return_tensor, replacing an
  earlier kind-follows-input rule that could not express "a host float that defers".
  Defaults unchanged, so existing callers and tests/unit_tests/test_utils.py are untouched.
• training_log() now returns a tuple (report_memory_flag, exit_duration_done);
  tasks/finetune_utils.py updated.
• Timer reporting is lagged. Timer.elapsed() returns intervals resolved from completed
  CUDA events, so a log line reports slightly older windows than it used to. Every interval
  is still reported exactly once and in order; open windows are drained on stop() and
  bounded at 1024 pending events. Aggregates over many iterations remain exact.
• Exit granularity is now --log-interval rather than every iteration, since the
  --exit-duration-in-mins verdict rides the batched read. For a limit measured in minutes
  this is not a behavioral regression, and it reuses a knob that already exists rather than
  inventing one.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.