fix(flux): let the first real step recalibrate the DDP grad buffers

[olehtika]

Targets fix/flux-eval-correctness (#1055), because the failure is introduced by that branch and
#1069 on it already fixes the neighbouring half.

What breaks

Megatron's gradient buckets calibrate on their first batch: _ParamAndGradBucketGroup.reset()
records how many times each parameter registered a ready gradient as
golden_per_param_grad_ready_counts, and from the second batch on register_grad_ready() issues
the reduce-scatter only once that count recurs.

The boundary warmup added in 8e1c22a runs batches like any other, so it consumes that
calibration and the golden counts end up describing a synthetic step rather than the first real
one.

Under gradient accumulation the warmup counts never recur, and the bucket then never dispatches at
all — every parameter reports in and the collective is still missing:

AssertionError: Communication call has not been issued for this bucket (21/21 params have grad available)

The 21/21 is the diagnostic detail: the key set is complete, so the dict equality failed on the
counts, not on a parameter that never arrived.

Why #1069 does not cover it

#1069 installs finalize_model_grads_func for the warmup steps so their reduce-scatters are
awaited. That fixes the accumulation-1 failure (Should not have multiple communication calls outstanding at once) but leaves the calibration untouched.

Measured on 8×MI355X, Flux 12B MXFP6, warmup_train_steps: 2, 300-iteration arms, single node,
overlap_grad_reduce + distributed optimizer:

Warmup state	micro 64, GBS 512, accum 1	micro 32, GBS 512, accum 2
before #1069	dies, multiple communication calls outstanding	dies, has not been issued (21/21)
#1069 only	300/300 steps	dies at step 0, has not been issued (21/21), all 26 ranks
#1069 + this PR	300/300 steps	runs; stops at 200 on the 0.586 gate as expected

The middle row is the point: #1069 is necessary and not sufficient.

The fix

Restore is_first_batch and clear both count dicts after warmup, so the first real step calibrates
— what would have happened had warmup not run. It sits alongside the existing steps that undo
warmup's effects on parameters, the optimizer and the LR scheduler; this is the same class of
cleanup, for state those steps missed.

The outstanding-handle drain in the helper is a guard, not the fix. With #1069 applied it reports
0:

[MLPerf_WARMUP] Reset DDP grad-ready calibration on 43 bucket groups (0 outstanding collectives drained)

Before #1069 the same line reported 43 of 43 bucket groups still in flight, which is how that
half was originally found. Draining is still correct here because the handle would belong to a
synthetic step whose gradients are about to be discarded.

Scope

Both install paths share _run_warmup_and_restore, so _install_train_step_warmup gets the reset
too. That path calibrated on warmup batches before this branch existed and has not been observed
failing, so the reset there is defensive — whether it was safe by design or by coincidence of
matching microbatch schedules, I have not established.

Tests

Two unit tests alongside the ones #1069 added: that the calibration is restored, and that an
outstanding handle is awaited and cleared. Both pass against the container used for the runs above.

Convergence impact of the accumulation-1 path is unchanged: validation trajectories match the
pre-#1055 stack (1.2666 / 0.8609 / 0.7992 versus 1.2662 / 0.8955 / 0.7966 at steps 100/200/300),
and throughput is 77.21 vs 77.34 images/s/GPU.

Made with Cursor

[jasainio]

Approving. The fix is correct, the tests are real, and it solves the problem. Everything below is non-blocking — one finding that I think belongs in the description or a follow-up, and three smaller points.

What I checked

is_first_batch = True plus clearing both dicts is exactly the constructor's initial state for those fields (param_and_grad_buffer.py:224-227), and it composes correctly with what runs next: the first real zero_grad_buffer() calls reset(), which with empty counts records no golden and leaves is_first_batch set, so the first real step calibrates itself as though warmup had never run. models are the DDP-wrapped chunks, so bucket_groups / expert_parallel_bucket_groups resolve. Both new tests fail when the call is neutered.

I also drove the real state machine rather than trusting my reading of it — a _ParamAndGradBucketGroup built without its constructor, with Megatron's genuine reset / register_grad_ready / finish_grad_sync:

accumulation 1, no fix                golden={p0:1, p1:1, p2:1} -> trains
accumulation 2, no fix                golden={p0:2, p1:2, p2:2} -> AssertionError: Communication call
                                                                   has not been issued for this bucket
                                                                   (3/3 params have grad available)
accumulation 2, with this PR's reset  golden={p0:2, p1:2, p2:2} -> trains

Your table and your assertion, from the real code.

The cause is one level more specific, and it is the same omission as #1069

The description says warmup "runs batches like any other, so it consumes that calibration". True, but that alone does not explain the asymmetry the table reports — register_grad_ready only counts when is_last_microbatch is set, so under accumulation the counts should be identical either way.

They differ because config.no_sync_func is assigned in train() at training.py:2636, eleven lines above finalize_model_grads_func at :2647, and the boundary warmup fires at build_train_valid_test_data_iterators (:1032) — before train() at :1073. With no_sync_func unset, schedules.py:570-572 falls back to nullcontext, so every warmup microbatch looks like the last one and registers a grad. Warmup at accumulation 2 records golden counts of 2 per parameter; a real step produces 1; they never match.

So this is not adjacent to #1069, it is the same defect: two callbacks installed side by side in train(), both missing at the boundary, each surfacing as a differently-worded DDP assertion. Worth stating in the description, because it turns the accumulation table from an observation into a prediction.

It also settles the scope question you left open. The in-train_step warmup path runs after train() has installed no_sync_func, so its warmup batches produce the same counts as real ones and the calibration it records is already correct. That path is safe by design rather than by coincidence of matching microbatch schedules, and the reset there is genuinely defensive.

Suggestion, for this PR or a follow-up: install no_sync_func for the warmup alongside finalize_model_grads_func in step 3c, with the same save/restore. The counterfactual arm of the run above confirms it brings golden out at 1 per parameter and the accumulation-2 case trains with no reset at all. Two reasons beyond redundancy: warmup currently never exercises the accumulation path it exists to warm, since no microbatch ever runs with grad sync disabled; and train() at :2632 asserts config.no_sync_func is None when overlap_grad_reduce is on, so the restore has to be in a finally — better done deliberately than discovered. Keep the reset regardless, as the backstop for any other way warmup's control flow can diverge.

Smaller points

The drain half-restores. It mirrors finish_grad_sync's handle handling (:687-688) but not its param_gather_dispatched = False (:667). Harmless while #1069 is on the path, since finish_grad_sync runs every warmup step — but the drain exists precisely for the case where nothing awaited warmup's collectives, and there it leaves param_gather_dispatched True, which makes finish_param_sync skip its dispatch at :446. If it is a guard, it should restore the same fields the real path does.

The reset cannot fail. getattr(m, "bucket_groups", []) and the hasattr(group, "is_first_batch") skip both degrade to a silent no-op, and the function only logs. Correct for FSDP, which has no bucket groups — but it means an attribute rename in a Megatron bump would silently reintroduce a bug whose whole character is being invisible until it asserts. Since the description proposes the log line as a regression canary, consider failing hard when a chunk has bucket_groups but no group was reset.

Test gaps. Nothing covers expert_parallel_bucket_groups, and neither test asserts the reported counts, which is the canary the description leans on.

One non-defect worth recording

After the reset the first real step runs as a first batch, so its reduce-scatter is dispatched from finish_grad_sync rather than overlapped. Identical to a run without warmup, but it now falls inside the measured region. Negligible against a 16000-step budget; noting it so it is not mistaken for a regression later.

Line references are against the pinned Megatron-LM submodule on this branch.