perf(checkpoint): load standard HF safetensors with DCP

[yuhezhang-ai]

What changes

• Standard Hugging Face or external models whose checkpoint tensors already match their model weights now load safetensors with DCP.
• Falcon H1 is an example of this path: it uses the Transformers model implementation and has no NeMo state-dict adapter.
• Dense Llama, Qwen2, and Qwen3 also opt into low-memory DCP on one GPU. Their adapters pass checkpoint tensors through unchanged, apart from tied-weight handling.
• A shared passthrough adapter contains that common identity behavior.

DCP reads checkpoint tensors into model-owned weight storage. It avoids first building the complete checkpoint as a CPU dictionary and then installing that complete state into the model. This is ordinary DCP loading, not grouped streaming.

What still uses full CPU loading

This table describes base-checkpoint initialization after this PR:

Checkpoint case	One GPU	Distributed
Legacy PyTorch .bin checkpoint	Full CPU	Full CPU
Safetensors, standard HF/external model, no dequantization	DCP	DCP
Safetensors, standard HF/external model, dequantization requested	Full CPU	Full CPU
Safetensors, custom adapter with zero or small temporary tensors	DCP	DCP
Safetensors, custom adapter requiring model-sized transformed tensors	Full CPU	DCP with rank-local tensors
Safetensors, custom adapter with dequantization requested	Full CPU	DCP with rank-local conversion
Legacy generic Transformers MoE with separate checkpoint experts and grouped model experts (for example, Mixtral or Qwen1.5 MoE without a NeMo adapter)	Full CPU	Full CPU

The two custom-adapter fallbacks are specific to one GPU. Distributed custom-model loading already uses DCP, because each rank transforms only its local checkpoint shard.

The remaining full-CPU routes at every world size are legacy .bin loading, standard-HF dequantization, and the legacy generic Transformers MoE conversion. I did not find a checked-in recipe that exercises the standard-HF dequantization row: current recipes that request checkpoint dequantization use registered custom models and adapters instead. The final MoE row covers older model families; Mixtral is deprecated for 26.10, and the Qwen1.5 MoE recipe uses the separate Transformers/BitsAndBytes path. We may leave this compatibility fallback unchanged unless a current workload needs it.

Stack and scope

This PR is stacked on #3610. #3610 defines and validates the low-memory contract for custom MoE adapters. This PR covers ordinary HF/external models and the dense Llama, Qwen2, and Qwen3 passthrough adapters.

Performance

Measured HF-model speedup

Falcon H1 7B is the ordinary HF/external-model example for this PR. It uses the Transformers implementation and has no custom NeMo adapter. Before this change, every rank read the complete checkpoint into CPU memory. With this change, DCP reads the tensors needed by each rank directly into model weights.

The same 14.13 GB checkpoint and recipe were measured twice before and twice with this change:

Checkpoint path	Sample 1	Sample 2	Mean
Full CPU baseline	99.41 s	97.69 s	98.55 s
This PR: DCP	9.71 s	12.10 s	10.91 s

The mean checkpoint-load time is 9.04x faster, an 88.9% reduction. On the two DCP runs, destination setup was 0.06 s and final installation was 0.10 s; nearly all remaining time was the storage read.

Baseline jobs: 402367452, 402371059. Treatment jobs: 406906717, 406918235.

The stock 50-step Falcon recipe reproduced its existing step-12 3.13 GiB OOM after loading and training correctly. The established one-step form passed with finite step-0 loss 26.0887, gradient norm 225.1654, and validation loss 11.6139.

Llama, Qwen2, and Qwen3 coverage

The scoped jobs measured these DCP checkpoint loads:

Model	This PR load	This PR job	Matching release job
Llama 3.2 1B	4.44 s for 2.30 GB	406906716	398442178
Qwen2.5 0.5B	3.49 s for 0.92 GB	406906713	398442156
Qwen3 0.6B	3.87 s for 1.11 GB	406906714	398442157

These distributed jobs are correctness coverage, not additional speedup measurements: distributed custom adapters already used DCP before this PR. This PR changes the one-GPU route for these three adapters. The August 15 release jobs also predate the checkpoint-only timer, so their logs do not provide a comparable baseline load duration. No speedup ratio is claimed for these small models.

This CI did not measure controlled peak host RSS, so no numerical memory reduction is claimed.

Correctness and coverage

Exact-SHA scoped CI covered the affected paths:

• Falcon external HF: checkpoint loaded and the one-step job passed.
• Llama 3.2 1B: training and every checkpoint-robustness phase passed. AutoModel reload max KL was 0.0; vanilla-HF and cross-TP reload stayed below their thresholds.
• Qwen2.5 0.5B: complete 50-step CP2 job passed.
• Qwen3 0.6B: complete 50-step job passed.

Pipelines: full matrix, green Falcon rerun, August 15 release pipeline.

The GPU jobs ran on the routing implementation before the final code-only passthrough-base cleanup. That cleanup only moved the identical identity methods from three adapters into one shared base. On the current head:

• 208 checkpoint tests passed; 5 CUDA-only tests skipped.
• 36 adapter capability/storage tests passed.
• Real local safetensors reads cover matching keys, mapped HF keys, and all three passthrough adapters.
• The quantized-HF guard remains covered and retains the documented full-CPU conversion path.

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