Fix Diffusers export dtype resolution

modelopt/torch/export/unified_export_hf.py

@@ -826,11 +826,15 @@ 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."""
+    configured_dtype = getattr(model.config, "torch_dtype", None)
     if dtype is None:
-        return model.config.torch_dtype
-    if dtype != model.config.torch_dtype:
+        if configured_dtype is not None:
+            return configured_dtype
+        first_parameter = next(model.parameters(), None)
+        return first_parameter.dtype if first_parameter is not None else torch.float16
+    if configured_dtype is not None and dtype != configured_dtype:
         warnings.warn(
-            f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype "
+            f"Model's original dtype ({configured_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,9 +31,44 @@
     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
 
 
+@pytest.mark.parametrize(
+    ("configured_dtype", "dtype", "expected_dtype", "warning_count"),
+    [
+        (None, None, torch.float32, 0),
+        (None, torch.float16, torch.float16, 0),
+        (torch.bfloat16, None, torch.bfloat16, 0),
+        (torch.bfloat16, torch.bfloat16, torch.bfloat16, 0),
+        (torch.bfloat16, torch.float16, torch.float16, 1),
+    ],
+)
+def test_resolve_export_dtype(configured_dtype, dtype, expected_dtype, warning_count, recwarn):
+    model = torch.nn.Linear(1, 1)
+    model.config = (
+        SimpleNamespace(torch_dtype=configured_dtype) if configured_dtype is not None else object()
+    )
+
+    assert _resolve_export_dtype(model, dtype) == expected_dtype
+    assert len(recwarn) == warning_count
+    if warning_count:
+        assert str(recwarn[0].message) == (
+            "Model's original dtype (torch.bfloat16) differs from target dtype "
+            "(torch.float16), which may lead to numerical errors."
+        )
+
+
+def test_resolve_export_dtype_with_empty_diffusers_config():
+    # Import locally so Diffusers stays optional during torch-only test collection.
+    frozen_dict = pytest.importorskip("diffusers.configuration_utils").FrozenDict()
+    model = torch.nn.Linear(1, 1)
+    model.config = frozen_dict
+
+    assert _resolve_export_dtype(model, None) == torch.float32
+
+
 def test_hf_all_tied_weights_keys_contract():
     """Pin the transformers API we build tied_map from, so a version bump fails loud here.