From 1e9d346c6f672d00910b1734155f6fa3b81c9cd7 Mon Sep 17 00:00:00 2001 From: realAsma Date: Thu, 20 Aug 2026 23:17:09 +0000 Subject: [PATCH 1/5] Fix Diffusers export dtype resolution Signed-off-by: realAsma --- modelopt/torch/export/unified_export_hf.py | 5 +++-- tests/unit/torch/export/test_unified_export_hf.py | 9 +++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 77429b1cfaf..2bf6b429145 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -828,9 +828,10 @@ def _resolve_export_dtype(model: nn.Module, dtype: torch.dtype | None) -> torch. """Return the export dtype, defaulting to the model's own and warning on a mismatch.""" if dtype is None: return model.config.torch_dtype - if dtype != model.config.torch_dtype: + configured_dtype = getattr(model.config, "torch_dtype", None) + 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..3e70918f05e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -29,9 +29,18 @@ 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 +def test_resolve_export_dtype_uses_explicit_dtype_without_configured_dtype(recwarn): + model = torch.nn.Module() + model.config = object() + + assert _resolve_export_dtype(model, torch.float16) == torch.float16 + assert not recwarn + + def test_hf_all_tied_weights_keys_contract(): """Pin the transformers API we build tied_map from, so a version bump fails loud here. From cc2c21837faa2e58f597d69e0f850bcef0e05567 Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 21 Aug 2026 20:41:10 +0000 Subject: [PATCH 2/5] Fix implicit Diffusers export dtype resolution Signed-off-by: realAsma --- modelopt/torch/export/unified_export_hf.py | 7 +++++-- tests/unit/torch/export/test_unified_export_hf.py | 8 +++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 2bf6b429145..4dae93b5aa7 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -826,9 +826,12 @@ 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.""" - if dtype is None: - return model.config.torch_dtype configured_dtype = getattr(model.config, "torch_dtype", None) + if dtype is None: + 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 ({configured_dtype}) differs from target dtype " diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 3e70918f05e..464d00d3ea9 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -33,11 +33,13 @@ from modelopt.torch.quantization.nn import TensorQuantizer -def test_resolve_export_dtype_uses_explicit_dtype_without_configured_dtype(recwarn): - model = torch.nn.Module() +@pytest.mark.parametrize("dtype", [None, torch.float16]) +def test_resolve_export_dtype_without_configured_dtype(dtype, recwarn): + model = torch.nn.Linear(1, 1) model.config = object() - assert _resolve_export_dtype(model, torch.float16) == torch.float16 + expected_dtype = model.weight.dtype if dtype is None else dtype + assert _resolve_export_dtype(model, dtype) == expected_dtype assert not recwarn From 8b9f72367d89d42b1852c6bb609357c84154fa1a Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 21 Aug 2026 21:13:09 +0000 Subject: [PATCH 3/5] Add export dtype resolution coverage Signed-off-by: realAsma --- .../torch/export/test_unified_export_hf.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 464d00d3ea9..1535b67a03b 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -21,6 +21,7 @@ make_tied_linear_pair, wrap_in_parent_with_tied_keys, ) +from diffusers.configuration_utils import FrozenDict import modelopt.torch.quantization as mtq from modelopt.torch.export.model_utils import TiedWeightMap @@ -33,14 +34,27 @@ from modelopt.torch.quantization.nn import TensorQuantizer -@pytest.mark.parametrize("dtype", [None, torch.float16]) -def test_resolve_export_dtype_without_configured_dtype(dtype, recwarn): +@pytest.mark.parametrize( + ("config", "dtype", "expected_dtype", "warning_count"), + [ + (object(), None, torch.float32, 0), + (object(), torch.float16, torch.float16, 0), + (FrozenDict(torch_dtype=torch.bfloat16), None, torch.bfloat16, 0), + (FrozenDict(torch_dtype=torch.bfloat16), torch.bfloat16, torch.bfloat16, 0), + (FrozenDict(torch_dtype=torch.bfloat16), torch.float16, torch.float16, 1), + ], +) +def test_resolve_export_dtype(config, dtype, expected_dtype, warning_count, recwarn): model = torch.nn.Linear(1, 1) - model.config = object() + model.config = config - expected_dtype = model.weight.dtype if dtype is None else dtype assert _resolve_export_dtype(model, dtype) == expected_dtype - assert not recwarn + 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_hf_all_tied_weights_keys_contract(): From aaac495d40a6a06ba7ab236a74c1413c70dcb4bf Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 21 Aug 2026 22:04:23 +0000 Subject: [PATCH 4/5] Keep export dtype tests optional-dependency safe Signed-off-by: realAsma --- .../torch/export/test_unified_export_hf.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 1535b67a03b..0b14acf97d0 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -15,13 +15,14 @@ """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 ( make_tied_linear_pair, wrap_in_parent_with_tied_keys, ) -from diffusers.configuration_utils import FrozenDict import modelopt.torch.quantization as mtq from modelopt.torch.export.model_utils import TiedWeightMap @@ -35,18 +36,20 @@ @pytest.mark.parametrize( - ("config", "dtype", "expected_dtype", "warning_count"), + ("configured_dtype", "dtype", "expected_dtype", "warning_count"), [ - (object(), None, torch.float32, 0), - (object(), torch.float16, torch.float16, 0), - (FrozenDict(torch_dtype=torch.bfloat16), None, torch.bfloat16, 0), - (FrozenDict(torch_dtype=torch.bfloat16), torch.bfloat16, torch.bfloat16, 0), - (FrozenDict(torch_dtype=torch.bfloat16), torch.float16, torch.float16, 1), + (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(config, dtype, expected_dtype, warning_count, recwarn): +def test_resolve_export_dtype(configured_dtype, dtype, expected_dtype, warning_count, recwarn): model = torch.nn.Linear(1, 1) - model.config = config + 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 @@ -57,6 +60,14 @@ def test_resolve_export_dtype(config, dtype, expected_dtype, warning_count, recw ) +def test_resolve_export_dtype_with_empty_diffusers_config(): + 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. From 22a10db8fd3596219873802b663ef360b6747761 Mon Sep 17 00:00:00 2001 From: realAsma Date: Fri, 21 Aug 2026 22:34:58 +0000 Subject: [PATCH 5/5] Explain optional Diffusers test import Signed-off-by: realAsma --- tests/unit/torch/export/test_unified_export_hf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 0b14acf97d0..469ffd9227e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -61,6 +61,7 @@ def test_resolve_export_dtype(configured_dtype, dtype, expected_dtype, warning_c 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