Repository navigation
[fix] GTP+DDP potential data-race: publish DDP params before GTP prefetch reads them - #6388
Conversation
BUG: GTP prefetches ahead of the module owning the weight, and its all-gather input is param.data -- a view into the DDP bucket buffer. With --overlap-param-gather that bucket may still be mid-gather, so GTP silently gathers the PREVIOUS iteration's weights. Wrong results, no crash. FIX: backends attach a callable under PARAM_READY_ATTR; consumers call ensure_params_ready() before reading. DDP registers one per bucket group and stays GTP-agnostic. GTP calls it on forward gathers only. NOT COVERED: CUDA-graph capture. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
/claude strict-review |
|
/ok to test d1bee25 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Light review — the fix itself looks correct. I traced _BucketParamReadiness.__call__ against _ParamAndGradBucketGroup.finish_param_sync/start_param_sync and the hot-path / owns_schedule / in-flight branches all hold; the assert not bucket_group.param_gather_dispatched in the elif is implied by the hot-path guard above it, and the lazy dedup in ensure_params_ready is correct for A,B,A, A,B,B, and A,B,C orderings. Two things worth addressing, neither a correctness bug in the fix:
1. The "recompute is excluded" guarantee only covers the FP8/FP4 path
generalized_tensor_parallelism.py:1235 gates on in_fp8_activation_recompute_phase(), a TE flag set only by activation_recompute_forward(...) — i.e. only when recompute routes through te_checkpoint. megatron/core/recompute.py:133 picks that path only under config.fp8 or config.fp4; BF16/FP16/FP32 recompute goes through tensor_parallel.checkpoint (random.py:585), which sets Megatron's own IS_CHECKPOINTING flag and leaves the TE flag False.
So under BF16 + --recompute-granularity + GTP, ensure_params_ready() does run inside backward — contradicting the code comment ("FORWARD ONLY … recompute is a forward replayed inside backward") and the docstring of _worker_recompute_skips_readiness ("recompute must NOT ask DDP to publish parameters").
This is not a new hazard: _make_forward_pre_hook fires on every module forward including the recompute replay, so DDP already calls finish_param_sync there today, and the specific grad-aliasing hazard cited cannot occur in BF16 — reuse_grad_buf_for_mxfp8_param_ag asserts fp8_param_gather (distributed_data_parallel_config.py:259). But note finish_grad_sync resets param_gather_dispatched = False (param_and_grad_buffer.py:804), so on the BF16 path a mid-backward readiness call can take the "never dispatched => publishing STARTS a gather" branch — exactly what the comment says is being avoided.
Either widen the guard, or narrow the claim. To actually make it forward-only:
from megatron.core.tensor_parallel.random import is_checkpointing
...
if fwd and not in_fp8_activation_recompute_phase() and not is_checkpointing():
ensure_params_ready(weights)(is_checkpointing() covers the BF16 tensor_parallel.checkpoint window; it is set for both the forward and the recompute leg, so if you only want to skip the replay you would need a recompute-phase-specific flag instead.) If the BF16 behaviour is intentional and benign, please say so in the comment and the test docstring rather than stating recompute never publishes — as written both read as a guarantee the code does not provide, and test_recompute_forward_does_not_request_publication will keep passing regardless since it monkeypatches the TE flag directly.
2. The GPU repro tests skip in CI
test_gtp_ddp_param_sync_race.py calls _run_distributed(worker, WORLD) with WORLD = 4, and _run_distributed does pytest.skip when torch.distributed.get_world_size() != required_world_size (gtp_test_utils.py:52). CI launches unit tests at --nproc_per_node 8 (tests/unit_tests/run_ci_test.sh:135), so all four tests in this file skip in the tests/unit_tests/**/*.py bucket. That is the existing convention for the whole generalized_tensor_parallel/ directory (all 31 _run_distributed call sites there pass 4), so this PR did not introduce it — but it does mean the real 4-rank stale-value repro, the strongest part of this PR, will not guard against regression in CI. The CPU tests in tests/unit_tests/distributed/test_param_readiness.py do run, so the protocol and callback branches stay covered. Worth confirming you are happy with that, or adding a sub-group variant so the repro runs on the 8-GPU runner.
| # Re-wrapping a model chunk: a previous DDP may have left a marker | ||
| # pointing at ITS bucket group. This DDP owns the parameter now and | ||
| # publishes nothing, so the stale callback must go. | ||
| delattr(param, PARAM_READY_ATTR) |
There was a problem hiding this comment.
Do we need to consider the case that rewraps the same model chunk after changing a parameter’s requires_grad value? The current cleanup only visits parameters included in the new DDP bucket lists. A parameter managed by the previous DDP but excluded from the new buckets after requires_grad=False would retain its old PARAM_READY_ATTR.
There was a problem hiding this comment.
sorry but when should it happen for the case that rewraps the same model chunk after changing a parameter’s requires_grad value?
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test be0a522 |
… reads (NVIDIA#6388) DDP's overlapped param all-gather publishes a bucket group only in the owning module's forward pre-hook; GTP's one-block-ahead prefetch reads param.data before that hook runs, racing the gather. DDP now stamps a _BucketParamReadyCallback on every param it publishes asynchronously (PARAM_READY_CALLBACK_ATTR), and GTP calls ensure_params_ready(weights) before its forward consume/prefetch. Also aliases in_fp8_activation_recompute_phase as in_activation_recompute_phase at the GTP call sites, since the flag is set for BF16 recompute too. (squashed cherry-pick of unmerged PR NVIDIA#6388 @ be0a522; residual deltas vs the PR head are dev-vs-main context only)
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31722899402 |
…etch reads them (NVIDIA#6388) Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com> (cherry picked from commit 87cc590)
…etch reads them (NVIDIA#6388) Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com> (cherry picked from commit 87cc590)
Summary
[fix] GTP+DDP potential data-race: publish DDP params before GTP prefetch reads them
BUG: GTP prefetches ahead of the module owning the weight, and its all-gather input is
param.data-- a view into the DDP bucket buffer. With--overlap-param-gatherthat bucket may still be mid-gather, so GTP silently gathers the PREVIOUS iteration's weights. Wrong results, no crash.FIX: backends attach a callable under
PARAM_READY_ATTR; consumers callensure_params_ready()before reading. DDP registers one per bucket group and stays GTP-agnostic. GTP calls it on forward gathers only.Note
Originally reported by @JF-D and has one preliminary fix in #6232; both sign-offs are on the commit.
This MR steps further from #6232:
DDP is GTP agnostic. distributed_data_parallel.py has zero GTP references (grep-verified).disable_forward_pre_hook(param_sync=False)— the baseline returns early and never drains an in-flight gather; ours drains it (with skip_next=True).Stale marker on re-wrap— the baseline's persists and can point at the wrong bucket group; ours clears it.Weakrefs— the baseline holds bucket_group strongly; ours holds both weakly.Hot path ~2.7× cheaper— 190 ns vs ~517 ns (the baseline rebuilds a list+set per call and checks graph-capture before the published check). ≈ 5 ms → 1.8 ms/iter at 600 weights × 16 microbatches.Pre-checks