perf(checkpoint): avoid full CPU loads for single-GPU MoE

[yuhezhang-ai]

Summary

This replaces #3601, which GitHub automatically closed when its head branch was renamed. The old draft had no human reviews or discussion to migrate.

This PR lets two real single-GPU MoE configurations use AutoModel's existing DCP checkpoint loader instead of first building the complete checkpoint on CPU:

Before:
safetensors -> complete CPU checkpoint -> model-specific conversion -> GPU model

This PR:
safetensors -> DCP -> existing GPU model weights

Ling still needs a small exception for its fused attention weights:

fused QKV checkpoint tensor -> temporary fused QKV tensor -> final Q, K, and V model weights

There is no new streaming executor, load-group API, or per-adapter iteration protocol. The implementation reuses one standard DCP load and the adapters' existing to_hf / from_hf conversion.

Adapters opt into this route with one capability: supports_low_memory_dcp_load. It means that most checkpoint tensors load directly into the model's existing weight memory, while any remaining temporary tensors are small enough to be safe on one GPU. Adapters that need another model-sized set of tensors keep the CPU fallback.

This is stacked on #3580. That PR makes the existing MoE conversion path faster and more memory-efficient. This PR changes which single-GPU configurations can skip the full-CPU-checkpoint fallback entirely.

Why Nemotron can load directly

Nemotron Nano V3's one-GPU recipe is configured with TE experts and a DeepEP dispatcher. However, the MoE layer cannot use expert parallelism with one process, so it deliberately constructs ordinary GroupedExperts instead.

The adapter previously looked only at the configured backend. It therefore assumed the expert tensors were TE temporary stacks and selected the full CPU fallback, even though the actual one-GPU model has ordinary grouped parameters that DCP can fill directly.

This PR makes the shared MoE adapter check match the MoE constructor:

• TE + EP dispatcher + one process: ordinary grouped parameters; direct load is safe.
• TE + EP dispatcher + multiple processes: TE grouped storage; keep the existing conversion path.
• Non-EP dispatcher: ordinary grouped parameters; direct load is safe.
• MoK: keep its existing special handling.

Nemotron's expert checkpoint tensors are non-contiguous transposed views of the grouped model weights. DCP writes through those views into the model's existing memory.

Why Ling no longer needs load groups

Ling's expert and ordinary weights already have safe DCP destinations. Only its checkpoint QKV layout differs: the checkpoint stores one fused tensor while the model owns three Q, K, and V parameters.

The normal adapter path already handles this correctly:

1. Create the fused-QKV destinations (about 0.23 GiB total for Ling-mini-2.0).
2. Let DCP load the checkpoint once.
3. Split the loaded fused tensors into the final Q, K, and V weights.

The adapter now declares that this small temporary state is safe, so the loader does not materialize the complete 30.28 GiB checkpoint on CPU.

The discarded prototype loaded one dependency group at a time and reduced Ling peak host RSS further, from 11.27 GiB to 2.75 GiB. It required more than 1,000 lines of generic executor, validation, and adapter-specific planning code. The standard DCP implementation is slightly faster in the matched measurements below, keeps host memory 80.6% below the parent, and removes that framework/API complexity.

Small shared cleanup

AutoModel asks adapters for their Hugging Face checkpoint key names before sharding. The base implementation previously performed this keys-only operation with real model tensors, so an allocating adapter could accidentally convert real weights just to discover names.

The base adapter now performs key discovery with shape-only meta tensors. This also replaces Gemma4's model-specific version of the same workaround.

Performance

Nemotron Nano V3, one GPU

The 58.82 GiB checkpoint files were not already in the compute node's RAM cache (mincore: 0.0 GiB resident before construction). The run used an H100. The timed loader phase includes DCP planning, storage access, destination writes, and synchronization; it is not a pure filesystem-read timer.

Code state	Job	Checkpoint loader	Full construction	Peak host RAM (RSS)	Peak GPU allocation
#3580 parent: full CPU fallback	16225281	157.92 s	160.814 s	114.858 GiB	58.831 GiB
This PR (1933190d5)	16275130	28.64 s	30.005 s	6.037 GiB	58.840 GiB

Compared with the parent, this is:

• 5.51x faster in the checkpoint loader.
• 5.36x faster through full model construction.
• 108.82 GiB / 94.7% lower peak host RSS.
• About 9.3 MiB of GPU peak above the final 58.831 GiB model.

The run loaded all parameters, left no meta parameters, and reproduced the established pretrained embedding, expert gate/up, expert down, and shared-expert fingerprints.

Ling-mini-2.0, one GPU

Job 16275108 used the same 30.28 GiB Hugging Face checkpoint files, H100, staged image, and alternating fresh Python processes. The initial treatment process placed the files in the compute node's RAM cache and is excluded from the means.

Code state	Checkpoint loader, two samples	Mean loader	Mean construction	Peak host RAM (RSS)	Peak GPU allocation
#3580 parent (d6f363166)	26.77 s, 25.66 s	26.22 s	27.63 s	58.218 GiB	32.378 GiB
This PR (1933190d5)	12.81 s, 12.65 s	12.73 s	13.66 s	11.273 GiB	32.378 GiB

Compared with the parent, this is:

• 2.06x faster in the checkpoint loader.
• 2.02x faster through full model construction.
• 46.95 GiB / 80.6% lower peak host RSS.
• No increase in peak GPU allocation.

Every sample produced identical Q/K/V, expert, and routing-bias probes, the same forward checksum (176.02343750) and top token (670), finite logits, and no meta parameters.

Scope

The measurable route change is for single-device, non-quantized custom-model safetensors initialization.

• Ling and Nemotron Nano V3 now skip the full CPU checkpoint described above.
• The shared one-process TE/EP correction also applies to MoE adapters whose other tensors already satisfy the low-memory DCP requirement: Qwen3-MoE, Qwen3-Next, Qwen3-Omni-MoE, ERNIE-MoE, GLM-MoE, HY-V3, HY-MT2, and Laguna.
• Multi-GPU Ling and Nemotron already use the standard rank-local DCP path. Their distributed behavior is unchanged by this PR.
• Non-TE grouped backends that support direct loading were already on the standard DCP path.
• Quantized initialization, legacy .bin checkpoints, and adapters with full-size transformations keep their existing fallback and remain follow-up work under Roadmap: checkpoint I/O architecture, correctness, performance, and memory #3576.
• Pretraining from config is unchanged because it does not load a base checkpoint.

Validation

• Focused checkpoint, adapter, Ling, Nemotron, and shared MoE tests on the renamed capability: 322 passed, 9 skipped.
• Model infrastructure and Gemma4 adapter tests: 39 passed, 30 skipped (CUDA-only cases skipped on the CPU test host).
• Ruff format/check and git diff --check pass.
• Nemotron cold-start job 16275130: exact rebased SHA, 0.0 GiB initial page residency, expected fingerprints, no meta parameters, and 9.3 MiB GPU peak overhead.
• Ling matched A/B job 16275108: exact rebased SHA, identical parameter/forward probes across parent and treatment, finite logits, and no meta parameters.
• Scoped CI on the behavior-equivalent benchmark head (1933190d5): pipeline 63890795 passed both nemotron_nano_v3_singlegpu_lora and the distributed nemotron_nano_v3_hellaswag checkpoint-robustness job. The restored AutoModel logits matched exactly (max KL 0.0), and the consolidated checkpoint loaded by vanilla HF remained within tolerance (max KL 0.01525, threshold 0.07). The subsequent change only consolidates and renames the capability flag.

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