Single gpu disk offload PTQ for DSR1/Ultra

[Fridah-nv]

What does this PR do?

Type of change: New feature, bug fix, new tests

Enables single-GPU PTQ for models too large to fit in VRAM (e.g. Nemotron-Ultra-550B at 1.1 TB BF16, DeepSeek-R1 at 642 GB BF16) by adding accelerate disk/CPU offload support to the HF PTQ example and fixing the export path to correctly handle offloaded models.

G1 Offload-aware unified HF export (modelopt/torch/export/unified_export_hf.py)

The existing _export_transformers_checkpoint removed accelerate hooks before materializing weights, silently writing meta tensors (empty weights) to the checkpoint. Fix:

• _has_accelerate_offload(model)detects any disk/CPU-offload accelerate hook in the model tree.
• _process_quantized_modules_offloaded(model, dtype) new export path for offloaded models: materializes one decoder layer at a time via enable_weight_access_and_writeback, dispatches export handlers inside the context window, and snapshots the layer state dict before hooks re-offload the weights. A second pass collects non-decoder modules that are also disk-offloaded (embed, norm, lm_head) to avoid meta tensors in the returned state dict. Hooks are removed only after the full state dict is assembled.
• Meta-tensor guard in _export_quantized_weight raises RuntimeError on meta input instead of silently corrupting the checkpoint.

G2 Disk-offload CLI (examples/hf_ptq/hf_ptq.py, example_utils.py)

Three new arguments to hf_ptq.py:

• --offload_folder PATH  enable accelerate disk offload; shards spill here.
• --max_gpu_memory_gb N VRAM budget for the accelerate device map.
• --max_cpu_memory_gb N  CPU RAM budget for the accelerate device map.

Validation: --offload_folder is incompatible with --low_memory_mode and --use_seq_device_map.

G3 Streaming shard writer for 80 GB CPU RAM (modelopt/torch/export/unified_export_hf.py)

The G1 path accumulated the entire quantized state dict in CPU RAM before writing
(~764 GiB for Ultra 550B), blocking the 80 GB target.

New streaming path writes shard files layer-by-layer. Peak memory = 1 decoder layer +
1 shard buffer instead of the full checkpoint:

Model	Old peak CPU RAM	New peak CPU RAM
Ultra NemotronH 550B	~764 GiB	~57 GB
DeepSeek-R1	~630 GiB	~55 GB

Key pieces:

• _StreamingShardWriter(export_dir, max_shard_size) buffers tensors up to max_shard_size bytes, flushes to numbered temp files (__shard_part_NNNNN.safetensors), renames to canonical shard names at finalize(), writes
  model.safetensors.index.json. Single-shard exports produce model.safetensors with no index file.
• _postprocess_single_tensor(key, value, ...) per-tensor extraction of postprocess_state_dict logic (KV amax scale, skip/rename, squeeze) for streaming use.
• _parse_shard_size(size) converts "10GB" / "500MB" strings to bytes.
• _export_transformers_checkpoint_streaming(model, dtype, export_dir, max_shard_size) streams decoder layers via enable_weight_access_and_writeback, applies per-tensor postprocessing + name reversal + tied-alias filter, writes shard files directly. Non-decoder offloaded modules and GPU-resident tensors are handled in separate passes.
• export_hf_checkpoint dispatches to the streaming path when _has_accelerate_offload(model) is true; hf_quant_config.json, quant-config name reversal, and config.json update are shared between paths.

export_hf_checkpoint accepts a new max_shard_size parameter (default "10GB") that controls the shard size for both paths.

Supporting changes

• modelopt/torch/quantization/plugins/huggingface.py  get_nemotron_h_decoder_layers now checks both model.backbone.layers (remote-code variant) and model.model.layers (native HF variant), fixing layer discovery for NemotronH when loaded without trust_remote_code.
• modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload.yaml  new recipe combining NVFP4 W4A4 on MoE experts, FP8 KV cache, and layerwise calibration with calib_mutates_weights: false (required for disk-offload compatibility).
• example_utils.py  _FP8BF16Fallback shim: dequantizes block-scaled FP8 expert weights to BF16 for calibration forward passes when the kernels package is unavailable (e.g. DSR1 on nodes without finegrained FP8 kernel support).

