Fix HF PTQ empty-init dtype kwargs (#1857)

## Summary

Fixes NVBug 6359821: `hf_ptq.py` can fail for remote/custom
architectures like `DeciLMForCausalLM` when dtype-related kwargs are
forwarded into model construction paths that do not accept them.

This change keeps the fix scoped to the observed DeciLM/Llama Nemotron
path. It resolves the init config used for empty-weight construction,
derives dtype consistently from the resolved config, forwards the
supported dtype kwarg for the DeciLM empty-weight probe, and drops
unsupported dtype forwarding from the DeciLM real `from_pretrained()`
load.

NVBug: https://nvbugspro.nvidia.com/bug/6359821

## Validation

- `pre-commit run --files examples/hf_ptq/example_utils.py
tests/examples/hf_ptq/test_example_utils.py`
- `pytest_pwd tests/examples/hf_ptq/test_example_utils.py -q -x` (15
passed)
- Actual `Llama-3_3-Nemotron-Super-49B-v1` end-to-end `hf_ptq.py` export
on one node with 6 GPUs, Transformers 4.48.3:
#1857 (comment)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved model loading for Hugging Face remote-code scenarios by
safely re-deriving the initialization configuration when needed, with a
warning-based fallback.
* Ensured precision is derived consistently from the resolved config
(including dtype name handling) with a safe default when unspecified.
* Tightened forwarding of precision-related kwargs and
`trust_remote_code`, and avoided passing `max_memory` during config
loading.

* **Tests**
* Added unit coverage for initialization config resolution (including
failure fallback).
* Extended integration-style coverage to validate dtype/kwarg
forwarding, `trust_remote_code` behavior, and eval-mode initialization.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

examples/hf_ptq/example_utils.py

@@ -750,18 +750,29 @@ def has_pack_quantized_config(config):
             auto_model_module = getattr(transformers, architecture)
             from_config = auto_model_module._from_config
 
+        is_decilm = "DeciLM" in architecture
         config_for_init = _resolve_init_config(
             hf_config, auto_model_module, ckpt_path, config_kwargs
         )
 
         with init_empty_weights(include_buffers=True):
             # When computing the device_map, assuming bfloat16 precision by default,
             # unless specified by the hf_config.
-            torch_dtype = getattr(config_for_init, "torch_dtype", torch.bfloat16)
+            config_dtype = (
+                getattr(config_for_init, "dtype", None)
+                or getattr(config_for_init, "torch_dtype", None)
+                or torch.bfloat16
+            )
+            if isinstance(config_dtype, str):
+                config_dtype = getattr(torch, config_dtype)
             model_kwargs2 = model_kwargs.copy()
             if auto_model_module not in [AutoModelForCausalLM, AutoModel]:
                 model_kwargs2.pop("trust_remote_code", None)
-            model_kwargs2["dtype"] = torch_dtype
+            if is_decilm:
+                model_kwargs2["torch_dtype"] = config_dtype
+                model_kwargs2.pop("dtype", None)
+            else:
+                model_kwargs2["dtype"] = config_dtype
             model_kwargs2.pop("max_memory", None)
             model = from_config(config_for_init, **model_kwargs2)
 
@@ -783,10 +794,13 @@ def has_pack_quantized_config(config):
             )
             model_kwargs["max_memory"] = max_memory
 
+        model_kwargs2 = model_kwargs.copy()
+        if is_decilm:
+            model_kwargs2.pop("dtype", None)
         model = auto_model_module.from_pretrained(
             ckpt_path,
             device_map=device_map,
-            **model_kwargs,
+            **model_kwargs2,
         )
     model.eval()
     if has_pack_quantized_config(hf_config):

tests/examples/hf_ptq/test_example_utils.py

@@ -19,9 +19,11 @@
 """
 
 import json
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import patch
 
+import pytest
 import torch
 from _test_utils.examples.hf_ptq_example_utils import example_utils
 from safetensors.torch import save_file
@@ -229,3 +231,88 @@ def test_resolve_init_config_falls_back_when_rederive_raises():
     cfg = _remote_config()
     with patch.object(example_utils.AutoConfig, "from_pretrained", side_effect=ValueError()):
         assert example_utils._resolve_init_config(cfg, object, "/ckpt", {}) is cfg
+
+
+@pytest.mark.parametrize(
+    (
+        "architecture",
+        "model_class_name",
+        "expected_config_dtype_kwarg",
+        "unexpected_config_dtype_kwarg",
+    ),
+    [
+        ("DeciLMForCausalLM", "AutoModelForCausalLM", "torch_dtype", "dtype"),
+        ("LlamaForCausalLM", "LlamaForCausalLM", "dtype", "torch_dtype"),
+    ],
+)
+def test_get_model_uses_expected_dtype_kwarg(
+    monkeypatch,
+    architecture,
+    model_class_name,
+    expected_config_dtype_kwarg,
+    unexpected_config_dtype_kwarg,
+):
+    calls = {}
+    hf_config = SimpleNamespace(
+        architectures=[architecture],
+        dtype=torch.float16,
+        model_type="llama",
+        torch_dtype=torch.bfloat16,
+    )
+
+    class FakeModel:
+        def eval(self):
+            calls["eval"] = True
+
+    class FakeAutoModelForCausalLM:
+        @staticmethod
+        def from_config(config, **kwargs):
+            calls["from_config"] = kwargs
+            assert config is hf_config
+            assert kwargs[expected_config_dtype_kwarg] is torch.float16
+            assert unexpected_config_dtype_kwarg not in kwargs
+            assert "max_memory" not in kwargs
+            return FakeModel()
+
+        @staticmethod
+        def from_pretrained(*args, **kwargs):
+            calls["from_pretrained"] = kwargs
+            assert "dtype" not in kwargs
+            assert "torch_dtype" not in kwargs
+            return FakeModel()
+
+    class FakeLlamaForCausalLM(FakeAutoModelForCausalLM):
+        _from_config = FakeAutoModelForCausalLM.from_config
+
+        @staticmethod
+        def from_pretrained(*args, **kwargs):
+            calls["from_pretrained"] = kwargs
+            assert kwargs["dtype"] == "auto"
+            assert "torch_dtype" not in kwargs
+            return FakeModel()
+
+    monkeypatch.setattr(
+        example_utils.AutoConfig,
+        "from_pretrained",
+        lambda *args, **kwargs: hf_config,
+    )
+    if model_class_name == "AutoModelForCausalLM":
+        monkeypatch.setattr(example_utils, "AutoModelForCausalLM", FakeAutoModelForCausalLM)
+        monkeypatch.delattr(example_utils.transformers, architecture, raising=False)
+    else:
+        monkeypatch.setattr(example_utils.transformers, model_class_name, FakeLlamaForCausalLM)
+    monkeypatch.setattr(example_utils, "is_nemotron_vl", lambda config: False)
+    monkeypatch.setattr(example_utils, "is_speculative", lambda config: False)
+    monkeypatch.setattr(example_utils, "init_empty_weights", lambda include_buffers: nullcontext())
+    monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 1024})
+    monkeypatch.setattr(example_utils, "infer_auto_device_map", lambda model, max_memory: {"": 0})
+
+    model = example_utils.get_model("checkpoint", device="cpu", trust_remote_code=True)
+
+    assert isinstance(model, FakeModel)
+    assert calls["eval"]
+    if expected_config_dtype_kwarg == "torch_dtype":
+        assert calls["from_config"]["trust_remote_code"] is True
+    else:
+        assert "trust_remote_code" not in calls["from_config"]
+    assert calls["from_pretrained"]["trust_remote_code"] is True