Repository navigation
Keep fixed-batch GRPO gradients and optimizer state consistent across replicas - #19
Merged
kaiwenKevinn merged 2 commits intoOct 23, 2026
Merged
Conversation
This was referenced Oct 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 2 of 3; depends on #18 and is followed by #20. Merge in order: #18 → #19 → #20. This PR targets
main, so its current diff includes the publication/request changes until #18 is merged.Review only this part: optimizer and fixed-batch changes.
Adding a Hybrid replica must divide one fixed global GRPO batch while preserving complete prompt groups and each trajectory's weight. This change partitions that batch across Core and Hybrid, carries actual sample counts through tasks and gradient packets, and weights each contribution by its share of the batch. Empty Core partitions still participate in synchronization. Asynchronous gradient packets own their CPU storage.
The shared objective averages valid-token losses within each trajectory, then averages trajectories. PG and KL statistics use the same sample weighting. Zero-mask trajectories contribute zero and remain in the batch. For variable-length trajectories with microbatches larger than one, this intentionally changes the previous numerical objective, including Core-only runs. Clipping, the KL coefficient and Adam remain unchanged. Hybrid statistics travel through the existing gradient transport metadata; the following EHP lifecycle PR incorporates the statistics of gradients actually consumed by Core.
Hybrid optimizer restoration preserves model values, FP32 master parameters, Adam moments, independent per-parameter steps and DP-zero shard contents. The automatic non-fused Torch Adam/AdamW fallback bounds update temporaries per parameter; the supported BF16 distributed-optimizer path also defers gradient preparation. Transient logits and loss references are released before subsequent allocation peaks. A one-time cache release before the first Hybrid optimizer update retains the mitigation introduced after an observed NCCL allocation failure; its necessity has not been isolated by an ablation.
Validation covers the B12/n4/MB3 partition-gradient counterexample, variable-length and zero-mask losses, native MCore/Gloo DP2-to-DP1 optimizer restoration and the next update, task counts, transport metadata and output lifetimes. All 15 targeted cases passed again at the newly consolidated fixed-batch commit. The complete three-PR series passed a fresh 50-update Qwen3-4B run on four GPUs, with 200 complete groups, 800 trajectories, 12 successful Hybrid training batches, matching final publication receipts and clean completion. Production code at each PR boundary retains its previously validated content. A follow-up test-only correction uses asyncio.TimeoutError so the shared serving test also works on Python 3.10.
This PR builds on the publication/request lifecycle PR. Live state activation, resource ownership and consumed-Hybrid statistics integration are completed by the EHP lifecycle PR.
This three-part series supersedes #16 and #17.