Skip to content

fix(megatron): synchronize complete accumulation windows - #2146

Draft
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fix-megatron-overlap-accumulation
Draft

fix(megatron): synchronize complete accumulation windows#2146
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fix-megatron-overlap-accumulation

Conversation

@bvolpato

@bvolpato bvolpato commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Megatron gradient overlap can reduce an incomplete optimizer window when the number of microbatches changes. SkyRL defers finalize_model_grads until optim_step, but DDP backward hooks can still dispatch earlier. The pinned Megatron implementation learns the first window's ready counts: after a two-microbatch window, a one-microbatch window can fail with Communication call has not been issued, while a three-microbatch window can dispatch before its final backward.

This holds DDP synchronization disabled throughout each training schedule, including multiple forward_backward requests in one optimizer window, and explicitly starts asynchronous reductions at optim_step. Schedule callbacks are temporarily cleared so nested schedule contexts cannot re-enable the hooks; callbacks and DDP state are restored on exceptions. Forward-only execution is unchanged.

Related to #2008; this addresses the overlap-enabled accumulation path.

Reproduction and validation

CPU tests cover changing request sizes, shared model configuration, callback restoration, exceptions, and synchronous reduction ownership:

uv run --isolated --extra skyrl-train --extra dev pytest \
  tests/backends/skyrl_train/distributed/test_megatron_grad_sync.py -q
  • 12 CPU tests passed in an isolated Python 3.12/PyTorch 2.11 CPU environment.
  • Ruff, Black, and gitleaks pre-commit hooks passed.
  • Added a four-GPU regression using real Megatron workers and two independent DP=2 groups. It compares gradient norms and updated logprobs for overlap enabled/disabled across 2, 1, and 3 microbatches and a window spanning multiple requests:
uv run --isolated --extra dev --extra megatron pytest -s \
  tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_grad_sync_accumulation.py

The GPU regression has not run on this host, which has no GPUs. Keeping this draft pending distributed validation.

Downsides

Reductions still use the asynchronous DDP path, but they start after all backward calls in the optimizer window. This gives up backward/communication overlap for correctness. Restoring that overlap needs an explicit final-request boundary in the accumulation API.

Risk and rollback

This changes Megatron training synchronization timing, including Tinker accumulation. The GPU parity test is the remaining validation gate. If distributed parity fails, keep overlap_grad_reduce=false and revert the change before continuing training.

Signed-off-by: bvolpato <brunocvcunha@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant