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.

End-to-end validation

Two production-scale checkpoints were quantized end-to-end using the new disk-offload PTQ path on a single GB200 GPU (189 GiB VRAM).

DeepSeek-R1 (671B, MoE)

Checkpoint	DeepseekV3ForCausalLM, 671B params, 61 decoder layers
Input size	642 GB BF16
Recipe	nvfp4_experts_only-kv_fp8_layerwise_offload
--max_gpu_memory_gb	80
--max_cpu_memory_gb	80
--calib_size / --batch_size	8 / 1
--trust_remote_code	no (built-in transformers)
Wall-clock	40 min 12 s (load ~14 min, calib ~12 min, export ~14 min)
Peak GPU memory	88.9 GB
Peak process RSS	376 GB
Output	40 shards x ~10 GB = 403 GB (~37% compression)

NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 (550B, NemotronH MoE + Mamba)

Checkpoint	NemotronHForCausalLM, ~550B params, 108 decoder layers
Input size	~1.1 TB BF16
Recipe	nvfp4_experts_only-kv_fp8_layerwise_offload
--max_gpu_memory_gb / --max_cpu_memory_gb	170 / 500 and 80 / 80
--calib_size / --batch_size	8 / 1
--trust_remote_code	yes (NemotronHForCausalLM)

	170 GB GPU / 500 GB CPU	80 GB GPU / 80 GB CPU
Wall-clock	41 min 13 s	47 min 16 s
Peak GPU memory	165.8 GB	76.7 GB
Peak RSS (load)	789 GB transient	345 GB transient
Steady-state RSS	~454-496 GB	~50 GB
Output	34 shards x ~11 GB = 365 GB	34 shards x ~11 GB = 365 GB

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: da5e57c8-b99c-49be-bcb3-354d41a6bb8f

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 14.52145% with 259 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.13%. Comparing base (a5117fa) to head (e519853).
⚠️ Report is 2 commits behind head on fridah/omniml-4947-mem-monitor.

Files with missing lines	Patch %	Lines
modelopt/torch/export/unified_export_hf.py	11.11%	224 Missing ⚠️
modelopt/torch/export/quant_utils.py	35.71%	27 Missing ⚠️
modelopt/torch/quantization/plugins/huggingface.py	0.00%	6 Missing ⚠️
modelopt/torch/export/registry.py	33.33%	2 Missing ⚠️

Additional details and impacted files

@@                        Coverage Diff                         @@
##           fridah/omniml-4947-mem-monitor    #2008      +/-   ##
==================================================================
- Coverage                           66.22%   66.13%   -0.10%     
==================================================================
  Files                                 519      519              
  Lines                               60578    60831     +253     
==================================================================
+ Hits                                40120    40230     +110     
- Misses                              20458    20601     +143     

