Fix diffusers HF export crash on configs without torch_dtype

`_resolve_export_dtype` accessed `model.config.torch_dtype` directly, which
raises `AttributeError` for a diffusers `FrozenDict` config that has no
`torch_dtype` key. This broke `export_hf_checkpoint` for diffusers models
(UNet/DiT/Flux) under the minimum transformers/diffusers versions, failing the
`tf_min` unit matrix (tests/unit/torch/export/test_export_diffusers.py).

Read it via `getattr(..., None)` and fall back to the model's parameter dtype
(`next(model.parameters()).dtype`), matching how `diffusers_utils` already
resolves a diffusers model's dtype. Behavior is unchanged when the config
carries `torch_dtype`. Adds version-independent unit tests for the helper.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Author: shengliangxu

modelopt/torch/export/unified_export_hf.py

@@ -826,11 +826,17 @@ def _dispatch_export_handler(name: str, sub_module: nn.Module, ctx: ExportContex
 
 def _resolve_export_dtype(model: nn.Module, dtype: torch.dtype | None) -> torch.dtype:
     """Return the export dtype, defaulting to the model's own and warning on a mismatch."""
+    # Diffusers models expose their config as a ``FrozenDict`` that need not carry a
+    # ``torch_dtype`` key, so bare attribute access can raise ``AttributeError``. Fall back
+    # to the model's parameter dtype in that case (matching ``diffusers_utils``).
+    model_dtype = getattr(model.config, "torch_dtype", None)
+    if model_dtype is None:
+        model_dtype = next(model.parameters()).dtype
     if dtype is None:
-        return model.config.torch_dtype
-    if dtype != model.config.torch_dtype:
+        return model_dtype
+    if dtype != model_dtype:
         warnings.warn(
-            f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype "
+            f"Model's original dtype ({model_dtype}) differs from target dtype "
             f"({dtype}), which may lead to numerical errors."
         )
     return dtype

tests/unit/torch/export/test_unified_export_hf.py

@@ -15,6 +15,8 @@
 
 """Tests for tied-weight helpers in unified_export_hf."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 from _test_utils.torch.quantization.tied_modules import (
@@ -29,6 +31,7 @@
     postprocess_state_dict,
     sync_tied_input_amax,
 )
+from modelopt.torch.export.unified_export_hf import _resolve_export_dtype
 from modelopt.torch.quantization.nn import TensorQuantizer
 
 
@@ -484,3 +487,37 @@ def test_fuse_prequant_layernorm_fuses_and_removes_pre_quant_scale():
     for module in modules:
         assert not hasattr(module.input_quantizer, "_pre_quant_scale")
         assert module.fused_with_prequant
+
+
+def _module_with_config(config, param_dtype):
+    """A minimal ``nn.Module`` exposing a ``.config`` and a single parameter."""
+
+    class _M(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(2, 2).to(param_dtype)
+            self.config = config
+
+    return _M()
+
+
+def test_resolve_export_dtype_uses_config_torch_dtype():
+    """A config carrying ``torch_dtype`` is used verbatim."""
+    model = _module_with_config(SimpleNamespace(torch_dtype=torch.bfloat16), torch.float32)
+    assert _resolve_export_dtype(model, None) == torch.bfloat16
+
+
+def test_resolve_export_dtype_falls_back_to_param_dtype_without_torch_dtype():
+    """Diffusers-style configs (e.g. a ``FrozenDict``) need not carry ``torch_dtype``.
+
+    Regression test: bare ``model.config.torch_dtype`` used to raise ``AttributeError``
+    during diffusers export under minimum transformers/diffusers versions.
+    """
+    model = _module_with_config(SimpleNamespace(), torch.float16)  # no torch_dtype
+    assert _resolve_export_dtype(model, None) == torch.float16
+
+
+def test_resolve_export_dtype_returns_explicit_dtype_without_config_dtype():
+    """An explicit target dtype is returned, and the mismatch path does not crash."""
+    model = _module_with_config(SimpleNamespace(), torch.float32)
+    assert _resolve_export_dtype(model, torch.bfloat16) == torch.bfloat16