diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py index a9fa5a1ee..d1abfbd43 100644 --- a/primus/backends/megatron/patches/mlperf_warmup_patches.py +++ b/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 diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py index 6f257cf8a..509269ae1 100644 --- a/tests/unit_tests/backends/megatron/test_mlperf_patches.py +++ b/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