Usage

# Single-GPU PTQ for a model too large to fit in VRAM, using disk offload
python examples/hf_ptq/hf_ptq.py \
    --pyt_ckpt_path /path/to/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 \
    --recipe general/ptq/nvfp4_experts_only-kv_fp8_layerwise_offload \
    --export_path /path/to/output \
    --offload_folder /path/to/offload \
    --max_gpu_memory_gb 170 \
    --max_cpu_memory_gb 500 \
    --trust_remote_code \
    --calib_size 8 --batch_size 1 --skip_generate

Testing

Unit tests (tests/unit/torch/export/test_offload_export.py, 7 tests, CPU-only):

• _has_accelerate_offload detection (true/false/nested-module cases)
• _export_quantized_weight meta-tensor guard (raises on meta, passes on real)
• _process_quantized_modules_offloaded with disk-offloaded embed + GPU-resident decoder layer: verifies no meta tensor in returned state dict

GPU integration tests (tests/gpu/torch/export/test_offload_export.py, 2 tests):

• Tiny 2-layer LLaMA with CPU offload: FP8 quantization + export, asserts no meta tensors and valid hf_quant_config.json
• Same with layerwise FP8 (calib_mutates_weights=False): disk-offload path end-to-end

Validation runs (on GB200, single GPU):

• Nemotron-Ultra-550B (1.1 TB BF16): 108 layers, 33.2 min wall time, 161.94 GB peak VRAM, 15 safetensors shards. Export previously crashed with AttributeError: NemotronHForCausalLM has no attribute 'backbone' â fixed by the non-decoder offloaded tensor materialization pass.
• DeepSeek-R1 (642 GB BF16): validated on 4 GPUs with the layerwise offload recipe; single-GPU path blocked by kernels package dependency.

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?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ / ❌ / N/A
• Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

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

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (3)

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5ddc07e9-591d-4b06-91f0-14981a169d26

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fridah/single-gpu-disk-offload-ptq

---

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

[Fridah-nv]

/claude review

[codecov]

Codecov Report

❌ Patch coverage is 23.18339% with 222 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.51%. Comparing base (a77e58a) to head (f762313).

Files with missing lines	Patch %	Lines
modelopt/torch/export/unified_export_hf.py	22.48%	193 Missing ⚠️
modelopt/torch/export/quant_utils.py	32.35%	23 Missing ⚠️
modelopt/torch/quantization/plugins/huggingface.py	0.00%	6 Missing ⚠️

Additional details and impacted files

@@                        Coverage Diff                         @@
##           fridah/omniml-4947-mem-monitor    #2008      +/-   ##
==================================================================
- Coverage                           75.73%   69.51%   -6.22%     
==================================================================
  Files                                 518      518              
  Lines                               58544    58787     +243     
==================================================================
- Hits                                44336    40867    -3469     
- Misses                              14208    17920    +3712     

Flag	Coverage Δ
examples	43.17% <23.18%> (-0.13%)	⬇️
gpu	32.72% <6.57%> (-18.25%)	⬇️
regression	15.10% <6.57%> (+0.04%)	⬆️

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.

[realAsma]

RB: Can you explain to me the overall flow of the changes in this file?

[realAsma]

🤖 Bot comment.

🐝 I’m tracing the file-level flow and its before/after behavior now.

[realAsma]

🤖 Bot comment.

The high-level change is that export_hf_checkpoint() now has two Transformers paths:

• Normal model: keeps the existing flow: process all quantized modules → build/post-process one full state dict → call save_pretrained().
• Accelerate CPU/disk-offloaded model: _has_accelerate_offload() detects the hooks and routes to _export_transformers_checkpoint_streaming().

The new offloaded flow is:

1. Run the same model-level preparation as before (MoE preparation, resmooth/requantize, build quant config, and synchronize relevant amax values).
2. Discover decoder layers, then materialize one layer at a time inside enable_weight_access_and_writeback(..., writeback=False). While the real weights are available, dispatch the existing export handlers to pack quantized weights and reconstruct fused MoE weights.
3. Immediately post-process each tensor (drop quantizer internals, convert KV-cache amax to scales, squeeze scales, reverse Transformers conversion names, and skip tied aliases), move it to CPU, and feed it to _StreamingShardWriter. Exiting the context lets Accelerate offload that layer again.
4. Repeat the same materialize-and-stream treatment for offloaded non-decoder modules such as embeddings, norms, and lm_head, then stream any remaining resident parameters/buffers.
5. _StreamingShardWriter flushes at max_shard_size, renames temporary parts to canonical safetensor shard names, and writes the index. This bounds peak CPU memory to roughly one materialized layer plus one shard buffer instead of a full-model state dict.
6. Call save_pretrained(state_dict={}) only to emit config/custom-code artifacts, then the outer function writes hf_quant_config.json and updates/sanitizes config.json as before.

Two safety pieces support this: Accelerate hooks are no longer removed before offloaded weights are materialized, and _export_quantized_weight() now fails loudly if a meta tensor reaches it. _process_quantized_modules_offloaded() is the analogous full-state-dict helper for direct lower-level calls, but the public top-level offload path uses streaming to avoid the full CPU-RAM accumulation.

[realAsma]

RB: This PR should treat offloading as the special case and consider not offloading as the general case. I feel that this PR does not respect this assumption. Example there are many places where regular path is route to if not _offloaded: ... (is it possible to have a cleaner seperation which does not need to move general model export code under if not _offloaded?)

[realAsma]

🤖 Bot comment.

🐝 I’m checking the branch structure and shared-vs-offload responsibilities now.

[realAsma]

🤖 Bot comment.

Yes—your read is correct. The regular path should remain the straight-line default, with one early special-case dispatch:

if _has_accelerate_offload(model):
    return _export_offloaded_transformers_checkpoint(...)

# existing regular export flow, unchanged
post_state_dict, hf_quant_config = _export_transformers_checkpoint(...)
...
model.save_pretrained(...)

The offload helper should own its unsupported-option warnings, layer materialization/streaming, inline name conversion, and save_pretrained(state_dict={}) artifact workaround. Only genuinely common finalization (quant config + sanitized config.json) needs a small shared helper.

There is also duplicated dispatch today: export_hf_checkpoint() routes offloaded models directly to _export_transformers_checkpoint_streaming(), while _export_transformers_checkpoint() independently detects _offloaded and branches to _process_quantized_modules_offloaded(). That lower-level branch is bypassed by the public path and makes ordinary callers offload-aware. Removing that branch from the general exporter and keeping it inside the dedicated offload path would restore the intended separation and eliminate the repeated if not _offloaded guards.

[realAsma]

@Fridah-nv could you please address my concern - https://github.com/NVIDIA/Model-Optimizer/pull/2008/changes#r3648000772

Here is bot comment: https://github.com/NVIDIA/Model-Optimizer/pull/2008/changes#r3648016031

[Fridah-nv]

Addressed in f762313: the offloaded path now early-returns after its own block, and a _write_hf_export_config helper owns the shared hf_quant_config.json + config.json finalization. The regular export path is now the straight-line default with no guards; the offloaded block is a clean early exit at the top.

[realAsma]

RB: how does this code reach the _process_quantized_modules_offloaded? Can you show me a flowchart of the code flow?

[realAsma]

🤖 Bot comment.

🐝 I’m tracing the offloaded export call path now and will follow up with a flowchart.

[realAsma]

🤖 Bot comment.

Short answer: this branch does not reach _process_quantized_modules_offloaded.

flowchart TD
    A[export_hf_checkpoint] --> B[_has_accelerate_offload]
    B -->|true| C[_export_transformers_checkpoint_streaming]
    C --> D[discover decoder layers]
    D --> E[materialize one layer<br/>enable_weight_access_and_writeback]
    E --> F[_dispatch_export_handler for submodules]
    F --> G[reconstruct fused MoE]
    G --> H[_postprocess_single_tensor]
    H --> I[_StreamingShardWriter.add / finalize]

    B -->|false| J[_export_transformers_checkpoint]
    J --> K[its own _has_accelerate_offload check]
    K -->|false in this call path| L[_process_quantized_modules]

Loading

The offloaded helper is called at line 1354 only inside _export_transformers_checkpoint when its offload check is true. But export_hf_checkpoint calls that function only from the outer _offloaded == false branch, so it cannot select the helper on this entry path.

It is still reachable from direct callers of _export_transformers_checkpoint (for example the speculative exporters) and directly from its unit test. The streaming function instead implements the helper's layer-by-layer materialize/dispatch behavior inline at lines 1171–1211. So for this path, _process_quantized_modules_offloaded is effectively bypassed/redundant.

[realAsma]

RB: "It is still reachable from direct callers of _export_transformers_checkpoint (for example the speculative exporters) and directly from its unit test."

This does not look clean to me. Why can't speculative exporters call the _export_transformers_checkpoint_streaming for offload models?

[realAsma]

🤖 Bot comment.

🐝 I’m checking the speculative exporter contracts and streaming path now; I’ll follow up with the concrete compatibility constraints.

[realAsma]

@Fridah-nv could you also take a look?

[realAsma]

🤖 Bot comment.

Yes—the speculative exporters should dispatch to a streaming-capable path when the model is offloaded. They cannot call this helper verbatim only because its current contract is wrong for them: _export_transformers_checkpoint_streaming writes the entire base+draft checkpoint and HF config directly to export_dir and returns no state dict, while Eagle/DFlash must select and rename only eagle_module.* / dflash_module.*, borrow lm_head.weight for Eagle, validate the draft keys, and emit their own config (plus optional LoRA artifacts). Calling it as-is would first write the wrong full-model artifacts and then collide with the speculative model.safetensors.

The current fallback through _export_transformers_checkpoint is not a good low-memory solution either: it assembles the full state dict and defeats the reason for offload. The clean fix is to factor the layer-by-layer materialization/postprocessing into a reusable tensor iterator/sink, let normal HF export and speculative export provide different key filters/remappers/artifact writers, and dispatch on _has_accelerate_offload in both speculative exporters. Then _process_quantized_modules_offloaded can be removed if it has no remaining production caller; its unit test alone is not a reason to retain it.

[Fridah-nv]

_process_quantized_modules_offloaded is not reachable. export_hf_checkpoint dispatches offloaded models to _export_transformers_checkpoint_streaming and non-offloaded models to _export_transformers_checkpoint, so the if _offloaded: branch inside the latter was unreachable. f762313 removes _process_quantized_modules_offloaded entirely and replaces that branch with NotImplementedError.

[Fridah-nv]

"Why can't speculative exporters call the streaming path?"

They can't call _export_transformers_checkpoint_streaming verbatim. it writes the full base-model artifacts (model-NNNNN-of-NNNNN.safetensors, config.json, generation_config.json) and returns no state dict, while Eagle/DFlash need to filter and rename keys (eagle_module., dflash_module.), borrow lm_head.weight for Eagle, and emit their own config artifacts. A verbatim call would write the wrong key namespace and collide with the speculative checkpoint.

The right fix is to extract the layer-by-layer materialization into a reusable tensor iterator that spec exporters can plug their own key remapper and sink into. We've removed _process_quantized_modules_offloaded (it accumulated a full state dict, defeating offload) and left a NotImplementedError in _export_transformers_checkpoint for the offloaded case. The reusable iterator for spec exporters is scoped as a follow-up happy to open a tracking issue.

[shengliangxu]

It seems this line may remove existing safetensors files because save_pretrained with an empty state_dict will think the already written safetensors files obsolete and so remove them. could you double check?

[Fridah-nv]

This is correct. I forgot to push the local fix e4d69c7. It calls with explicit model.config.save_pretrained() + generation_config.save_pretrained() + shutil.copy2(*.py). This also avoids the MoE shared-tensor crash we hit on DSR1, where save_pretrained fired safetensors' shared-tensor validation on the live model parameters even with an empty state dict.