perf(checkpoint): reduce allocating grouped MoE load overhead

[yuhezhang-ai]

What this PR changes

This PR improves three grouped-MoE checkpoint paths left after #3574:

1. MoE reconstruction: write each expert directly into the final grouped tensor instead of building one tensor per expert and stacking them afterward. Use generation-0 garbage collection instead of repeatedly scanning the full Python heap.
2. Transformer Engine and MoK destinations: return the checkpoint-layout view that DCP should fill. TE's views are contiguous; some MoK views are not, but DCP supports both. This avoids allocating a second set of expert tensors and avoids full GC or torch.cuda.empty_cache() while those views remain in use.
3. Gemma4 grouped checkpoints: load expert tensors through transposed views of the model weights. With expert parallelism, each rank reads only its own experts instead of materializing all experts on every rank.

The loader also reports destination preparation, storage read, adapter conversion, and model installation separately.

This PR optimizes the existing load pipeline. General sequential checkpoint streaming remains follow-up work in #3576.

Why the model paths differ

Qwen TE and Gemma4 both transform expert weights, but their checkpoint layouts are different:

• Qwen TE: the HF checkpoint has separate keys for each expert. TE exposes a temporary grouped stack. The adapter turns that stack back into per-expert HF tensors before DCP reads them, then combines the loaded values into TE parameters.
• Gemma4: the HF checkpoint already stores all 128 experts in two grouped tensors per layer. The model stores transposed grouped tensors and also absorbs a small per-expert scale into the down projection.

For Gemma4 EP8, the old path expanded every rank's 16 experts into full 128-expert CPU tensors, read the full grouped checkpoint, and then sliced back to 16 experts. The new path gives DCP transposed, expert-sharded views of the model weights, so each rank reads only its own 16-expert slice. The small scale vector is applied in place after the read.

Performance

Gemma4 26B EP8

Both runs use the same 8×H100 gemma4_26b_a4b_moe_peft release recipe.

Load phase	Before this optimization	Current	Reduction
Destination preparation	165.90 s	0.07 s	99.96%
Checkpoint storage read	275.96 s	26.41 s	90.4%
Adapter conversion	9.37 s	0.01 s	99.9%
Model installation	0.13 s	0.12 s	—
Total model load	451.36 s	26.60 s	94.1% / 17.0x

• Before: job 405235261
• Current: job 405305252, pipeline 63712014

A matched local 8-H100 A/B measured 190.44 s before versus 24.46 s on the first optimized read and 13.87 s warm. Per-rank peak process RSS fell from 86.85 GiB to 7.41 GiB. Rank-by-rank expert probes matched, every rank owned 16 experts, and no meta parameters remained.

For single-device Gemma4, matched warm runs reduced peak process RSS from 64.30 GiB to 47.62 GiB (16.68 GiB / 25.9%). Model-load time itself was approximately neutral; EP is the large speed improvement.

Qwen3 30B TE: corrected speed and memory explanation

The earlier explanation was incomplete. TE creates its virtual grouped tensor with torch.stack(...).transpose(...). After the adapter slices and transposes an expert back to HF layout, that tensor is already contiguous. Therefore the old .contiguous() call reused the TE stack; it did not allocate a copy.

The large destination-setup speedup comes from avoiding repeated full gc.collect() calls. Full GC took roughly 19–23 seconds across the 46 expert projections. Generation-0 collection or no collection takes less than one second. torch.cuda.empty_cache() was not responsible for the old path's lower peak.

The initially proposed blank destinations were fast but always allocated a second set of expert tensors. The final implementation instead reuses TE's already-contiguous checkpoint views:

Code state	Load-only elapsed	Final CUDA allocation / rank	Peak CUDA allocation / rank	What it shows
Before this PR: cat then stack	59.25 s	7.109 GiB	14.093 GiB	Original reconstruction and full GC
Direct fill + from_hf generation-0 GC	35.43 s	7.109 GiB	13.999 GiB	Direct fill saves 96 MiB/rank; cleanup provides most of the time gain
Rejected blank TE destinations	15.21 s	7.109 GiB	20.609 GiB	Fast setup, but duplicates 6.61 GiB/rank
Current: reuse TE views	20.36 s	7.109 GiB	13.999 GiB	Keeps the fast setup without the memory regression