Flag	Coverage Δ
examples	40.42% <14.52%> (+22.08%)	⬆️
gpu	32.71% <6.60%> (-18.26%)	⬇️
regression	15.10% <6.60%> (+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.

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

[Edwardf0t1]

Inline notes on the five items I consider merge-blocking or worth resolving before merge. Overall the streaming-export architecture is sound and the Ultra-550B validation run (161 GB peak VRAM, 15 shards) is convincing evidence the happy path works — these are about CI, an aliasing edge case, and two scope questions.

[Edwardf0t1]

The streaming path drops the tied-weight safety net, so save_file can crash.

The batch path ends postprocess_state_dict with a data_ptr()-based dedup (quant_utils.py ~L1115-1135) that removes any aliased tensor. That's what guarantees safetensors.save_file never sees two keys sharing storage. The streaming path replaces it with name matching against _tied_weights_keys, gated on tie_word_embeddings=True.

_stream_tensor passes tensor.detach().contiguous().cpu() — for an already-contiguous CPU tensor all three are no-ops returning the same object, so aliasing survives into self._buffer and save_file raises RuntimeError: Some tensors share memory.

Concrete failure: a model with tie_word_embeddings=False in config but genuinely shared lm_head.weight / embed_tokens.weight storage, or any aliasing not literally listed in _tied_weights_keys (note transformers v5 changed _tied_weights_keys to a dict of patterns for several architectures, so exact-string matching can silently miss). Both keys land in the same shard buffer → crash after all prior shards are already written.

The comment above is right that data_ptr() is unreliable across the whole export — but it is reliable within a single shard buffer, because the writer holds a strong reference to every buffered tensor until _flush(). A per-buffer check restores parity without the cross-layer staleness problem:

def add(self, key: str, tensor: torch.Tensor) -> None:
    ptr = tensor.data_ptr()
    if ptr in self._buffer_ptrs:   # reset alongside self._buffer in _flush()
        return                      # (or .clone() if you'd rather keep both keys)
    self._buffer_ptrs[ptr] = key
    self._buffer[key] = tensor
    ...

[Fridah-nv]

I adopted the suggested change. In the streaming export case, the tied word embedding case is handled by name filter in _stream_tensor, while the storage need to guard on alias because of these cases:
ExportContext.tied_cache / moe_tied_cache are live dicts, and _export_quantized_weight deliberately re-points weight / weight_scale / weight_scale_2 / input_scale at a previously-processed module sharing source memory — so the downstream data_ptr dedup can collapse them. This is postprocess_state_dict, batch-path only; streaming has no equivalent, so such aliases would reach save_file directly.

Let me know if this makes sense. I'm actually not famliar with the moe_tied_cache part.

[Edwardf0t1]

Unexplained behavior change: dropping or "Deepseek" in architecture from this condition.

Previously any DeepSeek architecture took the AutoModelForCausalLM + assert trust_remote_code path even when transformers exposed the class. Now DeepseekV3ForCausalLM (present since transformers 4.53) falls through to getattr(transformers, architecture)._from_config, and further down model_kwargs2.pop("trust_remote_code", None) fires.

That's a real behavior change for the exact model family this PR targets, it isn't mentioned in the description, and there's no test covering it. It also appears to pull in the opposite direction from _install_transformers_compat_shims(), which exists specifically to make remote-code DSR1 load.

Could you explain the intent, or split this into its own PR so it can be evaluated independently?

[Fridah-nv]

I clean this up as well but need a full run on DSR1 to verify

[realAsma]

Looks great to me!!

Could you please benchmark DSR1/Ultra single GPU PTQ with max calibration (PTQ + export time, GPU memory, GPU utilization, CPU memory etc.) and add that information to the PR description?

[Edwardf0t1]

Re-reviewed at e519853 — all five of my earlier findings are resolved, and I verified the first two by running the checks (ruff format/ruff check clean; recipe docs now 21/21 with no stale rows). One remaining test-coverage gap below.

[Edwardf0t1]

Missing: a streaming-vs-batch key-parity test.

The unit tests added for the alias guard and the data_ptr-under-offload cases are good, and this test is a solid guard against the original meta-tensor bug. But every assertion here iterates only the keys that were written — so it can confirm the written tensors are sane, and can't detect anything missing.

Concretely: if a refactor caused _export_transformers_checkpoint_streaming to skip an entire decoder layer (a seen_keys collision, a _tied_weights_keys pattern over-matching and dropping too much, a _dispatch_export_handler early-return, or the non-decoder pass missing an offloaded lm_head), this test still passes — the surviving shards are non-empty and non-meta.

That's the highest-risk failure mode for this PR, because the streaming path reimplements a lot of _export_transformers_checkpoint + postprocess_state_dict per-tensor rather than sharing it, and the two can drift silently.

The cheap version is to export the same tiny model twice — once offloaded, once not — and diff the key sets:

def test_streaming_export_matches_batch_export_keys(tmp_path):
    # non-offloaded reference
    ref_model = ...  # same tiny llama, plain .cuda(), same quant_cfg + forward_loop
    export_hf_checkpoint(ref_model, export_dir=str(ref_dir))

    # offloaded, exercises the streaming path
    off_model, *_ = _make_cpu_offloaded_model(...)
    export_hf_checkpoint(off_model, export_dir=str(off_dir))

    assert _all_keys(off_dir) == _all_keys(ref_dir)

where _all_keys unions safe_open(...).keys() across every shard. Comparing values too would be even better, but key parity alone catches the whole "silently dropped tensors" class and is what the current assertions can't see.

Two smaller gaps in the same spirit, if you want them in the same pass:

• Multi-shard is never integration-tested. Tiny LLaMA fits in one shard, so model.safetensors.index.json generation inside the real export path is untested end-to-end — only _StreamingShardWriter is tested in isolation. Parametrizing this test with max_shard_size="1MB" would cover it.
• _parse_shard_size has no test.

[kevalmorabia97]

You need to also enable tools unit tests in https://github.com/NVIDIA/Model-Optimizer/blob/main/.github/workflows/unit_tests.yml