diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index c1d1058a92..1a85164ffa 100644 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -38,6 +38,7 @@ ) from QEfficient.compile.qnn_compiler import compile as qnn_compile from QEfficient.exporter.weight_free.export import embed_weight_spec_as_metadata, link_prepared_checkpoint_dir +from QEfficient.exporter.weight_free.mxfp6 import MXFP6_ONNX_OPSET from QEfficient.generation.cloud_infer import QAICInferenceSession from QEfficient.transformers.models.pytorch_transforms import ( BlockingAttentionTransform, @@ -250,6 +251,7 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None: # Flag for checking if weights are offloaded self._is_weights_offloaded: bool = False self._weight_free: bool = kwargs.get("weight_free", False) + self._mxfp6_config = kwargs.get("mxfp6_config", None) # Flag for checking if model has been transformed yet self.is_transformed: bool = False @@ -609,6 +611,12 @@ def _resolve_pkv_names(layer_idx, layer_state): active_transforms = [ transform for transform in self._onnx_transforms if transform not in excluded_transforms ] + if ( + getattr(getattr(self, "_mxfp6_config", None), "enabled", False) + and CustomOpTransform not in active_transforms + and CustomOpTransform not in excluded_transforms + ): + active_transforms.append(CustomOpTransform) needs_external_tensor_data = any( transform in active_transforms for transform in (FP16ClipTransform, SplitTensorsTransform) ) @@ -616,7 +624,11 @@ def _resolve_pkv_names(layer_idx, layer_state): "onnx_base_dir": str(export_dir) if needs_external_tensor_data else None, "model_name": self.model_name, "dynamic_axes": None if dynamo else dynamic_axes, - "onnx_export_opset": constants.get_onnx_export_opset(dynamo), + "onnx_export_opset": ( + MXFP6_ONNX_OPSET + if getattr(getattr(self, "_mxfp6_config", None), "enabled", False) + else constants.get_onnx_export_opset(dynamo) + ), } if onnx_transform_kwargs is not None: transform_kwargs.update(onnx_transform_kwargs) @@ -1073,6 +1085,14 @@ def _compile( """ layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False) + if getattr(getattr(self, "_mxfp6_config", None), "enabled", False): + compiler_mxfp6_keys = ("mxfp6_matmul", "mxfp6-matmul", "mxfp6") + if any(compiler_options.get(key, False) for key in compiler_mxfp6_keys): + raise ValueError( + "`mxfp6_matmul=True`/`mxfp6=True` cannot be used when QEff-owned `mxfp6=True` is active." + ) + if enable_qnn: + raise ValueError("QNN compilation is not supported when QEff-owned `mxfp6=True` is active.") for removed_option in ("compile_only", "compile-only"): if removed_option in compiler_options: diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py index 05f8773025..e443c02e6b 100644 --- a/QEfficient/base/onnx_transforms.py +++ b/QEfficient/base/onnx_transforms.py @@ -49,7 +49,7 @@ CtxScatterFuncCB3D, ) from QEfficient.customop.onnxscript_utils import get_onnxscript_func -from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func +from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func, UnpackMxfp6, UnpackMxfp6Func from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc from QEfficient.utils import constants from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, SIZE_THRESHOLD_DEFAULT @@ -116,6 +116,7 @@ class CustomOpTransform(BaseOnnxTransform): "CtxScatterFuncCB": (CtxScatterFuncCB, CtxScatterCB), "CtxGatherFuncCB": (CtxGatherFuncCB, CtxGatherCB), "CastToUInt4": (CastToUInt4Func, CastToUInt4), + "UnpackMxfp6": (UnpackMxfp6Func, UnpackMxfp6), "CtxChunkScatterBatchFunc": (CtxChunkScatterBatchFunc, CtxChunkScatterBatch), "CtxGatherFuncBlockedKVBatch": (CtxGatherFuncBlockedKVBatch, CtxGatherBlockedKVBatch), } @@ -136,8 +137,16 @@ def apply(cls, model: ModelProto, onnx_export_opset: int = constants.ONNX_LEGACY # Add function prototypes to model existing = {f.name for f in model.functions} - for func_name, onnxscript_func in cls._custom_ops.values(): - proto = get_onnxscript_func(onnxscript_func, onnx_export_opset).to_function_proto() + for op_name, (_, onnxscript_func) in cls._custom_ops.items(): + if onnxscript_func is None: + if op_name in used_op_types: + cls._ensure_opset_imports(model, "com.qti.aisw.onnx", 1) + op_applied = True + continue + try: + proto = get_onnxscript_func(onnxscript_func, onnx_export_opset).to_function_proto() + except AttributeError: + proto = onnxscript_func.to_function_proto() if proto.name not in used_op_types: continue if proto.name not in existing: diff --git a/QEfficient/customop/onnxscript_utils.py b/QEfficient/customop/onnxscript_utils.py index e28a34b773..ca7f3a6d06 100644 --- a/QEfficient/customop/onnxscript_utils.py +++ b/QEfficient/customop/onnxscript_utils.py @@ -68,6 +68,6 @@ def get_dynamo_onnxscript_func(onnxscript_func): def get_onnxscript_func(onnxscript_func, onnx_export_opset: int): """Return the ONNXScript variant matching the requested export opset.""" - if onnx_export_opset == constants.ONNX_DYNAMO_EXPORT_OPSET: + if onnx_export_opset >= constants.ONNX_DYNAMO_EXPORT_OPSET: return get_dynamo_onnxscript_func(onnxscript_func) return onnxscript_func diff --git a/QEfficient/customop/quantization_ops.py b/QEfficient/customop/quantization_ops.py index 3878804c7d..dc248fd3b0 100644 --- a/QEfficient/customop/quantization_ops.py +++ b/QEfficient/customop/quantization_ops.py @@ -67,6 +67,57 @@ def CastToUInt4(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8: return ops.Cast(reshaped, to=int(TensorProto.UINT4)) +@onnxscript.script(onnxscript.values.Opset("com.qti.aisw.onnx", 1)) +def UnpackMxfp6(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8: + """ + Unpack packed MXFP6 bytes into logical FLOAT6E2M3 values. + + Supports N-D input: all leading dimensions are preserved; the last + dimension is expanded by 4/3. Each 3-byte group stores four 6-bit codes: + + code0 = byte0[5:0] + code1 = byte0[7:6] | byte1[3:0] << 2 + code2 = byte1[7:4] | byte2[1:0] << 4 + code3 = byte2[7:2] + """ + three = ops.Constant(value_ints=[3]) + four = ops.Constant(value_ints=[4]) + sixteen = ops.CastLike(ops.Constant(value_ints=[16]), weight_packed) + sixty_four = ops.CastLike(ops.Constant(value_ints=[64]), weight_packed) + + packed_shape = ops.Shape(weight_packed) + leading_dims = ops.Slice(packed_shape, starts=[0], ends=[-1], axes=[0]) + last_dim = ops.Slice(packed_shape, starts=[-1], ends=[2147483647], axes=[0]) + group_count = ops.Div(last_dim, three) + grouped_shape = ops.Concat(leading_dims, group_count, three, axis=0) + grouped = ops.Reshape(weight_packed, grouped_shape) + + byte0 = ops.Slice(grouped, starts=[0], ends=[1], axes=[-1]) + byte1 = ops.Slice(grouped, starts=[1], ends=[2], axes=[-1]) + byte2 = ops.Slice(grouped, starts=[2], ends=[3], axes=[-1]) + + shift2 = ops.CastLike(ops.Constant(value_ints=[2]), weight_packed) + shift4 = ops.CastLike(ops.Constant(value_ints=[4]), weight_packed) + shift6 = ops.CastLike(ops.Constant(value_ints=[6]), weight_packed) + + code0 = ops.Mod(byte0, sixty_four) + code1_low = ops.BitShift(byte0, shift6, direction="RIGHT") + code1_high = ops.BitShift(ops.Mod(byte1, sixteen), shift2, direction="LEFT") + code1 = ops.Mod(ops.Add(code1_low, code1_high), sixty_four) + code2_low = ops.BitShift(byte1, shift4, direction="RIGHT") + code2_masked_high = ops.Mod(byte2, ops.CastLike(ops.Constant(value_ints=[4]), weight_packed)) + code2_high = ops.BitShift(code2_masked_high, shift4, direction="LEFT") + code2 = ops.Mod(ops.Add(code2_low, code2_high), sixty_four) + code3 = ops.Mod(ops.BitShift(byte2, shift2, direction="RIGHT"), sixty_four) + + stacked = ops.Concat(code0, code1, code2, code3, axis=-1) + last_dim_unpacked = ops.Div(ops.Mul(last_dim, four), three) + new_shape = ops.Concat(leading_dims, last_dim_unpacked, axis=0) + reshaped = ops.Reshape(stacked, new_shape) + + return ops.Cast(reshaped, to=int(TensorProto.FLOAT6E2M3)) + + class CastToUInt4Func(torch.autograd.Function): """ Custom op: unpacks packed uint8 → uint8 (values 0-15) in PyTorch. @@ -100,6 +151,39 @@ def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value: return output +class UnpackMxfp6Func(torch.autograd.Function): + """ + Custom op: unpacks packed MXFP6 bytes into logical 6-bit code positions. + + PyTorch forward returns uint8 code values for unit tests: + (..., in_features * 3 // 4) UINT8 -> (..., in_features) UINT8 + + ONNX symbolic emits an ONNXScript custom op. The subgraph ends with + Cast -> FLOAT6E2M3. + """ + + @staticmethod + def forward(weight_packed: torch.Tensor) -> torch.Tensor: + if weight_packed.shape[-1] % 3 != 0: + raise ValueError("Packed MXFP6 final dimension must be divisible by 3") + grouped = weight_packed.to(torch.uint8).reshape(*weight_packed.shape[:-1], weight_packed.shape[-1] // 3, 3) + code0 = grouped[..., 0] & 0x3F + code1 = ((grouped[..., 0] >> 6) | ((grouped[..., 1] & 0x0F) << 2)) & 0x3F + code2 = ((grouped[..., 1] >> 4) | ((grouped[..., 2] & 0x03) << 4)) & 0x3F + code3 = (grouped[..., 2] >> 2) & 0x3F + return torch.stack((code0, code1, code2, code3), dim=-1).reshape( + *weight_packed.shape[:-1], weight_packed.shape[-1] * 4 // 3 + ) + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value: + return g.onnxscript_op(UnpackMxfp6, weight_packed) + + class DequantizeLinearFunc(torch.autograd.Function): """ Emits a standard ONNX DequantizeLinear node (ai.onnx domain, not custom). diff --git a/QEfficient/exporter/weight_free/checkpoint_key_resolver.py b/QEfficient/exporter/weight_free/checkpoint_key_resolver.py index 62d255a8f8..665a2da2ff 100644 --- a/QEfficient/exporter/weight_free/checkpoint_key_resolver.py +++ b/QEfficient/exporter/weight_free/checkpoint_key_resolver.py @@ -180,6 +180,7 @@ def promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name ] backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model promoted_inputs: List[WeightSpecInput] = [] + spec_input_names = set() for name, init_value in list(model_ir.graph.initializers.items()): if name not in model_names: @@ -211,6 +212,31 @@ def promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name location=WeightSpecLocation(file=checkpoint_files.index(checkpoint_file), key=checkpoint_key), ) ) + spec_input_names.add(name) + + for value in model_ir.graph.inputs: + name = value.name + if name in spec_input_names or name not in model_names: + continue + + onnx_name = tied_weight_map.get(name, name) + checkpoint_key = find_checkpoint_key(onnx_name, checkpoint_index, backbone) + if checkpoint_key is None: + if _is_computed_initializer(onnx_name): + continue + raise ValueError( + f"Could not resolve model graph input '{name}' to a safetensors checkpoint key " + f"(resolved name: '{onnx_name}', model: '{model_ref}')." + ) + + checkpoint_file = checkpoint_index[checkpoint_key] + promoted_inputs.append( + WeightSpecInput( + name=name, + location=WeightSpecLocation(file=checkpoint_files.index(checkpoint_file), key=checkpoint_key), + ) + ) + spec_input_names.add(name) return WeightSpec( model_name=model_name, diff --git a/QEfficient/exporter/weight_free/export.py b/QEfficient/exporter/weight_free/export.py index db5b8c20ae..abdaa1cb56 100644 --- a/QEfficient/exporter/weight_free/export.py +++ b/QEfficient/exporter/weight_free/export.py @@ -15,6 +15,7 @@ from accelerate import init_empty_weights from QEfficient.exporter.weight_free.checkpoint_key_resolver import promote_initializers_and_build_spec +from QEfficient.exporter.weight_free.mxfp6 import finalize_mxfp6_export from QEfficient.exporter.weight_free.weight_spec import load_weight_spec, resolve_weight_spec_path, save_weight_spec from QEfficient.utils import load_json from QEfficient.utils.checkpoint_utils import resolve_checkpoint_dir @@ -47,6 +48,8 @@ def _run_quantizer_for_wf(qeff_model, target_dtype: torch.dtype): quant_config = getattr(qeff_model.model.config, "quantization_config", None) if quant_config is not None: + if getattr(getattr(qeff_model, "_mxfp6_config", None), "enabled", False): + raise ValueError("`mxfp6=True` does not currently support source HF quantization configs.") # For quantized models the meta model must use the same quantized layer types as the # checkpoint so that ONNX initializer names match the checkpoint's storage keys. # qeff_model.model was built via from_config (no real quantized checkpoint load), so @@ -133,6 +136,9 @@ def _prepare_checkpoint_for_weight_free_export( source_dir = resolve_checkpoint_dir(model_ref) dtype_suffix = str(target_dtype).replace("torch.", "") + mxfp6_config = getattr(qeff_model, "_mxfp6_config", None) + if getattr(mxfp6_config, "enabled", False): + dtype_suffix += f"-mxfp6-{mxfp6_config.scale_dtype}" # TODO(wf): For different flavours of the model that expect different checkpoint weight layouts, # we end up overriding old one. We need to add support of hashing/caching here. prepared_name = source_dir.name + f"-qeff-prepared-{dtype_suffix}" @@ -236,7 +242,13 @@ def export_weight_free_onnx( ) _prune_unused_fake_initializers(onnx_program) onnx_program.save(str(onnx_path)) - save_weight_spec(resolve_weight_spec_path(onnx_path), spec) + weight_spec_path = save_weight_spec(resolve_weight_spec_path(onnx_path), spec) + finalize_mxfp6_export( + onnx_path=onnx_path, + weight_spec_path=weight_spec_path, + prepared_model_ref=prepared_model_ref, + config=getattr(qeff_model, "_mxfp6_config", None), + ) return meta_qeff_model, onnx_transform_kwargs diff --git a/QEfficient/exporter/weight_free/mxfp6.py b/QEfficient/exporter/weight_free/mxfp6.py new file mode 100644 index 0000000000..4b9224142b --- /dev/null +++ b/QEfficient/exporter/weight_free/mxfp6.py @@ -0,0 +1,791 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ---------------------------------------------------------------------------- + +"""QEff-owned MXFP6 preparation for weight-free ONNX export.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Mapping, Optional, Sequence, Tuple, Union + +import onnx +import torch +from onnx import TensorProto, helper +from safetensors import safe_open + +from QEfficient.utils.checkpoint_utils import atomic_save, read_weight_map, write_index + +MXFP6_BLOCK_SIZE = 32 +MXFP6_ONNX_OPSET = 28 +MXFP6_QUANTIZER_VERSION = 2 +MXFP6_SUBFUNCTION_TOPOLOGY_VERSION = 1 +MXFP6_MAX_FINITE = 7.5 +MXFP6_PREPARED_SENTINEL = ".mxfp6_prepared" +MXFP6_SCALE_SUFFIX = ".mxfp6_scale" +MXFP6_PACKED_SUFFIX = ".mxfp6_packed" + +_SCALE_DTYPE_ALIASES = { + "float16": "float16", + "fp16": "float16", + "half": "float16", + "float32": "float32", + "fp32": "float32", + "float": "float32", + "bfloat16": "bfloat16", + "bf16": "bfloat16", + "fp8": "e8m0", + "e8m0": "e8m0", + "float8e8m0": "e8m0", + "float8_e8m0fnu": "e8m0", +} + +_SCALE_TORCH_DTYPES = { + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, +} + + +@dataclass(frozen=True) +class Mxfp6Config: + """Immutable normalized MXFP6 export settings.""" + + enabled: bool = False + scale_dtype: str = "float16" + + +@dataclass(frozen=True) +class _FlatMxfp6Target: + insert_before_node: onnx.NodeProto + + +@dataclass(frozen=True) +class _FunctionMxfp6Target: + function: onnx.FunctionProto + call_node: onnx.NodeProto + formal_index: int + formal_name: str + + +_Mxfp6Target = Union[_FlatMxfp6Target, _FunctionMxfp6Target] + + +@dataclass(frozen=True) +class _Mxfp6Replacement: + original_name: str + logical_key: str + value_info: onnx.ValueInfoProto + target: _Mxfp6Target + packed: torch.Tensor + scale: torch.Tensor + + +def normalize_mxfp6_config(enabled: bool, scale_dtype: str = "float16") -> Mxfp6Config: + """Normalize public MXFP6 options into the internal immutable config.""" + if not enabled: + return Mxfp6Config() + normalized = _SCALE_DTYPE_ALIASES.get(str(scale_dtype).lower()) + if normalized is None: + raise ValueError( + "`mxfp6_scale_dtype` must be one of: " + "float16/fp16/half, float32/fp32/float, bfloat16/bf16, fp8/e8m0/float8e8m0." + ) + return Mxfp6Config(enabled=True, scale_dtype=normalized) + + +def _tensorproto_enum(name: str) -> Optional[int]: + return getattr(TensorProto, name, None) + + +def _dql_schema_supports_opset28() -> bool: + try: + schema = onnx.defs.get_schema("DequantizeLinear", 28, "") + except Exception: + return False + return schema.since_version >= 28 + + +def _dql_schema_supports_output_dtype() -> bool: + try: + schema = onnx.defs.get_schema("DequantizeLinear", 28, "") + except Exception: + return False + return "output_dtype" in schema.attributes + + +def validate_mxfp6_capabilities(config: Mxfp6Config, feature_name: str = "mxfp6=True") -> None: + """Validate Python and ONNX capabilities needed for QEff-owned MXFP6 export.""" + if not config.enabled: + return + if sys.version_info < (3, 10): + raise AssertionError(f"{feature_name} requires Python >= 3.10") + missing = [] + if _tensorproto_enum("FLOAT6E2M3") is None: + missing.append("onnx.TensorProto.FLOAT6E2M3") + if config.scale_dtype == "e8m0" and _tensorproto_enum("FLOAT8E8M0") is None: + missing.append("onnx.TensorProto.FLOAT8E8M0") + if not _dql_schema_supports_opset28(): + missing.append("DequantizeLinear schema >= 28") + if missing: + raise AssertionError( + f"{feature_name} requires ONNX support for QEff-owned MXFP6 export, but this environment is missing: " + + ", ".join(missing) + ) + + +def _fp6_positive_codebook() -> List[Tuple[int, float]]: + values = [(0, 0.0)] + for exp_bits in range(4): + for mant in range(8): + if exp_bits == 0 and mant == 0: + continue + if exp_bits == 0: + value = mant / 8.0 + else: + value = (1.0 + mant / 8.0) * (2.0 ** (exp_bits - 1)) + values.append(((exp_bits << 3) | mant, value)) + return values + + +_POSITIVE_CODEBOOK = _fp6_positive_codebook() +_POSITIVE_CODES = torch.tensor([code for code, _ in _POSITIVE_CODEBOOK], dtype=torch.uint8) +_POSITIVE_VALUES = torch.tensor([value for _, value in _POSITIVE_CODEBOOK], dtype=torch.float32) +_POSITIVE_MIDPOINTS = (_POSITIVE_VALUES[:-1] + _POSITIVE_VALUES[1:]) / 2 +_LOWER_TIE_IS_ODD = (_POSITIVE_CODES[:-1].to(torch.int16) % 2) == 1 + + +def pack_fp6_codes(codes: torch.Tensor) -> torch.Tensor: + """Pack four 6-bit codes into three bytes, LSB-first.""" + flat = codes.detach().cpu().to(torch.uint8).flatten() + if flat.numel() % 4 != 0: + raise ValueError("FP6 code count must be divisible by 4") + grouped = flat.reshape(-1, 4).to(torch.int32) + packed0 = grouped[:, 0] | ((grouped[:, 1] & 0x03) << 6) + packed1 = ((grouped[:, 1] >> 2) & 0x0F) | ((grouped[:, 2] & 0x0F) << 4) + packed2 = ((grouped[:, 2] >> 4) & 0x03) | (grouped[:, 3] << 2) + return torch.stack((packed0, packed1, packed2), dim=1).to(torch.uint8).flatten() + + +def unpack_fp6_codes(packed: torch.Tensor) -> torch.Tensor: + """Unpack LSB-first ONNX FP6 bytes into 6-bit codes.""" + flat = packed.detach().cpu().to(torch.uint8).flatten() + if flat.numel() % 3 != 0: + raise ValueError("Packed FP6 byte count must be divisible by 3") + grouped = flat.reshape(-1, 3).to(torch.int32) + code0 = grouped[:, 0] & 0x3F + code1 = ((grouped[:, 0] >> 6) | ((grouped[:, 1] & 0x0F) << 2)) & 0x3F + code2 = ((grouped[:, 1] >> 4) | ((grouped[:, 2] & 0x03) << 4)) & 0x3F + code3 = (grouped[:, 2] >> 2) & 0x3F + return torch.stack((code0, code1, code2, code3), dim=1).to(torch.uint8).flatten() + + +def quantize_to_mxfp6(weight: torch.Tensor, scale_dtype: str) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a dense tensor to per-final-axis-block E2M3 codes. + + The E2M3 value table is defined locally from sign, two exponent bits, + and three mantissa bits, then packed using ONNX's LSB-first FP6 byte + layout. Runtime parity is gated on ONNX FLOAT6/DQL-28 support and is not + inferred from this helper alone. + """ + if not weight.is_floating_point(): + raise TypeError("MXFP6 quantization only supports floating-point weights") + if weight.shape[-1] % MXFP6_BLOCK_SIZE != 0: + raise ValueError( + f"MXFP6 requires the final weight dimension to be divisible by {MXFP6_BLOCK_SIZE}; got {tuple(weight.shape)}" + ) + fp32 = weight.detach().cpu().to(torch.float32) + if not torch.isfinite(fp32).all(): + raise ValueError("MXFP6 quantization does not support NaN or Inf weights") + + block_shape = (*fp32.shape[:-1], fp32.shape[-1] // MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE) + blocks = fp32.reshape(block_shape) + amax = blocks.abs().amax(dim=-1) + scales = torch.where( + amax == 0, + torch.ones_like(amax), + torch.pow(torch.tensor(2.0, dtype=torch.float32), torch.ceil(torch.log2(amax / MXFP6_MAX_FINITE))), + ) + normalized = blocks / scales.unsqueeze(-1) + abs_values = normalized.abs().clamp(max=MXFP6_MAX_FINITE).flatten() + lower_indices = torch.bucketize(abs_values, _POSITIVE_MIDPOINTS, right=False) + upper_indices = torch.bucketize(abs_values, _POSITIVE_MIDPOINTS, right=True) + tie = lower_indices != upper_indices + use_upper = tie & _LOWER_TIE_IS_ODD[lower_indices.clamp(max=_LOWER_TIE_IS_ODD.numel() - 1)] + indices = torch.where(use_upper, upper_indices, lower_indices) + quantized = _POSITIVE_CODES[indices] + quantized = quantized | (normalized.flatten() < 0).to(torch.uint8) * 0x20 + + packed_shape = (*fp32.shape[:-1], fp32.shape[-1] * 3 // 4) + packed = pack_fp6_codes(quantized.flatten()).reshape(packed_shape) + if scale_dtype == "e8m0": + scales_out = _encode_e8m0_scales(scales) + else: + scales_out = scales.to(_SCALE_TORCH_DTYPES[scale_dtype]) + return packed, scales_out.contiguous() + + +def _encode_e8m0_scales(scales: torch.Tensor) -> torch.Tensor: + """Encode positive power-of-two scales as E8M0 exponent bytes.""" + exponents = torch.round(torch.log2(scales.to(torch.float32))).to(torch.int32) + 127 + if torch.any((exponents < 0) | (exponents > 254)): + raise ValueError("MXFP6 E8M0 scale exponent is out of encodable range") + return exponents.to(torch.uint8) + + +def _graph_inputs_by_name(graph) -> Dict[str, onnx.ValueInfoProto]: + return {value_info.name: value_info for value_info in graph.input} + + +def _node_consumers(graph) -> Dict[str, List[onnx.NodeProto]]: + consumers: Dict[str, List[onnx.NodeProto]] = {} + for node in graph.node: + for name in node.input: + consumers.setdefault(name, []).append(node) + return consumers + + +def _is_final_axis_transpose(node: onnx.NodeProto) -> bool: + if node.op_type != "Transpose": + return False + perm_attr = next((attr for attr in node.attribute if attr.name == "perm"), None) + if perm_attr is None: + return True + perm = list(perm_attr.ints) + return len(perm) >= 2 and perm[:-2] == list(range(len(perm) - 2)) and perm[-2:] == [len(perm) - 1, len(perm) - 2] + + +def _function_input_index(function: onnx.FunctionProto, call_node: onnx.NodeProto, input_name: str) -> Optional[int]: + input_indices = [idx for idx, name in enumerate(call_node.input) if name == input_name] + if len(input_indices) != 1: + return None + input_index = input_indices[0] + if input_index >= len(function.input): + return None + return input_index + + +def _mxfp6_flat_insert_before_node( + graph, + input_name: str, +) -> Optional[onnx.NodeProto]: + consumers = _node_consumers(graph) + direct = consumers.get(input_name, []) + if len(direct) != 1: + return None + node = direct[0] + if node.op_type == "MatMul" and len(node.input) > 1 and node.input[1] == input_name: + return node + if not _is_final_axis_transpose(node) or not node.output: + return None + transpose_consumers = consumers.get(node.output[0], []) + if ( + len(transpose_consumers) == 1 + and transpose_consumers[0].op_type == "MatMul" + and len(transpose_consumers[0].input) > 1 + and transpose_consumers[0].input[1] == node.output[0] + ): + return node + return None + + +def _mxfp6_target( + graph, + input_name: str, + function_lookup: Mapping[Tuple[str, str], onnx.FunctionProto], +) -> Optional[_Mxfp6Target]: + insert_before_node = _mxfp6_flat_insert_before_node(graph, input_name) + if insert_before_node is not None: + return _FlatMxfp6Target(insert_before_node=insert_before_node) + + consumers = _node_consumers(graph).get(input_name, []) + if len(consumers) != 1: + return None + call_node = consumers[0] + function = function_lookup.get((call_node.domain, call_node.op_type)) + if function is None: + return None + formal_index = _function_input_index(function, call_node, input_name) + if formal_index is None: + return None + formal_name = function.input[formal_index] + if _mxfp6_flat_insert_before_node(function, formal_name) is None: + return None + return _FunctionMxfp6Target( + function=function, + call_node=call_node, + formal_index=formal_index, + formal_name=formal_name, + ) + + +def _is_mxfp6_candidate_topology( + graph, + input_name: str, + function_lookup: Optional[Mapping[Tuple[str, str], onnx.FunctionProto]] = None, +) -> bool: + consumers = _node_consumers(graph).get(input_name, []) + if len(consumers) != 1: + return False + node = consumers[0] + if node.op_type in {"MatMul", "Transpose"}: + return True + if function_lookup is None: + return False + function = function_lookup.get((node.domain, node.op_type)) + if function is None: + return False + formal_index = _function_input_index(function, node, input_name) + if formal_index is None: + return False + return _is_mxfp6_candidate_topology(function, function.input[formal_index]) + + +def _is_lm_head_weight_name(name: str) -> bool: + components = [component for component in name.replace("/", ".").split(".") if component] + return len(components) >= 2 and components[-2:] == ["lm_head", "weight"] + + +def _is_lm_head_weight(spec_input) -> bool: + return _is_lm_head_weight_name(spec_input.name) or _is_lm_head_weight_name(spec_input.location.key) + + +def _insert_before_node(graph, target_node: onnx.NodeProto, new_node: onnx.NodeProto) -> None: + for idx, node in enumerate(graph.node): + if node is target_node: + graph.node.insert(idx, new_node) + return + graph.node.insert(0, new_node) + + +def _append_value_info(container, value_info: onnx.ValueInfoProto) -> None: + names = {existing.name for existing in container.value_info} + if value_info.name not in names: + container.value_info.append(value_info) + + +def _remove_graph_input(graph, name: str) -> onnx.ValueInfoProto: + for idx, value_info in enumerate(graph.input): + if value_info.name == name: + removed = graph.input[idx] + del graph.input[idx] + return removed + raise ValueError(f"ONNX graph input '{name}' not found") + + +def _tensor_value_info(name: str, elem_type: int, shape: Sequence[int]) -> onnx.ValueInfoProto: + return helper.make_tensor_value_info(name, elem_type, list(shape)) + + +def _set_default_opset(model: onnx.ModelProto, version: int) -> None: + for opset in model.opset_import: + if opset.domain == "": + opset.version = max(opset.version, version) + return + model.opset_import.append(helper.make_opsetid("", version)) + + +def _ensure_opset(model: onnx.ModelProto, domain: str, version: int) -> None: + if any(opset.domain == domain for opset in model.opset_import): + return + model.opset_import.append(helper.make_opsetid(domain, version)) + + +def _function_key(function: onnx.FunctionProto) -> Tuple[str, str]: + return function.domain, function.name + + +def _group_function_replacements( + replacements: Sequence[_Mxfp6Replacement], +) -> Dict[Tuple[str, str], Dict[int, List[_Mxfp6Replacement]]]: + grouped: Dict[Tuple[str, str], Dict[int, List[_Mxfp6Replacement]]] = {} + for replacement in replacements: + if not isinstance(replacement.target, _FunctionMxfp6Target): + continue + function_group = grouped.setdefault(_function_key(replacement.target.function), {}) + function_group.setdefault(replacement.target.formal_index, []).append(replacement) + return grouped + + +def _validate_function_replacements( + model: onnx.ModelProto, + replacements: Sequence[_Mxfp6Replacement], +) -> None: + grouped = _group_function_replacements(replacements) + replacement_by_name = {replacement.original_name: replacement for replacement in replacements} + for (domain, name), replacements_by_formal in grouped.items(): + formal_indices = sorted(replacements_by_formal) + calls = [node for node in model.graph.node if node.domain == domain and node.op_type == name] + for call_node in calls: + for formal_index in formal_indices: + if formal_index >= len(call_node.input): + raise NotImplementedError( + "mxfp6=True cannot rewrite shared ONNX subfunctions when a call is missing a converted " + f"formal input: function '{domain}::{name}', formal index {formal_index}." + ) + actual_name = call_node.input[formal_index] + replacement = replacement_by_name.get(actual_name) + if replacement is None or not isinstance(replacement.target, _FunctionMxfp6Target): + raise NotImplementedError( + "mxfp6=True cannot partially rewrite shared ONNX subfunctions. " + f"Function '{domain}::{name}' call input '{actual_name}' at formal index {formal_index} " + "does not have a matching MXFP6 replacement." + ) + if _function_key(replacement.target.function) != (domain, name): + raise NotImplementedError( + f"mxfp6=True found an inconsistent ONNX subfunction replacement for input '{actual_name}'." + ) + if replacement.target.formal_index != formal_index: + raise NotImplementedError( + "mxfp6=True found an inconsistent formal input mapping for " + f"function '{domain}::{name}' input '{actual_name}'." + ) + + reference_by_formal = { + formal_index: replacements_for_formal[0] + for formal_index, replacements_for_formal in replacements_by_formal.items() + } + for formal_index, reference in reference_by_formal.items(): + reference_packed_shape = list(reference.packed.shape) + reference_scale_shape = list(reference.scale.shape) + reference_logical_shape = list(_load_tensor_shape(reference.value_info)) + for replacement in replacements_by_formal[formal_index][1:]: + if ( + list(replacement.packed.shape) != reference_packed_shape + or list(replacement.scale.shape) != reference_scale_shape + or list(_load_tensor_shape(replacement.value_info)) != reference_logical_shape + ): + raise NotImplementedError( + "mxfp6=True cannot rewrite a shared ONNX subfunction when converted actual inputs for " + f"formal index {formal_index} have different shapes." + ) + + +def _load_tensor_shape(value_info: onnx.ValueInfoProto) -> List[int]: + return [dim.dim_value for dim in value_info.type.tensor_type.shape.dim] + + +def _rewrite_function_replacements( + model: onnx.ModelProto, + replacements: Sequence[_Mxfp6Replacement], + scale_dtype: str, +) -> None: + grouped = _group_function_replacements(replacements) + if not grouped: + return + replacement_by_name = {replacement.original_name: replacement for replacement in replacements} + function_lookup = {_function_key(function): function for function in model.functions} + scale_elem_type = _scale_tensorproto_dtype(scale_dtype) + + for function_key, replacements_by_formal in grouped.items(): + function = function_lookup[function_key] + formal_indices = sorted(replacements_by_formal) + for formal_index in formal_indices: + replacement = replacements_by_formal[formal_index][0] + formal_name = function.input[formal_index] + if formal_name != replacement.target.formal_name: + raise NotImplementedError( + "mxfp6=True found a changed ONNX FunctionProto signature while rewriting " + f"function '{function.domain}::{function.name}'." + ) + logical_shape = list(replacement.packed.shape) + logical_shape[-1] = logical_shape[-1] * 4 // 3 + packed_shape = list(replacement.packed.shape) + scale_shape = list(replacement.scale.shape) + packed_formal_name = formal_name + MXFP6_PACKED_SUFFIX + scale_formal_name = formal_name + MXFP6_SCALE_SUFFIX + unpacked_name = formal_name + ".mxfp6_unpacked" + insert_before_node = _mxfp6_flat_insert_before_node(function, formal_name) + if insert_before_node is None: + raise NotImplementedError( + "mxfp6=True currently supports ONNX subfunction weights only as MatMul RHS, " + f"optionally through a sole final-axis Transpose. Unsupported formal '{formal_name}'." + ) + + function.input[formal_index] = packed_formal_name + _append_value_info(function, _tensor_value_info(packed_formal_name, TensorProto.UINT8, packed_shape)) + _append_value_info(function, _tensor_value_info(scale_formal_name, scale_elem_type, scale_shape)) + _append_value_info(function, _tensor_value_info(unpacked_name, TensorProto.FLOAT6E2M3, logical_shape)) + _append_value_info( + function, + _tensor_value_info(formal_name, _output_tensorproto_dtype(replacement.value_info), logical_shape), + ) + unpack_node = helper.make_node( + "UnpackMxfp6", + inputs=[packed_formal_name], + outputs=[unpacked_name], + name=formal_name + "_mxfp6_unpack", + domain="com.qti.aisw.onnx", + ) + dq_node = helper.make_node( + "DequantizeLinear", + inputs=[unpacked_name, scale_formal_name], + outputs=[formal_name], + name=formal_name + "_mxfp6_dq", + axis=-1, + block_size=MXFP6_BLOCK_SIZE, + **( + {"output_dtype": _output_tensorproto_dtype(replacement.value_info)} + if _dql_schema_supports_output_dtype() + else {} + ), + ) + _insert_before_node(function, insert_before_node, unpack_node) + _insert_before_node(function, insert_before_node, dq_node) + + for formal_index in formal_indices: + function.input.append(function.input[formal_index].removesuffix(MXFP6_PACKED_SUFFIX) + MXFP6_SCALE_SUFFIX) + + for call_node in [ + node for node in model.graph.node if node.domain == function.domain and node.op_type == function.name + ]: + for formal_index in formal_indices: + actual_name = call_node.input[formal_index] + replacement = replacement_by_name[actual_name] + call_node.input[formal_index] = replacement.original_name + MXFP6_PACKED_SUFFIX + for formal_index in formal_indices: + packed_actual_name = call_node.input[formal_index] + call_node.input.append(packed_actual_name.removesuffix(MXFP6_PACKED_SUFFIX) + MXFP6_SCALE_SUFFIX) + + _set_default_opset(function, MXFP6_ONNX_OPSET) + _ensure_opset(function, "com.qti.aisw.onnx", 1) + + +def _checkpoint_file_for_key(prepared_dir: Path, weight_map: Dict[str, str], key: str) -> Path: + shard_name = weight_map.get(key) + if shard_name is None: + raise ValueError(f"Could not find checkpoint key '{key}' in prepared checkpoint index") + return prepared_dir / shard_name + + +def _load_checkpoint_tensor(prepared_dir: Path, weight_map: Dict[str, str], key: str) -> torch.Tensor: + with safe_open(str(_checkpoint_file_for_key(prepared_dir, weight_map, key)), framework="pt") as handle: + return handle.get_tensor(key) + + +def _write_mxfp6_tensors( + prepared_dir: Path, + tensors: Dict[str, Tuple[torch.Tensor, torch.Tensor]], + weight_map: Dict[str, str], +) -> Dict[str, str]: + by_shard: Dict[str, Dict[str, torch.Tensor]] = {} + for key, (packed, scale) in tensors.items(): + shard_name = weight_map[key] + by_shard.setdefault(shard_name, {})[key + MXFP6_PACKED_SUFFIX] = packed + by_shard[shard_name][key + MXFP6_SCALE_SUFFIX] = scale + + for shard_name, replacements in by_shard.items(): + existing = {} + shard_path = prepared_dir / shard_name + with safe_open(str(shard_path), framework="pt") as handle: + for key in handle.keys(): + existing[key] = handle.get_tensor(key) + for key, tensor in replacements.items(): + if key not in existing: + existing[key] = tensor + atomic_save(existing, shard_path) + + new_weight_map = dict(weight_map) + for key in tensors: + new_weight_map[key + MXFP6_PACKED_SUFFIX] = new_weight_map[key] + new_weight_map[key + MXFP6_SCALE_SUFFIX] = new_weight_map[key] + write_index(prepared_dir, new_weight_map) + return new_weight_map + + +def finalize_mxfp6_export( + onnx_path: Path, weight_spec_path: Path, prepared_model_ref: str, config: Optional[Mxfp6Config] +) -> None: + """Rewrite saved weight-free ONNX and checkpoint files for QEff-owned MXFP6.""" + if config is None or not config.enabled: + return + from QEfficient.exporter.weight_free.weight_spec import ( + WeightSpecInput, + WeightSpecLocation, + load_weight_spec, + save_weight_spec, + ) + + validate_mxfp6_capabilities(config) + prepared_dir = Path(prepared_model_ref) + sentinel = prepared_dir / MXFP6_PREPARED_SENTINEL + spec = load_weight_spec(weight_spec_path) + model = onnx.load(str(onnx_path), load_external_data=False) + inputs_by_name = _graph_inputs_by_name(model.graph) + function_lookup = {(function.domain, function.name): function for function in model.functions} + weight_map = read_weight_map(prepared_dir) + + replacements = {} + mxfp6_replacements = [] + updated_inputs = [] + for spec_input in spec.inputs: + original_name = spec_input.name + logical_key = spec_input.location.key + value_info = inputs_by_name.get(original_name) + if value_info is None: + updated_inputs.append(spec_input) + continue + target = _mxfp6_target(model.graph, original_name, function_lookup) + if target is None: + if _is_mxfp6_candidate_topology(model.graph, original_name, function_lookup): + raise NotImplementedError( + "mxfp6=True currently supports only dense weights consumed as MatMul RHS, " + f"optionally through a sole final-axis Transpose. Unsupported topology for ONNX input " + f"'{original_name}'." + ) + updated_inputs.append(spec_input) + continue + if _is_lm_head_weight(spec_input): + updated_inputs.append(spec_input) + continue + tensor = _load_checkpoint_tensor(prepared_dir, weight_map, logical_key) + if tensor.ndim < 2: + updated_inputs.append(spec_input) + continue + packed, scale = quantize_to_mxfp6(tensor, config.scale_dtype) + replacements[logical_key] = (packed, scale) + mxfp6_replacements.append( + _Mxfp6Replacement( + original_name=original_name, + logical_key=logical_key, + value_info=value_info, + target=target, + packed=packed, + scale=scale, + ) + ) + + packed_input_name = original_name + MXFP6_PACKED_SUFFIX + scale_input_name = original_name + MXFP6_SCALE_SUFFIX + packed_key = logical_key + MXFP6_PACKED_SUFFIX + scale_key = logical_key + MXFP6_SCALE_SUFFIX + logical_shape = list(tensor.shape) + packed_shape = list(packed.shape) + scale_shape = list(scale.shape) + updated_inputs.append( + WeightSpecInput( + name=packed_input_name, + location=WeightSpecLocation(file=spec_input.location.file, key=packed_key), + role="mxfp6_weight", + metadata={ + "source_dtype": _value_info_dtype_name(value_info), + "logical_dtype": "float6e2m3", + "storage_dtype": "uint8", + "packing": "onnx_lsb_first_6bit", + "logical_shape": logical_shape, + "packed_shape": packed_shape, + "block_size": MXFP6_BLOCK_SIZE, + "axis": -1, + "scale_input": scale_input_name, + "unpack_output": original_name + ".mxfp6_unpacked", + }, + ) + ) + updated_inputs.append( + WeightSpecInput( + name=scale_input_name, + location=WeightSpecLocation(file=spec_input.location.file, key=scale_key), + role="mxfp6_scale", + weight_input=packed_input_name, + metadata={ + "logical_dtype": config.scale_dtype, + "storage_dtype": "uint8" if config.scale_dtype == "e8m0" else config.scale_dtype, + "logical_shape": scale_shape, + "packed_shape": scale_shape, + "block_size": MXFP6_BLOCK_SIZE, + "axis": -1, + }, + ) + ) + + if not replacements: + raise NotImplementedError( + "mxfp6=True currently supports only dense weights consumed as MatMul RHS, " + "optionally through a sole final-axis Transpose." + ) + + _validate_function_replacements(model, mxfp6_replacements) + scale_elem_type = _scale_tensorproto_dtype(config.scale_dtype) + for replacement in mxfp6_replacements: + original_name = replacement.original_name + packed_input_name = original_name + MXFP6_PACKED_SUFFIX + scale_input_name = original_name + MXFP6_SCALE_SUFFIX + unpacked_name = original_name + ".mxfp6_unpacked" + logical_shape = list(replacement.packed.shape) + logical_shape[-1] = logical_shape[-1] * 4 // 3 + packed_shape = list(replacement.packed.shape) + scale_shape = list(replacement.scale.shape) + _remove_graph_input(model.graph, original_name) + model.graph.input.append(_tensor_value_info(packed_input_name, TensorProto.UINT8, packed_shape)) + model.graph.input.append(_tensor_value_info(scale_input_name, scale_elem_type, scale_shape)) + if isinstance(replacement.target, _FlatMxfp6Target): + _append_value_info(model.graph, _tensor_value_info(unpacked_name, TensorProto.FLOAT6E2M3, logical_shape)) + unpack_node = helper.make_node( + "UnpackMxfp6", + inputs=[packed_input_name], + outputs=[unpacked_name], + name=original_name + "_mxfp6_unpack", + domain="com.qti.aisw.onnx", + ) + dq_node = helper.make_node( + "DequantizeLinear", + inputs=[unpacked_name, scale_input_name], + outputs=[original_name], + name=original_name + "_mxfp6_dq", + axis=-1, + block_size=MXFP6_BLOCK_SIZE, + **( + {"output_dtype": _output_tensorproto_dtype(replacement.value_info)} + if _dql_schema_supports_output_dtype() + else {} + ), + ) + _insert_before_node(model.graph, replacement.target.insert_before_node, unpack_node) + _insert_before_node(model.graph, replacement.target.insert_before_node, dq_node) + + _rewrite_function_replacements(model, mxfp6_replacements, config.scale_dtype) + _write_mxfp6_tensors(prepared_dir, replacements, weight_map) + spec.inputs = updated_inputs + spec.version = 7 + _set_default_opset(model, 28) + _ensure_opset(model, "com.qti.aisw.onnx", 1) + onnx.save(model, str(onnx_path)) + save_weight_spec(weight_spec_path, spec) + sentinel.write_text( + json.dumps( + {"quantizer_version": MXFP6_QUANTIZER_VERSION, "scale_dtype": config.scale_dtype}, + indent=2, + sort_keys=True, + ) + ) + + +def _scale_tensorproto_dtype(scale_dtype: str) -> int: + if scale_dtype == "float16": + return TensorProto.FLOAT16 + if scale_dtype == "float32": + return TensorProto.FLOAT + if scale_dtype == "bfloat16": + return TensorProto.BFLOAT16 + fp8 = _tensorproto_enum("FLOAT8E8M0") + if fp8 is None: + raise AssertionError("onnx.TensorProto.FLOAT8E8M0 is required for mxfp6_scale_dtype='e8m0'") + return fp8 + + +def _output_tensorproto_dtype(value_info: onnx.ValueInfoProto) -> int: + elem_type = value_info.type.tensor_type.elem_type + if elem_type: + return elem_type + return TensorProto.FLOAT16 + + +def _value_info_dtype_name(value_info: onnx.ValueInfoProto) -> str: + elem_type = value_info.type.tensor_type.elem_type + return TensorProto.DataType.Name(elem_type) if elem_type else "UNKNOWN" diff --git a/QEfficient/exporter/weight_free/ort_weight_injection.py b/QEfficient/exporter/weight_free/ort_weight_injection.py index cb56645d8a..fe7d766470 100644 --- a/QEfficient/exporter/weight_free/ort_weight_injection.py +++ b/QEfficient/exporter/weight_free/ort_weight_injection.py @@ -95,6 +95,11 @@ def load_weight_free_ort_inputs( """ weight_spec_path = Path(weight_spec_path) spec = load_weight_spec(weight_spec_path) + if any(spec_input.role.startswith("mxfp6") for spec_input in spec.inputs): + raise NotImplementedError( + "ONNX Runtime weight injection for QEff-owned MXFP6 weight-free exports is not supported. " + "Use the QAIC compile/runtime path with an ONNX build that supports FLOAT6E2M3 and DequantizeLinear-28." + ) candidate_roots = [] if weights_root is not None: candidate_roots.append(Path(weights_root).expanduser().resolve()) diff --git a/QEfficient/exporter/weight_free/weight_spec.py b/QEfficient/exporter/weight_free/weight_spec.py index b88b0c1871..d2e9c919b3 100644 --- a/QEfficient/exporter/weight_free/weight_spec.py +++ b/QEfficient/exporter/weight_free/weight_spec.py @@ -42,10 +42,17 @@ class WeightSpecLocation: @dataclass class WeightSpecInput: - """Mapping from an ONNX input name to its external checkpoint tensor.""" + """Mapping from an ONNX graph value to its external checkpoint tensor. + + Entries use the same file/key location contract regardless of whether the + checkpoint tensor is consumed directly or through an export-time transform. + """ name: str location: WeightSpecLocation # required: every spec entry must point to a file + role: str = "weight" + weight_input: str | None = None + metadata: Dict[str, Any] = field(default_factory=dict) @dataclass @@ -53,7 +60,7 @@ class WeightSpec: """Serializable weight-free export metadata. The spec records the checkpoint files shipped beside the ONNX model and - maps promoted ONNX weight inputs back to tensor keys in those files. + maps ONNX weight graph values back to tensor keys in those files. """ model_name: str @@ -66,6 +73,13 @@ def to_dict(self) -> Dict[str, Any]: """Return a JSON-serializable representation of the weight spec.""" data = asdict(self) data["model_id"] = str(data["model_id"]) + for entry in data["inputs"]: + if entry.get("role") == "weight": + entry.pop("role", None) + if entry.get("weight_input") is None: + entry.pop("weight_input", None) + if not entry.get("metadata"): + entry.pop("metadata", None) return data @@ -129,6 +143,9 @@ def load_weight_spec(path: Path) -> WeightSpec: WeightSpecInput( name=entry["name"], location=_load_location(entry["location"]), + role=entry.get("role", "weight"), + weight_input=entry.get("weight_input"), + metadata=entry.get("metadata", {}), ) for entry in data["inputs"] if entry.get("location") is not None # backward compat: skip old buffer-only entries diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index 213986d41e..07a246d56d 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -40,6 +40,7 @@ MoEExpertStackingCheckpointTransform, MoEFusedExpertSplitCheckpointTransform, ) +from QEfficient.exporter.weight_free.mxfp6 import normalize_mxfp6_config, validate_mxfp6_capabilities from QEfficient.generation.cloud_infer import QAICInferenceSession, is_retained_state_name from QEfficient.generation.text_generation_inference import ( CloudAI100ExecInfoNew, @@ -3641,6 +3642,8 @@ def from_pretrained( max_seq_len_cached: Optional[int] = None, layerwise: bool = False, weight_free: bool = False, + mxfp6: bool = False, + mxfp6_scale_dtype: str = "float16", *args, **kwargs, ): @@ -3685,6 +3688,14 @@ def from_pretrained( (dtype conversion, MoE expert restacking, etc.) is saved under the ``QEFF_CHECKPOINT_HOME`` environment variable if set, otherwise next to the source checkpoint under the Hugging Face cache. + mxfp6 : bool, optional + If True, quantizes dense weight-free MatMul weights to QEff-owned + MXFP6 and emits standard-domain DequantizeLinear nodes. Requires + ``weight_free=True``. Default is False. + mxfp6_scale_dtype : str, optional + Scale tensor dtype for MXFP6. Accepted aliases are float16/fp16/half, + float32/fp32/float, bfloat16/bf16, and fp8/e8m0/float8e8m0/float8_e8m0fnu. + Default is "float16". *args : Positional arguments passed directly to `cls._hf_auto_class.from_pretrained`. @@ -3699,10 +3710,19 @@ def from_pretrained( QEFFAutoModelForCausalLM An instance initialized with the pretrained weights. """ + mxfp6_config = normalize_mxfp6_config(mxfp6, mxfp6_scale_dtype) + if mxfp6_config.enabled and not weight_free: + raise ValueError("`mxfp6=True` requires `weight_free=True`.") if layerwise and weight_free: raise ValueError( "`layerwise=True` and `weight_free=True` are mutually exclusive; weight_free replaces layerwise mode." ) + if layerwise and mxfp6_config.enabled: + raise ValueError("`layerwise=True` and `mxfp6=True` are mutually exclusive.") + if mxfp6_config.enabled: + if "quantization_config" in kwargs or kwargs.get("load_in_4bit") or kwargs.get("load_in_8bit"): + raise ValueError("`mxfp6=True` does not currently support source HF quantization configs.") + validate_mxfp6_capabilities(mxfp6_config) if weight_free: validate_dynamo_export_requirements("weight_free=True") @@ -3753,9 +3773,14 @@ def from_pretrained( if qaic_config is not None: qaic_config["pretrained_model_name_or_path"] = pretrained_model_name_or_path + if mxfp6_config.enabled and getattr(model.config, "quantization_config", None) is not None: + raise ValueError("`mxfp6=True` does not currently support source HF quantization configs.") + # This is support models that should be classified to in a different auto class but transformers load them via this class kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {}) if model.__class__.__name__ in MISCLASSIFIED_CAUSAL_LM_TO_QEFF_AUTO_CLASS_MAP: + if mxfp6_config.enabled: + raise ValueError("`mxfp6=True` is currently supported only for text CausalLM weight-free export.") return MISCLASSIFIED_CAUSAL_LM_TO_QEFF_AUTO_CLASS_MAP[model.__class__.__name__]( model, kv_offload=kv_offload, @@ -3771,6 +3796,7 @@ def from_pretrained( pretrained_model_name_or_path=pretrained_model_name_or_path, max_seq_len_cached=max_seq_len_cached, weight_free=weight_free, + mxfp6_config=mxfp6_config, **kwargs, ) if layerwise: diff --git a/QEfficient/utils/export_utils.py b/QEfficient/utils/export_utils.py index 78bf2675be..bb67d8a6e4 100644 --- a/QEfficient/utils/export_utils.py +++ b/QEfficient/utils/export_utils.py @@ -25,6 +25,11 @@ RenameRepeatedSubgraphTransform, RenameWsubNodesTransform, ) +from QEfficient.exporter.weight_free.mxfp6 import ( + MXFP6_BLOCK_SIZE, + MXFP6_QUANTIZER_VERSION, + MXFP6_SUBFUNCTION_TOPOLOGY_VERSION, +) from QEfficient.transformers.cache_utils import InvalidIndexProvider from QEfficient.utils.cache import QEFF_HOME from QEfficient.utils.constants import ( @@ -433,6 +438,14 @@ def _generate_export_hash(qeff_model, args, kwargs, func): ) if getattr(qeff_model, "_weight_free", False): copy_of_hash_params["weight_free"] = True + mxfp6_config = getattr(qeff_model, "_mxfp6_config", None) + if getattr(mxfp6_config, "enabled", False): + copy_of_hash_params["mxfp6"] = True + copy_of_hash_params["mxfp6_scale_dtype"] = mxfp6_config.scale_dtype + copy_of_hash_params["mxfp6_block_size"] = MXFP6_BLOCK_SIZE + copy_of_hash_params["mxfp6_quantizer_version"] = MXFP6_QUANTIZER_VERSION + if getattr(qeff_model, "_use_onnx_subfunctions", False): + copy_of_hash_params["mxfp6_subfunction_topology_version"] = MXFP6_SUBFUNCTION_TOPOLOGY_VERSION if getattr(qeff_model, "_use_onnx_subfunctions", False): copy_of_hash_params["onnx_subfunction_version"] = 3 # Generate hash from relevant parameters diff --git a/QEfficient/utils/hash_utils.py b/QEfficient/utils/hash_utils.py index 4765f14d48..8c581fc102 100644 --- a/QEfficient/utils/hash_utils.py +++ b/QEfficient/utils/hash_utils.py @@ -20,7 +20,7 @@ def json_serializable(obj): if is_dataclass(obj): # Convert dataclass to dict for serialization return asdict(obj) - if obj.__class__.__name__ == "Dim": + if obj.__class__.__name__ in {"Dim", "_Dim"}: return str(obj) if obj.__class__.__name__ == "_DimHint": return str(obj) diff --git a/examples/dynamo/causal_lm/local_mxfp6_weight_free_export.py b/examples/dynamo/causal_lm/local_mxfp6_weight_free_export.py new file mode 100644 index 0000000000..d69c6ded89 --- /dev/null +++ b/examples/dynamo/causal_lm/local_mxfp6_weight_free_export.py @@ -0,0 +1,174 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +"""Local QEff-owned MXFP6 weight-free export smoke test. + +This script creates a tiny Llama checkpoint locally, exports it through the +weight-free Dynamo path with QEff-owned MXFP6 enabled, and prints the resulting +ONNX/WeightSpec structure. It does not download from Hugging Face. + +Example: + python examples/dynamo/causal_lm/local_mxfp6_weight_free_export.py \ + --work-dir /tmp/qeff_mxfp6_smoke +""" + +from __future__ import annotations + +import argparse +import json +import tempfile +from pathlib import Path + +import onnx +from transformers import LlamaConfig, LlamaForCausalLM + +from QEfficient.exporter.weight_free.weight_spec import load_weight_spec, resolve_weight_spec_path +from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM + + +def _make_config(args: argparse.Namespace) -> LlamaConfig: + return LlamaConfig( + num_hidden_layers=args.num_hidden_layers, + num_attention_heads=args.num_attention_heads, + num_key_value_heads=args.num_key_value_heads, + hidden_size=args.hidden_size, + intermediate_size=args.intermediate_size, + vocab_size=args.vocab_size, + max_position_embeddings=args.max_position_embeddings, + ) + + +def _load_onnx(onnx_path: Path): + return onnx.load(str(onnx_path), load_external_data=False) + + +def _count_nodes(model, op_type: str) -> int: + return sum(1 for node in model.graph.node if node.op_type == op_type) + + +def _count_function_body_nodes(model, op_type: str) -> int: + return sum(1 for function in model.functions for node in function.node if node.op_type == op_type) + + +def _count_functions(model, op_type: str) -> int: + return sum(1 for function in model.functions if function.name == op_type) + + +def _decoder_function_signatures(model) -> list[dict]: + return [ + { + "domain": function.domain, + "name": function.name, + "inputs": list(function.input), + } + for function in model.functions + if function.domain != "com.qti.aisw.onnx" + ] + + +def _mxfp6_graph_input_summary(model) -> dict: + graph_inputs = {value_info.name for value_info in model.graph.input} + return { + "packed_graph_inputs": sum(1 for name in graph_inputs if name.endswith(".mxfp6_packed")), + "scale_graph_inputs": sum(1 for name in graph_inputs if name.endswith(".mxfp6_scale")), + } + + +def _has_lm_head_component(name: str) -> bool: + return "lm_head" in [component for component in name.replace("/", ".").split(".") if component] + + +def _lm_head_summary(model, spec) -> dict: + graph_inputs = {value_info.name for value_info in model.graph.input} + spec_inputs = [entry for entry in spec.inputs if _has_lm_head_component(entry.name)] + return { + "graph_inputs": sorted(name for name in graph_inputs if _has_lm_head_component(name)), + "mxfp6_graph_inputs": sorted( + name + for name in graph_inputs + if _has_lm_head_component(name) and name.endswith((".mxfp6_packed", ".mxfp6_scale")) + ), + "spec_inputs": [ + { + "name": entry.name, + "key": entry.location.key, + "role": entry.role, + } + for entry in spec_inputs + ], + } + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Create and export a tiny local Llama checkpoint with QEff-owned MXFP6 weight-free export.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--work-dir", type=Path, default=None, help="Directory for local checkpoint/export artifacts") + parser.add_argument("--scale-dtype", default="float16", help="MXFP6 scale dtype alias") + parser.add_argument("--use-onnx-subfunctions", action="store_true", help="Enable ONNX subfunction export") + parser.add_argument("--num-hidden-layers", type=int, default=1) + parser.add_argument("--num-attention-heads", type=int, default=4) + parser.add_argument("--num-key-value-heads", type=int, default=4) + parser.add_argument("--hidden-size", type=int, default=32) + parser.add_argument("--intermediate-size", type=int, default=64) + parser.add_argument("--vocab-size", type=int, default=128) + parser.add_argument("--max-position-embeddings", type=int, default=64) + args = parser.parse_args() + + root = args.work_dir or Path(tempfile.mkdtemp(prefix="qeff_mxfp6_local_")) + root = root.expanduser().resolve() + model_dir = root / "model" + export_dir = root / "export" + model_dir.mkdir(parents=True, exist_ok=True) + export_dir.mkdir(parents=True, exist_ok=True) + + model = LlamaForCausalLM(_make_config(args)).eval() + model.save_pretrained(model_dir, safe_serialization=True) + + qeff_model = QEFFAutoModelForCausalLM.from_pretrained( + str(model_dir), + weight_free=True, + mxfp6=True, + mxfp6_scale_dtype=args.scale_dtype, + ) + qeff_model.model.eval() + export_result = qeff_model.export( + export_dir, + use_onnx_subfunctions=args.use_onnx_subfunctions, + offload_pt_weights=False, + ) + onnx_path = Path(export_result[-1] if isinstance(export_result, (list, tuple)) else export_result) + weight_spec_path = resolve_weight_spec_path(onnx_path) + spec = load_weight_spec(weight_spec_path) + onnx_model = _load_onnx(onnx_path) + + role_counts = {} + for spec_input in spec.inputs: + role_counts[spec_input.role] = role_counts.get(spec_input.role, 0) + 1 + + summary = { + "work_dir": str(root), + "onnx_path": str(onnx_path), + "weight_spec_path": str(weight_spec_path), + "weight_spec_version": spec.version, + "weight_spec_role_counts": role_counts, + "top_level_dequantize_linear_nodes": _count_nodes(onnx_model, "DequantizeLinear"), + "top_level_unpack_mxfp6_nodes": _count_nodes(onnx_model, "UnpackMxfp6"), + "function_body_dequantize_linear_nodes": _count_function_body_nodes(onnx_model, "DequantizeLinear"), + "function_body_unpack_mxfp6_nodes": _count_function_body_nodes(onnx_model, "UnpackMxfp6"), + "unpack_mxfp6_functions": _count_functions(onnx_model, "UnpackMxfp6"), + "decoder_function_signatures": _decoder_function_signatures(onnx_model), + "matmul_nodes": _count_nodes(onnx_model, "MatMul"), + "mxfp6_graph_input_summary": _mxfp6_graph_input_summary(onnx_model), + "lm_head_summary": _lm_head_summary(onnx_model, spec), + } + print(json.dumps(summary, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/examples/dynamo/causal_lm/requirements.txt b/examples/dynamo/causal_lm/requirements.txt index 61786599f1..176a5fe2c4 100644 --- a/examples/dynamo/causal_lm/requirements.txt +++ b/examples/dynamo/causal_lm/requirements.txt @@ -18,4 +18,8 @@ torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.28.0%2Bcpu-cp31 # Additional dynamo dependencies accelerate==1.9.0 compressed-tensors==0.17.0 +ml_dtypes>=0.5.4 +onnx @ git+https://github.com/onnx/onnx.git@3ab9181bbe2bb73cfd9ab88d4bdb8fad539d59b3 +onnx-ir==1.0.0 onnxscript==0.6.2 +protobuf>=6.31.1 diff --git a/tests/weight_free/conftest.py b/tests/weight_free/conftest.py index e427173e6e..2796a3b88e 100644 --- a/tests/weight_free/conftest.py +++ b/tests/weight_free/conftest.py @@ -34,6 +34,7 @@ def _parse_torch_version() -> tuple: def pytest_configure(config): config.addinivalue_line("markers", "weight_free: mark a test as part of the weight-free export test suite") + config.addinivalue_line("markers", "weight_free_unit: mark weight-free unit tests that do not run Dynamo export") config.addinivalue_line("markers", "weight_free_export: CPU-only weight-free export smoke and parity tests") @@ -50,7 +51,7 @@ def pytest_collection_modifyitems(config, items): for item in items: if not (item.fspath.parts and "weight_free" in str(item.fspath)): continue - if torch_version < (2, 13): + if torch_version < (2, 13) and not item.get_closest_marker("weight_free_unit"): item.add_marker( pytest.mark.skip(reason=f"Weight-free tests require torch >= 2.13; running {torch.__version__}") ) diff --git a/tests/weight_free/test_transforms.py b/tests/weight_free/test_transforms.py index b813f2337b..2cd5bf7625 100644 --- a/tests/weight_free/test_transforms.py +++ b/tests/weight_free/test_transforms.py @@ -22,6 +22,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import onnx import onnx_ir as ir import pytest import torch @@ -32,10 +33,12 @@ from QEfficient.base.checkpoint_transforms import CHECKPOINT_PREPARED_MANIFEST, CheckpointTransformPipeline from QEfficient.base.onnx_transforms import ( + CustomOpTransform, PreserveNestedCacheRetainedStateTransform, PruneFakeInitializersTransform, RenameRepeatedSubgraphTransform, ) +from QEfficient.customop.quantization_ops import UnpackMxfp6Func from QEfficient.exporter.weight_free import checkpoint_key_resolver from QEfficient.exporter.weight_free.checkpoint_key_resolver import find_checkpoint_key from QEfficient.exporter.weight_free.checkpoint_transforms import ( @@ -44,6 +47,27 @@ MoEExpertStackingCheckpointTransform, MoEFusedExpertSplitCheckpointTransform, ) +from QEfficient.exporter.weight_free.mxfp6 import ( + MXFP6_BLOCK_SIZE, + MXFP6_PACKED_SUFFIX, + MXFP6_SCALE_SUFFIX, + Mxfp6Config, + _is_lm_head_weight_name, + finalize_mxfp6_export, + normalize_mxfp6_config, + pack_fp6_codes, + quantize_to_mxfp6, + unpack_fp6_codes, +) +from QEfficient.exporter.weight_free.ort_weight_injection import load_weight_free_ort_inputs +from QEfficient.exporter.weight_free.weight_spec import ( + ExternalDataFile, + WeightSpec, + WeightSpecInput, + WeightSpecLocation, + load_weight_spec, + save_weight_spec, +) from QEfficient.transformers.models.llama.modeling_llama import QEffLlamaDecoderLayer from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM from QEfficient.utils import runtime_requirements @@ -51,6 +75,8 @@ from QEfficient.utils.runtime_requirements import validate_runtime_requirements from QEfficient.utils.torch_patches import temporarily_enable_nested_compile_regions +ONNX_FLOAT6E2M3 = int(TensorProto.FLOAT6E2M3) + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -463,6 +489,896 @@ def __init__(self): assert spec.inputs[0].location.key == "model.embed_tokens.weight" +@pytest.mark.weight_free_unit +class TestWeightFreeMxfp6: + @pytest.mark.parametrize( + "alias,normalized", + [ + ("float16", "float16"), + ("fp16", "float16"), + ("half", "float16"), + ("float32", "float32"), + ("fp32", "float32"), + ("float", "float32"), + ("bfloat16", "bfloat16"), + ("bf16", "bfloat16"), + ("fp8", "e8m0"), + ("e8m0", "e8m0"), + ("float8e8m0", "e8m0"), + ("float8_e8m0fnu", "e8m0"), + ], + ) + def test_mxfp6_scale_dtype_aliases(self, alias, normalized): + assert normalize_mxfp6_config(True, alias) == Mxfp6Config(enabled=True, scale_dtype=normalized) + + def test_mxfp6_scale_dtype_defaults_to_float16(self): + assert normalize_mxfp6_config(True) == Mxfp6Config(enabled=True, scale_dtype="float16") + + def test_mxfp6_rejects_unknown_scale_dtype(self): + with pytest.raises(ValueError, match="mxfp6_scale_dtype"): + normalize_mxfp6_config(True, "int8") + + def test_mxfp6_requires_weight_free(self): + with pytest.raises(ValueError, match="requires `weight_free=True`"): + QEFFAutoModelForCausalLM.from_pretrained("dummy-model", mxfp6=True) + + def test_mxfp6_rejects_layerwise(self): + with pytest.raises(ValueError, match="layerwise=True"): + QEFFAutoModelForCausalLM.from_pretrained("dummy-model", layerwise=True, weight_free=True, mxfp6=True) + + def test_mxfp6_rejects_source_quantization_config(self, monkeypatch): + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto.validate_mxfp6_capabilities", lambda config: None + ) + with pytest.raises(ValueError, match="source HF quantization configs"): + QEFFAutoModelForCausalLM.from_pretrained( + "dummy-model", + weight_free=True, + mxfp6=True, + quantization_config={"quant_method": "gptq"}, + ) + + def test_mxfp6_config_stored_on_wrapper(self, monkeypatch): + model_hf, _ = make_tiny_llama() + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto.validate_mxfp6_capabilities", lambda config: None + ) + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto.validate_dynamo_export_requirements", lambda name: None + ) + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto._build_meta_model", lambda *args, **kwargs: model_hf + ) + + qeff_model = QEFFAutoModelForCausalLM.from_pretrained( + "dummy-model", + weight_free=True, + mxfp6=True, + mxfp6_scale_dtype="half", + ) + + assert qeff_model._mxfp6_config == Mxfp6Config(enabled=True, scale_dtype="float16") + + def test_mxfp6_rejects_misclassified_vlm_redirect(self, monkeypatch): + class InternVLChatModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace() + + model_hf = InternVLChatModel() + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto.validate_mxfp6_capabilities", lambda config: None + ) + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto.validate_dynamo_export_requirements", lambda name: None + ) + monkeypatch.setattr( + "QEfficient.transformers.models.modeling_auto._build_meta_model", lambda *args, **kwargs: model_hf + ) + + with pytest.raises(ValueError, match="text CausalLM weight-free export"): + QEFFAutoModelForCausalLM.from_pretrained("dummy-model", weight_free=True, mxfp6=True) + + def test_mxfp6_rejects_compiler_owned_mxfp6_matmul(self, tmp_path): + model_hf, _ = make_tiny_llama() + qeff_model = QEFFAutoModelForCausalLM( + model_hf, + weight_free=True, + mxfp6_config=Mxfp6Config(enabled=True, scale_dtype="float16"), + ) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"placeholder") + + for option_name in ("mxfp6_matmul", "mxfp6-matmul", "mxfp6"): + with pytest.raises(ValueError, match="mxfp6"): + qeff_model._compile(onnx_path=str(onnx_path), compile_dir=str(tmp_path), **{option_name: True}) + + with pytest.raises(ValueError, match="QNN compilation"): + qeff_model._compile(onnx_path=str(onnx_path), compile_dir=str(tmp_path), enable_qnn=True) + + def test_fp6_pack_unpack_roundtrip(self): + codes = torch.arange(64, dtype=torch.uint8) + packed = pack_fp6_codes(codes) + assert packed.numel() == 48 + torch.testing.assert_close(unpack_fp6_codes(packed), codes) + + def test_fp6_pack_uses_onnx_lsb_first_layout(self): + codes = torch.tensor([0x00, 0x01, 0x02, 0x03], dtype=torch.uint8) + + torch.testing.assert_close( + pack_fp6_codes(codes), + torch.tensor([0x40, 0x20, 0x0C], dtype=torch.uint8), + ) + + def test_mxfp6_quantization_uses_block_scales_and_rejects_nonfinite(self): + weight = torch.zeros(2, MXFP6_BLOCK_SIZE * 2, dtype=torch.float32) + weight[0, :MXFP6_BLOCK_SIZE] = 15.0 + packed, scales = quantize_to_mxfp6(weight, "float16") + + assert packed.dtype == torch.uint8 + assert packed.shape == (2, MXFP6_BLOCK_SIZE * 2 * 3 // 4) + assert packed.numel() == weight.numel() * 6 // 8 + assert UnpackMxfp6Func.apply(packed).shape == weight.shape + assert scales.dtype == torch.float16 + torch.testing.assert_close(scales[0, 0], torch.tensor(2.0, dtype=torch.float16)) + torch.testing.assert_close(scales[0, 1], torch.tensor(1.0, dtype=torch.float16)) + + weight[0, 0] = float("nan") + with pytest.raises(ValueError, match="NaN or Inf"): + quantize_to_mxfp6(weight, "float16") + + @pytest.mark.parametrize( + "name", + [ + "lm_head.weight", + "model.lm_head.weight", + "base_model.lm_head.weight", + "base_model/lm_head/weight", + ], + ) + def test_lm_head_weight_name_predicate_matches_path_component_suffix(self, name): + assert _is_lm_head_weight_name(name) + + @pytest.mark.parametrize( + "name", + [ + "prefix_lm_head.weight", + "lm_head_projection.weight", + "model.lm_head.weight.extra", + "model.layers.0.lm_head_adapter.weight", + ], + ) + def test_lm_head_weight_name_predicate_rejects_partial_components(self, name): + assert not _is_lm_head_weight_name(name) + + def test_mxfp6_finalizer_writes_v7_location_metadata(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6._dql_schema_supports_output_dtype", lambda: True) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + weight = torch.arange(4 * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape(4, MXFP6_BLOCK_SIZE) + _write_safetensors_checkpoint(prepared, {"linear.weight": weight}) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]) + w = helper.make_tensor_value_info("linear.weight", TensorProto.FLOAT, [4, MXFP6_BLOCK_SIZE]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + graph = helper.make_graph( + [helper.make_node("MatMul", ["x", "linear.weight"], ["y"])], + "mxfp6_test", + [x, w], + [y], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="linear.weight", + location=WeightSpecLocation(file=0, key="linear.weight"), + ) + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + spec = load_weight_spec(spec_path) + raw_spec = json.loads(spec_path.read_text()) + assert spec.version == 7 + assert [entry.role for entry in spec.inputs] == ["mxfp6_weight", "mxfp6_scale"] + assert spec.inputs[0].name == "linear.weight" + MXFP6_PACKED_SUFFIX + assert spec.inputs[0].location.key == "linear.weight" + MXFP6_PACKED_SUFFIX + assert spec.inputs[0].metadata["logical_dtype"] == "float6e2m3" + assert spec.inputs[0].metadata["storage_dtype"] == "uint8" + assert spec.inputs[0].metadata["packing"] == "onnx_lsb_first_6bit" + assert spec.inputs[0].metadata["scale_input"] == "linear.weight" + MXFP6_SCALE_SUFFIX + assert spec.inputs[0].metadata["unpack_output"] == "linear.weight.mxfp6_unpacked" + assert spec.inputs[0].metadata["packed_shape"] == [4, MXFP6_BLOCK_SIZE * 3 // 4] + assert spec.inputs[1].location.key == "linear.weight" + MXFP6_SCALE_SUFFIX + assert raw_spec["files"] == [{"format": "safetensors", "path": "model.safetensors"}] + assert {entry["location"]["key"] for entry in raw_spec["inputs"]} == { + "linear.weight" + MXFP6_PACKED_SUFFIX, + "linear.weight" + MXFP6_SCALE_SUFFIX, + } + assert all(set(entry["location"]) == {"file", "key"} for entry in raw_spec["inputs"]) + + def json_keys(value): + if isinstance(value, dict): + yield from value + for nested_value in value.values(): + yield from json_keys(nested_value) + elif isinstance(value, list): + for nested_value in value: + yield from json_keys(nested_value) + + assert not {"base64", "packed_bytes", "payload", "scale_values", "values"} & set(json_keys(raw_spec)) + + tensors = _load_prepared_tensors(prepared) + assert tensors["linear.weight"].dtype == torch.float32 + assert tensors["linear.weight" + MXFP6_PACKED_SUFFIX].dtype == torch.uint8 + assert tensors["linear.weight" + MXFP6_SCALE_SUFFIX].dtype == torch.float16 + rewritten = onnx.load(str(onnx_path), load_external_data=False) + assert rewritten.opset_import[0].version == 28 + graph_inputs_by_name = {value_info.name: value_info for value_info in rewritten.graph.input} + graph_input_names = set(graph_inputs_by_name) + initializer_by_name = {initializer.name: initializer for initializer in rewritten.graph.initializer} + assert "linear.weight" not in graph_input_names + assert "linear.weight" + MXFP6_PACKED_SUFFIX in graph_input_names + assert "linear.weight" + MXFP6_SCALE_SUFFIX in graph_input_names + assert "linear.weight" + MXFP6_PACKED_SUFFIX not in initializer_by_name + assert "linear.weight" + MXFP6_SCALE_SUFFIX not in initializer_by_name + + def input_shape(value_info): + return [dim.dim_value for dim in value_info.type.tensor_type.shape.dim] + + packed_input = graph_inputs_by_name["linear.weight" + MXFP6_PACKED_SUFFIX] + scale_input = graph_inputs_by_name["linear.weight" + MXFP6_SCALE_SUFFIX] + assert packed_input.type.tensor_type.elem_type == TensorProto.UINT8 + assert input_shape(packed_input) == [4, MXFP6_BLOCK_SIZE * 3 // 4] + assert scale_input.type.tensor_type.elem_type == TensorProto.FLOAT16 + assert input_shape(scale_input) == [4, 1] + + value_info_by_name = {value_info.name: value_info for value_info in rewritten.graph.value_info} + unpack_value_info = value_info_by_name["linear.weight.mxfp6_unpacked"] + assert unpack_value_info.type.tensor_type.elem_type == TensorProto.FLOAT6E2M3 + assert input_shape(unpack_value_info) == [4, MXFP6_BLOCK_SIZE] + assert CustomOpTransform.apply(rewritten, onnx_export_opset=28) + unpack_functions = [ + fn for fn in rewritten.functions if fn.domain == "com.qti.aisw.onnx" and fn.name == "UnpackMxfp6" + ] + assert len(unpack_functions) == 1 + unpack_cast_nodes = [node for node in unpack_functions[0].node if node.op_type == "Cast"] + assert any( + any(attr.name == "to" and helper.get_attribute_value(attr) == ONNX_FLOAT6E2M3 for attr in node.attribute) + for node in unpack_cast_nodes + ) + unpack_nodes = [node for node in rewritten.graph.node if node.op_type == "UnpackMxfp6"] + dq_nodes = [node for node in rewritten.graph.node if node.op_type == "DequantizeLinear"] + assert len(unpack_nodes) == 1 + assert unpack_nodes[0].domain == "com.qti.aisw.onnx" + assert unpack_nodes[0].input == ["linear.weight" + MXFP6_PACKED_SUFFIX] + assert unpack_nodes[0].output == ["linear.weight.mxfp6_unpacked"] + assert len(dq_nodes) == 1 + assert dq_nodes[0].domain == "" + assert dq_nodes[0].input == ["linear.weight.mxfp6_unpacked", "linear.weight" + MXFP6_SCALE_SUFFIX] + assert dq_nodes[0].input[0] != "linear.weight" + MXFP6_PACKED_SUFFIX + assert [node.op_type for node in rewritten.graph.node[:2]] == ["UnpackMxfp6", "DequantizeLinear"] + attrs = {attr.name: helper.get_attribute_value(attr) for attr in dq_nodes[0].attribute} + assert attrs["axis"] == -1 + assert attrs["block_size"] == MXFP6_BLOCK_SIZE + assert attrs["output_dtype"] == TensorProto.FLOAT + + def test_mxfp6_finalizer_keeps_lm_head_weight_dense(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + hidden_size = 4 + out_features = MXFP6_BLOCK_SIZE + linear_weight = torch.arange(hidden_size * out_features, dtype=torch.float32).reshape(hidden_size, out_features) + lm_head_weight = linear_weight + 1000 + _write_safetensors_checkpoint( + prepared, + { + "linear.weight": linear_weight, + "base_model.lm_head.weight": lm_head_weight, + }, + ) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, hidden_size]) + h = helper.make_tensor_value_info("h", TensorProto.FLOAT, [1, hidden_size]) + linear = helper.make_tensor_value_info("linear.weight", TensorProto.FLOAT, [hidden_size, out_features]) + lm_head = helper.make_tensor_value_info("model.lm_head.weight", TensorProto.FLOAT, [hidden_size, out_features]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, out_features]) + logits = helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, out_features]) + graph = helper.make_graph( + [ + helper.make_node("MatMul", ["x", "linear.weight"], ["y"]), + helper.make_node("MatMul", ["h", "model.lm_head.weight"], ["logits"]), + ], + "mxfp6_lm_head_test", + [x, h, linear, lm_head], + [y, logits], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="linear.weight", + location=WeightSpecLocation(file=0, key="linear.weight"), + ), + WeightSpecInput( + name="model.lm_head.weight", + location=WeightSpecLocation(file=0, key="base_model.lm_head.weight"), + ), + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + spec = load_weight_spec(spec_path) + rewritten = onnx.load(str(onnx_path), load_external_data=False) + graph_input_names = {value_info.name for value_info in rewritten.graph.input} + spec_by_name = {entry.name: entry for entry in spec.inputs} + tensors = _load_prepared_tensors(prepared) + + assert "linear.weight" not in graph_input_names + assert "linear.weight" + MXFP6_PACKED_SUFFIX in graph_input_names + assert "linear.weight" + MXFP6_SCALE_SUFFIX in graph_input_names + assert spec_by_name["linear.weight" + MXFP6_PACKED_SUFFIX].role == "mxfp6_weight" + assert spec_by_name["linear.weight" + MXFP6_SCALE_SUFFIX].role == "mxfp6_scale" + + assert "model.lm_head.weight" in graph_input_names + assert spec_by_name["model.lm_head.weight"].role == "weight" + assert spec_by_name["model.lm_head.weight"].location.key == "base_model.lm_head.weight" + assert "model.lm_head.weight" + MXFP6_PACKED_SUFFIX not in graph_input_names + assert "model.lm_head.weight" + MXFP6_SCALE_SUFFIX not in graph_input_names + assert "model.lm_head.weight" + MXFP6_PACKED_SUFFIX not in spec_by_name + assert "model.lm_head.weight" + MXFP6_SCALE_SUFFIX not in spec_by_name + assert "base_model.lm_head.weight" + MXFP6_PACKED_SUFFIX not in tensors + assert "base_model.lm_head.weight" + MXFP6_SCALE_SUFFIX not in tensors + + unpack_nodes = [node for node in rewritten.graph.node if node.op_type == "UnpackMxfp6"] + dq_nodes = [node for node in rewritten.graph.node if node.op_type == "DequantizeLinear"] + assert len(unpack_nodes) == 1 + assert len(dq_nodes) == 1 + assert "lm_head" not in unpack_nodes[0].name + assert "lm_head" not in dq_nodes[0].name + + def test_mxfp6_finalizer_supports_final_axis_transpose_topology(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + weight = torch.arange(4 * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape(4, MXFP6_BLOCK_SIZE) + _write_safetensors_checkpoint(prepared, {"linear.weight": weight}) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + w = helper.make_tensor_value_info("linear.weight", TensorProto.FLOAT, [4, MXFP6_BLOCK_SIZE]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4]) + graph = helper.make_graph( + [ + helper.make_node("Transpose", ["linear.weight"], ["linear.weight.t"], perm=[1, 0]), + helper.make_node("MatMul", ["x", "linear.weight.t"], ["y"]), + ], + "mxfp6_transpose_test", + [x, w], + [y], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="linear.weight", + location=WeightSpecLocation(file=0, key="linear.weight"), + ) + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + rewritten = onnx.load(str(onnx_path), load_external_data=False) + unpack_nodes = [node for node in rewritten.graph.node if node.op_type == "UnpackMxfp6"] + dq_nodes = [node for node in rewritten.graph.node if node.op_type == "DequantizeLinear"] + transpose_nodes = [node for node in rewritten.graph.node if node.op_type == "Transpose"] + matmul_nodes = [node for node in rewritten.graph.node if node.op_type == "MatMul"] + + assert len(unpack_nodes) == 1 + assert unpack_nodes[0].output == ["linear.weight.mxfp6_unpacked"] + assert len(dq_nodes) == 1 + assert dq_nodes[0].input[0] == "linear.weight.mxfp6_unpacked" + assert dq_nodes[0].input[1] == "linear.weight" + MXFP6_SCALE_SUFFIX + assert dq_nodes[0].output == ["linear.weight"] + assert transpose_nodes[0].input == ["linear.weight"] + assert matmul_nodes[0].input[1] == "linear.weight.t" + + def test_mxfp6_finalizer_supports_function_transpose_matmul_topology(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + weight = torch.arange(4 * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape(4, MXFP6_BLOCK_SIZE) + _write_safetensors_checkpoint(prepared, {"model.layers.0.linear.weight": weight}) + + function_domain = "pkg.torch.__subgraph__" + function = helper.make_function( + domain=function_domain, + fname="repeated_subgraph0", + inputs=["hidden_states", "arg2_1"], + outputs=["fn_y"], + nodes=[ + helper.make_node("Transpose", ["arg2_1"], ["arg2_1_t"], perm=[1, 0]), + helper.make_node("MatMul", ["hidden_states", "arg2_1_t"], ["fn_y"]), + ], + opset_imports=[helper.make_opsetid("", 18)], + ) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + w = helper.make_tensor_value_info("layer.weight", TensorProto.FLOAT, [4, MXFP6_BLOCK_SIZE]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4]) + call_node = helper.make_node( + "repeated_subgraph0", + ["x", "layer.weight"], + ["y"], + domain=function_domain, + ) + graph = helper.make_graph([call_node], "mxfp6_function_test", [x, w], [y]) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid(function_domain, 1)], + ) + model.functions.append(function) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="layer.weight", + location=WeightSpecLocation(file=0, key="model.layers.0.linear.weight"), + ) + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + spec = load_weight_spec(spec_path) + rewritten = onnx.load(str(onnx_path), load_external_data=False) + graph_input_names = {value_info.name for value_info in rewritten.graph.input} + initializer_names = {initializer.name for initializer in rewritten.graph.initializer} + unpack_nodes = [node for node in rewritten.graph.node if node.op_type == "UnpackMxfp6"] + dq_nodes = [node for node in rewritten.graph.node if node.op_type == "DequantizeLinear"] + call_nodes = [node for node in rewritten.graph.node if node.op_type == "repeated_subgraph0"] + + assert [entry.role for entry in spec.inputs] == ["mxfp6_weight", "mxfp6_scale"] + assert "layer.weight" not in graph_input_names + assert "layer.weight" + MXFP6_PACKED_SUFFIX in graph_input_names + assert "layer.weight" + MXFP6_SCALE_SUFFIX in graph_input_names + assert "layer.weight" + MXFP6_PACKED_SUFFIX not in initializer_names + assert "layer.weight" + MXFP6_SCALE_SUFFIX not in initializer_names + assert len(unpack_nodes) == 0 + assert len(dq_nodes) == 0 + assert len(call_nodes) == 1 + assert call_nodes[0].input == ["x", "layer.weight" + MXFP6_PACKED_SUFFIX, "layer.weight" + MXFP6_SCALE_SUFFIX] + + rewritten_function = next( + fn for fn in rewritten.functions if fn.domain == function_domain and fn.name == "repeated_subgraph0" + ) + assert list(rewritten_function.input) == ["hidden_states", "arg2_1.mxfp6_packed", "arg2_1.mxfp6_scale"] + assert [node.op_type for node in rewritten_function.node] == [ + "UnpackMxfp6", + "DequantizeLinear", + "Transpose", + "MatMul", + ] + assert rewritten_function.node[0].input == ["arg2_1.mxfp6_packed"] + assert rewritten_function.node[0].output == ["arg2_1.mxfp6_unpacked"] + assert rewritten_function.node[1].input == ["arg2_1.mxfp6_unpacked", "arg2_1.mxfp6_scale"] + assert rewritten_function.node[1].output == ["arg2_1"] + assert rewritten_function.node[2].input == ["arg2_1"] + assert rewritten_function.node[3].input[1] == "arg2_1_t" + assert any(opset.domain == "" and opset.version == 28 for opset in rewritten_function.opset_import) + assert any( + opset.domain == "com.qti.aisw.onnx" and opset.version == 1 for opset in rewritten_function.opset_import + ) + assert CustomOpTransform.apply(rewritten, onnx_export_opset=28) + unpack_functions = [ + fn for fn in rewritten.functions if fn.domain == "com.qti.aisw.onnx" and fn.name == "UnpackMxfp6" + ] + assert len(unpack_functions) == 1 + + def test_mxfp6_finalizer_supports_function_direct_matmul_topology(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + weight = torch.arange(MXFP6_BLOCK_SIZE * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape( + MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE + ) + _write_safetensors_checkpoint(prepared, {"model.layers.0.linear.weight": weight}) + + function_domain = "pkg.torch.__subgraph__" + function = helper.make_function( + domain=function_domain, + fname="repeated_subgraph0", + inputs=["hidden_states", "arg2_1"], + outputs=["fn_y"], + nodes=[helper.make_node("MatMul", ["hidden_states", "arg2_1"], ["fn_y"])], + opset_imports=[helper.make_opsetid("", 18)], + ) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + w = helper.make_tensor_value_info("layer.weight", TensorProto.FLOAT, [MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + graph = helper.make_graph( + [helper.make_node("repeated_subgraph0", ["x", "layer.weight"], ["y"], domain=function_domain)], + "mxfp6_function_direct_test", + [x, w], + [y], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid(function_domain, 1)], + ) + model.functions.append(function) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="layer.weight", + location=WeightSpecLocation(file=0, key="model.layers.0.linear.weight"), + ) + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + rewritten = onnx.load(str(onnx_path), load_external_data=False) + call_node = next(node for node in rewritten.graph.node if node.op_type == "repeated_subgraph0") + rewritten_function = next( + fn for fn in rewritten.functions if fn.domain == function_domain and fn.name == "repeated_subgraph0" + ) + + assert [ + node.op_type for node in rewritten.graph.node if node.op_type in {"UnpackMxfp6", "DequantizeLinear"} + ] == [] + assert call_node.input == ["x", "layer.weight" + MXFP6_PACKED_SUFFIX, "layer.weight" + MXFP6_SCALE_SUFFIX] + assert list(rewritten_function.input) == ["hidden_states", "arg2_1.mxfp6_packed", "arg2_1.mxfp6_scale"] + assert [node.op_type for node in rewritten_function.node] == ["UnpackMxfp6", "DequantizeLinear", "MatMul"] + assert rewritten_function.node[2].input == ["hidden_states", "arg2_1"] + + def test_mxfp6_finalizer_preserves_multiple_function_formal_ordering(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + weight0 = torch.arange(MXFP6_BLOCK_SIZE * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape( + MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE + ) + weight1 = weight0 + 1000 + _write_safetensors_checkpoint( + prepared, + { + "model.layers.0.linear0.weight": weight0, + "model.layers.0.linear1.weight": weight1, + }, + ) + + function_domain = "pkg.torch.__subgraph__" + function = helper.make_function( + domain=function_domain, + fname="repeated_subgraph0", + inputs=["hidden_states", "arg1_1", "arg2_1"], + outputs=["fn_y"], + nodes=[ + helper.make_node("MatMul", ["hidden_states", "arg1_1"], ["hidden_1"]), + helper.make_node("MatMul", ["hidden_1", "arg2_1"], ["fn_y"]), + ], + opset_imports=[helper.make_opsetid("", 18)], + ) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + w0 = helper.make_tensor_value_info("layer.weight0", TensorProto.FLOAT, [MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE]) + w1 = helper.make_tensor_value_info("layer.weight1", TensorProto.FLOAT, [MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + graph = helper.make_graph( + [ + helper.make_node( + "repeated_subgraph0", ["x", "layer.weight0", "layer.weight1"], ["y"], domain=function_domain + ) + ], + "mxfp6_function_order_test", + [x, w0, w1], + [y], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid(function_domain, 1)], + ) + model.functions.append(function) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="layer.weight0", + location=WeightSpecLocation(file=0, key="model.layers.0.linear0.weight"), + ), + WeightSpecInput( + name="layer.weight1", + location=WeightSpecLocation(file=0, key="model.layers.0.linear1.weight"), + ), + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + rewritten = onnx.load(str(onnx_path), load_external_data=False) + call_node = next(node for node in rewritten.graph.node if node.op_type == "repeated_subgraph0") + rewritten_function = next( + fn for fn in rewritten.functions if fn.domain == function_domain and fn.name == "repeated_subgraph0" + ) + + assert call_node.input == [ + "x", + "layer.weight0" + MXFP6_PACKED_SUFFIX, + "layer.weight1" + MXFP6_PACKED_SUFFIX, + "layer.weight0" + MXFP6_SCALE_SUFFIX, + "layer.weight1" + MXFP6_SCALE_SUFFIX, + ] + assert list(rewritten_function.input) == [ + "hidden_states", + "arg1_1.mxfp6_packed", + "arg2_1.mxfp6_packed", + "arg1_1.mxfp6_scale", + "arg2_1.mxfp6_scale", + ] + assert [node.op_type for node in rewritten_function.node] == [ + "UnpackMxfp6", + "DequantizeLinear", + "MatMul", + "UnpackMxfp6", + "DequantizeLinear", + "MatMul", + ] + + def test_mxfp6_finalizer_rejects_incomplete_shared_function_calls(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + weight = torch.arange(MXFP6_BLOCK_SIZE * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape( + MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE + ) + _write_safetensors_checkpoint(prepared, {"model.layers.0.linear.weight": weight}) + + function_domain = "pkg.torch.__subgraph__" + function = helper.make_function( + domain=function_domain, + fname="repeated_subgraph0", + inputs=["hidden_states", "arg2_1"], + outputs=["fn_y"], + nodes=[helper.make_node("MatMul", ["hidden_states", "arg2_1"], ["fn_y"])], + opset_imports=[helper.make_opsetid("", 18)], + ) + + x0 = helper.make_tensor_value_info("x0", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + x1 = helper.make_tensor_value_info("x1", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + w0 = helper.make_tensor_value_info("layer0.weight", TensorProto.FLOAT, [MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE]) + w1 = helper.make_tensor_value_info("layer1.weight", TensorProto.FLOAT, [MXFP6_BLOCK_SIZE, MXFP6_BLOCK_SIZE]) + y0 = helper.make_tensor_value_info("y0", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + y1 = helper.make_tensor_value_info("y1", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + graph = helper.make_graph( + [ + helper.make_node("repeated_subgraph0", ["x0", "layer0.weight"], ["y0"], domain=function_domain), + helper.make_node("repeated_subgraph0", ["x1", "layer1.weight"], ["y1"], domain=function_domain), + ], + "mxfp6_function_incomplete_test", + [x0, x1, w0, w1], + [y0, y1], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid(function_domain, 1)], + ) + model.functions.append(function) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="layer0.weight", + location=WeightSpecLocation(file=0, key="model.layers.0.linear.weight"), + ) + ], + ), + ) + + with pytest.raises(NotImplementedError, match="partially rewrite shared ONNX subfunctions"): + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + def test_mxfp6_finalizer_skips_unsupported_function_formal_arg(self, tmp_path, monkeypatch): + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.validate_mxfp6_capabilities", lambda config: None) + monkeypatch.setattr("QEfficient.exporter.weight_free.mxfp6.TensorProto.FLOAT6E2M3", 999, raising=False) + + prepared = tmp_path / "prepared" + prepared.mkdir() + linear_weight = torch.arange(4 * MXFP6_BLOCK_SIZE, dtype=torch.float32).reshape(4, MXFP6_BLOCK_SIZE) + norm_weight = torch.arange(MXFP6_BLOCK_SIZE, dtype=torch.float32) + _write_safetensors_checkpoint( + prepared, + { + "model.layers.0.linear.weight": linear_weight, + "model.layers.0.input_layernorm.weight": norm_weight, + }, + ) + + function_domain = "pkg.torch.__subgraph__" + linear_function = helper.make_function( + domain=function_domain, + fname="repeated_subgraph0", + inputs=["hidden_states", "arg2_1"], + outputs=["linear_y"], + nodes=[ + helper.make_node("Transpose", ["arg2_1"], ["arg2_1_t"], perm=[1, 0]), + helper.make_node("MatMul", ["hidden_states", "arg2_1_t"], ["linear_y"]), + ], + opset_imports=[helper.make_opsetid("", 18)], + ) + rmsnorm_function = helper.make_function( + domain=function_domain, + fname="repeated_subgraph1", + inputs=["hidden_states", "arg5_1"], + outputs=["norm_y"], + nodes=[ + helper.make_node( + "CustomRMSNorm", + ["hidden_states", "arg5_1"], + ["norm_y"], + domain="com.qti.aisw.onnx", + epsilon_f=1e-5, + ) + ], + opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("com.qti.aisw.onnx", 1)], + ) + + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + linear_w = helper.make_tensor_value_info("linear.weight", TensorProto.FLOAT, [4, MXFP6_BLOCK_SIZE]) + norm_w = helper.make_tensor_value_info("norm.weight", TensorProto.FLOAT, [MXFP6_BLOCK_SIZE]) + linear_y = helper.make_tensor_value_info("linear_y", TensorProto.FLOAT, [1, 4]) + norm_y = helper.make_tensor_value_info("norm_y", TensorProto.FLOAT, [1, MXFP6_BLOCK_SIZE]) + graph = helper.make_graph( + [ + helper.make_node("repeated_subgraph0", ["x", "linear.weight"], ["linear_y"], domain=function_domain), + helper.make_node("repeated_subgraph1", ["x", "norm.weight"], ["norm_y"], domain=function_domain), + ], + "mxfp6_unsupported_function_test", + [x, linear_w, norm_w], + [linear_y, norm_y], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 18), + helper.make_opsetid(function_domain, 1), + helper.make_opsetid("com.qti.aisw.onnx", 1), + ], + ) + model.functions.extend([linear_function, rmsnorm_function]) + onnx_path = tmp_path / "model.onnx" + onnx.save(model, str(onnx_path)) + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id=str(prepared), + files=[ExternalDataFile(path="model.safetensors", format="safetensors")], + inputs=[ + WeightSpecInput( + name="linear.weight", + location=WeightSpecLocation(file=0, key="model.layers.0.linear.weight"), + ), + WeightSpecInput( + name="norm.weight", + location=WeightSpecLocation(file=0, key="model.layers.0.input_layernorm.weight"), + ), + ], + ), + ) + + finalize_mxfp6_export(onnx_path, spec_path, str(prepared), Mxfp6Config(enabled=True, scale_dtype="float16")) + + spec = load_weight_spec(spec_path) + rewritten = onnx.load(str(onnx_path), load_external_data=False) + graph_input_names = {value_info.name for value_info in rewritten.graph.input} + spec_by_name = {entry.name: entry for entry in spec.inputs} + + assert "linear.weight" not in graph_input_names + assert "linear.weight" + MXFP6_PACKED_SUFFIX in graph_input_names + assert "linear.weight" + MXFP6_SCALE_SUFFIX in graph_input_names + assert spec_by_name["linear.weight" + MXFP6_PACKED_SUFFIX].role == "mxfp6_weight" + assert spec_by_name["linear.weight" + MXFP6_SCALE_SUFFIX].role == "mxfp6_scale" + + assert "norm.weight" in graph_input_names + assert "norm.weight" + MXFP6_PACKED_SUFFIX not in graph_input_names + assert "norm.weight" + MXFP6_SCALE_SUFFIX not in graph_input_names + assert spec_by_name["norm.weight"].role == "weight" + assert "norm.weight" + MXFP6_PACKED_SUFFIX not in spec_by_name + assert "norm.weight" + MXFP6_SCALE_SUFFIX not in spec_by_name + + def test_ort_loader_rejects_mxfp6_weight_spec(self, tmp_path): + spec_path = save_weight_spec( + tmp_path / "weight_spec.json", + WeightSpec( + model_name="tiny", + model_id="tiny", + inputs=[ + WeightSpecInput( + name="linear.weight.mxfp6_packed", + location=WeightSpecLocation(file=0, key="linear.weight"), + role="mxfp6_weight", + ) + ], + ), + ) + + with pytest.raises(NotImplementedError, match="MXFP6"): + load_weight_free_ort_inputs(spec_path, {}) + + def _fake_export( self, example_inputs, @@ -516,6 +1432,38 @@ def test_weight_free_export_hash_differs_from_regular_dynamo(self): assert "weight_free" not in regular_params assert weight_free_params["weight_free"] is True + def test_mxfp6_export_hash_separates_scale_dtypes(self): + config = SimpleNamespace(to_diff_dict=lambda: {"model_type": "llama"}) + common_model = SimpleNamespace( + model=SimpleNamespace(config=config), + hash_params={"pretrained_model_name_or_path": "tiny"}, + _use_onnx_subfunctions=False, + _weight_free=True, + ) + fp16_model = SimpleNamespace( + **common_model.__dict__, + _mxfp6_config=Mxfp6Config(enabled=True, scale_dtype="float16"), + ) + bf16_model = SimpleNamespace( + **common_model.__dict__, + _mxfp6_config=Mxfp6Config(enabled=True, scale_dtype="bfloat16"), + ) + common_kwargs = { + "example_inputs": {"input_ids": torch.ones(1, 2, dtype=torch.int64)}, + "output_names": ["logits"], + "dynamic_axes": {"input_ids": {0: "batch_size"}}, + "dynamo": True, + } + + common_hash, _ = _generate_export_hash(common_model, (), dict(common_kwargs), _fake_export) + fp16_hash, fp16_params = _generate_export_hash(fp16_model, (), dict(common_kwargs), _fake_export) + bf16_hash, bf16_params = _generate_export_hash(bf16_model, (), dict(common_kwargs), _fake_export) + + assert len({common_hash, fp16_hash, bf16_hash}) == 3 + assert fp16_params["mxfp6"] is True + assert fp16_params["mxfp6_scale_dtype"] == "float16" + assert bf16_params["mxfp6_scale_dtype"] == "bfloat16" + class TestRuntimeRequirements: def test_validate_runtime_requirements_accepts_matching_requirements(self, monkeypatch):