The first three rows are the matched EP8 job 16206220. The final row is job 16219024, using the same cached Qwen3-30B checkpoint, EP8 topology, staged image, and load-only harness. Storage timing varies with filesystem state, so CUDA peak is the controlled comparison. In the final run, the loader's own phase timer reported 18.26 s total: 0.40 s destination setup, 17.42 s storage read, 0.31 s adapter conversion, and 0.13 s installation.

All four states completed with zero meta parameters and matching rank-local expert fingerprints. The current path reduces the rejected implementation's peak by 6.610 GiB/rank (32.1%) and returns exactly to the direct-fill peak.

Nemotron Nano V3 single-device fallback

Matched one-GPU runs used the same 58.82 GB checkpoint, recipe, image, and one-step workload.

Load phase	Before this PR	Current	Reduction
Adapter conversion	246.81 s	146.69 s	40.6%
Total model load	265.78 s	162.65 s	38.8%

The time reduction comes from replacing 46 full Python-heap scans—two expert projections across 23 MoE layers—with generation-0 collection. Direct fill separately removes one final-size reconstruction scratch buffer. The run completed successfully with the same finite first-step metrics as the parent. This older job did not record a clean load-only CUDA peak; that measurement is tracked with the sequential-loading follow-up.

Correctness boundaries

• Pretraining from config is unchanged. Without a checkpoint, load_model() is not called and model initialization remains unchanged.
• Checkpoint views are load-only. DCP fully overwrites the returned views. Save/export conversion still creates contiguous tensors that preserve their current values.
• No model-sized blank destinations are created. DCP can fill contiguous and non-contiguous views, so both TE and MoK reuse their existing checkpoint-layout storage.
• Partial loads preserve initialization. Missing allowed keys are removed before DCP runs, and incomplete expert groups fail validation.
• Quantized conversion is unchanged. Quantized paths keep value-preserving conversion and the existing rebuild behavior.
• Gemma4 EP direct load is topology-gated. It requires materialized expert DTensors sharded only on the expert axis across EP. Unsupported sharding and quantized conversion use the existing fallback.

Validation

• New TE storage test proves checkpoint-layout outputs reuse TE's temporary stack without allocating a second buffer.
• New MoK storage test proves its non-contiguous checkpoint-layout outputs reuse the existing grouped storage.
• New CUDA allocator regression test limits TE destination construction to 70 MiB for a 64 MiB stack; the rejected blank path would require 128 MiB.
• New CUDA allocator regression test limits direct fill to 40 MiB for 16 MiB of inputs plus one 16 MiB output; the old per-expert concatenations would raise the peak to 48 MiB.
• Focused GPU tests: 2 passed in Slurm job 16219024.
• MoE mixin CPU suite: 51 passed, 4 GPU-only skipped, 1 unrelated Gloo test deselected.
• Qwen3 30B TE load-only EP8: job 16219024, peak 13.999 GiB/rank, zero meta parameters, matching expert fingerprints.
• Qwen3 30B MoE checkpoint robustness: pipeline 63370169.
• Gemma4 26B EP 50-step run: job 405305252.
• Gemma4 full checkpoint robustness: job 405377241, pipeline 63721419 — base run plus all four save/reload phases passed.
• Ruff format/check and git diff --check pass.

Builds on #3574. Part of #3576.

[copy-pr-bot]

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.

[yuhezhang-ai]

/ok to test da03590

[yuhezhang-ai]

/ok to test 74f3539

[yuhezhang-ai]

/ok to test 8949e6b

[akoumpa]

/ok to test 2320d5a

[yuhezhang-ai]

/ok to test d6f3631