Skip to content

perf(checkpoint): avoid full GC scans during MoE export - #3621

Open
yuhezhang-ai wants to merge 1 commit into
yuhez/perf/moe-checkpoint-direct-fillfrom
yuhez/perf/moe-export-gc
Open

perf(checkpoint): avoid full GC scans during MoE export#3621
yuhezhang-ai wants to merge 1 commit into
yuhez/perf/moe-checkpoint-direct-fillfrom
yuhez/perf/moe-export-gc

Conversation

@yuhezhang-ai

@yuhezhang-ai yuhezhang-ai commented Aug 22, 2026

Copy link
Copy Markdown
Contributor

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 retaining torch.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).

Run #3580 head This PR Reduction
Pair 1, median across ranks 30.334 s 8.368 s 72.4% / 3.62x
Pair 2, median across ranks 31.480 s 8.587 s 72.7% / 3.67x

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).

Run #3580 head This PR Reduction
Pair 1, median across ranks 21.419 s 0.01643 s ~1300x
Pair 2, median across ranks 21.928 s 0.01646 s ~1330x

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

  • Added a focused regression test for both gate/up and down conversions. It verifies generation-zero collection,
    CUDA cache cleanup, output keys, values, and ordering.
  • MoE conversion class: 6 passed.
  • MoE mixin CPU suite: 52 passed, 4 GPU-only skipped, 1 pre-existing Gloo test deselected because the local
    sandbox cannot resolve its loopback address.
  • Qwen3-30B-A3B TE EP8 complete sharded model-save A/B: job 16310166 completed successfully; every run wrote
    the same number of files and bytes.
  • Qwen3-30B-A3B TE EP8 save-preparation A/B: job 16308055 completed successfully.
  • Ruff format/check and git diff --check pass.

Depends on #3580.

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
@yuhezhang-ai
yuhezhang-ai requested a review from a team as a code owner August 22, 2026 01:38
@copy-pr-bot

copy-pr-bot Bot commented Aug 22, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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