diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 77429b1cfaf..4dae93b5aa7 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/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 diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index b9fa29238d7..469ffd9227e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/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.