[NVBug 6779760] Test NVFP4 export of CPU-offloaded models

[cjluo-nv]

What does this PR do?

Type of change: new tests

Adds NVFP4 to the CPU-offloaded export GPU test (test_export_hf_checkpoint_cpu_offloaded). Until now that test covered only FP8.

NVBug 6779760: on a 64 GB DGX Spark, the NVFP4 playbook (ModelOpt 0.35.0) crashed during export with RuntimeError: Expected all tensors to be on the same device ... cuda:0 and cpu. On GB10, cudaMemGetInfo does not count reclaimable page cache as free memory, so hf_ptq CPU-offloads part of the model.

• Calibration: accelerate moves the offloaded layers onto the GPU, so the quantizers' amax is created on CUDA.
• Export (0.35.0): removing the hooks moves the offloaded weights back to CPU, while amax stays on CUDA.
• Failure: the NVFP4 scale math divides a CPU per-block amax by a CUDA global scale.

Main is already fixed. Device alignment came in #432 (0.39.0), and offloaded models have used the streaming exporter since #2008 (0.46.0). This PR adds a regression guard for that exact configuration.

The fixture also needed a fix. It mapped the parent modules ("", "model") to GPU 0, which makes accelerate load the "offloaded" layer onto the GPU before offloading it. After hook removal, that layer's weight came back on GPU instead of CPU, so the fixture never produced the bug's layout. from_pretrained(device_map=...), which is how hf_ptq loads, does not do this. The fixture now leaves the parents unmapped.

Follow-up (not in this PR): tests/_test_utils/torch/quantization/offload.py::make_cpu_offloaded_model builds the same root-mapped device_map. It is used by test_vllm_fakequant_hf_export.py and test_accelerate_gpu.py.

Testing

All runs were on a DGX Spark (GB10) in nvcr.io/nvidia/tensorrt-llm/release:spark-single-gpu-dev, with torch 2.8.0a0, transformers 5.18.0 and accelerate 1.10.1.

• pytest tests/gpu/torch/export/test_offload_export.py: 5 passed.
• Without the fix: I reintroduced 0.35.0's export behaviour, i.e. batch export with hook removal for offloaded models and no device alignment in get_weights_scaling_factor or get_weight_scaling_factor.
  • New test: only [quant_cfg2-NVFP4] fails, with the bug's RuntimeError ... cuda:0 and cpu.
  • Unmodified test file, same reverted code: 5 passed, so it could not catch this.
• Fixture layout: layer 0 now ends up with weight on cpu and amax on cuda:0 after hook removal, matching from_pretrained(device_map=...). Before the fix, both were on cuda:0.
• End to end (playbook, not this PR's code):
  • 0.35.0 reproduces the crash.
  • main and 0.41.0 export DeepSeek-R1-Distill-Llama-8B NVFP4 successfully under the same CPU-offload split. Dequantized weights of offloaded and GPU-resident layers match BF16 (weight_scale_2 exact, ~0.09 relative error on both).

Not run: other GPU architectures, and CI's environment.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: N/A
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: ❌

Additional Information

NVBug 6779760. The playbook fix (bump the Spark section from 0.35.0 to 0.41.0) is tracked on the bug.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Tests
  • Expanded validation of CPU-offloaded exports across default FP8, layerwise FP8, and NVFP4 configurations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: ee33bd25-f94a-44d6-9fd0-6342cbfef377

📥 Commits

Reviewing files that changed from the base of the PR and between aa4b0d9 and 0696809.

📒 Files selected for processing: 1

• tests/gpu/torch/export/test_offload_export.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Walkthrough

The CPU-offloaded export test now covers default FP8, layerwise FP8, and NVFP4 configurations. It adjusts device-map construction and checks that the exported quantization algorithm matches the selected configuration.

Changes: CPU-Offloaded Export Test

Layer / File(s)	Summary
Parameterize offloaded export test tests/gpu/torch/export/test_offload_export.py	The test description includes NVFP4 calibration. The device map skips root and model parent modules. The test covers two FP8 configurations and NVFP4, and checks the exported algorithm against the selected quant_algo.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Other

Suggested reviewers: fridah-nv

Merge Risk: ⚪ Minimal · up to 06968

The expanded test covers FP8 and NVFP4 CPU-offloaded exports and checks their serialized algorithms; no concrete merge-blocking issue is evident from the inspected changes.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly identifies the main change: adding an NVFP4 export test for CPU-offloaded models. The bug reference is relevant and the wording is concise.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The PR changes only tests/gpu/torch/export/test_offload_export.py; it does not change modelopt package code, examples, pyproject.toml, or requirements files. Added lines contain no torch.load(..…

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.65%. Comparing base (f3678b6) to head (0696809).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2825      +/-   ##
==========================================
+ Coverage   69.33%   78.65%   +9.31%     
==========================================
  Files         654      655       +1     
  Lines       72665    72777     +112     
==========================================
+ Hits        50385    57242    +6857     
+ Misses      22280    15535    -6745     

Flag	Coverage Δ
gpu	59.44% <ø> (+37.61%)	⬆️
unit	59.93% <ø> (+0.06%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[meenchen]

Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.

LGTM: this is a well-motivated regression test for NVBug 6779760, but it edits an existing test file, so a human needs to sign off.

Needs action:

• Sign off on the edits to tests/gpu/torch/export/test_offload_export.py. The _make_cpu_offloaded_model fix and the parametrized quant_algo check are justified, and the existing FP8 cases and their assertions are unchanged.

No action needed:

• The fixture change also affects test_streaming_export_matches_batch_export. The author reports all 5 tests passing, which includes it.
• The author checked that the test catches the bug: with 0.35.0's export behaviour put back, only the NVFP4 case fails, with the original cuda:0 and cpu error. The unmodified test file passes, which shows the old test missed it. I inspected the tests but didn't run them.
• The PR body says the same root-mapped device_map in _test_utils/torch/quantization/offload.py will be fixed in a follow-up PR. That's fine.