perf(checkpoint): avoid full GC scans during MoE export - #3621
Open
yuhezhang-ai wants to merge 1 commit into
Open
perf(checkpoint): avoid full GC scans during MoE export#3621yuhezhang-ai wants to merge 1 commit into
yuhezhang-ai wants to merge 1 commit into
Conversation
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
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.
What this PR changes
Grouped-MoE native-to-Hugging-Face conversion cleans up temporary expert views after both gate/up and down
projections. The allocating CUDA path currently runs a full Python heap scan for every projection. This PR limits
those scans to newly created objects with
gc.collect(0), while retainingtorch.cuda.empty_cache().This is intentionally a small follow-up to #3580. It does not change checkpoint routing,
loading, tensor layouts, or serialization. The affected path is used while preparing TE and other allocating grouped
expert layouts for save/export; allocating quantized conversions can also use the same cleanup sites.
Part of #3576.
Performance
Complete sharded model checkpoint save
A matched local A/B measured a complete synchronous sharded model-weight save with Qwen3-30B-A3B, Transformer
Engine grouped experts, DeepEP, EP8, and 8 H100 80GB GPUs. The timer includes model state collection,
native-to-HF conversion, metadata/planning, the full DCP safetensors write, and the final rank wait. Model loading,
optimizer state, and optional consolidated HF export are excluded. Two alternating pairs ran on the same node
(Slurm job
16310166).Every run wrote 61,069,868,639 bytes in 12 files, including 8 safetensors shards. File counts and total output
size matched between the base and this PR. Peak CUDA allocation was also unchanged at 13.952607 GiB per rank.
The DCP write itself remained about 7.75-8.19 seconds; the end-to-end improvement comes from removing roughly
22 seconds of full-heap GC during conversion.
Isolated native-to-HF preparation
A matched local A/B used Qwen3-30B-A3B, Transformer Engine grouped experts, DeepEP, EP8, and 8 H100 80GB GPUs.
It timed the real native-to-HF conversion and contiguous-view preparation performed before a checkpoint save. Two
alternating pairs ran on the same node (Slurm job
16308055).The old path made 96 full-heap scans per rank and spent 21.38-21.89 seconds in GC. The new path made the same 96
generation-zero scans in about 0.00038 seconds. Peak CUDA allocation was unchanged at 13.858858 GiB per rank.
Both sides produced 2,739 HF keys per rank, including 2,304 expert keys. Rank-local expert fingerprints and CUDA
peaks matched in both A/B pairs.
Validation
CUDA cache cleanup, output keys, values, and ordering.
sandbox cannot resolve its loopback address.
16310166completed successfully; every run wrotethe same number of files and bytes.
16308055completed successfully.git diff --checkpass.Depends on #3580.