Skip to content

Support final normalization of streamed Megatron gradients - #2173

Draft
oran-goodman wants to merge 2 commits into
NovaSky-AI:mainfrom
oran-goodman:codex/streaming-gradient-scale
Draft

Support final normalization of streamed Megatron gradients#2173
oran-goodman wants to merge 2 commits into
NovaSky-AI:mainfrom
oran-goodman:codex/streaming-gradient-scale

Conversation

@oran-goodman

@oran-goodman oran-goodman commented Sep 7, 2026

Copy link
Copy Markdown
  • Allow an optimizer request to scale accumulated gradients after distributed reduction and before clipping or Adam updates. This lets GSPO begin backward computation before the full batch denominator is known.
  • Advertise streaming_accumulation_v1 only for non-colocated Megatron; default requests keep their existing behavior, and unsupported runtimes reject non-default scaling.

Testing: 14 CPU tests passed, including Adam parameters/moments and expert buffers; Ruff 0.11.9 and Black 24.10.0 passed. Qwen3.5-4B GPU qualification passed three updates with an interleaved adapter: gradient norms and outputs matched exactly, as did all 55.1M final main parameter entries and both Adam moment arrays with matched request packing. Different packing changes frozen-model numerics.

The Trajectory integration completed 20 steps on both configurations: streamed run 1050962 versus reference 1050942 reduced median step time 16.5% and weight-publication intervals 19.0%, with 9.8% higher accepted action-token throughput (initial step excluded). Both had zero stale-group rejections. This is one stochastic pair, not bitwise full-batch parity. The temporary runtime was shut down and checkpoints preserved.

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