Fix hf_ptq.py discarding completed PTQ run on sanity-generate() failure (#2480)

### What does this PR do? Type of change: Bug fix `post_quantize()` in `examples/hf_ptq/hf_ptq.py` ran the optional post-quantization sanity-check `full_model.generate()` unguarded, directly before `export_quantized()`. Any exception raised there aborted the whole run and discarded a completed calibration without exporting a checkpoint. Root cause (traced from [NVBug 6752977](https://nvbugspro.nvidia.com/bug/6752977), DGX Spark GB10 / DeepSeek-R1-Distill-Llama-8B / NVFP4): `get_model()` loads with `device_map="auto"`, relying on `accelerate`'s `infer_auto_device_map`/`get_max_memory()` to decide GPU vs. CPU placement. On DGX Spark's unified-memory single-GPU host, that memory probe under-reports GPU capacity, so part of even an 8B model can land on CPU — and the existing fallback shrinks the GPU budget further (`* gpu_mem_percentage`), compounding it. Calibration survives this because it never invokes the real fake-quant kernel, but the post-PTQ sanity `generate()` does, and NVFP4's dynamic-block-quantize op (`modelopt/torch/quantization/tensor_quant.py`) hard-asserts `amax.is_cuda` with no CPU fallback, so any CPU-offloaded layer crashes there — after ~5.8 hours of calibration, before export. This PR does not attempt to fix the underlying `device_map`/memory-probing behavior (unverified without the actual hardware/logs, which weren't reachable from this environment). Instead it makes the failure mode safe: a failure in the optional sanity check now only skips that check and warns, and export always proceeds, regardless of why `generate()` failed. ### Usage No new API. Behavior change only: `examples/hf_ptq/hf_ptq.py` now completes export even if the post-quantization sanity `generate()` call raises. ### Testing - Added `tests/examples/hf_ptq/test_hf_ptq_args.py::test_post_quantize_export_survives_a_failed_sanity_generate`, which drives `post_quantize()` with a `full_model.generate()` that raises and asserts `export_quantized()` still runs. - Ran `pytest tests/examples/hf_ptq/test_hf_ptq_args.py` (48 passed). - Ran `pre-commit` on the changed files (`ruff-format` reformatted line wrapping only). ### Before your PR is "*Ready for review*" - 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`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌  ### Additional Information Fixes NVBug 6752977. Linked JIRA: OMNIML-5932. 🤖 Generated with [Claude Code](https://claude.com/claude-code)  ## Summary by CodeRabbit - **Bug Fixes** - Quantized checkpoint export now continues when the optional post-quantization generation check fails. - A warning is shown when the generation check cannot complete.  Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
2026-10-02 03:14:52 +08:00 · 2026-09-19 01:21:07 +00:00
parent ed5c5ed369
commit 0a8e70804a
3 changed files with 57 additions and 5 deletions
@@ -69,6 +69,7 @@ Changelog

 - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute.
 - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted.
+- Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless.
 - Fix ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` storing the MoE router at Megatron's ``moe_router_dtype``, which is a routing *compute* dtype, not a storage one. The router now exports at the export ``dtype`` like every other unquantized weight, matching what ``hf_ptq.py`` and the released NVFP4 checkpoints contain; pass ``moe_router_dtype`` to ``export_mcore_gpt_to_hf`` explicitly if you want the old fp32 storage.
 - Fix unified Megatron export writing a second, unreferenced copy of the vocab embedding when a model with MTP layers is exported with pipeline parallelism. The duplicate was never loaded but inflated the checkpoint by the size of the embedding (about 1 GB for Qwen3.6-35B-A3B); re-export to reclaim the space.
 - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings.
@@ -1162,11 +1162,18 @@ def post_quantize(
        pass
    elif model_type != "llama4" and not is_nemotron_vl_model:
        # Our fake quantizer may not be fully compatible with torch.compile.
-        generated_ids_after_ptq = full_model.generate(
-            preview_input_ids,
-            attention_mask=preview_attention_mask,
-            max_new_tokens=100,
-        )
+        # This is a best-effort sanity check: e.g. a `device_map="auto"` load that offloads
+        # part of the model to CPU (seen on unified-memory single-GPU hosts) can make a
+        # quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block
+        # quantization) don't support. Don't let that discard the completed calibration.
+        try:
+            generated_ids_after_ptq = full_model.generate(
+                preview_input_ids,
+                attention_mask=preview_attention_mask,
+                max_new_tokens=100,
+            )
+        except Exception as e:
+            warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}")
    elif is_nemotron_vl_model and tokenizer is not None:
        generated_ids_after_ptq = run_nemotron_vl_preview(
            full_model,
@@ -872,3 +872,47 @@ def test_recipe_superseded_action_is_wired_to_both_flags(monkeypatch):
        assert isinstance(by_dest[dest], RecipeSupersededAction), (
            f"--{dest} lost its deprecation action"
        )
+
+
+# --- post-quantization sanity-check generate() must not block export ----------------------------
+
+
+def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch):
+    """A device-placement issue (e.g. `device_map="auto"` offloading a layer to CPU on a
+    unified-memory single-GPU host) can make the post-PTQ sanity `generate()` raise. That must
+    not discard the completed calibration: export should still run. Regression test for
+    NVBug 6752977."""
+    hf_ptq = _import_hf_ptq(monkeypatch)
+
+    full_model = SimpleNamespace(
+        generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))
+    )
+    export_calls = []
+    monkeypatch.setattr(
+        hf_ptq,
+        "export_quantized",
+        lambda *a, **k: export_calls.append((a, k)),
+    )
+    monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
+
+    args = SimpleNamespace(specdec_offline_dataset=None, verbose=False)
+
+    with pytest.warns(UserWarning, match="Post-quantization generation sanity check failed"):
+        hf_ptq.post_quantize(
+            args=args,
+            full_model=full_model,
+            language_model=full_model,
+            model_type="llama",
+            tokenizer=None,
+            processor=None,
+            preview_input_ids=torch.zeros(1, 4, dtype=torch.long),
+            preview_attention_mask=None,
+            generated_ids_before_ptq=torch.zeros(1, 4, dtype=torch.long),
+            is_nemotron_vl_model=False,
+            first_text_speech_dataset=None,
+            default_padding_side="right",
+            default_pad_token=None,
+            calib_dataloader=None,
+        )
+
+    assert len(export_calls) == 1