Skip to content

Keep fixed-batch GRPO gradients and optimizer state consistent across replicas - #19

Merged
kaiwenKevinn merged 2 commits into
NetX-lab:mainfrom
tcyjohn:codex/final-optimizer-20261022
Oct 23, 2026
Merged

kaiwenKevinn merged 2 commits into
NetX-lab:mainfrom
tcyjohn:codex/final-optimizer-20261022

Conversation

@tcyjohn

@tcyjohn tcyjohn commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

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.

@kaiwenKevinn
kaiwenKevinn merged commit 26f721c into NetX-lab:main Oct 23, 2026
3 checks passed
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.

2 participants