From 24ff1fb7de1bb1bfedc6547d62ee82d786170280 Mon Sep 17 00:00:00 2001 From: Paresh Shukla <1964694+pareshs44@users.noreply.github.com> Date: Fri, 10 Jul 2026 15:30:03 -0700 Subject: [PATCH] Enable fp4 quantization for 4D tensors --- src/coreai_opt/_utils/torch_utils.py | 26 ++++- src/coreai_opt/quantization/_export_utils.py | 36 +++--- .../test_eager_mlir_export_embedding.py | 98 ++++++++++++++++ tests/models/simple.py | 64 +++++++++++ tests/quantization/test_eager_quant.py | 106 +++++++++++------- .../quantization/test_graph_mode_quantizer.py | 6 +- tests/test_utils/test_torch_utils.py | 51 +++++++++ 7 files changed, 326 insertions(+), 61 deletions(-) create mode 100644 tests/export/test_eager_mlir_export_embedding.py diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 915891c..b0c214f 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -426,12 +426,23 @@ def mmap_module_state_dict(module: torch.nn.Module, path: str | PathLike[str]) - reload it via mmap, replacing the module's parameters/buffers with mmap views so the in-RAM tensors can be released. + Handles ``SubbyteTensor`` subclasses (e.g. ``Float4Tensor``) that safetensors + cannot serialize directly by unwrapping to plain uint8 for storage and + re-wrapping after reload. + Requires all tensors in ``module.state_dict()`` to be on CPU. Raises ``ValueError`` otherwise — mmap is a CPU-only mechanism """ + from coreai_torch._compression._floatx import SubbyteTensor as _SubbyteTensor # noqa: PLC0415 from safetensors.torch import load_file, save_file # noqa: PLC0415 state_dict = module.state_dict() + + # Keys whose tensors are SubbyteTensor wrappers (Float4Tensor, etc.) that + # safetensors cannot serialize. Track their class for re-wrapping after load. + subbyte_keys: dict[str, type] = {} + tensors_to_save: dict[str, torch.Tensor] = {} + for name, tensor in state_dict.items(): if not isinstance(tensor, torch.Tensor): continue @@ -439,10 +450,17 @@ def mmap_module_state_dict(module: torch.nn.Module, path: str | PathLike[str]) - raise ValueError( f"mmap_module_state_dict requires CPU tensors; '{name}' is on {tensor.device}." ) + if isinstance(tensor, _SubbyteTensor): + subbyte_keys[name] = type(tensor) + tensors_to_save[name] = tensor.elem.contiguous() + else: + tensors_to_save[name] = tensor.contiguous() - save_file( - {k: v.contiguous() for k, v in state_dict.items() if isinstance(v, torch.Tensor)}, - path, - ) + save_file(tensors_to_save, path) mmap_sd = load_file(path, device="cpu") + + # Re-wrap SubbyteTensor keys from their underlying uint8 representation. + for name, tensor_cls in subbyte_keys.items(): + mmap_sd[name] = tensor_cls(mmap_sd[name]) + module.load_state_dict(mmap_sd, assign=True) diff --git a/src/coreai_opt/quantization/_export_utils.py b/src/coreai_opt/quantization/_export_utils.py index eb0e7d4..c705819 100644 --- a/src/coreai_opt/quantization/_export_utils.py +++ b/src/coreai_opt/quantization/_export_utils.py @@ -313,12 +313,17 @@ def validate_fp4_export( ) -> None: """Validate that FP4 quantization config is supported for MLIR export. - FP4 export only supports weight-only quantization on 2D weight tensors - (Linear layers) with PerBlockGranularity and block_size=32. + FP4 export requires weight-only quantization with ``PerBlockGranularity`` + that resolves to blocks of size 32 along the weight's last axis and no + blocking along any other axis (i.e. resolved block sizes + ``(1, ..., 1, 32)``). This covers 2D Linear weights, which resolve to + ``(1, 32)``, and higher-rank weights such as MoE experts that resolve to + ``(1, ..., 1, 32)``. Args: fake_quant_mod (FakeQuantizeImplBase): The fake quantization module to validate. - quantized_data (torch.Tensor | None): The quantized weight tensor to validate shape. + quantized_data (torch.Tensor | None): The quantized weight tensor used + to resolve per-axis block sizes against the actual weight shape. Raises: ValueError: If the FP4 configuration is not supported for export. @@ -329,20 +334,21 @@ def validate_fp4_export( f"Got quantization_target={fake_quant_mod.quantization_target}." ) - if quantized_data is not None and quantized_data.ndim != 2: + granularity = fake_quant_mod._granularity + if not isinstance(granularity, PerBlockGranularity): raise ValueError( - "FP4 weight quantization export is only supported for " - f"2D weight tensors (Linear layers). Got {quantized_data.ndim}D tensor." + f"FP4 quantization requires PerBlockGranularity. Got {type(granularity).__name__}." ) - granularity = fake_quant_mod._granularity - if ( - not isinstance(granularity, PerBlockGranularity) - or granularity.block_size != _FP4_EXPORT_BLOCK_SIZE - ): + if quantized_data is None: + return + + resolved_block_size = granularity.get_block_size(quantized_data.shape) + expected_block_size = (1,) * (quantized_data.ndim - 1) + (_FP4_EXPORT_BLOCK_SIZE,) + if resolved_block_size != expected_block_size: raise ValueError( - f"FP4 quantization requires PerBlockGranularity with " - f"block_size={_FP4_EXPORT_BLOCK_SIZE}. " - f"Got {type(granularity).__name__} with " - f"block_size={getattr(granularity, 'block_size', 'N/A')}." + f"FP4 export requires per-axis block sizes {expected_block_size} for a " + f"{quantized_data.ndim}D weight (blocks of {_FP4_EXPORT_BLOCK_SIZE} along " + f"the last axis, no blocking elsewhere). Got resolved block sizes " + f"{resolved_block_size} from granularity={granularity!r}." ) diff --git a/tests/export/test_eager_mlir_export_embedding.py b/tests/export/test_eager_mlir_export_embedding.py new file mode 100644 index 0000000..671b3f3 --- /dev/null +++ b/tests/export/test_eager_mlir_export_embedding.py @@ -0,0 +1,98 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Eager Core AI export tests for FP4-quantized nn.Embedding variants. + +These mirror the embedding patterns seen in common LLMs: + - plain : standard nn.Embedding + - scaled : embedding output scaled by sqrt(hidden_size) + - tied : embedding weight shared with an lm_head Linear + +All variants use FP4 (float4_e2m1fn) weight-only quantization with symmetric +per-block granularity (block_size=32), which is the only granularity supported +for FP4 MLIR export. +""" + +import pytest +import torch +from torch import nn + +from coreai_opt import ExportBackend +from coreai_opt.quantization import ( + ModuleQuantizerConfig, + QuantizationSpec, + Quantizer, + QuantizerConfig, +) +from coreai_opt.quantization.spec import PerBlockGranularity, QuantizationScheme +from tests.models.simple import ( + PlainEmbeddingModel, + ScaledEmbeddingModel, + TiedEmbeddingModel, +) + +from . import export_utils + +_VOCAB_SIZE = 256 +_EMBED_DIM = 64 +_SEQ_LEN = 8 + + +def _fp4_config() -> QuantizerConfig: + """FP4 weight-only config: symmetric, per-block (block_size=32) along axis 1.""" + return QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={ + "weight": QuantizationSpec( + dtype="float4_e2m1fn", + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerBlockGranularity(axis=1, block_size=32), + ), + }, + op_input_spec=None, + op_output_spec=None, + ), + execution_mode="eager", + ) + + +def _run_fp4_embedding_export(model: nn.Module, expected_shift_scale_ops: int) -> None: + """Quantize the embedding model with FP4, finalize, and export/run on Core AI.""" + model = model.eval().to(dtype=torch.float16) + input_ids = torch.randint(0, _VOCAB_SIZE, (1, _SEQ_LEN), dtype=torch.int32) + + quantizer = Quantizer(model, _fp4_config()) + prepared_model = quantizer.prepare((input_ids,)) + + with torch.no_grad(): + prepared_model_output = prepared_model(input_ids) + + finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI) + + export_utils.convert_and_verify( + finalized_model=finalized_model, + input_data=input_ids, + expected_ops={"constexpr_blockwise_shift_scale": expected_shift_scale_ops}, + export_backend=ExportBackend.CoreAI, + prepared_model_output=prepared_model_output, + ) + + +@pytest.mark.parametrize( + "model_factory, expected_shift_scale_ops", + [ + pytest.param(PlainEmbeddingModel, 1, id="plain"), + pytest.param(ScaledEmbeddingModel, 1, id="scaled"), + pytest.param(TiedEmbeddingModel, 2, id="tied"), + ], +) +def test_fp4_embedding_export( + model_factory: type[nn.Module], expected_shift_scale_ops: int +) -> None: + """Eager Core AI export with FP4-quantized nn.Embedding variants.""" + _run_fp4_embedding_export( + model_factory(vocab_size=_VOCAB_SIZE, embed_dim=_EMBED_DIM), + expected_shift_scale_ops, + ) diff --git a/tests/models/simple.py b/tests/models/simple.py index 18919ab..46b7375 100644 --- a/tests/models/simple.py +++ b/tests/models/simple.py @@ -171,3 +171,67 @@ def simple_mha_model(): def simple_mha_model_input(): """Fixture providing example input tensor for MHA model.""" return torch.randn(1, 10, 64) + + +class PlainEmbeddingModel(nn.Module): + """Simple model with standard nn.Embedding.""" + + def __init__(self, vocab_size: int = 256, embed_dim: int = 64) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, embed_dim) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + +class ScaledEmbeddingModel(nn.Module): + """Simple model with embedding output scaled by sqrt(hidden_size).""" + + def __init__(self, vocab_size: int = 256, embed_dim: int = 64) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, embed_dim) + self.embed_scale = embed_dim**0.5 + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) * torch.tensor( + self.embed_scale, dtype=self.embed_tokens.weight.dtype + ) + + +class TiedEmbeddingModel(nn.Module): + """Simple model with embedding weight tied with an lm_head Linear.""" + + def __init__(self, vocab_size: int = 256, embed_dim: int = 64) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, embed_dim) + self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False) + # Tie weights. + self.lm_head.weight = self.embed_tokens.weight + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + hidden = self.embed_tokens(input_ids) + return self.lm_head(hidden) + + +@pytest.fixture +def plain_embedding_model(): + """Fixture providing a plain nn.Embedding model.""" + return PlainEmbeddingModel() + + +@pytest.fixture +def scaled_embedding_model(): + """Fixture providing a sqrt(hidden_size)-scaled embedding model.""" + return ScaledEmbeddingModel() + + +@pytest.fixture +def tied_embedding_model(): + """Fixture providing an embedding model with weight tied to an lm_head.""" + return TiedEmbeddingModel() + + +@pytest.fixture +def embedding_model_input(): + """Fixture providing example token-id input for the embedding models.""" + return torch.randint(0, 256, (1, 8), dtype=torch.int32) diff --git a/tests/quantization/test_eager_quant.py b/tests/quantization/test_eager_quant.py index 3d57bb9..a6a10a1 100644 --- a/tests/quantization/test_eager_quant.py +++ b/tests/quantization/test_eager_quant.py @@ -442,51 +442,79 @@ def test_finalize_mmap_files_exist(self, use_mmap, basic_config, tmp_path): else: assert files == [] - @pytest.mark.parametrize( - "shared_weight_model", - [True, False], - ids=["shared_weights", "no_sharing"], - ) - def test_finalize_mmap_matches_non_mmap_output( - self, basic_config, tmp_path, shared_weight_model - ): - """Finalize with and without ``mmap_dir`` produces numerically identical - outputs, including for models with shared (weight-tied) layers.""" - if shared_weight_model: - model_cls = SharedWeightModel - example_input_shape = (1, 4) - else: - model_cls = SimpleLinearModel - example_input_shape = (1, 10) - - model_no_mmap = model_cls() - model_with_mmap = copy.deepcopy(model_no_mmap) - example_input = torch.rand(*example_input_shape) + def _finalize_with_and_without_mmap(self, model, config, example_input, tmp_path): + """Finalize ``model`` for the CoreAI backend both without and with + ``mmap_dir``, returning the ``(no_mmap, with_mmap)`` finalized models.""" + model_with_mmap = copy.deepcopy(model) finalized_no_mmap = self._quantize_model( - model_no_mmap, basic_config, example_input, None, ExportBackend.CoreAI + model, config, example_input, None, ExportBackend.CoreAI ) finalized_with_mmap = self._quantize_model( - model_with_mmap, basic_config, example_input, str(tmp_path), ExportBackend.CoreAI + model_with_mmap, config, example_input, str(tmp_path), ExportBackend.CoreAI ) + return finalized_no_mmap, finalized_with_mmap - if shared_weight_model: - # Tied-weight modules should share a single dequant parametrization - # post-finalize, both with and without mmap. - assert ( - finalized_no_mmap.linear1.parametrizations.weight[0] - is finalized_no_mmap.linear2.parametrizations.weight[0] - ), "no-mmap finalize did not preserve sharing for weight-tied modules" + @staticmethod + def _assert_forward_outputs_equal(model_a, model_b, example_input): + """Assert two models produce numerically identical forward outputs.""" + with torch.no_grad(): + out_a = model_a(example_input) + out_b = model_b(example_input) + assert torch.equal(out_a, out_b) + + def test_finalize_mmap_matches_non_mmap_output(self, basic_config, tmp_path): + """Standard per-tensor quantization: finalize with and without + ``mmap_dir`` produces numerically identical outputs.""" + example_input = torch.rand(1, 10) + finalized_no_mmap, finalized_with_mmap = self._finalize_with_and_without_mmap( + SimpleLinearModel(), basic_config, example_input, tmp_path + ) + self._assert_forward_outputs_equal(finalized_no_mmap, finalized_with_mmap, example_input) + + def test_finalize_mmap_preserves_weight_sharing(self, basic_config, tmp_path): + """Weight-tied modules keep a single shared dequant parametrization after + finalize, both with and without ``mmap_dir`` and ensure outputs still match.""" + example_input = torch.rand(1, 4) + finalized_no_mmap, finalized_with_mmap = self._finalize_with_and_without_mmap( + SharedWeightModel(), basic_config, example_input, tmp_path + ) + self._assert_forward_outputs_equal(finalized_no_mmap, finalized_with_mmap, example_input) + + for label, finalized in (("no-mmap", finalized_no_mmap), ("mmap", finalized_with_mmap)): assert ( - finalized_with_mmap.linear1.parametrizations.weight[0] - is finalized_with_mmap.linear2.parametrizations.weight[0] - ), "mmap finalize did not preserve sharing for weight-tied modules" + finalized.linear1.parametrizations.weight[0] + is finalized.linear2.parametrizations.weight[0] + ), f"{label} finalize did not preserve sharing for weight-tied modules" + + def test_finalize_mmap_matches_non_mmap_output_fp4_per_block(self, tmp_path): + """FP4 per-block weights stored as Float4Tensor finalize identically + with and without ``mmap_dir``, and the mmap path emits a safetensors + file for the weight.""" + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={ + "weight": QuantizationSpec( + dtype="float4_e2m1fn", + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerBlockGranularity(axis=1, block_size=32), + ), + }, + op_input_spec=None, + op_output_spec=None, + ), + execution_mode="eager", + ) - with torch.no_grad(): - out_no_mmap = finalized_no_mmap(example_input) - out_with_mmap = finalized_with_mmap(example_input) + example_input = torch.randn(1, 64) + finalized_no_mmap, finalized_with_mmap = self._finalize_with_and_without_mmap( + nn.Linear(64, 32, bias=False), config, example_input, tmp_path + ) + self._assert_forward_outputs_equal(finalized_no_mmap, finalized_with_mmap, example_input) - assert torch.equal(out_no_mmap, out_with_mmap) + # FP4 weights are stored as Float4Tensor; a safetensors file must be + # emitted for the mmap-backed weight. + assert any(f.endswith(".safetensors") for f in os.listdir(tmp_path)) def test_finalize_state_dict_safetensors_roundtrip(self, basic_config, tmp_path): """An mmap-finalized model survives a state_dict save → load_file → @@ -3315,7 +3343,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32), PerTensorGranularity(), False, - "FP4 quantization requires PerBlockGranularity with block_size=32", + "FP4 quantization requires PerBlockGranularity", id="per_tensor_granularity_rejected", ), pytest.param( @@ -3323,7 +3351,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32), PerBlockGranularity(axis=1, block_size=16), False, - "FP4 quantization requires PerBlockGranularity with block_size=32", + r"FP4 export requires per-axis block sizes \(1, 32\) for a 2D weight", id="wrong_block_size_rejected", ), pytest.param( @@ -3339,7 +3367,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32, 5, 5), PerBlockGranularity(axis=1, block_size=32), False, - "FP4 weight quantization export is only supported for 2D weight tensors", + r"FP4 export requires per-axis block sizes \(1, 1, 1, 32\) for a 4D weight", id="conv_layer_rejected", ), ], diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index 81d3d3a..d88fa8d 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -1328,7 +1328,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32), PerTensorGranularity(), False, - "FP4 quantization requires PerBlockGranularity with block_size=32", + "FP4 quantization requires PerBlockGranularity", id="per_tensor_granularity_rejected", ), pytest.param( @@ -1336,7 +1336,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32), PerBlockGranularity(axis=1, block_size=16), False, - "FP4 quantization requires PerBlockGranularity with block_size=32", + r"FP4 export requires per-axis block sizes \(1, 32\) for a 2D weight", id="wrong_block_size_rejected", ), pytest.param( @@ -1352,7 +1352,7 @@ class TestFP4MLIRExportValidation: torch.randn(1, 32, 5, 5), PerBlockGranularity(axis=1, block_size=32), False, - "FP4 weight quantization export is only supported for 2D weight tensors", + r"FP4 export requires per-axis block sizes \(1, 1, 1, 32\) for a 4D weight", id="conv_layer_rejected", ), ], diff --git a/tests/test_utils/test_torch_utils.py b/tests/test_utils/test_torch_utils.py index 5afa6a6..4511bbe 100644 --- a/tests/test_utils/test_torch_utils.py +++ b/tests/test_utils/test_torch_utils.py @@ -7,10 +7,13 @@ import pytest import torch +from coreai_torch._compression._floatx import Float4Tensor +from torch import nn from torchao.quantization.pt2e import allow_exported_model_train_eval from coreai_opt._utils.fx_utils import normalize_module_fqn from coreai_opt._utils.torch_utils import ( + mmap_module_state_dict, move_model_to_eval, move_model_to_train, ) @@ -88,3 +91,51 @@ class TestNormalizeModuleFqn: def test_normalize_module_fqn(raw: str, expected: str) -> None: """Verify various path formats are normalized correctly.""" assert normalize_module_fqn(raw) == expected + + +class TestMmapModuleStateDict: + """Test mmap_module_state_dict serialization and reload.""" + + @staticmethod + def test_standard_tensors_roundtrip(tmp_path): + """Standard (non-subbyte) tensors are saved and reloaded correctly.""" + model = nn.Linear(8, 4, bias=True) + original_weight = model.weight.data.clone() + original_bias = model.bias.data.clone() + + mmap_module_state_dict(model, tmp_path / "model.safetensors") + + assert torch.equal(model.weight.data, original_weight) + assert torch.equal(model.bias.data, original_bias) + + @staticmethod + def test_mixed_standard_and_float4(tmp_path): + """Module with both standard and Float4Tensor parameters roundtrips.""" + + module = nn.Module() + module.register_buffer("normal", torch.randn(4, 4)) + uint8_data = torch.randint(0, 255, (2, 8), dtype=torch.uint8) + module.register_buffer("compressed", Float4Tensor(uint8_data)) + + original_normal = module.normal.clone() + + mmap_module_state_dict(module, tmp_path / "model.safetensors") + + assert torch.equal(module.normal, original_normal) + assert isinstance(module.compressed, Float4Tensor) + assert torch.equal(module.compressed.elem, uint8_data) + + @staticmethod + def test_raises_on_non_cpu_tensor(tmp_path): + """Raises ValueError when a tensor is not on CPU.""" + if torch.cuda.is_available(): + device = "cuda" + elif torch.backends.mps.is_available(): + device = "mps" + else: + pytest.skip("No non-CPU device available") + + model = nn.Linear(4, 4).to(device) + + with pytest.raises(ValueError, match="requires CPU tensors"): + mmap_module_state_dict(model, tmp_path / "model.safetensors")