perf(checkpoint): bound GPT-OSS MXFP4 loading

[yuhezhang-ai]

What this changes

GPT-OSS MXFP4 checkpoints currently take the full-CPU fallback on one GPU:

1. Read the complete compressed checkpoint into CPU memory.
2. Convert all expert weights to BF16 on CPU.
3. Copy the complete BF16 model into GPU weight memory.

This PR removes that fallback for the maintained GPT-OSS adapter on one GPU:

• Ordinary tensors load directly into their final model weight memory.
• DCP reads one decoder layer's compressed MXFP4 expert tensors at a time.
• That layer is converted on GPU and copied into its final model weights.
• The compressed tensors are released before the next layer is read.
• Backend-only _extra_state entries remain initialized because they are not stored in Hugging Face checkpoints.

The result is 25 load parts: one part for ordinary tensors and one part for each of the 24 decoder layers. The largest
compressed temporary allocation is 0.39 GB. Final expert weights may be BF16 or FP32.

This uses the bounded-load interface added by the parent PR (#3619). It does not add another checkpoint framework or
change distributed GPT-OSS loading. Partial decoders, distributed expert tensors, and unsupported model dtypes keep
their existing path.

Why

gpt_oss_20b_single_gpu_peft.yaml is an active one-GPU recipe targeting DGX Spark. Loading a complete BF16 CPU model
before filling the GPU model is especially expensive on a system where CPU and GPU allocations share physical memory.
It is also slow because MXFP4 conversion runs on CPU.

Real GPT-OSS 20B result

Slurm job 16311724 loaded the official openai/gpt-oss-20b 13.76 GB MXFP4 checkpoint on one H100. Branch, parent,
and branch ran as fresh processes in the same non-exclusive one-GPU allocation. The checkpoint was already present in
a local snapshot and the harness used offline mode, so downloads are not included.

Code path	Model load	Peak process RSS	Peak CUDA during load	Simultaneous RSS + CUDA
Parent #3619	35.21 s	50.14 GiB	40.58 GiB	89.10 GiB
This PR, first run	7.46 s	3.49 GiB	44.36 GiB	45.66 GiB
This PR, second run	5.50 s	3.47 GiB	44.85 GiB	46.17 GiB

Using the more conservative first branch run, checkpoint loading is 4.7x faster. Peak process RSS falls by
46.65 GiB (93%), and simultaneous process RSS plus CUDA allocation falls by 43.44 GiB (49%).

CUDA allocation rises by 3.78 GiB during loading because one layer's compressed tensors and conversion workspace now
live on GPU. That is much smaller than the removed CPU model copy, so the combined footprint is still roughly halved.

The parent spent 28.56 seconds converting MXFP4 on CPU. This PR completed all GPU-side conversion in 0.36 seconds;
the remaining 6.76 seconds in the first run were DCP storage reads.

All three runs reported the same parameter L2 norm (256972.9989) and completed a finite one-token forward pass.

Scoped recipe CI

The active gpt_oss_20b_single_gpu_peft release recipe passed all 50 training steps on one GB10/DGX Spark GPU in
job 407428835. It loaded the checkpoint in 25
parts with a largest compressed temporary allocation of 0.39 GB, reported the same parameter norm
(256972.9989), and kept finite loss and gradients throughout training. Reported training memory stayed near 45 GiB.

Checkpoint loading took 401.90 seconds in this CI run, of which 387.50 seconds was reading from the S3-mounted CI
cache. This run is correctness and memory-safety evidence; the local-snapshot H100 A/B above is the speed comparison.

Validation

• ruff format --check and ruff check on the changed files
• 46 focused GPT-OSS adapter and checkpoint-routing tests
• Production safetensors test proving DCP is used and the full-CPU loader is not called
• BF16 and FP32 expert destination tests
• Exact key-coverage, _extra_state, partial-decoder, unsupported-dtype, and distributed-fallback tests
• Official GPT-OSS 20B one-GPU load and finite forward in Slurm job 16311724
• Active GPT-OSS 20B one-GPU release recipe on GB10, including 50 training steps

[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

[copy-pr-bot]

/ok to test

@yuhezhang-ai, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/