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

primus/backends/megatron/patches/mlperf_warmup_patches.py

@@ -222,6 +222,53 @@ def _build_synthetic_iterator(primus_args):
     return MegatronDataloaderWrapper(mock_loader)
 
 
+def _reset_ddp_grad_ready_calibration(models):
+    """Undo the DDP grad-ready calibration that the warmup steps consumed.
+
+    Megatron's gradient buffers 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 is reached again. Warmup steps are batches like
+    any other, so they consume the calibration -- the golden counts end up describing a synthetic
+    step, and the first real step is measured against them.
+
+    Under gradient accumulation the two do not agree, and the bucket then never dispatches at all.
+    Every parameter reports in and the collective is still missing, which surfaces at
+    ``finish_grad_sync`` as "Communication call has not been issued for this bucket (21/21 params
+    have grad available)" -- the full 21/21 being what distinguishes a count mismatch from a
+    parameter that never arrived. Restoring ``is_first_batch`` and clearing both dicts makes the
+    first real step calibrate, which is what would have happened had warmup not run.
+
+    The outstanding-handle drain is a guard rather than part of the fix. Since #1069 installs
+    ``finalize_model_grads_func`` around the warmup steps their reduce-scatters are awaited, so it
+    is expected to report 0; before #1069 it reported 43 of 43 bucket groups on Flux 12B. Draining
+    is still correct here because the handle would belong to a synthetic step whose gradients are
+    about to be discarded.
+    """
+    drained = groups_reset = 0
+    for m in models:
+        groups = list(getattr(m, "bucket_groups", [])) + list(
+            getattr(m, "expert_parallel_bucket_groups", [])
+        )
+        for group in groups:
+            if not hasattr(group, "is_first_batch"):
+                continue
+            handle = getattr(group, "grad_reduce_handle", None)
+            if handle is not None:
+                handle.wait()
+                group.grad_reduce_handle = None
+                drained += 1
+            group.is_first_batch = True
+            group.golden_per_param_grad_ready_counts = {}
+            group.per_param_grad_ready_counts = {}
+            groups_reset += 1
+    _log(
+        f"Reset DDP grad-ready calibration on {groups_reset} bucket groups "
+        f"({drained} outstanding collectives drained)"
+    )
+    return groups_reset
+
+
 def _run_warmup_and_restore(
     *,
     warmup_steps,
@@ -365,6 +412,9 @@ def _run_warmup_and_restore(
     except TypeError:
         optimizer.zero_grad()
 
+    # ---- 11b. Undo the DDP grad-ready calibration the warmup steps consumed ----
+    _reset_ddp_grad_ready_calibration(models)
+
     # ---- 12. Reset counters ----
     megatron_args.consumed_train_samples = 0
     megatron_args.skipped_train_samples = 0

tests/unit_tests/backends/megatron/test_mlperf_patches.py

@@ -697,6 +697,103 @@ def _train_step(fwd, data_iter, mdl, opt, sched, cfg, fwdbwd, iteration=None):
         assert seen == [preexisting]
         assert config.finalize_model_grads_func is preexisting
 
+    def test_warmup_restores_the_ddp_grad_ready_calibration(self, monkeypatch):
+        """Warmup must not leave its own calibration behind for the first real step.
+
+        Megatron's gradient buckets calibrate on their first batch, and from the second on issue
+        the reduce-scatter only once that golden count recurs. Warmup batches consume the
+        calibration, so the golden counts describe a synthetic step. Under gradient accumulation
+        they then never recur: every parameter reports in and the collective is still never
+        issued, which ``finish_grad_sync`` raises as "Communication call has not been issued for
+        this bucket". Installing the grad-finalize callback does not cover this -- it was
+        reproduced on 8x MI355X at accumulation 2 with that fix already applied.
+        """
+        _install_fake_grad_finalize(monkeypatch)
+        _, megatron_args = _install_fake_megatron(monkeypatch)
+        megatron_args.transformer_impl = "transformer_engine"
+        _silence_log_rank_0(monkeypatch)
+        model, optimizer, scheduler = _warmup_actors(monkeypatch)
+
+        param = next(model[0].parameters())
+        group = SimpleNamespace(
+            is_first_batch=False,
+            golden_per_param_grad_ready_counts={param: 1},
+            per_param_grad_ready_counts={param: 1},
+            grad_reduce_handle=None,
+        )
+        model[0].bucket_groups = [group]
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _run_warmup_and_restore,
+        )
+
+        _run_warmup_and_restore(
+            warmup_steps=1,
+            train_step_fn=lambda *a, **k: None,
+            forward_step_func=lambda *a, **k: None,
+            synthetic_iter=iter(()),
+            model=model,
+            optimizer=optimizer,
+            opt_param_scheduler=scheduler,
+            config=SimpleNamespace(finalize_model_grads_func=None),
+            forward_backward_func=lambda *a, **k: None,
+            iteration=0,
+        )
+
+        assert group.is_first_batch is True, "the first real step must calibrate, not inherit"
+        assert group.golden_per_param_grad_ready_counts == {}
+        assert group.per_param_grad_ready_counts == {}
+
+    def test_warmup_drains_an_outstanding_grad_reduce_handle(self, monkeypatch):
+        """A collective left in flight by warmup must be awaited, not handed on.
+
+        Expected to be a no-op now that the grad-finalize callback is installed for the warmup
+        steps, so this pins the guard rather than a live failure: the handle would belong to a
+        synthetic step whose gradients are about to be discarded, and the first real step cannot
+        see why its bucket is busy.
+        """
+        _install_fake_grad_finalize(monkeypatch)
+        _, megatron_args = _install_fake_megatron(monkeypatch)
+        megatron_args.transformer_impl = "transformer_engine"
+        _silence_log_rank_0(monkeypatch)
+        model, optimizer, scheduler = _warmup_actors(monkeypatch)
+
+        class _Handle:
+            def __init__(self):
+                self.waited = False
+
+            def wait(self):
+                self.waited = True
+
+        handle = _Handle()
+        group = SimpleNamespace(
+            is_first_batch=False,
+            golden_per_param_grad_ready_counts={},
+            per_param_grad_ready_counts={},
+            grad_reduce_handle=handle,
+        )
+        model[0].bucket_groups = [group]
+
+        from primus.backends.megatron.patches.mlperf_warmup_patches import (
+            _run_warmup_and_restore,
+        )
+
+        _run_warmup_and_restore(
+            warmup_steps=1,
+            train_step_fn=lambda *a, **k: None,
+            forward_step_func=lambda *a, **k: None,
+            synthetic_iter=iter(()),
+            model=model,
+            optimizer=optimizer,
+            opt_param_scheduler=scheduler,
+            config=SimpleNamespace(finalize_model_grads_func=None),
+            forward_backward_func=lambda *a, **k: None,
+            iteration=0,
+        )
+
+        assert handle.waited is True
+        assert group.grad_reduce_handle is None
+
 
 # ============================================================================
 # Run lifecycle