fix: keep GLM-5.2 weight broadcast within GPU memory at 4-node scale - #3199
Open
samsja wants to merge 1 commit into
Open
fix: keep GLM-5.2 weight broadcast within GPU memory at 4-node scale#3199samsja wants to merge 1 commit into
samsja wants to merge 1 commit into
Conversation
With 4 train nodes the resident FP32 master shards are ~87 GiB/rank and the per-layer broadcast (full expert gather + fp8 conversion) peaked ~49 GiB above that, OOMing the post-step broadcast once training transients filled the cache. - empty_cache() before broadcast so cached training pools are reusable - drop consumed expert tensors eagerly and quantize into preallocated stacks instead of list+torch.stack copies (~49 -> ~31 GiB transient per layer) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
samsja
marked this pull request as ready for review
August 6, 2026 03:58
samsja
force-pushed
the
fix/glm52-broadcast-memory
branch
from
August 6, 2026 04:03
0be264c to
7c5a576
Compare
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
1 participant
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.
Summary
At 4 train nodes the resident FP32 master shards for GLM-5.2 are ~87 GiB/rank. The NCCL weight broadcast gathers each layer's full routed-expert tensors and fp8-quantizes them, transiently peaking ~49 GiB above that — the startup broadcast fit with ~4 GiB to spare, but the post-step broadcast OOMed because ~42 GiB of training transients sat in allocator caches (incl. per-stream pools) the broadcast's default stream could not reuse.
Two changes:
torch.cuda.empty_cache()before each in-loop broadcast so cached training pools become reusable (measured: reserved drops 129.7 → 93.9 GiB entering the broadcast).convert_tt_layer_to_vllm_kernel: drop consumed expert tensors eagerly and quantize each expert into preallocated stacks instead of list +torch.stackcopies. Per-layer transient peak drops from ~49 GiB to ~31 GiB (measured 135.4 → 117.8 GiB reserved during broadcast).Validation
GLM-5.2 wordle RL on 8 nodes (4 trainer + 2 prefill + 2 decode disaggregated NIXL), seq_len 8192, cp1/ep8, fp8 trainer quantization, optimizer+gradient CPU offload (#3196), sign_sgd: 20/20 steps with in-loop NCCL broadcasts every step, peak reserved stable at 129.8 GiB, no OOM. Layer-by-layer memory probes confirmed the numbers above.
🤖 Generated with Claude Code