perf(checkpoint): bound quantized DCP load memory

[yuhezhang-ai]

What this changes

This PR makes FP8 checkpoint conversion use a small part of the model at a time instead of keeping temporary tensors for the complete checkpoint alive at once.

The parent PR (#3616) already chooses DCP according to whether dequantization is actually needed. Two expensive cases remain for Mistral3/Devstral FP8 checkpoints:

• Distributed 123B model: DCP reads rank-local FP8 tensors, but the adapter retains all temporary FP8 tensors until the complete read finishes. On TP=8, that peaks at 71.95 GiB of GPU memory per rank.
• Single-GPU 24B VLM: the safe fallback first builds and converts the complete checkpoint in CPU memory, then installs it into the model. This uses about 66.5 GiB of process memory in addition to the GPU model and makes loading much slower.

This PR adds an optional adapter hook for checkpoints that require dtype or layout conversion:

1. The adapter describes one self-contained part of the checkpoint.
2. DCP reads only that part.
3. Non-quantized tensors are written directly into model weights. FP8 tensors are converted directly into their final BF16 or FP32 model weights.
4. Temporary tensors are released before the next part is read.

Ordinary adapters do not implement the hook and keep their existing loading path.

The Mistral3 adapter groups at most eight text-decoder layers per part. It supports both the causal-LM layout and the untied full-VLM layout used by mistralai/Devstral-Small-2-24B-Instruct-2512. In the VLM layout, the vision encoder and multimodal projector load directly into their final model weights rather than becoming additional temporary copies.

FP32 master-weight configurations use the same bounded path. The model weights already have FP32 storage when loading begins, so each FP8 checkpoint part is converted directly into those FP32 destinations. The later FSDP mixed-precision policy can still use BF16 for compute.

Safety boundaries

The new path is considered only when all of the following are true:

• this is the initial model load from a safetensors checkpoint;
• dequantization is actually required;
• one complete model state is being loaded, without a requested key subset; and
• the model adapter explicitly provides the load parts.

The Mistral3 implementation additionally requires a complete decoder, BF16 or FP32 final weights, and the supported per-tensor FP8 checkpoint format. Pipeline-parallel ranks that own only part of the decoder and tied VLM checkpoints keep their existing fallback. Unsupported adapters and dtypes are unchanged.

The shared loader rejects missing checkpoint tensors, repeated checkpoint requests, repeated model destinations, and incomplete model coverage instead of allowing a partial load.

Measured results

Devstral 123B, TP=8 + LoRA

These runs used the same official mistralai/Devstral-2-123B-Instruct-2512 checkpoint, staged image, TP=8 + LoRA setup, and one-token forward pass.

Code path	Slurm job	Model load	Peak allocated GPU memory/rank	Peak reserved GPU memory/rank
Parent #3616: retain all rank-local FP8 destinations	16296055	90.08 s	71.95 GiB	73.35–73.36 GiB
Intermediate: one decoder layer per part	16295655	101.50 s	30.25 GiB	30.37 GiB
This PR: up to eight decoder layers per part	16296372	91.16 s	30.44 GiB	31.48 GiB

The final path reduces peak allocated GPU memory by 41.51 GiB/rank (57.7%). Model-load time changes by +1.08 seconds (1.2%), so the 123B result is a memory/OOM improvement with loading speed preserved, not a speed claim.

The final run used 12 DCP parts; its largest temporary part was 1.29 GiB/rank. It completed with 1,852 sharded parameters, trainable LoRA weights, and finite logits.

Devstral Small 2 24B VLM, one H100

This is the full Mistral3ForConditionalGeneration VLM checkpoint, including its vision encoder and multimodal projector. The repository recipe uses text-only SQuAD data and normally runs in CI on one multi-GPU node; this benchmark deliberately used one GPU to exercise the path that previously built the complete checkpoint in CPU memory.

Both measurements used the same checkpoint files after download. The comparison uses the fastest observed full-CPU run, making the speed comparison conservative.

Code path	Slurm job	Model load	Peak CPU/process memory	Peak GPU memory	Peak CPU + GPU memory
Parent #3616: full checkpoint converted on CPU first	16301212	56.57 s	66.48 GiB	44.73 GiB	111.21 GiB
This PR: DCP load and conversion in six parts	16301550	14.40 s	5.19 GiB	48.87 GiB	54.06 GiB

For this real single-GPU VLM case, the new path is 3.9x faster. Peak CPU/process memory falls by 61.29 GiB, and peak CPU + GPU memory falls by 57.15 GiB (51.4%). GPU memory increases by 4.14 GiB because one bounded FP8 decoder group is temporarily present alongside the BF16 model; that allocation is released before the next group.

The loader read the 24.02 GB checkpoint in 12.65 seconds, including 12.55 seconds of storage reads and 0.02 seconds of FP8 conversion. The final one-token forward produced finite logits.

CPU/process and GPU memory were sampled together on a discrete-memory H100 system. Their sum is useful for estimating total loading pressure, but it is not a direct DGX Spark unified-memory measurement.

Validation

• Real Devstral 123B TP=8 + LoRA load and forward: Slurm 16296372
• Real Devstral Small 2 24B full-VLM load and forward on one H100: Slurm 16301550
• Exact tiny TP=2 check: Slurm 16295582; every reconstructed parameter matched, including 16 real DTensor shards, and the forward was finite
• Focused checkpoint, Mistral3 adapter, and VLM model unit suites: passed
• BF16 and FP32 destination conversion parity and grouped-load unit coverage: passed
• Ruff format/check and git diff --check: passed

Addresses the Devstral FP8 loading-memory case discussed in #2114.

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