diff --git a/QEfficient/base/modeling_qeff.py b/QEfficient/base/modeling_qeff.py index 6be24bee44..3d89596c07 100755 --- a/QEfficient/base/modeling_qeff.py +++ b/QEfficient/base/modeling_qeff.py @@ -14,7 +14,7 @@ import warnings from abc import ABC, abstractmethod from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Type, Union import onnx import torch @@ -117,6 +117,7 @@ class QEFFBaseModel(ABC): _layerwise_active = False _pytorch_transforms: List[PytorchTransform] _onnx_transforms = [BaseOnnxTransform] + _checkpoint_transforms: List[Type] = [] def _transform_names(self) -> List[str]: return [x.__name__ for x in self._pytorch_transforms + self._onnx_transforms] @@ -160,6 +161,7 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None: self.onnx_path: Optional[str] = None self.qpc_path: Optional[str] = None self.qpc_session: Optional[QAICInferenceSession] = None + self.weight_spec_path: Optional[str] = None self.model_architecture = ( (arch := getattr(self.model.config, "architectures", None)) and len(arch) > 0 and arch[0] ) or None @@ -432,6 +434,78 @@ def _export_via_dynamo( else: os.environ["TORCH_INVOKE_ALLOW_CREATE_FALLBACK"] = prev_invoke_fallback + def _export_via_weightfree( + self, + tmp_onnx_path: Path, + example_inputs: Dict[str, torch.Tensor], + input_names: List[str], + output_names: List[str], + dynamic_axes: Dict, + export_kwargs: Dict, + onnx_transform_kwargs: Optional[Dict] = None, + ): + """Export through the weight-free dynamo path with checkpoint-backed weights.""" + from QEfficient.customop.dynamo_ops import DYNAMO_CUSTOM_OP_TABLE + from QEfficient.exporter.weight_free.core import export_weight_free_onnx + from QEfficient.utils.export_utils import convert_dynamic_axes_to_dynamic_shapes + + model_config = getattr(self.model, "config", None) + dynamic_shapes = convert_dynamic_axes_to_dynamic_shapes(dynamic_axes, model_config) + + sig_keys = list(inspect.signature(self.model.forward).parameters.keys()) + sig_key_set = set(sig_keys) + ordered_inputs, ordered_shapes = {}, {} + for key in sig_keys: + if key in example_inputs: + ordered_inputs[key] = example_inputs[key] + if key in dynamic_shapes: + ordered_shapes[key] = dynamic_shapes[key] + example_inputs = { + **ordered_inputs, + **{key: value for key, value in example_inputs.items() if key not in sig_key_set}, + } + dynamic_shapes = { + **ordered_shapes, + **{key: value for key, value in dynamic_shapes.items() if key not in sig_key_set}, + } + + wf_export_kwargs = dict(export_kwargs) + wf_export_kwargs.setdefault("report", False) + wf_export_kwargs.setdefault("optimize", False) + wf_export_kwargs["dynamo"] = True + wf_export_kwargs["opset_version"] = constants.ONNX_DYNAMO_EXPORT_OPSET + wf_export_kwargs["custom_translation_table"] = { + **(wf_export_kwargs.pop("custom_translation_table", None) or {}), + **DYNAMO_CUSTOM_OP_TABLE, + } + + export_func = export_weight_free_onnx + if self.model.__class__.__name__ in {"QEffKimiK25DecoderWrapper", "QEffKimiK25EncoderWrapper"}: + from QEfficient.exporter.weight_free.core import export_loaded_model_weight_free_onnx + + export_func = export_loaded_model_weight_free_onnx + + prev_invoke_fallback = os.environ.get("TORCH_INVOKE_ALLOW_CREATE_FALLBACK") + os.environ["TORCH_INVOKE_ALLOW_CREATE_FALLBACK"] = "1" + try: + _, updated_onnx_transform_kwargs, cleanup = export_func( + qeff_model=self, + tmp_onnx_path=tmp_onnx_path, + example_inputs=example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_shapes=dynamic_shapes, + export_kwargs=wf_export_kwargs, + onnx_transform_kwargs=onnx_transform_kwargs or {}, + ) + finally: + if prev_invoke_fallback is None: + os.environ.pop("TORCH_INVOKE_ALLOW_CREATE_FALLBACK", None) + else: + os.environ["TORCH_INVOKE_ALLOW_CREATE_FALLBACK"] = prev_invoke_fallback + + return updated_onnx_transform_kwargs, cleanup + @export_wrapper def _export( self, @@ -444,6 +518,7 @@ def _export( prefill_only: Optional[bool] = False, dynamo: bool = False, dynamic_shapes: Optional[Dict[str, Dict[int, Any]]] = None, + use_weight_free_export: bool = False, **export_kwargs, ) -> str: """ @@ -476,14 +551,21 @@ def _export( # TODO: Hack for retain_full_kv, handle this outside export_kwargs.pop("retain_full_kv", None) onnx_path = export_dir / f"{self.model_name}.onnx" + weight_spec_path = onnx_path.with_name("weight_spec.json") # Return early if ONNX already exists if onnx_path.is_file(): self.onnx_path = onnx_path + if weight_spec_path.is_file(): + self.weight_spec_path = str(weight_spec_path) return onnx_path # check if the model is in meta state or weights are offloaded - self._model_offloaded_check() + if not use_weight_free_export: + self._model_offloaded_check() + + if use_weight_free_export and not dynamo: + raise NotImplementedError("Weight-free export requires dynamo=True.") export_dir.mkdir(parents=True, exist_ok=True) @@ -556,8 +638,28 @@ def _resolve_pkv_names(layer_idx, layer_state): dynamic_axes = {rename_map.get(k, k): v for k, v in dynamic_axes.items()} input_names = aligned_input_names + cleanup_fn = None try: - if dynamo: + if use_weight_free_export: + tmp_onnx_dir = export_dir / "onnx_weightfree_tmp" + tmp_onnx_dir.mkdir(parents=True, exist_ok=True) + tmp_onnx_path = tmp_onnx_dir / onnx_path.name + onnx_transform_kwargs, cleanup_fn = self._export_via_weightfree( + tmp_onnx_path=tmp_onnx_path, + example_inputs=example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + export_kwargs=export_kwargs, + onnx_transform_kwargs=onnx_transform_kwargs, + ) + shutil.move(str(tmp_onnx_path), str(onnx_path)) + tmp_weight_spec = tmp_onnx_dir / "weight_spec.json" + if tmp_weight_spec.exists(): + shutil.move(str(tmp_weight_spec), str(weight_spec_path)) + self.weight_spec_path = str(weight_spec_path) + shutil.rmtree(tmp_onnx_dir, ignore_errors=True) + elif dynamo: self._export_via_dynamo( onnx_path, example_inputs, @@ -576,7 +678,8 @@ def _resolve_pkv_names(layer_idx, layer_state): export_kwargs, ) logger.info("PyTorch export successful") - self._offload_model_weights(offload_pt_weights) + if not use_weight_free_export: + self._offload_model_weights(offload_pt_weights) model = onnx.load(onnx_path, load_external_data=False) needs_external_tensor_data = any( @@ -585,13 +688,18 @@ def _resolve_pkv_names(layer_idx, layer_state): transform_kwargs = { "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, # dynamo uses dynamic_shapes, not axes - "onnx_export_opset": constants.get_onnx_export_opset(dynamo), + "dynamic_axes": None if (dynamo or use_weight_free_export) else dynamic_axes, + "onnx_export_opset": constants.get_onnx_export_opset(dynamo or use_weight_free_export), } if onnx_transform_kwargs is not None: transform_kwargs.update(onnx_transform_kwargs) - onnx_transforms = OnnxTransformPipeline(transforms=self._onnx_transforms) + active_transforms = [ + transform + for transform in self._onnx_transforms + if not (use_weight_free_export and transform is SplitTensorsTransform) + ] + onnx_transforms = OnnxTransformPipeline(transforms=active_transforms) model, transformed = onnx_transforms.apply(model, **transform_kwargs) # Keep this strictly layerwise-scoped so regular non-layerwise export @@ -603,6 +711,15 @@ def _resolve_pkv_names(layer_idx, layer_state): model.metadata_props.append( onnx.StringStringEntryProto(key="qeff_transforms", value=",".join(self._transform_names())) ) + if use_weight_free_export and self.weight_spec_path is not None: + import json as _json + + from QEfficient.exporter.weight_free.core import _upsert_metadata_prop + + weight_spec_json = _json.dumps( + load_json(Path(self.weight_spec_path)), separators=(",", ":"), sort_keys=True + ) + _upsert_metadata_prop(model, "com.qti.aisw.extdata", weight_spec_json) logger.info("ONNX transforms applied") onnx_path_tmp = onnx_path.with_suffix(onnx_path.suffix + ".tmp") @@ -615,8 +732,19 @@ def _resolve_pkv_names(layer_idx, layer_state): except Exception as e: logger.error(f"ONNX export or transforms failed: {e}") raise e + finally: + if cleanup_fn is not None: + cleanup_fn() self.onnx_path = onnx_path + if use_weight_free_export and self.weight_spec_path is not None: + from QEfficient.exporter.weight_free.spec import load_weight_spec + + spec = load_weight_spec(Path(self.weight_spec_path)) + prepared_out = Path(spec.model_id) + symlink = onnx_path.parent / prepared_out.name + if prepared_out.exists() and not symlink.exists(): + symlink.symlink_to(prepared_out) return onnx_path def get_onnx_path( @@ -631,6 +759,7 @@ def get_onnx_path( qaic_config: Optional[dict] = None, moe_prefill_packed_chunk_size: Optional[int] = None, kv_cache_prefix: Optional[str] = None, + use_weight_free_export: bool = False, **compiler_options, ): kwargs = { @@ -638,6 +767,7 @@ def get_onnx_path( "use_onnx_subfunctions": use_onnx_subfunctions, "dynamo": dynamo, "retain_full_kv": retain_full_kv, + "use_weight_free_export": use_weight_free_export, } layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False) if layerwise_cache_probe: @@ -1009,6 +1139,7 @@ def _compile( layerwise_cache_probe = compiler_options.pop("_layerwise_cache_probe", False) moe_prefill_packed_chunk_size = compiler_options.pop("moe_prefill_packed_chunk_size", None) + use_weight_free_export = compiler_options.pop("use_weight_free_export", False) for removed_option in ("compile_only", "compile-only"): if removed_option in compiler_options: @@ -1040,6 +1171,7 @@ def _compile( moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size, _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **compiler_options, ) if QEFFBaseModel._layerwise_active: diff --git a/QEfficient/base/onnx_transforms.py b/QEfficient/base/onnx_transforms.py index b24b4ebcb5..c35de5daf4 100644 --- a/QEfficient/base/onnx_transforms.py +++ b/QEfficient/base/onnx_transforms.py @@ -45,7 +45,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, update_cast_to_uint4_output_types 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 @@ -597,7 +597,7 @@ def apply( **kwargs, ) -> Tuple[ModelProto, bool]: if not self.transforms: - return model, False + return model, update_cast_to_uint4_output_types(model) # Same logic as before, but replace `transforms` with `self.transforms` mapping: Dict[str, Tuple[TensorProto, str]] = {} @@ -648,6 +648,8 @@ def _set_external_data(tensor, file_name): model, onnx_export_opset=kwargs.get("onnx_export_opset", constants.ONNX_LEGACY_EXPORT_OPSET) ) + cast_to_uint4_types_updated = update_cast_to_uint4_output_types(model) + if RenameFunctionOutputsTransform in requested: applied[RenameFunctionOutputsTransform] = RenameFunctionOutputsTransform.apply( model, layer_idx=kwargs.get("layer_idx", 0) @@ -665,4 +667,4 @@ def _set_external_data(tensor, file_name): for t, done in applied.items(): logger.info(f"Transform '{t.__name__}' applied={done}") - return model, any(applied.values()) + return model, any(applied.values()) or cast_to_uint4_types_updated diff --git a/QEfficient/customop/dynamo_ops.py b/QEfficient/customop/dynamo_ops.py index 81318efbbf..3b9a6152aa 100644 --- a/QEfficient/customop/dynamo_ops.py +++ b/QEfficient/customop/dynamo_ops.py @@ -23,6 +23,7 @@ CtxScatterCB3D, ) from QEfficient.customop.onnxscript_utils import get_dynamo_onnxscript_func +from QEfficient.customop.quantization_ops import CastToUInt4 from QEfficient.customop.rms_norm import CustomRMSNorm # noqa: E402 @@ -40,6 +41,21 @@ def _(hidden_states: torch.Tensor, weight: torch.Tensor, epsilon: float) -> torc return torch.empty_like(hidden_states) +@torch.library.custom_op("qefficient::cast_to_uint4", mutates_args=()) +def cast_to_uint4_op(weight_packed: torch.Tensor) -> torch.Tensor: + """Unpack packed uint8 nibbles into uint8 values for eager/fake execution.""" + lower = weight_packed & 0x0F + upper = (weight_packed >> 4) & 0x0F + output_shape = (*weight_packed.shape[:-1], weight_packed.shape[-1] * 2) + return torch.stack([lower, upper], dim=-1).reshape(output_shape) + + +@cast_to_uint4_op.register_fake +def _(weight_packed: torch.Tensor) -> torch.Tensor: + output_shape = (*weight_packed.shape[:-1], weight_packed.shape[-1] * 2) + return torch.empty(output_shape, dtype=weight_packed.dtype, device=weight_packed.device) + + @torch.library.custom_op("qefficient::ctx_scatter", mutates_args=()) def ctx_scatter_op(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> torch.Tensor: """Custom context scatter operation""" @@ -374,6 +390,7 @@ def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) -> DYNAMO_CUSTOM_OP_TABLE = { torch.ops.qefficient.rms_norm.default: get_dynamo_onnxscript_func(CustomRMSNorm), + torch.ops.qefficient.cast_to_uint4.default: get_dynamo_onnxscript_func(CastToUInt4), torch.ops.qefficient.ctx_scatter.default: get_dynamo_onnxscript_func(CtxScatter), torch.ops.qefficient.ctx_scatter_3d.default: get_dynamo_onnxscript_func(CtxScatter3D), torch.ops.qefficient.ctx_scatter_cb.default: get_dynamo_onnxscript_func(CtxScatterCB), diff --git a/QEfficient/customop/quantization_ops.py b/QEfficient/customop/quantization_ops.py index 3878804c7d..0c205abfbe 100644 --- a/QEfficient/customop/quantization_ops.py +++ b/QEfficient/customop/quantization_ops.py @@ -7,15 +7,18 @@ import onnxscript import torch -from onnx import TensorProto +from onnx import ModelProto, TensorProto +from onnxscript.onnx_types import UINT4 +from QEfficient.customop.onnxscript_utils import qeff_custom_op +from QEfficient.customop.utils import select_interface from QEfficient.utils import constants -ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET)) +ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET)) -@onnxscript.script(onnxscript.values.Opset("com.qti.aisw.onnx", 1)) -def CastToUInt4(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8: +@qeff_custom_op("com.qti.aisw.onnx", 1) +def CastToUInt4(weight_packed: onnxscript.UINT8) -> UINT4: """ Unpack packed uint8 weights into uint4 values and cast output to UINT4. Supports N-D input: all leading dimensions are preserved; only the last @@ -100,6 +103,39 @@ def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value: return output +def cast_to_uint4(weight_packed: torch.Tensor) -> torch.Tensor: + return select_interface(CastToUInt4Func.apply, torch.ops.qefficient.cast_to_uint4)(weight_packed) + + +def update_cast_to_uint4_output_types(model: ModelProto) -> bool: + """Correct exported CastToUInt4 value-info to logical UINT4. + + PyTorch fake tensors cannot carry UINT4 dtype, so Dynamo annotates the + custom-op node outputs as UINT8 even though the ONNXScript function returns + UINT4. QAIC consumes the graph-level value-info for quantized + dequantization, so keep the node metadata consistent with the custom-op + body. + """ + cast_outputs = { + output_name + for node in model.graph.node + if node.domain == "com.qti.aisw.onnx" and node.op_type == "CastToUInt4" + for output_name in node.output + } + if not cast_outputs: + return False + + transformed = False + for value in list(model.graph.value_info) + list(model.graph.output) + list(model.graph.input): + if value.name not in cast_outputs: + continue + tensor_type = value.type.tensor_type + if tensor_type.elem_type != TensorProto.UINT4: + tensor_type.elem_type = TensorProto.UINT4 + transformed = True + return transformed + + class DequantizeLinearFunc(torch.autograd.Function): """ Emits a standard ONNX DequantizeLinear node (ai.onnx domain, not custom). @@ -147,3 +183,18 @@ def symbolic( axis_i=2, block_size_i=block_size, ) + + +def dequantize_linear( + weight_unpacked: torch.Tensor, scale: torch.Tensor, zeros: torch.Tensor, block_size: int +) -> torch.Tensor: + if torch._dynamo.is_compiling(): + return torch.onnx.ops.symbolic( + "::DequantizeLinear", + (weight_unpacked, scale, zeros), + {"axis": 2, "block_size": block_size}, + dtype=scale.dtype, + shape=weight_unpacked.shape, + version=18, + ) + return DequantizeLinearFunc.apply(weight_unpacked, scale, zeros, block_size) diff --git a/QEfficient/exporter/weight_free/__init__.py b/QEfficient/exporter/weight_free/__init__.py new file mode 100644 index 0000000000..6397204542 --- /dev/null +++ b/QEfficient/exporter/weight_free/__init__.py @@ -0,0 +1,43 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +from QEfficient.exporter.weight_free.core import ( + export_weight_free_onnx as export_weight_free_onnx, +) +from QEfficient.exporter.weight_free.core import ( + load_weight_free_ort_inputs as load_weight_free_ort_inputs, +) +from QEfficient.exporter.weight_free.core import ( + log_weight_free_export as log_weight_free_export, +) +from QEfficient.exporter.weight_free.spec import ( + ExternalDataFile as ExternalDataFile, +) +from QEfficient.exporter.weight_free.spec import ( + load_weight_spec as load_weight_spec, +) +from QEfficient.exporter.weight_free.spec import ( + resolve_weight_spec_path as resolve_weight_spec_path, +) +from QEfficient.exporter.weight_free.spec import ( + save_weight_spec as save_weight_spec, +) +from QEfficient.exporter.weight_free.transforms import ( + BaseCheckpointTransform as BaseCheckpointTransform, +) +from QEfficient.exporter.weight_free.transforms import ( + CheckpointTransformPipeline as CheckpointTransformPipeline, +) +from QEfficient.exporter.weight_free.transforms import ( + DtypeConversionCheckpointTransform as DtypeConversionCheckpointTransform, +) +from QEfficient.exporter.weight_free.transforms import ( + MoEExpertStackingCheckpointTransform as MoEExpertStackingCheckpointTransform, +) +from QEfficient.exporter.weight_free.transforms import ( + MoEFusedExpertSplitCheckpointTransform as MoEFusedExpertSplitCheckpointTransform, +) diff --git a/QEfficient/exporter/weight_free/core.py b/QEfficient/exporter/weight_free/core.py new file mode 100644 index 0000000000..4a8dd71a1d --- /dev/null +++ b/QEfficient/exporter/weight_free/core.py @@ -0,0 +1,720 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import os +import sys +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +import numpy as np +import onnx_ir as ir +import torch +from accelerate import init_empty_weights +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file +from torch import nn + +from QEfficient.transformers.embeddings.embedding_utils import PooledModel +from QEfficient.transformers.models.pytorch_transforms import PoolingTransform +from QEfficient.utils.export_utils import ( + _cleanup_onnx_subfunctions, + _setup_onnx_subfunctions, + get_decoder_layer_classes_for_export, +) +from QEfficient.utils.logging_utils import logger +from QEfficient.utils.torch_patches import ( + temporarily_disable_nested_compile_regions, + temporarily_enable_nested_compile_regions, +) + +from .spec import ( + ExternalDataFile, + TiedWeightAlias, + WeightSpec, + WeightSpecInput, + WeightSpecLocation, + load_weight_spec, + resolve_weight_spec_path, + save_weight_spec, +) +from .transforms import CheckpointTransformPipeline + +# Memory profiler from scripts/memory_profiling — optional, degrades gracefully +# if the scripts directory is not on the path or matplotlib is missing. +sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "scripts")) +try: + from memory_profiling import QEffMemoryProfiler as _QEffMemoryProfiler + + _HAS_PROFILER = True +except ImportError: + _HAS_PROFILER = False + +# Last checkpoint-prep profiler stats — written after every export call so +# _runner.py can read them without text-parsing stdout. +_last_prep_peak_rss_mb: float = 0.0 +_last_prep_duration_seconds: float = 0.0 +_checkpoint_prep_ran: bool = False + + +def _to_meta(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return torch.empty_like(value, device="meta") + if isinstance(value, tuple): + return tuple(_to_meta(item) for item in value) + if isinstance(value, list): + return [_to_meta(item) for item in value] + if isinstance(value, dict): + return {key: _to_meta(item) for key, item in value.items()} + return value + + +@lru_cache(maxsize=None) +def _resolve_checkpoint_dir(model_id_or_path: str) -> Path: + candidate = Path(model_id_or_path).expanduser() + if candidate.exists(): + return candidate + + # Try safetensors first (preferred format for weight-free export) + snapshot_dir = snapshot_download( + repo_id=model_id_or_path, + allow_patterns=["*.safetensors", "*.json"], + ignore_patterns=["*.onnx", "*.ot", "*.md", "*.txt", "*.pdf", "*.msgpack", "*.h5", "*.pth"], + resume_download=True, + ) + snapshot_path = Path(snapshot_dir) + + # Check if any weight files were actually downloaded (not just .json config/tokenizer files) + has_weights = ( + bool(list(snapshot_path.glob("*.safetensors"))) or (snapshot_path / "model.safetensors.index.json").exists() + ) + + if not has_weights: + # Model has no safetensors on Hub — fall back to .bin files. + # CheckpointTransformPipeline.apply() will auto-convert them to safetensors on first use. + snapshot_dir = snapshot_download( + repo_id=model_id_or_path, + allow_patterns=["*.bin", "*.json"], + ignore_patterns=[ + "*.onnx", + "*.ot", + "*.md", + "*.txt", + "*.pdf", + "*.msgpack", + "*.h5", + "*.pth", + "flax_model*", + "tf_model*", + ], + resume_download=True, + ) + + return Path(snapshot_dir) + + +def _resolve_checkpoint_files(model_id_or_path: str) -> List[str]: + checkpoint_dir = _resolve_checkpoint_dir(model_id_or_path) + checkpoint_files = sorted(str(path) for path in checkpoint_dir.glob("*.safetensors")) + if not checkpoint_files: + raise FileNotFoundError(f"No safetensors checkpoint files found for {model_id_or_path}") + return checkpoint_files + + +def _module_name_map(model: nn.Module) -> Dict[int, str]: + return {id(module): name for name, module in model.named_modules()} + + +def _collect_tied_weights(model: nn.Module) -> List[TiedWeightAlias]: + if not getattr(model.config, "tie_word_embeddings", False): + return [] + + input_embeddings = model.get_input_embeddings() + output_embeddings = model.get_output_embeddings() + if input_embeddings is None or output_embeddings is None: + return [] + + module_names = _module_name_map(model) + canonical_name = module_names.get(id(input_embeddings)) + alias_name = module_names.get(id(output_embeddings)) + if not canonical_name or not alias_name or canonical_name == alias_name: + return [] + + return [TiedWeightAlias(alias=f"{alias_name}.weight", canonical=f"{canonical_name}.weight")] + + +def _build_meta_qeff_model(qeff_model): + model_ref = qeff_model.hash_params.get("pretrained_model_name_or_path") + if not model_ref: + raise ValueError( + "Weight-free export requires checkpoint metadata. " + "Pass `pretrained_model_name_or_path=...` when constructing the QEff model manually." + ) + + quant_config = getattr(qeff_model.model.config, "quantization_config", None) + + config = copy.deepcopy(qeff_model.model.config) + source_model = getattr(qeff_model.model, "model", qeff_model.model) + with init_empty_weights(): + try: + meta_model = qeff_model._hf_auto_class.from_config(config, attn_implementation="eager") + except ValueError: + meta_model = source_model.__class__(config) + + if qeff_model.__class__.__name__ == "QEffCausalLMForTextImageToTextModel" and hasattr(meta_model, "vision_tower"): + meta_model.vision_tower = nn.Identity() + + if quant_config is None: + target_dtype = getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.float32 + if target_dtype == torch.bfloat16: + target_dtype = torch.float16 + meta_model = meta_model.to(dtype=target_dtype) + + if not hasattr(meta_model, "get_qeff_language_decoder") and not hasattr(meta_model, "get_qeff_vision_encoder"): + from QEfficient.transformers.models.pytorch_transforms import KVCacheExternalModuleMapperTransform + + meta_model, _ = KVCacheExternalModuleMapperTransform.apply(meta_model) + + meta_qeff_model = qeff_model.__class__( + meta_model, + continuous_batching=getattr(qeff_model, "continuous_batching", False), + qaic_config=copy.deepcopy(getattr(qeff_model.model, "qaic_config", None)), + max_seq_len_cached=getattr(qeff_model.model.config, "max_seq_len_cached", None), + pretrained_model_name_or_path=model_ref, + enable_proxy=getattr(qeff_model, "_enable_proxy", False), + ) + meta_qeff_model.hash_params.update(copy.deepcopy(qeff_model.hash_params)) + + if isinstance(qeff_model.model, PooledModel): + meta_qeff_model.model, _ = PoolingTransform.apply(meta_qeff_model.model, qeff_model.model.pooling_fn) + + if quant_config is not None: + # 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. + # We apply the quantizer's architecture preprocessing (layer-type replacement only, + # no weight loading) AFTER __init__ so that Mxfp4GptOssExpertDequantizeTransform — + # which is part of _pytorch_transforms and targets QEffMxfp4GptOssExperts — has + # already run as a no-op and will not undo the replacement below. + from QEfficient.transformers.quantizers.auto import ( + QEFF_AUTO_QUANTIZATION_CONFIG_MAPPING, + QEFF_AUTO_QUANTIZER_MAPPING, + ) + + # quantization_config may be a plain dict (AutoConfig.from_pretrained) or a proper + # config object (QEFFAutoModelForCausalLM.from_pretrained). Normalise to an object. + if isinstance(quant_config, dict): + quant_type = quant_config.get("quant_method") or quant_config.get("quant_type") + config_cls = QEFF_AUTO_QUANTIZATION_CONFIG_MAPPING.get(quant_type) + if config_cls is None: + raise NotImplementedError( + f"Weight-free export is not implemented for quantization type '{quant_type}'. Supported: mxfp4" + ) + init_kwargs = {k: v for k, v in quant_config.items() if k != "quant_method"} + quant_config = config_cls(**init_kwargs) + else: + quant_method = getattr(quant_config, "quant_method", None) or getattr(quant_config, "quant_type", None) + quant_type = quant_method.value if hasattr(quant_method, "value") else quant_method + + quantizer_cls = QEFF_AUTO_QUANTIZER_MAPPING.get(quant_type) if quant_type else None + if quantizer_cls is None: + raise NotImplementedError( + f"Weight-free export is not implemented for quantization type '{quant_type}'. Supported: mxfp4" + ) + quantizer = quantizer_cls(quant_config) + # Run inside init_empty_weights so newly created quantized layer buffers stay on + # the meta device and are treated as weight-spec entries, not embedded constants. + with init_empty_weights(): + quantizer._process_model_before_weight_loading(meta_qeff_model.model) + + meta_qeff_model.model.eval() + return meta_qeff_model + + +def _checkpoint_root(model_id_or_path: str, checkpoint_files: Sequence[str]) -> Optional[Path]: + if not checkpoint_files: + return None + + candidate = Path(model_id_or_path).expanduser() + if candidate.exists(): + return candidate.parent + + first_checkpoint = Path(checkpoint_files[0]) + for parent in first_checkpoint.parents: + if parent.name.startswith("models--"): + return parent.parent + return first_checkpoint.parent + + +def _load_checkpoint_index(checkpoint_files: List[str]) -> Dict[str, str]: + tensor_to_file = {} + for checkpoint_file in checkpoint_files: + handle = safe_open(checkpoint_file, framework="pt") + for key in handle.keys(): + tensor_to_file[key] = checkpoint_file + return tensor_to_file + + +def _build_location( + checkpoint_files: Sequence[str], + checkpoint_file: Optional[str], + tensor_key: str, +) -> Optional[WeightSpecLocation]: + if checkpoint_file is None: + return None + + return WeightSpecLocation(file=list(checkpoint_files).index(checkpoint_file), key=tensor_key) + + +def _find_checkpoint_key( + onnx_name: str, + checkpoint_index: Dict[str, str], + backbone: nn.Module, +) -> Optional[str]: + """ + Resolve an ONNX initializer name to its key in the safetensors checkpoint. + + Four lookups are attempted in order: + + 1. Direct match — ONNX name == checkpoint key. + Covers decoder-only LLMs (Llama, GPT-OSS) and any model exported + without extra wrappers. + + 2. Strip our PooledModel prefix ("base_model.") and try the bare key. + Covers embedding models whose checkpoint was saved from the base class + directly (e.g. BAAI/bge-base saved from BertModel → bare keys). + + 3. Strip "base_model." then prepend backbone.base_model_prefix. + Covers embedding models whose checkpoint was saved from a task-specific + class (e.g. BAAI/bge-reranker saved from XLMRobertaForSequenceClassification + → "roberta." prefix). HuggingFace defines base_model_prefix on every + model class as exactly the attribute name the task class uses to store + the backbone, so this is generalized — no per-model-family list needed. + + 4. Strip backbone.base_model_prefix from the front of the ONNX name. + Covers ForCausalLM models (e.g. GPT2LMHeadModel) whose checkpoint was + saved from the base model class (e.g. GPT2Model) so keys lack the LM + wrapper prefix — ONNX name "transformer.wte.weight" maps to checkpoint + key "wte.weight" after stripping the "transformer." prefix. + """ + # 1. Direct match + if onnx_name in checkpoint_index: + return onnx_name + + # Strip PooledModel's "base_model." wrapper prefix + stripped = onnx_name[len("base_model.") :] if onnx_name.startswith("base_model.") else onnx_name + + # 2. Bare key — checkpoint saved from base class + if stripped in checkpoint_index: + return stripped + + if stripped.startswith("model."): + without_wrapper = stripped[len("model.") :] + if without_wrapper in checkpoint_index: + return without_wrapper + + prefix = getattr(backbone, "base_model_prefix", "") + + # 3. Task-class prefix — checkpoint saved from ForSequenceClassification etc. + if prefix: + prefixed = f"{prefix}.{stripped}" + if prefixed in checkpoint_index: + return prefixed + + # 4. Strip base_model_prefix from ONNX name — checkpoint saved from the + # base model class so keys lack the ForCausalLM wrapper prefix + # (e.g. GPT2: ONNX "transformer.wte.weight" → checkpoint "wte.weight"). + if prefix and stripped.startswith(f"{prefix}."): + without_prefix = stripped[len(f"{prefix}.") :] + if without_prefix in checkpoint_index: + return without_prefix + + # 5. MLP attribute rename: transformers 5.x renamed 'block_sparse_moe' → 'mlp' + # for Mixtral. The original checkpoint uses 'block_sparse_moe' but the model + # (and ONNX) uses 'mlp'. Try the reverse substitution. + if ".mlp." in stripped: + candidate = stripped.replace(".mlp.", ".block_sparse_moe.") + if candidate in checkpoint_index: + return candidate + + # 6. MoE router attribute rename: QEff wrappers may call the router 'gate' + # while the checkpoint stores it as 'router' (e.g. Qwen3-MoE). + if stripped.endswith(".mlp.gate.weight"): + candidate = stripped[: -len(".gate.weight")] + ".router.weight" + if candidate in checkpoint_index: + return candidate + # Reverse: checkpoint calls it 'gate', ONNX uses 'router' + if stripped.endswith(".mlp.router.weight"): + candidate = stripped[: -len(".router.weight")] + ".gate.weight" + if candidate in checkpoint_index: + return candidate + + return None + + +def _promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name: str, qeff_model) -> WeightSpec: + from QEfficient.transformers.embeddings.embedding_utils import PooledModel + + model_ir = onnx_program.model + model_names = {name for name, _ in qeff_model.model.named_parameters()} + model_names.update({name for name, _ in qeff_model.model.named_buffers()}) + tied_weights = _collect_tied_weights(qeff_model.model) + tied_weight_map = {entry.alias: entry.canonical for entry in tied_weights} + checkpoint_files = _resolve_checkpoint_files(model_ref) + checkpoint_root = _checkpoint_root(model_ref, checkpoint_files) + checkpoint_index = _load_checkpoint_index(checkpoint_files) + relative_checkpoint_files = [ + ExternalDataFile( + path=str(Path(checkpoint_file).relative_to(checkpoint_root)) + if checkpoint_root is not None + else Path(checkpoint_file).name, + format="safetensors", + ) + for checkpoint_file in checkpoint_files + ] + backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model + promoted_inputs: List[WeightSpecInput] = [] + + for name, init_value in list(model_ir.graph.initializers.items()): + onnx_name = tied_weight_map.get(name, name) + checkpoint_key = _find_checkpoint_key(onnx_name, checkpoint_index, backbone) + + if checkpoint_key is None and name in model_names: + checkpoint_key = _find_checkpoint_key(name, checkpoint_index, backbone) + + if checkpoint_key is None: + # Computed buffer (e.g. sin_cached, cos_cached) — leave as ONNX initializer. + # The compiler embeds it in the model; it is not loaded from a checkpoint file. + continue + + location = _build_location(checkpoint_files, checkpoint_index[checkpoint_key], checkpoint_key) + + model_ir.graph.inputs.append( + ir.Value( + name=name, + shape=init_value.shape, + type=ir.TensorType(init_value.dtype), + ) + ) + del model_ir.graph.initializers[name] + promoted_inputs.append(WeightSpecInput(name=name, location=location)) + + return WeightSpec( + model_name=model_name, + model_id=model_ref, + files=relative_checkpoint_files, + inputs=promoted_inputs, + ) + + +def _prune_unused_fake_initializers(onnx_program) -> None: + """Remove FakeTensor initializers not referenced by any graph node. + + During weight-free dynamo export, meta-device parameters that are not + actually consumed in the forward graph can end up as orphan initializers. + Serialising them fails (no data to write), so we drop them here before save. + """ + from torch._subclasses.fake_tensor import FakeTensor + + initializers = onnx_program.model.graph.initializers + used_names = {name for node in onnx_program.model.graph for name in node.inputs} + used_names.update(output.name for output in onnx_program.model.graph.outputs) + for name in list(initializers): + const_value = getattr(initializers[name], "const_value", None) + raw_value = getattr(const_value, "raw", None) + if isinstance(raw_value, FakeTensor) and name not in used_names: + del initializers[name] + + +def _upsert_metadata_prop(model, key: str, value: str) -> None: + """Insert or update a metadata_props entry on an ONNX model. + + Used to embed weight_spec.json into the ONNX so the QAIC compiler + can locate external weight files without a separate sidecar lookup. + """ + import onnx + + for entry in model.metadata_props: + if entry.key == key: + entry.value = value + return + model.metadata_props.append(onnx.StringStringEntryProto(key=key, value=value)) + + +def _set_module_tensor(module: nn.Module, tensor_name: str, tensor: torch.Tensor) -> None: + parts = tensor_name.split(".") + parent = module + for part in parts[:-1]: + parent = parent[int(part)] if part.isdigit() else getattr(parent, part) + leaf_name = parts[-1] + existing = getattr(parent, leaf_name) + if isinstance(existing, torch.nn.Parameter): + setattr(parent, leaf_name, torch.nn.Parameter(tensor, requires_grad=existing.requires_grad)) + else: + setattr(parent, leaf_name, tensor) + + +def _materialize_checkpoint_metadata_tensors(qeff_model, model_ref: str) -> None: + checkpoint_files = _resolve_checkpoint_files(model_ref) + checkpoint_index = _load_checkpoint_index(checkpoint_files) + backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model + tensors = list(qeff_model.model.named_parameters()) + list(qeff_model.model.named_buffers()) + for name, tensor in tensors: + if not name.endswith("weight_shape") or not getattr(tensor, "is_meta", False): + continue + checkpoint_key = _find_checkpoint_key(name, checkpoint_index, backbone) + if checkpoint_key is None: + continue + checkpoint_file = checkpoint_index[checkpoint_key] + with safe_open(checkpoint_file, framework="pt", device="cpu") as handle: + concrete_tensor = handle.get_tensor(checkpoint_key) + _set_module_tensor(qeff_model.model, name, concrete_tensor) + + +def export_loaded_model_weight_free_onnx( + qeff_model, + tmp_onnx_path: Path, + example_inputs: Dict[str, torch.Tensor], + input_names: List[str], + output_names: List[str], + dynamic_shapes: Dict[str, Any], + export_kwargs: Dict[str, Any], + onnx_transform_kwargs: Dict[str, Any], +): + """Export an already-transformed model, then externalize its initializers. + + This Kimi fallback avoids compressed-tensors meta tracing limitations while + still producing a weight-free ONNX plus weight_spec.json. + """ + onnx_program = torch.onnx.export( + qeff_model.model, + args=(), + f=None, + kwargs=example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_axes=None, + dynamic_shapes=dynamic_shapes, + **export_kwargs, + ) + if onnx_program is None: + raise RuntimeError("torch.onnx.export returned None for loaded-model weight-free export") + + checkpoint_dir = tmp_onnx_path.parent / "loaded_weight_checkpoint" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = checkpoint_dir / "model.safetensors" + state_dict = { + name: tensor.detach().cpu().contiguous().clone() + for name, tensor in qeff_model.model.state_dict().items() + if torch.is_tensor(tensor) and not tensor.is_meta + } + save_file(state_dict, str(checkpoint_path)) + + spec = _promote_initializers_and_build_spec( + onnx_program, + model_ref=str(checkpoint_dir), + model_name=qeff_model.model_name, + qeff_model=qeff_model, + ) + save_weight_spec(tmp_onnx_path.with_name("weight_spec.json"), spec) + onnx_program.save(str(tmp_onnx_path)) + + def _cleanup(): + return None + + return qeff_model, onnx_transform_kwargs, _cleanup + + +def export_weight_free_onnx( + qeff_model, + tmp_onnx_path: Path, + example_inputs: Dict[str, torch.Tensor], + input_names: List[str], + output_names: List[str], + dynamic_shapes: Dict[str, Any], + export_kwargs: Dict[str, Any], + onnx_transform_kwargs: Dict[str, Any], +): + meta_qeff_model = _build_meta_qeff_model(qeff_model) + cleanup_required = False + + if getattr(qeff_model, "_use_onnx_subfunctions", False): + _, subfunc_kwargs, _ = _setup_onnx_subfunctions( + meta_qeff_model, + (), + { + "use_dynamo": True, + "onnx_transform_kwargs": copy.deepcopy(onnx_transform_kwargs), + "output_names": list(output_names), + }, + ) + onnx_transform_kwargs = subfunc_kwargs.get("onnx_transform_kwargs", onnx_transform_kwargs) + cleanup_required = True + + decoder_layer_classes = get_decoder_layer_classes_for_export(meta_qeff_model.model) + if getattr(meta_qeff_model, "_use_onnx_subfunctions", False) and decoder_layer_classes: + export_context = temporarily_enable_nested_compile_regions(meta_qeff_model.model, decoder_layer_classes) + else: + export_context = temporarily_disable_nested_compile_regions(meta_qeff_model.model, decoder_layer_classes) + + meta_example_inputs = _to_meta(example_inputs) + model_ref = meta_qeff_model.hash_params["pretrained_model_name_or_path"] + _materialize_checkpoint_metadata_tensors(meta_qeff_model, model_ref) + + meta_qeff_model.model.requires_grad_(False) + with export_context: + onnx_program = torch.onnx.export( + meta_qeff_model.model, + args=(), + f=None, + kwargs=meta_example_inputs, + input_names=input_names, + output_names=output_names, + dynamic_axes=None, + dynamic_shapes=dynamic_shapes, + **export_kwargs, + ) + if onnx_program is None: + raise RuntimeError("torch.onnx.export returned None for weight-free dynamo export") + + # Prepare checkpoint: stack MoE experts (if needed) and convert dtype. + # Store next to the SOURCE checkpoint directory (not inside the hashed export dir) + # so any model config variant pointing at the same source reuses the prepared data. + # Include the dtype in the directory name to avoid collisions between fp16/fp32 exports. + target_dtype = ( + getattr(qeff_model.model.config, "dtype", None) + or getattr(qeff_model.model.config, "torch_dtype", None) + or torch.float32 + ) + if target_dtype == torch.bfloat16: + target_dtype = torch.float16 + dtype_suffix = str(target_dtype).replace("torch.", "") # "float16" or "float32" + prep_pipeline = CheckpointTransformPipeline(transforms=qeff_model._checkpoint_transforms) + source_dir = _resolve_checkpoint_dir(model_ref) + prepared_out = source_dir.parent / (source_dir.name + f"-qeff-prepared-{dtype_suffix}") + + _prep_profiler = _QEffMemoryProfiler(sampling_interval=0.05, verbose=False) + _prep_profiler.start_monitoring() + _prep_profiler.mark_operation("Checkpoint Prep") + + prepared_model_ref = str( + prep_pipeline.apply( + src=source_dir, + out=prepared_out, + target_dtype=target_dtype, + ) + ) + + _prep_profiler.stop_monitoring() + print(_prep_profiler.get_memory_report()) + logger.info(_prep_profiler.get_memory_report()) + + # Expose prep stats as module-level variables so external runners + # (e.g. compare_weightfree._runner.py) can read them without parsing stdout. + import QEfficient.exporter.weight_free.core as _self_module + + _self_module._last_prep_peak_rss_mb = _prep_profiler.peak_rss + _self_module._last_prep_duration_seconds = _prep_profiler.operation_durations.get("Checkpoint Prep", 0.0) + _self_module._checkpoint_prep_ran = True + + spec = _promote_initializers_and_build_spec( + onnx_program=onnx_program, + model_ref=prepared_model_ref, + model_name=qeff_model.model_name, + qeff_model=meta_qeff_model, + ) + _prune_unused_fake_initializers(onnx_program) + onnx_program.save(str(tmp_onnx_path)) + save_weight_spec(resolve_weight_spec_path(tmp_onnx_path), spec) + + def cleanup(): + if cleanup_required: + _cleanup_onnx_subfunctions(meta_qeff_model) + + return meta_qeff_model, onnx_transform_kwargs, cleanup + + +def _load_checkpoint_tensor(checkpoint_file: str, key: str) -> np.ndarray: + handle = safe_open(checkpoint_file, framework="pt") + tensor = handle.get_tensor(key).detach().cpu() + # numpy does not support bfloat16; cast to float32 for ORT compatibility + if tensor.dtype == torch.bfloat16: + tensor = tensor.to(torch.float32) + return tensor.numpy() + + +def _default_weights_roots(weight_spec_path: Path, spec) -> List[Path]: + roots = [] + ext_root = os.environ.get("AIC_EXTERNAL_DATA_ROOT") + if ext_root: + roots.append(Path(ext_root).expanduser()) + roots.append(weight_spec_path.parent) + candidate = Path(spec.model_id).expanduser() + if candidate.exists(): + roots.append(candidate.parent) + else: + checkpoint_dir = _resolve_checkpoint_dir(spec.model_id) + checkpoint_root = _checkpoint_root(spec.model_id, [str(path) for path in checkpoint_dir.glob("*.safetensors")]) + if checkpoint_root is not None: + roots.append(checkpoint_root) + + deduped_roots: List[Path] = [] + seen = set() + for root in roots: + resolved = root.resolve() + if resolved in seen: + continue + seen.add(resolved) + deduped_roots.append(resolved) + return deduped_roots + + +def _resolve_location_file( + location: WeightSpecLocation, + files: Sequence[ExternalDataFile], + candidate_roots: Sequence[Path], +) -> Path: + if isinstance(location.file, int): + location_path = Path(files[location.file].path) + else: + location_path = Path(location.file) + if location_path.is_absolute(): + return location_path + + for root in candidate_roots: + candidate = root / location_path + if candidate.exists(): + return candidate + + return candidate_roots[0] / location_path if candidate_roots else location_path + + +def load_weight_free_ort_inputs( + weight_spec_path: Path, + runtime_inputs: Dict[str, np.ndarray], + weights_root: Optional[Path] = None, +) -> Dict[str, np.ndarray]: + weight_spec_path = Path(weight_spec_path) + spec = load_weight_spec(weight_spec_path) + candidate_roots = [] + if weights_root is not None: + candidate_roots.append(Path(weights_root).expanduser().resolve()) + candidate_roots.extend(_default_weights_roots(weight_spec_path, spec)) + + ort_inputs = dict(runtime_inputs) + for spec_input in spec.inputs: + if spec_input.name in ort_inputs: + continue + checkpoint_file = _resolve_location_file(spec_input.location, spec.files, candidate_roots) + ort_inputs[spec_input.name] = _load_checkpoint_tensor(str(checkpoint_file), spec_input.location.key) + + return ort_inputs + + +def log_weight_free_export(onnx_path: Path) -> None: + logger.info(f"Weight-free ONNX exported to {onnx_path} with spec {resolve_weight_spec_path(onnx_path)}") diff --git a/QEfficient/exporter/weight_free/spec.py b/QEfficient/exporter/weight_free/spec.py new file mode 100644 index 0000000000..995ab328e0 --- /dev/null +++ b/QEfficient/exporter/weight_free/spec.py @@ -0,0 +1,98 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import json +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Union + +WEIGHT_SPEC_VERSION = 5 + + +@dataclass +class TiedWeightAlias: + alias: str + canonical: str + + +@dataclass +class ExternalDataFile: + path: str + format: str + + +CheckpointFile = ExternalDataFile + + +@dataclass +class WeightSpecLocation: + file: Union[int, str] + key: str + + +@dataclass +class WeightSpecInput: + name: str + location: WeightSpecLocation # required: every spec entry must point to a file + + +@dataclass +class WeightSpec: + model_name: str + model_id: str + files: List[ExternalDataFile] = field(default_factory=list) + inputs: List[WeightSpecInput] = field(default_factory=list) + version: int = WEIGHT_SPEC_VERSION + + def to_dict(self) -> Dict[str, Any]: + data = asdict(self) + data["model_id"] = str(data["model_id"]) + return data + + +def save_weight_spec(path: Path, spec: WeightSpec) -> Path: + with path.open("w", encoding="utf-8") as handle: + json.dump(spec.to_dict(), handle, indent=2, sort_keys=True) + return path + + +def _load_files(raw: list) -> List[ExternalDataFile]: + if not raw: + return [] + # Backward compat: old format stored plain strings + if isinstance(raw[0], str): + return [ExternalDataFile(path=entry, format="safetensors") for entry in raw] + return [ExternalDataFile(**entry) for entry in raw] + + +def _load_location(raw: dict) -> WeightSpecLocation: + # Backward compat: old format had a redundant "type" field on the location + return WeightSpecLocation(file=raw["file"], key=raw["key"]) + + +def load_weight_spec(path: Path) -> WeightSpec: + with path.open("r", encoding="utf-8") as handle: + data = json.load(handle) + + return WeightSpec( + model_name=data["model_name"], + model_id=data["model_id"], + files=_load_files(data.get("files", data.get("checkpoint_files", []))), + inputs=[ + WeightSpecInput( + name=entry["name"], + location=_load_location(entry["location"]), + ) + for entry in data["inputs"] + if entry.get("location") is not None # backward compat: skip old buffer-only entries + ], + version=data.get("version", WEIGHT_SPEC_VERSION), + ) + + +def resolve_weight_spec_path(onnx_path: Path) -> Path: + return onnx_path.with_name("weight_spec.json") diff --git a/QEfficient/exporter/weight_free/transforms.py b/QEfficient/exporter/weight_free/transforms.py new file mode 100644 index 0000000000..46c9e8db62 --- /dev/null +++ b/QEfficient/exporter/weight_free/transforms.py @@ -0,0 +1,905 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +""" +Checkpoint transforms for weight-free ONNX export. + +Two transforms are provided, selected in priority order by CheckpointTransformPipeline: + + MoEExpertStackingCheckpointTransform — stacks per-expert HF keys into batched + tensors AND converts dtype, all in a single read pass over the checkpoint. + Skipped automatically (is_applicable=False) for dense models. + + DtypeConversionCheckpointTransform — converts all floating-point tensors to + target_dtype. Used as the fallback path for dense models. + +Usage on model classes (mirrors _pytorch_transforms / _onnx_transforms):: + + _checkpoint_transforms = [ + MoEExpertStackingCheckpointTransform, # no-op for dense (is_applicable=False) + DtypeConversionCheckpointTransform, # fallback for dense models + ] +""" + +import json +import os +import re +import shutil +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Type + +import psutil +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from QEfficient.transformers.quantizers.quantizer_utils import convert_moe_packed_tensors +from QEfficient.utils.logging_utils import logger + +# --------------------------------------------------------------------------- +# System-state helpers — used to derive worker counts at runtime +# --------------------------------------------------------------------------- + + +def _available_ram_gb() -> float: + """Available (free + reclaimable) RAM on the current machine in GB.""" + return psutil.virtual_memory().available / 1024**3 + + +def _cpu_count() -> int: + """Logical CPU count with a safe fallback.""" + return os.cpu_count() or 8 + + +def _estimate_layer_stack_gb( + expert_entries: Dict[Tuple[int, int, str], Tuple[str, str]], + layer_idx: int, + num_experts: int, + src: Path, + target_dtype: torch.dtype = torch.float32, +) -> float: + """Estimate peak RAM (GB) required to stack one MoE layer's experts. + + At the moment stacker.stack(target_dtype) runs, five tensors exist in RAM: + + Inputs (checkpoint dtype, e.g. BF16): + gate [E, I, H] + up [E, I, H] + down [E, H, I] + + Outputs (target_dtype, e.g. FP32 — twice as large when converting BF16→FP32): + gate_up [E, 2I, H] cat(gate, up).to(target_dtype) + down_out [E, H, I] down.to(target_dtype) + + Using source dtype bytes for the outputs underestimates by ~45% when + converting BF16→FP32, causing too many parallel workers and OOM. + Returns 1.0 GB as a safe fallback if the shape cannot be read. + """ + sample = next( + (v for (li, ei, k), v in expert_entries.items() if li == layer_idx and k in ("gate_proj", "linear", "w1")), + None, + ) + if sample is None: + return 1.0 + + shard_name, orig_key = sample + try: + with safe_open(str(src / shard_name), framework="pt") as f: + sl = f.get_slice(orig_key) + shape = sl.get_shape() # [I, H] + dtype_str = sl.get_dtype() + except Exception: + return 1.0 + + src_bytes = {"F32": 4, "F16": 2, "BF16": 2, "I8": 1}.get(dtype_str, 2) + tgt_bytes = {torch.float32: 4, torch.float16: 2, torch.bfloat16: 2}.get(target_dtype, 4) + ffn_dim, hidden_dim = shape + + # Three input accumulators in source dtype + two output tensors in target dtype + input_elements = num_experts * ( + ffn_dim * hidden_dim + ffn_dim * hidden_dim + hidden_dim * ffn_dim + ) # gate + up + down + output_elements = num_experts * (2 * ffn_dim * hidden_dim + hidden_dim * ffn_dim) # gate_up + down_out + return (input_elements * src_bytes + output_elements * tgt_bytes) / 1024**3 + + +# --------------------------------------------------------------------------- +# Auxiliary file names copied alongside the prepared checkpoint +# --------------------------------------------------------------------------- +_AUX_FILES = [ + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "special_tokens_map.json", + "chat_template.jinja", + "vocab.json", + "merges.txt", +] + +_SENTINEL = ".checkpoint_prepared" + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _convert_bin_to_safetensors(src: Path) -> None: + """Load a .bin checkpoint via transformers and re-save as safetensors in place. + + Uses save_pretrained(safe_serialization=True) which correctly handles tied + weights and multi-shard layouts, writing model.safetensors (single file) or + model-NNNNN-of-MMMMM.safetensors + model.safetensors.index.json (multi-shard). + Idempotent — the caller already verified no safetensors files exist before calling. + """ + import gc + + from transformers import AutoConfig, AutoModelForCausalLM + + logger.info(f"No safetensors files found in {src}. Auto-converting .bin → safetensors (one-time).") + config = AutoConfig.from_pretrained(str(src), trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + str(src), + config=config, + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + model.save_pretrained(str(src), safe_serialization=True) + del model + gc.collect() + logger.info(f"Conversion complete — safetensors files written to {src}") + + +def _read_weight_map(src: Path) -> Dict[str, str]: + """Return {tensor_key: shard_filename} from model.safetensors.index.json, + or by scanning all *.safetensors for single-file checkpoints.""" + index_path = src / "model.safetensors.index.json" + if index_path.exists(): + return json.loads(index_path.read_text())["weight_map"] + shard_files = sorted(src.glob("*.safetensors")) + if not shard_files: + raise FileNotFoundError(f"No safetensors files found in {src}") + weight_map: Dict[str, str] = {} + for sf in shard_files: + with safe_open(str(sf), framework="pt") as f: + for k in f.keys(): + weight_map[k] = sf.name + return weight_map + + +def _atomic_save(tensors: Dict[str, torch.Tensor], dst: Path) -> None: + tmp = dst.with_suffix(dst.suffix + ".tmp") + save_file({k: v.contiguous() for k, v in tensors.items()}, str(tmp)) + tmp.replace(dst) + + +def _write_index(out: Path, weight_map: Dict[str, str]) -> None: + files = set(weight_map.values()) + total_size = sum((out / f).stat().st_size for f in files if (out / f).exists()) + index = { + "metadata": {"total_size": total_size}, + "weight_map": dict(sorted(weight_map.items())), + } + (out / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + + +def _copy_aux_files(src: Path, out: Path) -> None: + for name in _AUX_FILES: + src_file = src / name + if src_file.exists() and not (out / name).exists(): + shutil.copy2(str(src_file), str(out / name)) + + +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +class BaseCheckpointTransform: + """Base class for checkpoint file transforms. Not to be instantiated. + + Each subclass produces a *complete* prepared checkpoint directory in ``out``. + The pipeline picks the first applicable transform and stops — no chaining. + """ + + def __init__(self): + raise TypeError("Checkpoint transform classes are not to be instantiated.") + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + **kwargs, + ) -> bool: + """Transform checkpoint at ``src``, write result to ``out``. + Returns True if the checkpoint was prepared, False if skipped (idempotent).""" + raise NotImplementedError + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + """Return True if this transform should run for the given checkpoint.""" + return True + + +# --------------------------------------------------------------------------- +# Transform 1: dtype conversion only — dense model path +# --------------------------------------------------------------------------- + + +class DtypeConversionCheckpointTransform(BaseCheckpointTransform): + """Convert all floating-point tensors to ``target_dtype``. + + One pass per shard, shards processed in parallel via ThreadPoolExecutor. + Used as the dense-model fallback; for MoE checkpoints, + MoEExpertStackingCheckpointTransform handles dtype conversion as part + of its own single pass and this transform is never reached. + """ + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + max_workers: Optional[int] = None, + **kwargs, + ) -> bool: + sentinel = out / _SENTINEL + if sentinel.exists(): + logger.info("DtypeConversionCheckpointTransform: prepared checkpoint exists, skipping.") + return False + + out.mkdir(parents=True, exist_ok=True) + _copy_aux_files(src, out) + + weight_map = _read_weight_map(src) + shard_names = sorted(set(weight_map.values())) + new_name_for = { + shard: (f"model_{idx:04d}.safetensors" if len(shard_names) > 1 else "model.safetensors") + for idx, shard in enumerate(shard_names) + } + + # I/O-bound: one thread per shard, capped at 4× CPU count and hard-capped + # at 256 — beyond that OS scheduling overhead outweighs I/O parallelism gains. + n_workers = max_workers if max_workers is not None else min(len(shard_names), _cpu_count() * 4, 256) + + def _process_shard(shard_name: str) -> None: + tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in f.keys(): + t = f.get_tensor(key) + tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + _atomic_save(tensors, out / new_name_for[shard_name]) + + logger.info( + f"DtypeConversionCheckpointTransform: converting {len(shard_names)} shards " + f"→ {target_dtype} | workers={n_workers} (cpus={_cpu_count()})" + ) + with ThreadPoolExecutor(max_workers=n_workers) as ex: + futures = [ex.submit(_process_shard, s) for s in shard_names] + for fut in as_completed(futures): + fut.result() + + new_weight_map = {k: new_name_for[v] for k, v in weight_map.items()} + _write_index(out, new_weight_map) + sentinel.touch() + logger.info(f"DtypeConversionCheckpointTransform: done → {out}") + return True + + +# --------------------------------------------------------------------------- +# Internal stacker helper for MoE layers +# --------------------------------------------------------------------------- + + +class _LayerStacker: + """Accumulates per-expert tensors for one MoE layer and produces batched output.""" + + def __init__(self, prefix: str, num_experts: int): + self.prefix = prefix + self.num_experts = num_experts + self._gate: Optional[torch.Tensor] = None + self._up: Optional[torch.Tensor] = None + self._down: Optional[torch.Tensor] = None + + def add(self, expert_idx: int, kind: str, tensor: torch.Tensor) -> None: + # Accept qwen3-moe names (gate_proj/up_proj/down_proj), + # grok-1 names (linear/linear_v/linear_1), + # and Mixtral names (w1=gate, w3=up, w2=down) — map to the same accumulators. + if kind in ("gate_proj", "linear", "w1"): + ffn_dim, hidden_dim = tensor.shape + if self._gate is None: + self._gate = torch.empty(self.num_experts, ffn_dim, hidden_dim, dtype=tensor.dtype) + self._gate[expert_idx] = tensor + elif kind in ("up_proj", "linear_v", "w3"): + ffn_dim, hidden_dim = tensor.shape + if self._up is None: + self._up = torch.empty(self.num_experts, ffn_dim, hidden_dim, dtype=tensor.dtype) + self._up[expert_idx] = tensor + else: # down_proj / linear_1 / w2 — shape is [hidden_dim, ffn_dim] + hidden_dim, ffn_dim = tensor.shape + if self._down is None: + self._down = torch.empty(self.num_experts, hidden_dim, ffn_dim, dtype=tensor.dtype) + self._down[expert_idx] = tensor + + def stack(self, target_dtype: torch.dtype) -> Dict[str, torch.Tensor]: + # Output in the exact layout that model __qeff_init__ creates so + # _promote_initializers_and_build_spec finds an exact checkpoint key match. + # _gate [E, I, H] → transpose(1,2) → gate_proj [E, H, I] + # _up [E, I, H] → transpose(1,2) → up_proj [E, H, I] + # _down [E, H, I] → transpose(1,2) → down_proj_t [E, I, H] + gate_proj = self._gate.to(target_dtype).transpose(1, 2).contiguous() + up_proj = self._up.to(target_dtype).transpose(1, 2).contiguous() + down_proj_t = self._down.to(target_dtype).transpose(1, 2).contiguous() + return { + f"{self.prefix}.gate_proj": gate_proj, # [E, H, I] + f"{self.prefix}.up_proj": up_proj, # [E, H, I] + f"{self.prefix}.down_proj_t": down_proj_t, # [E, I, H] + } + + +# --------------------------------------------------------------------------- +# Transform 2: MoE expert stacking + dtype conversion — single pass +# --------------------------------------------------------------------------- + + +class MoEExpertStackingCheckpointTransform(BaseCheckpointTransform): + """Stack per-expert checkpoint keys into batched tensors AND convert dtype. + + Detects the HuggingFace per-expert layout:: + + *.experts.{E}.gate_proj.weight [I, H] x num_experts + *.experts.{E}.up_proj.weight [I, H] x num_experts + *.experts.{E}.down_proj.weight [H, I] x num_experts + + and produces:: + + *.experts.gate_proj [E, H, I] (gate weights, transposed) + *.experts.up_proj [E, H, I] (up weights, transposed) + *.experts.down_proj_t [E, I, H] (down weights, transposed) + + matching the derived parameter layout that QEff MoE model __qeff_init__ + creates, so _promote_initializers_and_build_spec finds an exact key match. + Non-expert keys receive dtype conversion in the same pass. + + Parallelism: + + - Phase 1 (scan): one thread per shard, reads keys only (I/O bound, cheap). + - Phase 2 (stack): one thread per layer, loads and stacks its experts. + - Phase 3 (base): one thread per shard, converts non-expert keys. + + Phases 2 and 3 run concurrently once phase 1 completes. + """ + + EXPERT_RE = re.compile( + r"^(.+\.layers\.(\d+)\..+?\.experts)\.(\d+)\.(gate_proj|up_proj|down_proj|linear|linear_v|linear_1|w1|w2|w3)\.weight$" + ) + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + return any(cls.EXPERT_RE.match(k) for k in weight_map) + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + max_workers_scan: Optional[int] = None, + max_workers_layers: Optional[int] = None, + max_workers_base: Optional[int] = None, + **kwargs, + ) -> bool: + sentinel = out / _SENTINEL + if sentinel.exists(): + logger.info("MoEExpertStackingCheckpointTransform: prepared checkpoint exists, skipping.") + return False + + out.mkdir(parents=True, exist_ok=True) + _copy_aux_files(src, out) + + weight_map = _read_weight_map(src) + shard_names = sorted(set(weight_map.values())) + + # ── Phase 1: parallel key scan — no tensor data loaded ──────────────── + # + # expert_entries[(layer_idx, expert_idx, kind)] = (shard_name, orig_key) + # layer_prefix[layer_idx] = prefix up to .experts + # base_entries[orig_key] = shard_name + expert_entries: Dict[Tuple[int, int, str], Tuple[str, str]] = {} + layer_prefix: Dict[int, str] = {} + base_entries: Dict[str, str] = {} + + def _scan(shard_name: str) -> Tuple[Dict, Dict, Dict]: + loc_e: Dict[Tuple[int, int, str], Tuple[str, str]] = {} + loc_p: Dict[int, str] = {} + loc_b: Dict[str, str] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in f.keys(): + m = cls.EXPERT_RE.match(key) + if m: + loc_e[(int(m.group(2)), int(m.group(3)), m.group(4))] = (shard_name, key) + loc_p[int(m.group(2))] = m.group(1) + else: + loc_b[key] = shard_name + return loc_e, loc_p, loc_b + + # Phase 1: I/O-bound — cap at 4× logical CPUs, no point exceeding shard count. + # Hard cap at 256: beyond that, OS scheduling overhead outweighs I/O gains. + n_workers_scan = ( + max_workers_scan if max_workers_scan is not None else min(len(shard_names), _cpu_count() * 4, 256) + ) + logger.info( + f"MoEExpertStackingCheckpointTransform: scanning {len(shard_names)} shards " + f"(workers={n_workers_scan}, cpus={_cpu_count()}, ram_avail={_available_ram_gb():.1f} GB)..." + ) + with ThreadPoolExecutor(max_workers=n_workers_scan) as ex: + for loc_e, loc_p, loc_b in ex.map(_scan, shard_names): + expert_entries.update(loc_e) + layer_prefix.update(loc_p) + base_entries.update(loc_b) + + experts_per_layer: Dict[int, set] = {} + for layer_idx, expert_idx, _ in expert_entries: + experts_per_layer.setdefault(layer_idx, set()).add(expert_idx) + layer_indices = sorted(experts_per_layer.keys()) + sample_n = len(next(iter(experts_per_layer.values()))) if experts_per_layer else 0 + logger.info(f" {len(layer_indices)} MoE layers × {sample_n} experts each.") + + new_weight_map: Dict[str, str] = {} + + # ── Phase 2: parallel layer stacking ────────────────────────────────── + # Each layer thread loads its own experts (grouped by shard to open each + # shard at most once per layer), stacks, converts dtype, writes atomically. + def _stack_layer(layer_idx: int) -> Tuple[str, List[str]]: + num_exp = len(experts_per_layer[layer_idx]) + stacker = _LayerStacker(layer_prefix[layer_idx], num_exp) + + # Detect which kind names are present (qwen3-moe: gate_proj/up_proj/down_proj; + # grok-1: linear/linear_v/linear_1). + kinds_present = {k for (li, _, k) in expert_entries if li == layer_idx} + + by_shard: Dict[str, List[Tuple[int, str, str]]] = {} + for exp_idx in range(num_exp): + for kind in kinds_present: + shard_name, orig_key = expert_entries[(layer_idx, exp_idx, kind)] + by_shard.setdefault(shard_name, []).append((exp_idx, kind, orig_key)) + + for shard_name, entries in by_shard.items(): + with safe_open(str(src / shard_name), framework="pt") as f: + for exp_idx, kind, orig_key in entries: + stacker.add(exp_idx, kind, f.get_tensor(orig_key)) + + stacked = stacker.stack(target_dtype) + out_name = f"experts-layer-{layer_idx:05d}.safetensors" + _atomic_save(stacked, out / out_name) + return out_name, list(stacked.keys()) + + # Phase 2: memory-bound — each layer holds all E×3 expert tensors + the + # stacked output in RAM simultaneously. Derive the worker count from + # available RAM so we never OOM: keep 20% headroom, compute RAM per layer + # from the actual tensor shapes in the checkpoint. + if max_workers_layers is not None: + n_workers_layers = max_workers_layers + elif layer_indices: + sample_layer = layer_indices[0] + layer_gb = _estimate_layer_stack_gb( + expert_entries, sample_layer, len(experts_per_layer[sample_layer]), src, target_dtype + ) + available_gb = _available_ram_gb() + usable_gb = available_gb * 0.8 + n_workers_layers = max(1, min(len(layer_indices), int(usable_gb / layer_gb))) + else: + n_workers_layers = 1 + layer_gb = 0.0 + + logger.info( + f" Stacking {len(layer_indices)} layers → {target_dtype} | " + f"workers={n_workers_layers} (~{layer_gb:.2f} GB/layer, " + f"{_available_ram_gb():.1f} GB available)..." + ) + with ThreadPoolExecutor(max_workers=n_workers_layers) as ex: + futures = {ex.submit(_stack_layer, li): li for li in layer_indices} + for fut in as_completed(futures): + li = futures[fut] + out_name, out_keys = fut.result() + for key in out_keys: + new_weight_map[key] = out_name + logger.info(f" layer {li:5d} → {out_name}") + + # ── Phase 3: parallel base shard conversion ──────────────────────────── + by_shard_base: Dict[str, List[str]] = {} + for key, shard_name in base_entries.items(): + by_shard_base.setdefault(shard_name, []).append(key) + + base_shard_list = sorted(by_shard_base) + new_base_name_for = {shard: f"base-{idx:04d}.safetensors" for idx, shard in enumerate(base_shard_list)} + + def _convert_base(shard_name: str, keys: List[str]) -> None: + tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in keys: + t = f.get_tensor(key) + tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + _atomic_save(tensors, out / new_base_name_for[shard_name]) + + # Phase 3: mixed I/O + memory — one thread per shard, capped at CPU count. + n_workers_base = max_workers_base if max_workers_base is not None else min(len(base_shard_list), _cpu_count()) + logger.info(f" Converting {len(base_shard_list)} base shards → {target_dtype} | workers={n_workers_base}...") + with ThreadPoolExecutor(max_workers=n_workers_base) as ex: + futures_base = [ex.submit(_convert_base, s, keys) for s, keys in by_shard_base.items()] + for fut in as_completed(futures_base): + fut.result() + + for key, shard_name in base_entries.items(): + new_weight_map[key] = new_base_name_for[shard_name] + + _write_index(out, new_weight_map) + sentinel.touch() + logger.info(f"MoEExpertStackingCheckpointTransform: done → {out}") + return True + + +# --------------------------------------------------------------------------- +# Transform 3: GptOss MXFP4 dequantize + split fused projections +# --------------------------------------------------------------------------- + + +class GptOssMxfp4ExpertDequantSplitCheckpointTransform(BaseCheckpointTransform): + """Dequantize MXFP4-packed stacked expert tensors and split fused gate_up_proj. + + Detects the GptOss MXFP4 checkpoint layout:: + + *.experts.gate_up_proj_blocks [E, 2*I, G, B] U8 + *.experts.gate_up_proj_scales [E, 2*I, G] U8 + *.experts.gate_up_proj_bias [E, 2*I] BF16 + *.experts.down_proj_blocks [E, I, G, B] U8 + *.experts.down_proj_scales [E, I, G] U8 + *.experts.down_proj_bias [E, H] BF16 + + and produces:: + + *.experts.gate_proj [E, H, I] (dequant gate_up_proj, first half) + *.experts.up_proj [E, H, I] (dequant gate_up_proj, second half) + *.experts.gate_proj_bias [E, I] (gate_up_proj_bias, first half) + *.experts.up_proj_bias [E, I] (gate_up_proj_bias, second half) + *.experts.down_proj [E, H, I] (dequant down_proj) + *.experts.down_proj_bias [E, H] (dtype-converted, unchanged key) + + matching the derived parameter layout that QEffGptOssExperts.__qeff_init__ + creates, so _promote_initializers_and_build_spec finds an exact key match. + Non-expert keys receive dtype conversion in the same pass. + + Parallelism mirrors MoEExpertStackingCheckpointTransform: + - Phase 1 (scan): one thread per shard — collect expert tensor locations. + - Phase 2 (dequant): one thread per layer — dequant, split, write. + - Phase 3 (base): one thread per shard — dtype-convert non-expert keys. + """ + + _BLOCKS_RE = re.compile(r"^(.+\.layers\.(\d+)\..+?\.experts)\.(gate_up_proj|down_proj)_blocks$") + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + return any(cls._BLOCKS_RE.match(k) for k in weight_map) + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + max_workers_scan: Optional[int] = None, + max_workers_layers: Optional[int] = None, + max_workers_base: Optional[int] = None, + **kwargs, + ) -> bool: + sentinel = out / _SENTINEL + if sentinel.exists(): + logger.info("GptOssMxfp4ExpertDequantSplitCheckpointTransform: prepared checkpoint exists, skipping.") + return False + + out.mkdir(parents=True, exist_ok=True) + _copy_aux_files(src, out) + + weight_map = _read_weight_map(src) + shard_names = sorted(set(weight_map.values())) + + # ── Phase 1: scan — collect expert tensor locations ────────────────── + # expert_locs[(layer_idx, kind)] = (blocks_shard, blocks_key, scales_shard, scales_key) + # bias_locs[(layer_idx, kind)] = (shard, key) for gate_up_proj_bias / down_proj_bias + # layer_prefix[layer_idx] = prefix up to .experts + # base_entries[orig_key] = shard_name + _SCALES_RE = re.compile(r"^(.+\.layers\.(\d+)\..+?\.experts)\.(gate_up_proj|down_proj)_scales$") + _BIAS_RE = re.compile(r"^(.+\.layers\.(\d+)\..+?\.experts)\.(gate_up_proj|down_proj)_bias$") + + expert_locs: Dict[Tuple[int, str], Dict] = {} # {(layer, kind): {blocks/scales: (shard, key)}} + bias_locs: Dict[Tuple[int, str], Tuple[str, str]] = {} + layer_prefix: Dict[int, str] = {} + base_entries: Dict[str, str] = {} + + def _scan(shard_name: str): + loc_e: Dict[Tuple[int, str], Dict] = {} + loc_b: Dict[Tuple[int, str], Tuple[str, str]] = {} + loc_p: Dict[int, str] = {} + loc_base: Dict[str, str] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in f.keys(): + m = cls._BLOCKS_RE.match(key) + if m: + li, kind = int(m.group(2)), m.group(3) + loc_e.setdefault((li, kind), {})["blocks"] = (shard_name, key) + loc_p[li] = m.group(1) + continue + m = _SCALES_RE.match(key) + if m: + li, kind = int(m.group(2)), m.group(3) + loc_e.setdefault((li, kind), {})["scales"] = (shard_name, key) + loc_p[li] = m.group(1) + continue + m = _BIAS_RE.match(key) + if m: + li, kind = int(m.group(2)), m.group(3) + loc_b[(li, kind)] = (shard_name, key) + loc_p[li] = m.group(1) + continue + loc_base[key] = shard_name + return loc_e, loc_b, loc_p, loc_base + + n_scan = max_workers_scan if max_workers_scan is not None else min(len(shard_names), _cpu_count() * 4, 256) + logger.info( + f"GptOssMxfp4ExpertDequantSplitCheckpointTransform: scanning {len(shard_names)} shards " + f"(workers={n_scan})..." + ) + with ThreadPoolExecutor(max_workers=n_scan) as ex: + for loc_e, loc_b, loc_p, loc_base in ex.map(_scan, shard_names): + for k, v in loc_e.items(): + expert_locs.setdefault(k, {}).update(v) + bias_locs.update(loc_b) + layer_prefix.update(loc_p) + base_entries.update(loc_base) + + layer_indices = sorted({li for li, _ in expert_locs}) + logger.info(f" Found {len(layer_indices)} MoE layers.") + + new_weight_map: Dict[str, str] = {} + + # ── Phase 2: per-layer dequant + split ──────────────────────────────── + def _process_layer(layer_idx: int) -> Tuple[str, List[str]]: + prefix = layer_prefix[layer_idx] + tensors: Dict[str, torch.Tensor] = {} + + def _load(shard: str, key: str) -> torch.Tensor: + with safe_open(str(src / shard), framework="pt") as f: + return f.get_tensor(key) + + # gate_up_proj: dequant → [E, H, 2*I], then split interleaved (gate=even cols, up=odd cols) + # HF _apply_gate uses gate_up[..., ::2] for gate and gate_up[..., 1::2] for up, + # so columns are interleaved: col 0=gate0, col 1=up0, col 2=gate1, col 3=up1, ... + gu_blocks_shard, gu_blocks_key = expert_locs[(layer_idx, "gate_up_proj")]["blocks"] + gu_scales_shard, gu_scales_key = expert_locs[(layer_idx, "gate_up_proj")]["scales"] + gu_blocks = _load(gu_blocks_shard, gu_blocks_key) + gu_scales = _load(gu_scales_shard, gu_scales_key) + gate_up = convert_moe_packed_tensors(gu_blocks, gu_scales, dtype=target_dtype) + tensors[f"{prefix}.gate_proj"] = gate_up[..., 0::2].contiguous() + tensors[f"{prefix}.up_proj"] = gate_up[..., 1::2].contiguous() + + # gate_up_proj_bias: split [E, 2*I] → [E, I] + [E, I] (same interleaved convention) + if (layer_idx, "gate_up_proj") in bias_locs: + bias_shard, bias_key = bias_locs[(layer_idx, "gate_up_proj")] + gu_bias = _load(bias_shard, bias_key).to(target_dtype) + tensors[f"{prefix}.gate_proj_bias"] = gu_bias[..., 0::2].contiguous() + tensors[f"{prefix}.up_proj_bias"] = gu_bias[..., 1::2].contiguous() + + # down_proj: dequant → [E, H, I] + dp_blocks_shard, dp_blocks_key = expert_locs[(layer_idx, "down_proj")]["blocks"] + dp_scales_shard, dp_scales_key = expert_locs[(layer_idx, "down_proj")]["scales"] + dp_blocks = _load(dp_blocks_shard, dp_blocks_key) + dp_scales = _load(dp_scales_shard, dp_scales_key) + tensors[f"{prefix}.down_proj"] = convert_moe_packed_tensors(dp_blocks, dp_scales, dtype=target_dtype) + + # down_proj_bias: pass through with dtype conversion + if (layer_idx, "down_proj") in bias_locs: + dp_bias_shard, dp_bias_key = bias_locs[(layer_idx, "down_proj")] + tensors[f"{prefix}.down_proj_bias"] = _load(dp_bias_shard, dp_bias_key).to(target_dtype) + + out_name = f"experts-layer-{layer_idx:05d}.safetensors" + _atomic_save(tensors, out / out_name) + return out_name, list(tensors.keys()) + + n_layers = ( + max_workers_layers if max_workers_layers is not None else max(1, min(len(layer_indices), _cpu_count())) + ) + logger.info(f" Dequantizing {len(layer_indices)} layers | workers={n_layers}...") + with ThreadPoolExecutor(max_workers=n_layers) as ex: + futures = {ex.submit(_process_layer, li): li for li in layer_indices} + for fut in as_completed(futures): + li = futures[fut] + out_name, out_keys = fut.result() + for key in out_keys: + new_weight_map[key] = out_name + logger.info(f" layer {li:5d} → {out_name}") + + # ── Phase 3: base shard dtype conversion ────────────────────────────── + by_shard_base: Dict[str, List[str]] = {} + for key, shard_name in base_entries.items(): + by_shard_base.setdefault(shard_name, []).append(key) + + base_shard_list = sorted(by_shard_base) + new_base_name_for = {shard: f"base-{idx:04d}.safetensors" for idx, shard in enumerate(base_shard_list)} + + def _convert_base(shard_name: str, keys: List[str]) -> None: + tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(src / shard_name), framework="pt") as f: + for key in keys: + t = f.get_tensor(key) + tensors[key] = t.to(target_dtype) if t.is_floating_point() else t + _atomic_save(tensors, out / new_base_name_for[shard_name]) + + n_base = max_workers_base if max_workers_base is not None else min(len(base_shard_list), _cpu_count()) + logger.info(f" Converting {len(base_shard_list)} base shards | workers={n_base}...") + with ThreadPoolExecutor(max_workers=n_base) as ex: + futures_base = [ex.submit(_convert_base, s, keys) for s, keys in by_shard_base.items()] + for fut in as_completed(futures_base): + fut.result() + + for key, shard_name in base_entries.items(): + new_weight_map[key] = new_base_name_for[shard_name] + + _write_index(out, new_weight_map) + sentinel.touch() + logger.info(f"GptOssMxfp4ExpertDequantSplitCheckpointTransform: done → {out}") + return True + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- + + +class MoEFusedExpertSplitCheckpointTransform(BaseCheckpointTransform): + """Split already-stacked MoE expert weights into the derived layout. + + Some MoE checkpoints (e.g. Mixtral transformers >= 5.x) store experts + as per-layer fused tensors rather than per-expert individual weights: + + *.experts.gate_up_proj [E, 2*I, H] (gate and up concatenated) + *.experts.down_proj [E, H, I] + + The QEff model wrappers create derived parameters that the ONNX + initializer names refer to: + + *.experts.gate_proj [E, H, I] = gate_up_proj[:, :ffn_dim, :].T(1,2) + *.experts.up_proj [E, H, I] = gate_up_proj[:, ffn_dim:, :].T(1,2) + *.experts.down_proj_t [E, I, H] = down_proj.T(1,2) + + is_applicable returns True only when the fused format is detected. + Old-format checkpoints with per-expert keys (e.g. experts.0.gate_proj.weight) + are handled by MoEExpertStackingCheckpointTransform instead. + Also handles dtype conversion in the same pass. + """ + + _FUSED_GATE_UP_RE = re.compile(r"^(.+\.experts)\.gate_up_proj$") + _FUSED_DOWN_RE = re.compile(r"^(.+\.experts)\.down_proj$") + + @classmethod + def is_applicable(cls, weight_map: Dict[str, str]) -> bool: + return any(cls._FUSED_GATE_UP_RE.match(k) for k in weight_map) + + @classmethod + def apply( + cls, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + **kwargs, + ) -> bool: + index_path = src / "model.safetensors.index.json" + if index_path.exists(): + weight_map: Dict[str, str] = json.loads(index_path.read_text())["weight_map"] + else: + shards = sorted(src.glob("*.safetensors")) + if not shards: + return False + weight_map = {} + for shard in shards: + with safe_open(str(shard), framework="pt") as f: + for k in f.keys(): + weight_map[k] = shard.name + + if not cls.is_applicable(weight_map): + return False + + out.mkdir(parents=True, exist_ok=True) + + new_weight_map: Dict[str, str] = {} + for shard_name in sorted(set(weight_map.values())): + shard_src = src / shard_name + if not shard_src.exists(): + continue + + out_tensors: Dict[str, torch.Tensor] = {} + with safe_open(str(shard_src), framework="pt") as f: + for key in f.keys(): + tensor = f.get_tensor(key).to(target_dtype) + gate_up_m = cls._FUSED_GATE_UP_RE.match(key) + down_m = cls._FUSED_DOWN_RE.match(key) + + if gate_up_m: + prefix = gate_up_m.group(1) + ffn_dim = tensor.shape[1] // 2 + # Split fused [E,2I,H] → gate/up each [E,H,I] + out_tensors[f"{prefix}.gate_proj"] = tensor[:, :ffn_dim, :].transpose(1, 2).contiguous() + out_tensors[f"{prefix}.up_proj"] = tensor[:, ffn_dim:, :].transpose(1, 2).contiguous() + new_weight_map[f"{prefix}.gate_proj"] = shard_name + new_weight_map[f"{prefix}.up_proj"] = shard_name + # Keep original for completeness + out_tensors[key] = tensor + new_weight_map[key] = shard_name + elif down_m: + prefix = down_m.group(1) + # Transpose [E,H,I] → [E,I,H] + out_tensors[f"{prefix}.down_proj_t"] = tensor.transpose(1, 2).contiguous() + new_weight_map[f"{prefix}.down_proj_t"] = shard_name + out_tensors[key] = tensor + new_weight_map[key] = shard_name + else: + out_tensors[key] = tensor + new_weight_map[key] = shard_name + + save_file({k: v.contiguous() for k, v in out_tensors.items()}, str(out / shard_name)) + + (out / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": new_weight_map}, indent=2) + ) + return True + + +class CheckpointTransformPipeline: + """Selects and runs the first applicable checkpoint transform. + + Transforms are priority-ordered. The first one whose ``is_applicable()`` + returns True is executed and the pipeline stops. Each transform produces a + complete prepared checkpoint — there is no chaining between transforms. + + Example:: + + pipeline = CheckpointTransformPipeline([ + MoEExpertStackingCheckpointTransform, # MoE models: stacks + converts + DtypeConversionCheckpointTransform, # dense models: converts only + ]) + prepared_dir = pipeline.apply(src, out, target_dtype=torch.float32) + """ + + def __init__(self, transforms: List[Type[BaseCheckpointTransform]]): + self.transforms = transforms + + def apply( + self, + src: Path, + out: Path, + target_dtype: torch.dtype = torch.float32, + **kwargs, + ) -> Path: + src, out = Path(src), Path(out) + + # Auto-convert .bin checkpoints to safetensors on first use. + # Idempotent: skipped on subsequent runs once safetensors files exist. + has_safetensors = bool(list(src.glob("*.safetensors"))) or (src / "model.safetensors.index.json").exists() + if not has_safetensors and list(src.glob("*.bin")): + _convert_bin_to_safetensors(src) + + weight_map = _read_weight_map(src) + for transform in self.transforms: + if transform.is_applicable(weight_map): + transform.apply(src, out, target_dtype=target_dtype, **kwargs) + return out + return src # no transform applicable — source is already usable as-is diff --git a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py index bfd03faf25..0103f42fe3 100644 --- a/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py +++ b/QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py @@ -25,7 +25,7 @@ CtxScatterFunc3DInt, ) from QEfficient.customop.matmulnbits import QMOE, QuantLinearTorchFunction -from QEfficient.customop.quantization_ops import CastToUInt4Func, DequantizeLinearFunc +from QEfficient.customop.quantization_ops import cast_to_uint4, dequantize_linear from QEfficient.customop.rms_norm import CustomRMSNormFunc from QEfficient.customop.utils import select_interface from QEfficient.transformers.cache_utils import QEffDynamicCache, QEffDynamicCompressedKVRopeCache @@ -1089,23 +1089,17 @@ def original_moe(self, x, topk_ids, topk_weight): return final_out def moe_waa_unpack(self, hidden_states, topk_indices, topk_weights): - gate_proj_unpacked = CastToUInt4Func.apply(self.all_gate_qweight) - gate_zeros_unpacked = CastToUInt4Func.apply(self.all_gate_qzeros) - gate_proj_dq = DequantizeLinearFunc.apply( - gate_proj_unpacked, self.all_gate_scales, gate_zeros_unpacked, self.group_size - ) + gate_proj_unpacked = cast_to_uint4(self.all_gate_qweight) + gate_zeros_unpacked = cast_to_uint4(self.all_gate_qzeros) + gate_proj_dq = dequantize_linear(gate_proj_unpacked, self.all_gate_scales, gate_zeros_unpacked, self.group_size) - up_proj_unpacked = CastToUInt4Func.apply(self.all_up_qweight) - up_zeros_unpacked = CastToUInt4Func.apply(self.all_up_qzeros) - up_proj_dq = DequantizeLinearFunc.apply( - up_proj_unpacked, self.all_up_scales, up_zeros_unpacked, self.group_size - ) + up_proj_unpacked = cast_to_uint4(self.all_up_qweight) + up_zeros_unpacked = cast_to_uint4(self.all_up_qzeros) + up_proj_dq = dequantize_linear(up_proj_unpacked, self.all_up_scales, up_zeros_unpacked, self.group_size) - down_proj_unpacked = CastToUInt4Func.apply(self.all_down_qweight) - down_zeros_unpacked = CastToUInt4Func.apply(self.all_down_qzeros) - down_proj_dq = DequantizeLinearFunc.apply( - down_proj_unpacked, self.all_down_scales, down_zeros_unpacked, self.group_size - ) + down_proj_unpacked = cast_to_uint4(self.all_down_qweight) + down_zeros_unpacked = cast_to_uint4(self.all_down_qzeros) + down_proj_dq = dequantize_linear(down_proj_unpacked, self.all_down_scales, down_zeros_unpacked, self.group_size) num_experts = self.all_gate_qweight.shape[0] expert_in = hidden_states.unsqueeze(0).expand(num_experts, -1, -1) @@ -1541,7 +1535,10 @@ def forward( residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + if num_q_ffn_blocks is not None and self.mlp.__class__.__name__ == "DeepseekV3MoE": + hidden_states = self.mlp(hidden_states, num_q_ffn_blocks) + else: + hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) diff --git a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py index 37168e2822..103d6f3cbc 100644 --- a/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/QEfficient/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -340,9 +340,15 @@ def __qeff_init__(self): new_blocks = [] for old_block in old_blocks: - new_block = MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) - new_block.load_state_dict(old_block.state_dict()) - new_blocks.append(new_block.to(device=old_block.wqkv.weight.device, dtype=old_block.wqkv.weight.dtype)) + old_weight = old_block.wqkv.weight + if old_weight.is_meta: + with torch.device("meta"): + new_block = MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) + new_block.load_state_dict(old_block.state_dict(), assign=True) + else: + new_block = MoonViTEncoderLayer(**self.block_cfg, use_deterministic_attn=False) + new_block.load_state_dict(old_block.state_dict()) + new_blocks.append(new_block.to(device=old_weight.device, dtype=old_weight.dtype)) self.blocks = nn.ModuleList(new_blocks) @@ -521,15 +527,12 @@ def forward( inputs_embeds = merged_inputs_embeds attention_mask = merged_attention_mask - merged_image_tokens = ( - torch._shape_as_tensor(vision_embeds_for_state)[:1] - .view(1, 1) - .to(device=image_idx.device, dtype=torch.int64) + merged_image_tokens = torch.full( + (1, 1), vision_embeds_for_state.shape[0], device=image_idx.device, dtype=torch.int64 ) default_image_idx = torch.clamp(merged_image_tokens - 1, min=0) - input_batch = torch._shape_as_tensor(input_ids)[:1].view(1, 1).to(device=image_idx.device) - image_idx_batch = torch._shape_as_tensor(image_idx)[:1].view(1, 1).to(device=image_idx.device) - + input_batch = torch.full((1, 1), input_ids.shape[0], device=image_idx.device, dtype=torch.int64) + image_idx_batch = torch.full((1, 1), image_idx.shape[0], device=image_idx.device, dtype=torch.int64) if position_ids is None: post_media_position = torch.zeros_like(image_idx, dtype=torch.bool) else: @@ -588,6 +591,8 @@ def forward( output_kvs = getattr(outputs, "compressed_kvs", None) else: output_kvs = getattr(outputs, "past_key_values", None) + if vision_embeds_for_state is not None: + vision_embeds_for_state = vision_embeds_for_state + torch.zeros_like(vision_embeds_for_state) image_idx_output = image_idx[:1, :1] if image_idx is not None else image_idx return logits, vision_embeds_for_state, image_idx_output, output_kvs @@ -613,8 +618,7 @@ def _qeff_merge_input_ids_with_image_features( target_device = inputs_embeds.device image_features = image_features.to(target_device) - image_shape = torch._shape_as_tensor(image_features).to(device=input_ids.device, dtype=input_ids.dtype) - num_image_tokens = image_shape[0] + num_image_tokens = torch.full((1, 1), image_features.shape[0], device=input_ids.device, dtype=input_ids.dtype) image_token_mask = input_ids == image_token_index non_image_mask = ~image_token_mask @@ -640,22 +644,21 @@ def _qeff_merge_input_ids_with_image_features( dim=1 ) - image_start_positions = torch.where( - image_token_mask, - new_token_positions - num_image_tokens.view(1, 1) + 1, - torch.zeros_like(new_token_positions), + image_features_for_batch = image_features.unsqueeze(0).expand(input_ids.shape[0], -1, -1) + media_embedding = torch.cat([image_features_for_batch, final_embedding[:, image_features.shape[0] :, :]], dim=1) + media_attention_mask = torch.cat( + [ + torch.ones( + (input_ids.shape[0], image_features.shape[0]), + dtype=attention_mask.dtype, + device=input_ids.device, + ), + final_attention_mask[:, image_features.shape[0] :], + ], + dim=1, ) - image_start = image_start_positions.max(dim=1, keepdim=True).values - image_positions = merged_positions.squeeze(1) - image_start - max_image_index = num_image_tokens.view(1, 1) - 1 - safe_image_positions = torch.minimum(torch.clamp(image_positions, min=0), max_image_index) - image_slots = torch.logical_and(image_positions >= 0, image_positions < num_image_tokens.view(1, 1)) - image_slots = torch.logical_and(image_slots, has_image) - image_slots = torch.logical_and(image_slots, torch.logical_not(text_position_one_hot.any(dim=1))) - - gathered_image_embeddings = image_features[safe_image_positions.to(torch.long)] - final_embedding = torch.where(image_slots.unsqueeze(-1), gathered_image_embeddings, final_embedding) - final_attention_mask = torch.logical_or(final_attention_mask.bool(), image_slots).to(final_attention_mask.dtype) + final_embedding = torch.where(has_image.unsqueeze(-1), media_embedding, final_embedding) + final_attention_mask = torch.where(has_image, media_attention_mask, final_attention_mask) position_ids = torch.cumsum(final_attention_mask, dim=1) - 1 position_ids = torch.where(final_attention_mask == 0, torch.full_like(position_ids, -1), position_ids) diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index e811c94f04..914583d6d6 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -32,6 +32,12 @@ from QEfficient.base.modeling_qeff import QEFFBaseModel from QEfficient.base.onnx_transforms import FP16ClipTransform, SplitTensorsTransform from QEfficient.base.pytorch_transforms import SplitGateUpWeightsTransform +from QEfficient.exporter.weight_free.transforms import ( + DtypeConversionCheckpointTransform, + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, +) from QEfficient.generation.cloud_infer import QAICInferenceSession, is_retained_state_name from QEfficient.generation.text_generation_inference import ( CloudAI100ExecInfoNew, @@ -1069,6 +1075,8 @@ class QEffVisionEncoderForTextImageToTextModel(QEFFBaseModel): of multimodal models for optimal performance on Cloud AI 100 hardware. """ + _hf_auto_class = AutoModelForImageTextToText + _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, @@ -1078,6 +1086,12 @@ class QEffVisionEncoderForTextImageToTextModel(QEFFBaseModel): KVCacheExternalModuleMapperTransform, ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] def __init__(self, model: nn.modules, **kwargs): """ @@ -1126,6 +1140,8 @@ def export(self, inputs, output_names, dynamic_axes, export_dir=None, offload_pt export_dir=export_dir, offload_pt_weights=offload_pt_weights, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), + dynamo=kwargs.get("dynamo", False), + use_weight_free_export=kwargs.get("use_weight_free_export", False), ) def compile( @@ -1204,6 +1220,8 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): of multimodal models for optimal performance on Cloud AI 100 hardware. """ + _hf_auto_class = AutoModelForImageTextToText + _pytorch_transforms = [ AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform, @@ -1216,6 +1234,12 @@ class QEffCausalLMForTextImageToTextModel(QEFFBaseModel): SplitGateUpWeightsTransform, ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] def __init__(self, model, qaic_config: Optional[dict] = None, **kwargs): """ @@ -1328,6 +1352,8 @@ def export( export_dir=export_dir, offload_pt_weights=offload_pt_weights, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), + dynamo=kwargs.get("dynamo", False), + use_weight_free_export=kwargs.get("use_weight_free_export", False), ) def compile( @@ -1547,6 +1573,7 @@ def export( layerwise_window_size: int = 1, kv_cache_prefix: Optional[str] = None, offload_pt_weights: Optional[bool] = None, + use_weight_free_export: bool = False, **kwargs, ) -> str: """ @@ -1581,6 +1608,7 @@ def export( enable_chunking=enable_chunking, layerwise_window_size=layerwise_window_size, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **kwargs, ) bs: int = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE @@ -1662,6 +1690,8 @@ def export( export_dir=export_dir, offload_pt_weights=False, use_onnx_subfunctions=use_onnx_subfunctions, + dynamo=use_weight_free_export, + use_weight_free_export=use_weight_free_export, ) # TODO: remove the current pt weight offload capability once CustomLoader is in place @@ -1686,6 +1716,8 @@ def export( prefill_seq_len=prefill_seq_len, _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, + dynamo=use_weight_free_export, + use_weight_free_export=use_weight_free_export, ) return self.onnx_path @@ -1870,6 +1902,7 @@ def compile( layerwise: bool = False, layerwise_window_size: int = 1, kv_cache_prefix: Optional[str] = None, + use_weight_free_export: bool = False, moe_prefill_packed_chunk_size: int = constants.MOE_PREFILL_PACKED_CHUNK_SIZE, **compiler_options, ) -> str: @@ -1958,6 +1991,7 @@ def compile( qaic_config=qaic_config, moe_prefill_packed_chunk_size=moe_prefill_packed_chunk_size, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **compiler_options, ) self.vision_model.onnx_path = vision_wrapper.vision_model.onnx_path @@ -1989,6 +2023,7 @@ def compile( qaic_config=qaic_config, layerwise_window_size=layerwise_window_size, kv_cache_prefix=kv_cache_prefix, + use_weight_free_export=use_weight_free_export, **compiler_options, ) @@ -2081,6 +2116,7 @@ def compile( _layerwise_cache_probe=layerwise_cache_probe, kv_cache_prefix=kv_cache_prefix, offload_pt_weights=offload_pt_weights, + use_weight_free_export=use_weight_free_export, ) if layerwise_cache_probe: return self.lang_model.onnx_path @@ -2127,7 +2163,7 @@ def compile( compiler_options["node_precision_info"] = self.model.get_npi_file(self.model.name_or_path) if not skip_lang: - custom_io_lang = {} + custom_io_lang = {"vision_embeds": CUSTOM_IO_DTYPE_MAP[target_dtype]} for output_name in output_names["lang"]: if output_name.endswith("_RetainedState"): dtype = ( @@ -2399,8 +2435,21 @@ def kv_offload_generate( } } - vision_inputs_fp16 = {"pixel_values", "image_masks"} - vision_inputs.update({k: vision_inputs[k].astype("float16") for k in vision_inputs_fp16 if k in vision_inputs}) + for input_name in ("pixel_values", "image_masks"): + if input_name not in vision_inputs or input_name not in vision_session.binding_index_map: + continue + binding = vision_session.bindings[vision_session.binding_index_map[input_name]] + vision_inputs[input_name] = vision_inputs[input_name].astype( + vision_session.aic_to_np_dtype_mapping[binding.type] + ) + + # Required for KIMI-K25 + grid_thws_val = inputs.pop("grid_thws", None) + if grid_thws_val is not None: + h_val = int(grid_thws_val[0, 1].item()) + w_val = int(grid_thws_val[0, 2].item()) + vision_inputs["h_shape"] = np.ones((h_val), dtype=np.int64) + vision_inputs["w_shape"] = np.ones((w_val), dtype=np.int64) # Required for KIMI-K25 grid_thws_val = inputs.pop("grid_thws", None) @@ -2520,6 +2569,9 @@ def kv_offload_generate( if self._write_io_dir is not None: write_io_files(lang_inputs, outputs, self._write_io_dir, "prefill", "aic_batch_io", True, False) + if "image_idx_output" in outputs: + lang_inputs["image_idx"] = chunk_inputs["image_idx"] + prefill_time = perf_counter() - lang_start + vision_end - vision_start # Skip inputs/outputs again lang_session.skip_buffers( @@ -3452,6 +3504,12 @@ class QEFFAutoModelForCausalLM(QEFFBaseModel): ] _onnx_transforms = [] + _checkpoint_transforms = [ + GptOssMxfp4ExpertDequantSplitCheckpointTransform, + MoEExpertStackingCheckpointTransform, + MoEFusedExpertSplitCheckpointTransform, + DtypeConversionCheckpointTransform, + ] def prefill( self, @@ -3728,7 +3786,9 @@ def get_seq_len_and_handle_specialized_prefill_model( self.hash_params["chunking"] = True if self.model.config.model_type in {"qwen3_moe", "gpt_oss", "glm4_moe"}: return max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) - return constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + seq_len = max(prefill_seq_len or 0, constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN) + self.hash_params["chunking_seq_len"] = seq_len + return seq_len num_q_blocks = ( self.hash_params["blocking_config"].num_q_blocks if self.hash_params.get("blocking_kwargs", None) else None @@ -3894,10 +3954,6 @@ def export( # TODO: move this to a DA Serving utility class if self.model.config.model_type in SPECIALIZED_DISAGG_SERVING_MODEL_ARCH: if prefill_only: - if not enable_chunking and self.continuous_batching: - raise NotImplementedError( - "Looks like you are trying to run prefix-caching without chunking, this feature is not available yet!" - ) self.__update_prefill_transform(enable=True, enable_chunking=enable_chunking) self.hash_params.pop("retain_full_kv", None) seq_len = self.get_seq_len_and_handle_specialized_prefill_model( @@ -4118,6 +4174,17 @@ def _legacyify_cache(obj): output_names = apply_kv_cache_prefix(output_names, kv_cache_prefix) self.hash_params["kv_cache_prefix"] = kv_cache_prefix + if prefill_only: + effective_prefill_seq_len = ( + prefill_seq_len if prefill_seq_len is not None else seq_len if enable_chunking else None + ) + assert effective_prefill_seq_len is not None, "prefill_seq_len must be provided when prefill_only is True" + num_q_blocks_ffn = effective_prefill_seq_len // constants.EXPERT_BLOCKING_PACKED_CHUNK_SIZE + num_q_blocks_ffn = num_q_blocks_ffn if num_q_blocks_ffn > 0 else 1 + model_body = getattr(self.model, "model", None) + if model_body is not None: + setattr(model_body, "num_q_blocks_ffn", num_q_blocks_ffn) + if QEFFBaseModel._layerwise_active: return self._export_layerwise( example_inputs, diff --git a/QEfficient/utils/constants.py b/QEfficient/utils/constants.py index cff175a70e..f26b09087b 100644 --- a/QEfficient/utils/constants.py +++ b/QEfficient/utils/constants.py @@ -148,6 +148,7 @@ def get_default_aic_hw_version() -> str: DEFAULT_AIC_HW_VERSION = get_default_aic_hw_version() ONNX_TRANSFORM_MEMORY_CLEANUP_INTERVAL = 100 +EXPERT_BLOCKING_PACKED_CHUNK_SIZE = int(os.environ.get("EXPERT_BLOCKING_PACKED_CHUNK_SIZE", "256")) # Generic config key aliases used across model families. ATTENTION_HEAD_CONFIG_KEYS = ("num_attention_heads", "n_head", "n_heads", "num_heads") KV_HEAD_CONFIG_KEYS = ("num_key_value_heads", "n_kv_heads", "num_kv_heads", "effective_n_kv_heads") diff --git a/QEfficient/utils/export_utils.py b/QEfficient/utils/export_utils.py index bc61f868a2..9218191e2a 100755 --- a/QEfficient/utils/export_utils.py +++ b/QEfficient/utils/export_utils.py @@ -69,6 +69,7 @@ def convert_dynamic_axes_to_dynamic_shapes( torch.onnx.export(dynamic_shapes=...). """ max_seq_len = getattr(model_config, "max_position_embeddings", 1024) + max_image_dim = max(max_seq_len, 65536) model_type = getattr(model_config, "model_type", None) batch_min = 1 if model_type == "gpt_oss" else 2 @@ -87,6 +88,12 @@ def resolve_dim(dim_name: str): dim_registry[dim_name] = Dim("comp_ctx_lengths", min=DYNAMO_DIM_MIN_COMP_CTX_LENGTHS, max=max_seq_len) elif "ctx_len" in dim_name: dim_registry[dim_name] = Dim("ctx_len", min=2, max=max_seq_len) + elif dim_name == "num_patches": + dim_registry[dim_name] = Dim("num_patches", min=1, max=max_image_dim) + elif dim_name == "num_image_tokens": + dim_registry[dim_name] = Dim("num_image_tokens", min=1, max=max_image_dim) + elif dim_name in {"grid_h", "grid_w"}: + dim_registry[dim_name] = Dim(dim_name, min=1, max=max_image_dim) elif "sliding_window" in dim_name: dim_registry[dim_name] = Dim( "sliding_window", diff --git a/QEfficient/utils/torch_patches.py b/QEfficient/utils/torch_patches.py index 18cab51231..8e705eedf2 100644 --- a/QEfficient/utils/torch_patches.py +++ b/QEfficient/utils/torch_patches.py @@ -26,6 +26,7 @@ - _translate_fx_graph / _convert_fx_arg_to_onnx_arg nested tensor constants """ +import inspect from contextlib import contextmanager import torch @@ -337,3 +338,50 @@ def temporarily_enable_nested_compile_regions(model, target_classes=None): delattr(module, "forward") else: setattr(module, "forward", previous_forward) + + +@contextmanager +def temporarily_disable_nested_compile_regions(model, target_classes=None): + """ + Replace nested_compile_region-wrapped ``forward`` methods with their original + underlying functions for the duration of plain dynamo export (flat graph path). + + Used when use_weight_free_export=True and use_onnx_subfunctions=False so that + @nested_compile_region boundaries on decoder layer forward() methods do not + create unwanted subgraph splits during tracing. + """ + target_classes = tuple(target_classes) if target_classes else None + patched_modules = [] + + try: + for module in model.modules(): + if target_classes and not isinstance(module, target_classes): + continue + + bound_forward = getattr(module, "forward", None) + if bound_forward is None: + continue + + wrapped_forward = getattr(bound_forward, "__func__", bound_forward) + if getattr(wrapped_forward, "__qualname__", "") != "mark_compile_region..wrap..inner": + continue + + closure = getattr(wrapped_forward, "__closure__", None) or () + original_forward = next( + (cell.cell_contents for cell in closure if inspect.isfunction(cell.cell_contents)), + None, + ) + if original_forward is None: + continue + + previous_forward = module.__dict__.get("forward", _MISSING_INSTANCE_ATTR) + setattr(module, "forward", original_forward.__get__(module, type(module))) + patched_modules.append((module, previous_forward)) + + yield + finally: + for module, previous_forward in reversed(patched_modules): + if previous_forward is _MISSING_INSTANCE_ATTR: + delattr(module, "forward") + else: + setattr(module, "forward", previous_forward) diff --git a/examples/kimi_k2/export_kimi_k25_dynamo.py b/examples/kimi_k2/export_kimi_k25_dynamo.py new file mode 100644 index 0000000000..f2b03bb1f1 --- /dev/null +++ b/examples/kimi_k2/export_kimi_k25_dynamo.py @@ -0,0 +1,366 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +"""Export, compile, and generate with a reduced Kimi K2.5 model slice. + +The ONNX export step uses the Dynamo path directly on the dual-QPC component +wrappers. The exported ONNX files are then passed to the normal QEfficient +compile and generation APIs. +""" + +import argparse +import importlib.util +import os +from io import BytesIO +from pathlib import Path + +import requests +import torch +from PIL import Image + +from QEfficient import QEFFAutoModelForImageTextToText + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +parse_expert_ids = load_kimi_utils.parse_expert_ids +set_deterministic = load_kimi_utils.set_deterministic +load_kimi_k25_layer_subset_model = load_kimi_utils.load_kimi_k25_layer_subset_model + +DEFAULT_EXPORT_DIR = Path.home() / ".cache" / "qeff" / "kimi_k25_dynamo_export" +DEFAULT_COMPILE_DIR = Path.home() / ".cache" / "qeff" / "kimi_k25_dynamo_compile" +DEFAULT_IMAGE_URL = "https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png" + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Export Kimi K2.5 ONNX with Dynamo, compile QPCs, and run generation.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--model-path", + type=Path, + default="/home/huggingface_hub/models--moonshotai--Kimi-K2.5/snapshots/4d01dfe0332d63057c186e0b262165819efb6611", + help="Local Kimi-K2.5 snapshot path. Uses the local HF cache/default snapshot when omitted.", + ) + parser.add_argument("--export-dir", type=Path, default=DEFAULT_EXPORT_DIR) + parser.add_argument("--compile-dir", type=Path, default=DEFAULT_COMPILE_DIR) + parser.add_argument("--vision-onnx-path", type=Path, default=None) + parser.add_argument("--lang-onnx-path", type=Path, default=None) + parser.add_argument("--vision-qpc-path", type=Path, default=None) + parser.add_argument("--lang-qpc-path", type=Path, default=None) + parser.add_argument("--component", choices=("lang", "vision", "both"), default="both") + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument("--prefill-seq-len", type=int, default=2) + parser.add_argument("--ctx-len", type=int, default=1024) + parser.add_argument("--num-devices", type=int, default=1) + parser.add_argument("--num-cores", type=int, default=16) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--image-url", type=str, default=DEFAULT_IMAGE_URL) + parser.add_argument("--image-path", type=Path, default=None) + parser.add_argument("--image-height", type=int, default=None) + parser.add_argument("--image-width", type=int, default=None) + parser.add_argument("--generation-len", type=int, default=10) + parser.add_argument( + "--device-ids", + type=lambda device_ids: [int(device_id) for device_id in device_ids.strip("[]").split(",")], + default=[0], + help="Device IDs for generation, e.g. [0] or [0,1].", + ) + parser.add_argument("--mxfp6-matmul", action="store_true") + parser.add_argument("--mxint8-kv-cache", action="store_true") + parser.add_argument("--mos", type=int, default=1) + parser.add_argument("--aic-enable-depth-first", action="store_true") + parser.add_argument("--skip-export", action="store_true") + parser.add_argument("--skip-compile", action="store_true") + parser.add_argument("--skip-generate", action="store_true") + parser.add_argument( + "--use-onnx-subfunctions", + action="store_true", + help="Enable repeated-subgraph ONNX functions during Dynamo export.", + ) + parser.add_argument( + "--keep-weights", + action="store_true", + help="Keep PyTorch weights resident after export instead of offloading them to meta tensors.", + ) + parser.add_argument( + "--use-weight-free-export", + action="store_true", + help="Export ONNX graphs without embedded weights and emit weight_spec.json metadata.", + ) + args = parser.parse_args() + if (args.image_height is None) != (args.image_width is None): + parser.error("--image-height and --image-width must be provided together.") + if not args.skip_generate and args.component != "both": + parser.error("Generation requires --component both so both vision and language QPCs are available.") + if args.skip_export and (args.vision_onnx_path is None or args.lang_onnx_path is None) and not args.skip_compile: + parser.error("--skip-export requires --vision-onnx-path and --lang-onnx-path unless --skip-compile is set.") + if args.skip_compile and not args.skip_generate and (args.vision_qpc_path is None or args.lang_qpc_path is None): + parser.error("--skip-compile generation requires --vision-qpc-path and --lang-qpc-path.") + return args + + +def validate_dynamo_torch_version(): + torch_version = torch.__version__ + major, minor = (int(part) for part in torch_version.split("+")[0].split(".")[:2]) + if (major, minor) < (2, 13): + raise RuntimeError( + f"Kimi K2.5 Dynamo export requires PyTorch >= 2.13, but {torch_version} is installed. " + "Install the Dynamo requirements before running this example." + ) + + +def configure_qaic_tool_path(): + qaic_exec_path = Path("/opt/qti-aic/exec") + if qaic_exec_path.exists(): + os.environ["PATH"] = f"{qaic_exec_path}{os.pathsep}{os.environ.get('PATH', '')}" + + +def build_qeff_model(args): + model_path = args.model_path + subset_model_path = None + model_ref = model_path + if args.use_weight_free_export: + subset_model_path = args.export_dir.expanduser().resolve() / "kimi_k25_weightfree_subset" + model_ref = subset_model_path + model, tokenizer, processor = load_kimi_k25_layer_subset_model( + model_path=model_path, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + seed=args.seed, + subset_model_path=subset_model_path, + ) + model.eval() + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText( + model, qaic_config=qaic_config, pretrained_model_name_or_path=str(model_ref) + ) + qeff_model.model.eval() + qeff_model.vision_model.model.eval() + qeff_model.lang_model.model.eval() + qeff_model.transform( + ctx_len=args.ctx_len, + seq_len=args.prefill_seq_len, + bs=1, + num_devices=args.num_devices, + qaic_config=qaic_config, + aic_num_cores=args.num_cores, + ) + precompute_vision_rope_cache(qeff_model) + return qeff_model, tokenizer, processor, qaic_config + + +def precompute_vision_rope_cache(qeff_model): + rope_2d = qeff_model.vision_model.model.model.vision_tower.encoder.rope_2d + rope_2d._ensure_precomputed_freqs(torch.device("cpu")) + + +def get_component_export_args(qeff_model, prefill_seq_len: int): + inputs = qeff_model.model.get_dummy_inputs( + kv_offload=True, + continuous_batching=qeff_model.continuous_batching, + prefill_seq_len=prefill_seq_len, + ) + normalize_nested_cache_inputs_for_dynamo(inputs) + dynamic_axes = qeff_model.model.get_onnx_dynamic_axes( + kv_offload=True, + continuous_batching=qeff_model.continuous_batching, + comp_ctx_lengths=qeff_model.comp_ctx_lengths_decode, + ) + output_names = qeff_model.model.get_output_names(kv_offload=True) + add_static_dynamic_axes_for_dynamo(inputs, dynamic_axes) + remove_vision_output_dynamic_axes_for_dynamo(dynamic_axes) + freeze_batch_axes_for_dynamo(dynamic_axes) + return inputs, output_names, dynamic_axes + + +def normalize_nested_cache_inputs_for_dynamo(inputs): + for component_inputs in inputs.values(): + if "compressed_kvs" in component_inputs: + component_inputs["compressed_kvs"] = [tuple(layer) for layer in component_inputs["compressed_kvs"]] + if "past_key_values" in component_inputs: + component_inputs["past_key_values"] = [list(layer) for layer in component_inputs["past_key_values"]] + + +def add_static_dynamic_axes_for_dynamo(inputs, dynamic_axes): + nested_cache_inputs = {"past_key_values", "compressed_kvs"} + for component_name, component_inputs in inputs.items(): + component_dynamic_axes = dynamic_axes[component_name] + for input_name in component_inputs: + if input_name not in nested_cache_inputs: + component_dynamic_axes.setdefault(input_name, {}) + + +def remove_vision_output_dynamic_axes_for_dynamo(dynamic_axes): + dynamic_axes["vision"].pop("vision_embeds", None) + + +def freeze_batch_axes_for_dynamo(dynamic_axes): + for component_dynamic_axes in dynamic_axes.values(): + for axes_map in component_dynamic_axes.values(): + for axis_idx, dim_name in tuple(axes_map.items()): + if dim_name in {"batch_size", "full_batch_size"}: + axes_map.pop(axis_idx) + + +def export_vision(qeff_model, inputs, output_names, dynamic_axes, args): + return qeff_model.vision_model._export( + inputs["vision"], + output_names=output_names["vision"], + dynamic_axes=dynamic_axes["vision"], + export_dir=args.export_dir, + offload_pt_weights=False, + dynamo=True, + use_onnx_subfunctions=args.use_onnx_subfunctions, + use_weight_free_export=args.use_weight_free_export, + ) + + +def export_language(qeff_model, inputs, output_names, dynamic_axes, args): + qeff_model.lang_model.hash_params["prefill_only"] = False + onnx_path = qeff_model.lang_model._export( + inputs["lang"], + output_names=output_names["lang"], + dynamic_axes=dynamic_axes["lang"], + export_dir=args.export_dir, + offload_pt_weights=not args.keep_weights, + dynamo=True, + use_onnx_subfunctions=args.use_onnx_subfunctions, + use_weight_free_export=args.use_weight_free_export, + ) + return onnx_path + + +def export_components(qeff_model, inputs, output_names, dynamic_axes, args): + exported_paths = {"vision": args.vision_onnx_path, "lang": args.lang_onnx_path} + if args.skip_export: + return {component_name: path for component_name, path in exported_paths.items() if path is not None} + if args.component in {"vision", "both"}: + exported_paths["vision"] = export_vision(qeff_model, inputs, output_names, dynamic_axes, args) + if args.component in {"lang", "both"}: + exported_paths["lang"] = export_language(qeff_model, inputs, output_names, dynamic_axes, args) + return exported_paths + + +def load_generation_image(args): + if args.image_path is not None: + image = Image.open(args.image_path).convert("RGB") + else: + response = requests.get(args.image_url, timeout=30) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + if args.image_height is not None: + image = image.resize((args.image_width, args.image_height)) + return image + + +def compile_components(qeff_model, exported_paths, image, qaic_config, args): + if args.skip_compile: + qpc_paths = {"vision_qpc_path": str(args.vision_qpc_path), "lang_qpc_path": str(args.lang_qpc_path)} + qeff_model.vision_model.qpc_path = qpc_paths["vision_qpc_path"] + qeff_model.lang_model.qpc_path = qpc_paths["lang_qpc_path"] + qeff_model.qpc_paths = qpc_paths + return qpc_paths + + qpc_paths = qeff_model.compile( + vision_onnx_path=str(exported_paths.get("vision")), + lang_onnx_path=str(exported_paths.get("lang")), + compile_dir=str(args.compile_dir), + qaic_config=qaic_config, + prefill_seq_len=args.prefill_seq_len, + ctx_len=args.ctx_len, + num_cores=args.num_cores, + num_devices=args.num_devices, + mxfp6_matmul=args.mxfp6_matmul, + mxint8_kv_cache=args.mxint8_kv_cache, + skip_vision=args.component == "lang", + skip_lang=args.component == "vision", + aic_enable_depth_first=args.aic_enable_depth_first, + mos=args.mos, + image_height=image.height if image is not None else args.image_height, + image_width=image.width if image is not None else args.image_width, + use_weight_free_export=args.use_weight_free_export, + ) + return qpc_paths + + +def build_generation_inputs(processor, qeff_model, image, args): + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + return inputs + + +def generate(qeff_model, tokenizer, processor, image, args): + inputs = build_generation_inputs(processor, qeff_model, image, args) + output = qeff_model.generate( + inputs=inputs, + device_ids=args.device_ids, + generation_len=args.generation_len, + image_height=image.height, + image_width=image.width, + ) + print(output.generated_ids) + print(tokenizer.batch_decode(output.generated_ids)) + print(output) + return output + + +def main(): + os.environ.setdefault("HF_HUB_CACHE", "/home/huggingface_hub") + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + + args = parse_args() + validate_dynamo_torch_version() + configure_qaic_tool_path() + qeff_model, tokenizer, processor, qaic_config = build_qeff_model(args) + inputs, output_names, dynamic_axes = get_component_export_args(qeff_model, args.prefill_seq_len) + + exported_paths = export_components(qeff_model, inputs, output_names, dynamic_axes, args) + + for component_name, onnx_path in exported_paths.items(): + print(f"{component_name} ONNX exported to: {onnx_path}") + + image = load_generation_image(args) if not args.skip_generate else None + if not args.skip_compile or not args.skip_generate: + qpc_paths = compile_components(qeff_model, exported_paths, image, qaic_config, args) + print(f"QPC paths: {qpc_paths}") + if not args.skip_generate: + generate(qeff_model, tokenizer, processor, image, args) + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k2/export_kimi_k25_vision.py b/examples/kimi_k2/export_kimi_k25_vision.py new file mode 100644 index 0000000000..93f46cd6b3 --- /dev/null +++ b/examples/kimi_k2/export_kimi_k25_vision.py @@ -0,0 +1,396 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import argparse +import copy +import importlib.util +from io import BytesIO +from pathlib import Path +from time import perf_counter + +import numpy as np +import requests +import torch +from PIL import Image + +from QEfficient import QEFFAutoModelForImageTextToText +from QEfficient.generation.cloud_infer import QAICInferenceSession + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +load_kimi_k25_class = load_kimi_utils.load_kimi_k25_class +load_layer_subset_model = load_kimi_utils.load_layer_subset_model +parse_expert_ids = load_kimi_utils.parse_expert_ids +prepare_config = load_kimi_utils.prepare_config +set_deterministic = load_kimi_utils.set_deterministic + +PREFILL_SEQ_LEN = 512 +CTX_LEN = 2048 +BS = 1 +GENERATION_LEN = 10 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Run Kimi K2.5 vision disaggregated vision -> prefill -> decode flow.") + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument( + "--full-model", + action="store_true", + help="Load the full model. By default, the script loads a small layer subset for faster startup.", + ) + parser.add_argument("--num-vision-layers", type=int, default=NUM_VISION_LAYERS) + parser.add_argument("--num-text-layers", type=int, default=NUM_TEXT_LAYERS) + parser.add_argument("--expert-ids", type=parse_expert_ids, default=LOADED_EXPERT_IDS) + parser.add_argument("--num-experts-per-token", type=int, default=NUM_EXPERTS_PER_TOKEN) + parser.add_argument( + "--image-url", + type=str, + default="https://huggingface.co/moonshotai/Kimi-K2.5/resolve/main/figures/kimi-logo.png", + ) + parser.add_argument( + "--image-height", + type=int, + default=None, + help="Image height in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image height.", + ) + parser.add_argument( + "--image-width", + type=int, + default=None, + help="Image width in pixels for Kimi-K2.5 vision compile. Defaults to the loaded image width.", + ) + parser.add_argument("--prompt", type=str, default="Describe this image.") + parser.add_argument("--prefill-seq-len", type=int, default=PREFILL_SEQ_LEN) + parser.add_argument("--ctx-len", type=int, default=CTX_LEN) + parser.add_argument("--generation-len", type=int, default=GENERATION_LEN) + parser.add_argument("--num-cores", type=int, default=16) + parser.add_argument("--vision-num-devices", type=int, default=1) + parser.add_argument("--lang-num-devices", type=int, default=4) + parser.add_argument("--mxfp6-matmul", action="store_true") + parser.add_argument("--mxint8-kv-cache", action="store_true") + args = parser.parse_args() + if (args.image_height is None) != (args.image_width is None): + parser.error("--image-height and --image-width must be provided together.") + return args + + +def _clone_inputs(inputs): + return {key: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for key, value in inputs.items()} + + +def _numpy(value): + if torch.is_tensor(value): + return value.detach().cpu().numpy() + return np.asarray(value) + + +def _session_input_names(session: QAICInferenceSession) -> set[str]: + input_names = set(session.input_names) + input_names.update(name.rsplit("/", 1)[-1] for name in session.input_names) + return input_names + + +def _cast_for_session(session: QAICInferenceSession, name: str, value: np.ndarray) -> np.ndarray: + binding_index = session.binding_index_map.get(name) + if binding_index is None: + return value + dtype = session.aic_to_np_dtype_mapping[session.bindings[binding_index].type] + return value.astype(dtype, copy=False) + + +def _filter_session_inputs(session: QAICInferenceSession, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + input_names = _session_input_names(session) + return {name: _cast_for_session(session, name, value) for name, value in inputs.items() if name in input_names} + + +def _resolve_qpc_path(qpc_paths, key: str): + if isinstance(qpc_paths, dict): + qpc_path = qpc_paths.get(key) + if qpc_path is None: + raise KeyError(f"Missing {key!r} in compile output keys: {list(qpc_paths.keys())}") + return qpc_path + return qpc_paths + + +def _update_retained_states(target_inputs: dict[str, np.ndarray], source_outputs: dict[str, np.ndarray]): + for output_name, value in source_outputs.items(): + output_basename = output_name.rsplit("/", 1)[-1] + if output_basename.endswith("_RetainedState"): + target_inputs[output_basename.removesuffix("_RetainedState")] = value + + +def _get_next_token_ids(logits: np.ndarray) -> np.ndarray: + logits = np.asarray(logits) + return logits[:, -1, :].argmax(axis=-1).astype(np.int64).reshape(BS, 1) + + +def _compile_disagg_qpcs(qeff_model: QEFFAutoModelForImageTextToText, args, image: Image.Image): + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + common_compile_kwargs = { + "qaic_config": qaic_config, + "batch_size": BS, + "ctx_len": args.ctx_len, + "image_height": image.height, + "image_width": image.width, + "num_cores": args.num_cores, + "mxfp6_matmul": args.mxfp6_matmul, + "mxint8_kv_cache": args.mxint8_kv_cache, + "split_model_io": True, + "mos": 1, + "aic_enable_depth_first": True, + "use_onnx_subfunctions": True, + "layerwise": False, + } + + print("Compiling vision QPC...") + vision_qpc_path = qeff_model.compile( + prefill_seq_len=args.prefill_seq_len, + skip_vision=False, + skip_lang=True, + num_devices=args.vision_num_devices, + **common_compile_kwargs, + ) + + print("Compiling prefill QPC...") + prefill_qpc_path = qeff_model.compile( + prefill_seq_len=args.prefill_seq_len, + prefill_only=True, + enable_chunking=True, + skip_vision=True, + skip_lang=False, + num_devices=args.lang_num_devices, + **common_compile_kwargs, + ) + + print("Compiling decode QPC...") + decode_qpc_path = qeff_model.compile( + prefill_seq_len=1, + prefill_only=False, + skip_vision=True, + skip_lang=False, + num_devices=args.lang_num_devices, + **common_compile_kwargs, + ) + + return vision_qpc_path, prefill_qpc_path, decode_qpc_path + + +def _run_disagg_generation( + inputs: dict[str, torch.Tensor], + vision_session: QAICInferenceSession, + prefill_session: QAICInferenceSession, + decode_session: QAICInferenceSession, + *, + prefill_seq_len: int, + generation_len: int, +) -> np.ndarray: + inputs = {name: _numpy(value) for name, value in _clone_inputs(inputs).items()} + input_ids_length = inputs["input_ids"].shape[1] + num_chunks = -(input_ids_length // -prefill_seq_len) + padded_len = num_chunks * prefill_seq_len + + inputs["input_ids"] = np.pad( + inputs["input_ids"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=1, + ) + inputs["attention_mask"] = np.pad( + inputs["attention_mask"], + ((0, 0), (0, padded_len - input_ids_length)), + constant_values=0, + ) + + grid_thws = inputs.pop("grid_thws").astype(np.int64) + h = int(grid_thws[0, 1]) + w = int(grid_thws[0, 2]) + vision_inputs = { + "pixel_values": inputs["pixel_values"], + "h_shape": np.ones((h,), dtype=np.int64), + "w_shape": np.ones((w,), dtype=np.int64), + } + + vision_start = perf_counter() + print("Running vision QPC...") + vision_outputs = vision_session.run(_filter_session_inputs(vision_session, vision_inputs)) + vision_session.deactivate() + vision_time = perf_counter() - vision_start + + vision_embeds = vision_outputs.get("vision_embeds") + if vision_embeds is None: + raise RuntimeError(f"Vision QPC did not return vision_embeds. Outputs: {vision_outputs.keys()}") + + lang_inputs = { + "input_ids": inputs["input_ids"].astype(np.int64), + "position_ids": np.where(inputs["attention_mask"] > 0, np.arange(padded_len), -1).astype(np.int64), + "vision_embeds": vision_embeds, + "image_idx": np.zeros((BS, 1), dtype=np.int64), + } + + prefill_start = perf_counter() + print("Running prefill QPC...") + prefill_session.set_buffers(vision_outputs) + chunk_inputs = lang_inputs.copy() + prefill_outputs = None + for chunk_idx in range(num_chunks): + start = chunk_idx * prefill_seq_len + end = (chunk_idx + 1) * prefill_seq_len + chunk_inputs["input_ids"] = lang_inputs["input_ids"][:, start:end] + chunk_inputs["position_ids"] = lang_inputs["position_ids"][:, start:end] + prefill_outputs = prefill_session.run(_filter_session_inputs(prefill_session, chunk_inputs)) + _update_retained_states(chunk_inputs, prefill_outputs) + if "image_idx_output" in prefill_outputs: + chunk_inputs["image_idx"] = prefill_outputs["image_idx_output"].astype(np.int64) + + prefill_session.deactivate() + if prefill_outputs is None: + raise RuntimeError("QAIC prefill did not execute.") + prefill_time = perf_counter() - prefill_start + vision_time + print(f"Prefill time, including vision: {prefill_time:.2f} secs") + + generated_ids = [_get_next_token_ids(prefill_outputs["logits"])] + decode_inputs = { + "input_ids": generated_ids[-1], + "position_ids": np.max(lang_inputs["position_ids"], axis=-1, keepdims=True).astype(np.int64) + 1, + "vision_embeds": chunk_inputs.get("vision_embeds", vision_embeds), + "image_idx": chunk_inputs.get("image_idx", np.zeros((BS, 1), dtype=np.int64)), + } + _update_retained_states(decode_inputs, prefill_outputs) + + print("Running decode QPC...") + decode_start = perf_counter() + for _ in range(1, generation_len): + decode_outputs = decode_session.run(_filter_session_inputs(decode_session, decode_inputs)) + generated_ids.append(_get_next_token_ids(decode_outputs["logits"])) + decode_inputs["input_ids"] = generated_ids[-1] + decode_inputs["position_ids"] = decode_inputs["position_ids"] + 1 + if "image_idx_output" in decode_outputs: + decode_inputs["image_idx"] = decode_outputs["image_idx_output"].astype(np.int64) + _update_retained_states(decode_inputs, decode_outputs) + + decode_time = perf_counter() - decode_start + if generation_len > 1: + print(f"Decode tok/sec: {(generation_len - 1) / decode_time:.2f}") + return np.concatenate(generated_ids, axis=1) + + +def _load_model(args): + set_deterministic(1234) + config = prepare_config(args.model_path) + kimi_cls = load_kimi_k25_class(args.model_path) + + model_kwargs = { + "config": config, + "trust_remote_code": True, + "attn_implementation": "eager", + "torch_dtype": torch.float32, + } + + if args.full_model: + model, tokenizer, processor = kimi_cls.from_pretrained(str(args.model_path), **model_kwargs) + elif args.num_vision_layers is not None and args.num_text_layers is not None: + model, tokenizer, processor = load_layer_subset_model( + model_path=args.model_path, + kimi_cls=kimi_cls, + config=config, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + ) + print( + "Loaded layer subset: " + f"vision={model.config.vision_config.vt_num_hidden_layers}, " + f"text={model.config.text_config.num_hidden_layers}, " + f"experts={model.config.text_config.n_routed_experts}" + ) + else: + raise ValueError("Pass both --num-vision-layers and --num-text-layers to load a layer subset.") + + model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" + return model.eval().to("cpu"), tokenizer, processor + + +def _prepare_inputs(processor, args): + image = Image.open(BytesIO(requests.get(args.image_url, timeout=30).content)).convert("RGB") + if args.image_height is not None: + image = image.resize((args.image_width, args.image_height)) + + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs = {name: (value.to("cpu") if torch.is_tensor(value) else value) for name, value in inputs.items()} + return image, inputs + + +def main(): + args = parse_args() + model, tokenizer, processor = _load_model(args) + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText( + model, + kv_offload=True, + config=model.config, + torch_dtype=torch.float32, + qaic_config=qaic_config, + layerwise=False, + ) + + image, inputs = _prepare_inputs(processor, args) + inputs["pixel_values"] = inputs["pixel_values"].to(qeff_model.model.config.torch_dtype) + + vision_qpc_path, prefill_qpc_path, decode_qpc_path = _compile_disagg_qpcs(qeff_model, args, image) + print(f"Vision QPC path: {vision_qpc_path}") + print(f"Prefill QPC path: {prefill_qpc_path}") + print(f"Decode QPC path: {decode_qpc_path}") + + sessions = [] + try: + vision_session = QAICInferenceSession(_resolve_qpc_path(vision_qpc_path, "vision_qpc_path")) + prefill_session = QAICInferenceSession(_resolve_qpc_path(prefill_qpc_path, "lang_prefill_qpc_path")) + decode_session = QAICInferenceSession(_resolve_qpc_path(decode_qpc_path, "lang_decode_qpc_path")) + sessions.extend([vision_session, prefill_session, decode_session]) + + generated_ids = _run_disagg_generation( + inputs, + vision_session, + prefill_session, + decode_session, + prefill_seq_len=args.prefill_seq_len, + generation_len=args.generation_len, + ) + finally: + for session in sessions: + session.deactivate() + + print(generated_ids) + print(tokenizer.batch_decode(torch.as_tensor(generated_ids), skip_special_tokens=True)) + + +if __name__ == "__main__": + main() diff --git a/examples/kimi_k2/test_kimi_k25_dynamo.py b/examples/kimi_k2/test_kimi_k25_dynamo.py new file mode 100644 index 0000000000..0ceb0da322 --- /dev/null +++ b/examples/kimi_k2/test_kimi_k25_dynamo.py @@ -0,0 +1,429 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import importlib.util +import os +import re +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import torch +from export_kimi_k25_dynamo import ( + DEFAULT_COMPILE_DIR, + DEFAULT_EXPORT_DIR, + DEFAULT_IMAGE_URL, + compile_components, + configure_qaic_tool_path, + export_components, + get_component_export_args, + load_generation_image, + precompute_vision_rope_cache, + validate_dynamo_torch_version, +) + +from QEfficient import QEFFAutoModelForImageTextToText + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +parse_expert_ids = load_kimi_utils.parse_expert_ids +set_deterministic = load_kimi_utils.set_deterministic +load_kimi_k25_layer_subset_model = load_kimi_utils.load_kimi_k25_layer_subset_model + +TEXT_PROMPT = "Describe this image." +NEW_GENERATION_TOKENS = 10 +CTX_LEN = 1024 +PREFILL_SEQ_LEN = 2 + + +def _has_qaic_runtime_access() -> bool: + try: + import qaicrt + + _ctx = qaicrt.Context() + return True + except (ImportError, OSError, RuntimeError, AttributeError): + return False + + +def _skip_test(reason: str): + try: + import pytest + + pytest.skip(reason) + except ImportError as exc: + raise RuntimeError(reason) from exc + + +def _env_bool(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + return default if value is None else int(value) + + +def _env_path(name: str) -> Path | None: + value = os.environ.get(name) + return None if not value else Path(value).expanduser().resolve() + + +def _parse_device_ids(value: str) -> list[int]: + return [int(device_id) for device_id in value.strip().strip("[]").split(",") if device_id.strip()] + + +def _find_free_qaic_device_id() -> int | None: + qaic_util_path = Path("/opt/qti-aic/tools/qaic-util") + if not qaic_util_path.exists(): + return None + + try: + result = subprocess.run( + [str(qaic_util_path), "-q"], + check=False, + text=True, + capture_output=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + + current_qid = None + for line in result.stdout.splitlines(): + qid_match = re.match(r"^QID\s+(\d+)", line.strip()) + if qid_match: + current_qid = int(qid_match.group(1)) + continue + + free_match = re.search(r"Nsp Free:\s*(\d+)", line) + if free_match and current_qid is not None and int(free_match.group(1)) > 0: + return current_qid + + return None + + +def _resolve_device_ids() -> list[int] | None: + value = os.environ.get("KIMI_K25_DYNAMO_DEVICE_IDS") or os.environ.get("KIMI_K25_DYNAMO_DEVICE_ID") + if value: + return _parse_device_ids(value) + + free_device_id = _find_free_qaic_device_id() + if free_device_id is None: + return None + return [free_device_id] + + +def _clone_inputs(inputs): + return {name: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for name, value in inputs.items()} + + +def _decode_tokens(tokenizer, token_ids) -> str: + decoded = tokenizer.batch_decode(torch.as_tensor(token_ids), skip_special_tokens=True) + return decoded[0] if decoded else "" + + +@torch.no_grad() +def _greedy_generate_hf(model, inputs, max_new_tokens: int) -> torch.Tensor: + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + eos_token_id = getattr(model.config, "eos_token_id", None) + if eos_token_id is None and hasattr(model.config, "text_config"): + eos_token_id = getattr(model.config.text_config, "eos_token_id", None) + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + if eos_token_id is not None and torch.all(next_token == eos_token_id): + break + + return torch.cat(new_tokens, dim=1) + + +@torch.no_grad() +def _greedy_generate_qeff(qeff_model, inputs, args) -> torch.Tensor: + prefill_seq_len = args.prefill_seq_len + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + grid_thws = inputs["grid_thws"].to(torch.long) + + h_shape = torch.ones(int(grid_thws[0, 1].item()), dtype=torch.int64) + w_shape = torch.ones(int(grid_thws[0, 2].item()), dtype=torch.int64) + vision_embeds = qeff_model.vision_model.model( + inputs["pixel_values"].to(qeff_model.model.config.torch_dtype), + h_shape, + w_shape, + ).detach() + + input_ids_length = generated_ids.shape[1] + num_chunks = -(input_ids_length // -prefill_seq_len) + padded_len = num_chunks * prefill_seq_len + generated_ids = torch.nn.functional.pad( + generated_ids, + (0, padded_len - input_ids_length), + "constant", + qeff_model.model.config.pad_token_id, + ) + attention_mask = torch.nn.functional.pad(attention_mask, (0, padded_len - input_ids_length), "constant", 0) + position_ids = torch.where( + attention_mask.bool(), + torch.arange(padded_len, dtype=torch.long).view(1, -1), + torch.full((generated_ids.shape[0], padded_len), -1, dtype=torch.long), + ) + + export_inputs, _, _ = get_component_export_args(qeff_model, prefill_seq_len) + compressed_kvs = [ + [cache_tensor.clone()[: generated_ids.shape[0]] for cache_tensor in layer_cache] + for layer_cache in export_inputs["lang"]["compressed_kvs"] + ] + image_idx = torch.zeros((generated_ids.shape[0], 1), dtype=torch.int64) + new_tokens = [] + + logits = None + for chunk_idx in range(num_chunks): + start = chunk_idx * prefill_seq_len + end = start + prefill_seq_len + logits, _, image_idx_output, compressed_kvs = qeff_model.lang_model.model( + input_ids=generated_ids[:, start:end], + position_ids=position_ids[:, start:end], + vision_embeds=vision_embeds, + image_idx=image_idx, + compressed_kvs=compressed_kvs, + ) + if image_idx_output is not None: + image_idx = image_idx_output + + next_token = logits.argmax(2) + if next_token.ndim == 2 and next_token.shape[1] > 1: + next_token = next_token[:, -1:] + new_tokens.append(next_token) + + decode_position_ids = position_ids.max(dim=-1, keepdim=True).values + 1 + for _ in range(1, args.generation_len): + logits, _, image_idx_output, compressed_kvs = qeff_model.lang_model.model( + input_ids=next_token, + position_ids=decode_position_ids, + vision_embeds=vision_embeds, + image_idx=image_idx, + compressed_kvs=compressed_kvs, + ) + if image_idx_output is not None: + image_idx = image_idx_output + next_token = logits.argmax(2) + if next_token.ndim == 2 and next_token.shape[1] > 1: + next_token = next_token[:, -1:] + new_tokens.append(next_token) + decode_position_ids = decode_position_ids + 1 + + return torch.cat(new_tokens, dim=1) + + +def _build_generation_inputs(processor, image, args, dtype): + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": image}, + {"type": "text", "text": args.prompt}, + ], + }, + ] + inputs = processor( + messages=messages, + add_generation_prompt=True, + tokenize=False, + return_tensors="pt", + ) + inputs["pixel_values"] = inputs["pixel_values"].to(dtype) + return inputs + + +def _load_hf_model(args): + model_path = args.model_path + model, tokenizer, processor = load_kimi_k25_layer_subset_model( + model_path=model_path, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + seed=args.seed, + ) + return model.eval(), tokenizer, processor + + +def _build_qeff_model_from_hf(model, args): + qaic_config = {"mla_absorption": {"cache_compressed": True, "absorption": False, "online": False}} + qeff_model = QEFFAutoModelForImageTextToText(model, qaic_config=qaic_config) + qeff_model.model.eval() + qeff_model.vision_model.model.eval() + qeff_model.lang_model.model.eval() + qeff_model.transform( + ctx_len=args.ctx_len, + seq_len=args.prefill_seq_len, + bs=1, + num_devices=args.num_devices, + qaic_config=qaic_config, + aic_num_cores=args.num_cores, + ) + precompute_vision_rope_cache(qeff_model) + return qeff_model, qaic_config + + +def _make_args(device_ids: list[int]) -> SimpleNamespace: + vision_qpc_path = _env_path("KIMI_K25_DYNAMO_VISION_QPC_PATH") + lang_qpc_path = _env_path("KIMI_K25_DYNAMO_LANG_QPC_PATH") + skip_compile = vision_qpc_path is not None and lang_qpc_path is not None + + vision_onnx_path = _env_path("KIMI_K25_DYNAMO_VISION_ONNX_PATH") + lang_onnx_path = _env_path("KIMI_K25_DYNAMO_LANG_ONNX_PATH") + skip_export = skip_compile or (vision_onnx_path is not None and lang_onnx_path is not None) + + return SimpleNamespace( + model_path=_env_path("KIMI_K25_DYNAMO_MODEL_PATH"), + export_dir=_env_path("KIMI_K25_DYNAMO_EXPORT_DIR") or DEFAULT_EXPORT_DIR, + compile_dir=_env_path("KIMI_K25_DYNAMO_COMPILE_DIR") or DEFAULT_COMPILE_DIR, + vision_onnx_path=vision_onnx_path, + lang_onnx_path=lang_onnx_path, + vision_qpc_path=vision_qpc_path, + lang_qpc_path=lang_qpc_path, + component="both", + num_vision_layers=_env_int("KIMI_K25_DYNAMO_NUM_VISION_LAYERS", NUM_VISION_LAYERS), + num_text_layers=_env_int("KIMI_K25_DYNAMO_NUM_TEXT_LAYERS", NUM_TEXT_LAYERS), + expert_ids=LOADED_EXPERT_IDS, + num_experts_per_token=NUM_EXPERTS_PER_TOKEN, + prefill_seq_len=_env_int("KIMI_K25_DYNAMO_PREFILL_SEQ_LEN", PREFILL_SEQ_LEN), + ctx_len=_env_int("KIMI_K25_DYNAMO_CTX_LEN", CTX_LEN), + num_devices=len(device_ids), + num_cores=_env_int("KIMI_K25_DYNAMO_NUM_CORES", 16), + seed=_env_int("KIMI_K25_DYNAMO_SEED", 1234), + prompt=os.environ.get("KIMI_K25_DYNAMO_PROMPT", TEXT_PROMPT), + image_url=os.environ.get("KIMI_K25_DYNAMO_IMAGE_URL", DEFAULT_IMAGE_URL), + image_path=_env_path("KIMI_K25_DYNAMO_IMAGE_PATH"), + image_height=None, + image_width=None, + generation_len=_env_int("KIMI_K25_DYNAMO_GENERATION_LEN", NEW_GENERATION_TOKENS), + device_ids=device_ids, + mxfp6_matmul=_env_bool("KIMI_K25_DYNAMO_MXFP6_MATMUL"), + mxint8_kv_cache=_env_bool("KIMI_K25_DYNAMO_MXINT8_KV_CACHE"), + mos=_env_int("KIMI_K25_DYNAMO_MOS", 1), + aic_enable_depth_first=_env_bool("KIMI_K25_DYNAMO_AIC_ENABLE_DEPTH_FIRST"), + skip_export=skip_export, + skip_compile=skip_compile, + skip_generate=False, + use_onnx_subfunctions=_env_bool("KIMI_K25_DYNAMO_USE_ONNX_SUBFUNCTIONS"), + keep_weights=_env_bool("KIMI_K25_DYNAMO_KEEP_WEIGHTS"), + ) + + +def check_kimi_k25_dynamo_hf_vs_qeff_vs_qaic(): + os.environ.setdefault("HF_HUB_CACHE", "/home/huggingface_hub") + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + + validate_dynamo_torch_version() + configure_qaic_tool_path() + if not _has_qaic_runtime_access(): + _skip_test("QAIC generation skipped: no QAIC runtime access.") + + device_ids = _resolve_device_ids() + if device_ids is None: + _skip_test("QAIC generation skipped: no QAIC device has free NSPs.") + + args = _make_args(device_ids) + set_deterministic(args.seed) + model, tokenizer, processor = _load_hf_model(args) + image = load_generation_image(args) + inputs = _build_generation_inputs(processor, image, args, model.config.torch_dtype) + + hf_tokens = _greedy_generate_hf(copy.deepcopy(model), _clone_inputs(inputs), args.generation_len).cpu() + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + + qeff_model, qaic_config = _build_qeff_model_from_hf(model, args) + qeff_tokens = _greedy_generate_qeff(qeff_model, _clone_inputs(inputs), args).cpu() + print("QEFF:", _decode_tokens(tokenizer, qeff_tokens), "\n", qeff_tokens) + + assert torch.equal(hf_tokens, qeff_tokens), ( + "HF and QEff PyTorch tokens do not match for the Dynamo export wrapper path: " + f"hf={hf_tokens.tolist()}, qeff={qeff_tokens.tolist()}" + ) + + export_inputs, output_names, dynamic_axes = get_component_export_args(qeff_model, args.prefill_seq_len) + exported_paths = export_components(qeff_model, export_inputs, output_names, dynamic_axes, args) + qpc_paths = compile_components(qeff_model, exported_paths, image, qaic_config, args) + print(f"Dynamo ONNX paths: {exported_paths}") + print(f"Dynamo QPC paths: {qpc_paths}") + + qaic_output = qeff_model.generate( + inputs=_clone_inputs(inputs), + device_ids=args.device_ids, + generation_len=args.generation_len, + image_height=image.height, + image_width=image.width, + ) + qaic_tokens = torch.as_tensor(qaic_output.generated_ids[:, : args.generation_len], dtype=hf_tokens.dtype) + print("QAIC:", _decode_tokens(tokenizer, qaic_tokens), "\n", qaic_tokens) + + if torch.equal(hf_tokens, qaic_tokens): + return + + mismatch_message = ( + "HF/QEff and QAIC tokens do not match for the Dynamo exported and compiled model: " + f"hf={hf_tokens.tolist()}, qeff={qeff_tokens.tolist()}, qaic={qaic_tokens.tolist()}" + ) + if _env_bool("KIMI_K25_DYNAMO_STRICT_HF_QAIC", False): + raise AssertionError(mismatch_message) + + print( + "QAIC token drift tolerated for this reduced Kimi K2.5 Dynamo smoke test. " + "Set KIMI_K25_DYNAMO_STRICT_HF_QAIC=1 to require exact HF-vs-QAIC token parity. " + f"{mismatch_message}" + ) + assert qaic_tokens.shape == hf_tokens.shape, "HF and QAIC generated token shapes do not match" + + +def test_kimi_k25_dynamo_hf_vs_qeff_vs_qaic(): + check_kimi_k25_dynamo_hf_vs_qeff_vs_qaic() + + +if __name__ == "__main__": + check_kimi_k25_dynamo_hf_vs_qeff_vs_qaic() diff --git a/examples/kimi_k2/test_kimi_k25_dynamo_weight_free.py b/examples/kimi_k2/test_kimi_k25_dynamo_weight_free.py new file mode 100644 index 0000000000..7b3d7be54b --- /dev/null +++ b/examples/kimi_k2/test_kimi_k25_dynamo_weight_free.py @@ -0,0 +1,522 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import copy +import importlib.util +import json +import os +import re +import subprocess +from pathlib import Path +from time import perf_counter +from types import SimpleNamespace + +import onnx +import torch +from export_kimi_k25_dynamo import ( + DEFAULT_COMPILE_DIR, + DEFAULT_EXPORT_DIR, + DEFAULT_IMAGE_URL, + build_generation_inputs, + build_qeff_model, + compile_components, + configure_qaic_tool_path, + export_components, + get_component_export_args, + load_generation_image, + validate_dynamo_torch_version, +) +from onnx import numpy_helper +from safetensors.torch import save_file + +from QEfficient.generation.cloud_infer import QAICInferenceSession + +LOAD_KIMI_UTILS_PATH = Path(__file__).resolve().parents[2] / "tests" / "utils" / "load_kimi_utils.py" +_load_kimi_spec = importlib.util.spec_from_file_location("load_kimi_utils", LOAD_KIMI_UTILS_PATH) +if _load_kimi_spec is None or _load_kimi_spec.loader is None: + raise ImportError(f"Unable to load Kimi helpers from {LOAD_KIMI_UTILS_PATH}") +load_kimi_utils = importlib.util.module_from_spec(_load_kimi_spec) +_load_kimi_spec.loader.exec_module(load_kimi_utils) + +LOADED_EXPERT_IDS = load_kimi_utils.LOADED_EXPERT_IDS +NUM_EXPERTS_PER_TOKEN = load_kimi_utils.NUM_EXPERTS_PER_TOKEN +NUM_TEXT_LAYERS = load_kimi_utils.NUM_TEXT_LAYERS +NUM_VISION_LAYERS = load_kimi_utils.NUM_VISION_LAYERS +parse_expert_ids = load_kimi_utils.parse_expert_ids +set_deterministic = load_kimi_utils.set_deterministic +load_kimi_k25_layer_subset_model = load_kimi_utils.load_kimi_k25_layer_subset_model + +TEXT_PROMPT = "Describe this image." +NEW_GENERATION_TOKENS = 10 +CTX_LEN = 1024 +PREFILL_SEQ_LEN = 2 + + +class _QEffHolder: + def __init__(self, model): + self.model = model + + +def _has_qaic_runtime_access() -> bool: + try: + import qaicrt + + _ctx = qaicrt.Context() + return True + except (ImportError, OSError, RuntimeError, AttributeError): + return False + + +def _skip_test(reason: str): + try: + import pytest + + pytest.skip(reason) + except ImportError as exc: + raise RuntimeError(reason) from exc + + +def _env_bool(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + return default if value is None else int(value) + + +def _env_path(name: str) -> Path | None: + value = os.environ.get(name) + return None if not value else Path(value).expanduser().resolve() + + +def _env_path_any(*names: str) -> Path | None: + for name in names: + path = _env_path(name) + if path is not None: + return path + return None + + +def _env_int_any(default: int, *names: str) -> int: + for name in names: + value = os.environ.get(name) + if value is not None: + return int(value) + return default + + +def _env_bool_any(default: bool, *names: str) -> bool: + for name in names: + value = os.environ.get(name) + if value is not None: + return value.strip().lower() in {"1", "true", "yes", "on"} + return default + + +def _parse_device_ids(value: str) -> list[int]: + return [int(device_id) for device_id in value.strip().strip("[]").split(",") if device_id.strip()] + + +def _find_free_qaic_device_id() -> int | None: + candidate_paths = [Path("/opt/qti-aic/tools/qaic-util"), Path("/opt/qti-aic/exec/qaic-util")] + qaic_util_path = next((path for path in candidate_paths if path.exists()), None) + if qaic_util_path is None: + return None + + try: + result = subprocess.run( + [str(qaic_util_path), "-q"], + check=False, + text=True, + capture_output=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return None + + current_qid = None + for line in result.stdout.splitlines(): + qid_match = re.match(r"^QID\s+(\d+)", line.strip()) + if qid_match: + current_qid = int(qid_match.group(1)) + continue + + free_match = re.search(r"Nsp Free:\s*(\d+)", line) + if free_match and current_qid is not None and int(free_match.group(1)) > 0: + return current_qid + + return None + + +def _device_can_activate(qpc_path: Path, device_id: int) -> bool: + session = None + try: + session = QAICInferenceSession(str(qpc_path), [device_id]) + session.deactivate() + return True + except (ImportError, OSError, RuntimeError, AttributeError, ValueError): + return False + finally: + del session + + +def _find_activatable_device_id(vision_qpc_path: Path, lang_qpc_path: Path, max_device_id: int) -> int | None: + for device_id in range(max_device_id + 1): + if _device_can_activate(vision_qpc_path, device_id) and _device_can_activate(lang_qpc_path, device_id): + return device_id + return None + + +def _resolve_device_ids(args: SimpleNamespace | None = None) -> list[int] | None: + value = ( + os.environ.get("KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_IDS") + or os.environ.get("KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_ID") + or os.environ.get("KIMI_K25_DYNAMO_DEVICE_IDS") + or os.environ.get("KIMI_K25_DYNAMO_DEVICE_ID") + ) + if value: + return _parse_device_ids(value) + + if args is not None and args.vision_qpc_path is not None and args.lang_qpc_path is not None: + device_id = _find_activatable_device_id( + args.vision_qpc_path, + args.lang_qpc_path, + _env_int("KIMI_K25_DYNAMO_WEIGHT_FREE_MAX_DEVICE_ID", 63), + ) + if device_id is not None: + return [device_id] + + free_device_id = _find_free_qaic_device_id() + if free_device_id is None: + return None + return [free_device_id] + + +def _has_explicit_device_ids() -> bool: + return any( + os.environ.get(name) + for name in ( + "KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_IDS", + "KIMI_K25_DYNAMO_WEIGHT_FREE_DEVICE_ID", + "KIMI_K25_DYNAMO_DEVICE_IDS", + "KIMI_K25_DYNAMO_DEVICE_ID", + ) + ) + + +def _clone_inputs(inputs): + return {name: (value.clone() if torch.is_tensor(value) else copy.deepcopy(value)) for name, value in inputs.items()} + + +def _decode_tokens(tokenizer, token_ids) -> str: + decoded = tokenizer.batch_decode(torch.as_tensor(token_ids), skip_special_tokens=True) + return decoded[0] if decoded else "" + + +@torch.no_grad() +def _greedy_generate_hf(model, inputs, max_new_tokens: int) -> torch.Tensor: + generated_ids = inputs["input_ids"].to(torch.long) + attention_mask = inputs["attention_mask"].to(torch.long) + pixel_values = inputs["pixel_values"] + grid_thws = inputs["grid_thws"] + new_tokens = [] + + for _ in range(max_new_tokens): + outputs = model( + input_ids=generated_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + grid_thws=grid_thws, + use_cache=False, + return_dict=True, + ) + logits = outputs[0] if isinstance(outputs, tuple) else outputs.logits + next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) + new_tokens.append(next_token) + + generated_ids = torch.cat([generated_ids, next_token], dim=1) + attention_mask = torch.cat( + [ + attention_mask, + torch.ones((attention_mask.shape[0], 1), dtype=attention_mask.dtype, device=attention_mask.device), + ], + dim=1, + ) + + return torch.cat(new_tokens, dim=1) + + +def _load_hf_model(args): + model_path = args.model_path + model, tokenizer, processor = load_kimi_k25_layer_subset_model( + model_path=model_path, + num_vision_layers=args.num_vision_layers, + num_text_layers=args.num_text_layers, + loaded_expert_ids=args.expert_ids, + num_experts_per_tok=args.num_experts_per_token, + dtype=torch.float32, + seed=args.seed, + subset_model_path=args.export_dir / "kimi_k25_weightfree_subset" if args.use_weight_free_export else None, + ) + return model.eval(), tokenizer, processor + + +def _make_args(device_ids: list[int] | None = None) -> SimpleNamespace: + device_ids = device_ids or [0] + vision_qpc_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_VISION_QPC_PATH", "KIMI_K25_DYNAMO_VISION_QPC_PATH") + lang_qpc_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_LANG_QPC_PATH", "KIMI_K25_DYNAMO_LANG_QPC_PATH") + skip_compile = vision_qpc_path is not None and lang_qpc_path is not None + + vision_onnx_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_VISION_ONNX_PATH", "KIMI_K25_DYNAMO_VISION_ONNX_PATH") + lang_onnx_path = _env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_LANG_ONNX_PATH", "KIMI_K25_DYNAMO_LANG_ONNX_PATH") + skip_export = skip_compile or (vision_onnx_path is not None and lang_onnx_path is not None) + + return SimpleNamespace( + model_path=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_MODEL_PATH", "KIMI_K25_DYNAMO_MODEL_PATH"), + export_dir=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_EXPORT_DIR", "KIMI_K25_DYNAMO_EXPORT_DIR") + or DEFAULT_EXPORT_DIR.with_name(f"{DEFAULT_EXPORT_DIR.name}_weight_free"), + compile_dir=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_COMPILE_DIR", "KIMI_K25_DYNAMO_COMPILE_DIR") + or DEFAULT_COMPILE_DIR.with_name(f"{DEFAULT_COMPILE_DIR.name}_weight_free"), + vision_onnx_path=vision_onnx_path, + lang_onnx_path=lang_onnx_path, + vision_qpc_path=vision_qpc_path, + lang_qpc_path=lang_qpc_path, + component="both", + num_vision_layers=_env_int_any( + NUM_VISION_LAYERS, "KIMI_K25_DYNAMO_WEIGHT_FREE_NUM_VISION_LAYERS", "KIMI_K25_DYNAMO_NUM_VISION_LAYERS" + ), + num_text_layers=_env_int_any( + NUM_TEXT_LAYERS, "KIMI_K25_DYNAMO_WEIGHT_FREE_NUM_TEXT_LAYERS", "KIMI_K25_DYNAMO_NUM_TEXT_LAYERS" + ), + expert_ids=LOADED_EXPERT_IDS, + num_experts_per_token=NUM_EXPERTS_PER_TOKEN, + prefill_seq_len=_env_int_any( + PREFILL_SEQ_LEN, "KIMI_K25_DYNAMO_WEIGHT_FREE_PREFILL_SEQ_LEN", "KIMI_K25_DYNAMO_PREFILL_SEQ_LEN" + ), + ctx_len=_env_int_any(CTX_LEN, "KIMI_K25_DYNAMO_WEIGHT_FREE_CTX_LEN", "KIMI_K25_DYNAMO_CTX_LEN"), + num_devices=len(device_ids), + num_cores=_env_int_any(16, "KIMI_K25_DYNAMO_WEIGHT_FREE_NUM_CORES", "KIMI_K25_DYNAMO_NUM_CORES"), + seed=_env_int_any(1234, "KIMI_K25_DYNAMO_WEIGHT_FREE_SEED", "KIMI_K25_DYNAMO_SEED"), + prompt=os.environ.get( + "KIMI_K25_DYNAMO_WEIGHT_FREE_PROMPT", os.environ.get("KIMI_K25_DYNAMO_PROMPT", TEXT_PROMPT) + ), + image_url=os.environ.get( + "KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_URL", os.environ.get("KIMI_K25_DYNAMO_IMAGE_URL", DEFAULT_IMAGE_URL) + ), + image_path=_env_path_any("KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_PATH", "KIMI_K25_DYNAMO_IMAGE_PATH"), + image_height=_env_int_any(None, "KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_HEIGHT", "KIMI_K25_DYNAMO_IMAGE_HEIGHT"), + image_width=_env_int_any(None, "KIMI_K25_DYNAMO_WEIGHT_FREE_IMAGE_WIDTH", "KIMI_K25_DYNAMO_IMAGE_WIDTH"), + generation_len=_env_int_any( + NEW_GENERATION_TOKENS, "KIMI_K25_DYNAMO_WEIGHT_FREE_GENERATION_LEN", "KIMI_K25_DYNAMO_GENERATION_LEN" + ), + device_ids=device_ids, + mxfp6_matmul=_env_bool_any(False, "KIMI_K25_DYNAMO_WEIGHT_FREE_MXFP6_MATMUL", "KIMI_K25_DYNAMO_MXFP6_MATMUL"), + mxint8_kv_cache=_env_bool_any( + False, "KIMI_K25_DYNAMO_WEIGHT_FREE_MXINT8_KV_CACHE", "KIMI_K25_DYNAMO_MXINT8_KV_CACHE" + ), + mos=_env_int_any(1, "KIMI_K25_DYNAMO_WEIGHT_FREE_MOS", "KIMI_K25_DYNAMO_MOS"), + aic_enable_depth_first=_env_bool_any( + False, "KIMI_K25_DYNAMO_WEIGHT_FREE_AIC_ENABLE_DEPTH_FIRST", "KIMI_K25_DYNAMO_AIC_ENABLE_DEPTH_FIRST" + ), + skip_export=skip_export, + skip_compile=skip_compile, + skip_generate=False, + use_onnx_subfunctions=_env_bool_any( + False, "KIMI_K25_DYNAMO_WEIGHT_FREE_USE_ONNX_SUBFUNCTIONS", "KIMI_K25_DYNAMO_USE_ONNX_SUBFUNCTIONS" + ), + keep_weights=True, + use_weight_free_export=True, + ) + + +def _update_onnx_weight_spec_metadata(onnx_path: Path, spec: dict): + model = onnx.load(str(onnx_path), load_external_data=False) + value = json.dumps(spec, separators=(",", ":"), sort_keys=True) + for prop in model.metadata_props: + if prop.key == "com.qti.aisw.extdata": + prop.value = value + break + else: + prop = model.metadata_props.add() + prop.key = "com.qti.aisw.extdata" + prop.value = value + tmp_path = onnx_path.with_suffix(onnx_path.suffix + ".tmp") + onnx.save(model, str(tmp_path)) + tmp_path.replace(onnx_path) + + +def _materialize_weight_free_extdata(component_model, component_dir: Path, onnx_name: str): + spec_path = component_dir / "weight_spec.json" + if not spec_path.is_file(): + raise FileNotFoundError(f"Missing weight-free spec: {spec_path}") + + spec = json.loads(spec_path.read_text()) + state_dict = component_model.state_dict() + tensors = {} + missing = [] + for entry in spec["inputs"]: + name = entry["name"] + tensor = state_dict.get(name) + if tensor is None: + missing.append(name) + continue + tensors[name] = tensor.detach().cpu().contiguous() + if missing: + raise RuntimeError( + f"Weight-free extdata tensors are missing from {component_model.__class__.__name__}: {missing}" + ) + + checkpoint_dir = component_dir / "loaded_weight_checkpoint" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = checkpoint_dir / "model.safetensors" + save_file(tensors, str(checkpoint_path)) + + spec["files"] = [{"format": "safetensors", "path": "loaded_weight_checkpoint/model.safetensors"}] + for entry in spec["inputs"]: + entry["location"]["file"] = 0 + entry["location"]["key"] = entry["name"] + + spec_path.write_text(json.dumps(spec, indent=2)) + _update_onnx_weight_spec_metadata(component_dir / onnx_name, spec) + + +def _embed_packed_uint8_weights(component_dir: Path, onnx_name: str): + spec_path = component_dir / "weight_spec.json" + checkpoint_path = component_dir / "loaded_weight_checkpoint" / "model.safetensors" + spec = json.loads(spec_path.read_text()) + packed_names = { + entry["name"] + for entry in spec["inputs"] + if entry["name"].endswith("_qweight") or entry["name"].endswith("_qzeros") + } + if not packed_names: + return + + from safetensors import safe_open + + model = onnx.load(str(component_dir / onnx_name), load_external_data=False) + existing_initializers = {initializer.name for initializer in model.graph.initializer} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle: + for name in sorted(packed_names): + if name not in existing_initializers: + model.graph.initializer.append(numpy_helper.from_array(handle.get_tensor(name).numpy(), name=name)) + + keep_inputs = [value for value in model.graph.input if value.name not in packed_names] + del model.graph.input[:] + model.graph.input.extend(keep_inputs) + + spec["inputs"] = [entry for entry in spec["inputs"] if entry["name"] not in packed_names] + value = json.dumps(spec, separators=(",", ":"), sort_keys=True) + for prop in model.metadata_props: + if prop.key == "com.qti.aisw.extdata": + prop.value = value + break + else: + prop = model.metadata_props.add() + prop.key = "com.qti.aisw.extdata" + prop.value = value + + onnx_path = component_dir / onnx_name + tmp_path = onnx_path.with_suffix(onnx_path.suffix + ".tmp") + onnx.save(model, str(tmp_path)) + tmp_path.replace(onnx_path) + spec_path.write_text(json.dumps(spec, indent=2)) + + +def _prepare_weight_free_artifacts(qeff_model, exported_paths: dict[str, Path]): + vision_onnx_path = Path(exported_paths["vision"]) + lang_onnx_path = Path(exported_paths["lang"]) + _materialize_weight_free_extdata(qeff_model.vision_model.model, vision_onnx_path.parent, vision_onnx_path.name) + _materialize_weight_free_extdata(qeff_model.lang_model.model, lang_onnx_path.parent, lang_onnx_path.name) + _embed_packed_uint8_weights(lang_onnx_path.parent, lang_onnx_path.name) + + +def check_kimi_k25_dynamo_weight_free_hf_vs_qaic(): + os.environ.setdefault("HF_HUB_CACHE", "/home/huggingface_hub") + os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + + validate_dynamo_torch_version() + configure_qaic_tool_path() + if not _has_qaic_runtime_access(): + _skip_test("QAIC generation skipped: no QAIC runtime access.") + + args = _make_args() + if args.skip_compile: + device_ids = _resolve_device_ids(args) + else: + device_ids = _resolve_device_ids() + if device_ids is None and not args.skip_compile: + device_ids = [0] + if device_ids is None: + _skip_test("QAIC generation skipped: no QAIC device has free NSPs.") + args = _make_args(device_ids) + + set_deterministic(args.seed) + hf_model, tokenizer, processor = _load_hf_model(args) + image = load_generation_image(args) + hf_inputs = build_generation_inputs(processor, _QEffHolder(hf_model), image, args) + hf_tokens = _greedy_generate_hf(hf_model, _clone_inputs(hf_inputs), args.generation_len).cpu() + print("HF:", _decode_tokens(tokenizer, hf_tokens), "\n", hf_tokens) + + qeff_model, tokenizer, processor, qaic_config = build_qeff_model(args) + qeff_inputs = build_generation_inputs(processor, qeff_model, image, args) + export_inputs, output_names, dynamic_axes = get_component_export_args(qeff_model, args.prefill_seq_len) + + export_start = perf_counter() + exported_paths = export_components(qeff_model, export_inputs, output_names, dynamic_axes, args) + export_end = perf_counter() + print(f"Weight-free Dynamo ONNX export time: {export_end - export_start:.2f} sec (skip_export={args.skip_export})") + + if not args.skip_compile: + prep_start = perf_counter() + _prepare_weight_free_artifacts(qeff_model, exported_paths) + prep_end = perf_counter() + print(f"Weight-free external-data prep time: {prep_end - prep_start:.2f} sec") + + compile_start = perf_counter() + qpc_paths = compile_components(qeff_model, exported_paths, image, qaic_config, args) + compile_end = perf_counter() + print( + f"Weight-free Dynamo QAIC compile time: {compile_end - compile_start:.2f} sec (skip_compile={args.skip_compile})" + ) + print(f"Weight-free Dynamo ONNX paths: {exported_paths}") + print(f"Weight-free Dynamo QPC paths: {qpc_paths}") + + if not _has_explicit_device_ids(): + activatable_device_id = _find_activatable_device_id( + Path(qpc_paths["vision_qpc_path"]), + Path(qpc_paths["lang_qpc_path"]), + _env_int("KIMI_K25_DYNAMO_WEIGHT_FREE_MAX_DEVICE_ID", 63), + ) + if activatable_device_id is not None: + args.device_ids = [activatable_device_id] + print(f"Using activatable QAIC device id: {activatable_device_id}") + + qaic_output = qeff_model.generate( + inputs=_clone_inputs(qeff_inputs), + device_ids=args.device_ids, + generation_len=args.generation_len, + ) + qaic_tokens = torch.as_tensor(qaic_output.generated_ids[:, : args.generation_len], dtype=hf_tokens.dtype) + print("QAIC:", _decode_tokens(tokenizer, qaic_tokens), "\n", qaic_tokens) + + assert torch.equal(hf_tokens, qaic_tokens), ( + "HF and QAIC tokens do not match for the weight-free Dynamo exported and compiled model: " + f"hf={hf_tokens.tolist()}, qaic={qaic_tokens.tolist()}" + ) + + +def test_kimi_k25_dynamo_weight_free_hf_vs_qaic(): + check_kimi_k25_dynamo_weight_free_hf_vs_qaic() + + +if __name__ == "__main__": + check_kimi_k25_dynamo_weight_free_hf_vs_qaic() diff --git a/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py b/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py index 9cb6d5f273..aca6aa3245 100644 --- a/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py +++ b/tests/transformers/qeff_classes/test_automodel_for_causal_lm.py @@ -171,7 +171,10 @@ def test_causal_lm_hash_creation(config, cb, subfunc, prefill_only, tmp_path): ) model = AutoModelForCausalLM.from_config(config, **model_kwargs) qeff_model = QEFFAutoModelForCausalLM(model, cb) - qeff_model.export(tmp_path, use_onnx_subfunctions=subfunc, prefill_only=prefill_only) + export_kwargs = {"use_onnx_subfunctions": subfunc, "prefill_only": prefill_only} + if prefill_only: + export_kwargs["prefill_seq_len"] = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + qeff_model.export(tmp_path, **export_kwargs) hash_params = {} hash_params["config"] = qeff_model.model.config.to_diff_dict() hash_params["peft_config"] = None diff --git a/tests/utils/load_kimi_utils.py b/tests/utils/load_kimi_utils.py index 0b9dbfcf5d..e6d4ea8705 100644 --- a/tests/utils/load_kimi_utils.py +++ b/tests/utils/load_kimi_utils.py @@ -327,6 +327,7 @@ def load_layer_subset_model( loaded_expert_ids, num_experts_per_tok: int, dtype, + subset_model_path: Path | None = None, ): checkpoint_index = json.loads((model_path / "model.safetensors.index.json").read_text()) weight_map = checkpoint_index["weight_map"] @@ -338,8 +339,17 @@ def load_layer_subset_model( num_experts_per_tok=num_experts_per_tok, ) - with tempfile.TemporaryDirectory() as tmpdir: - temp_model_path = Path(tmpdir) + tempdir_context = None + if subset_model_path is None: + tempdir_context = tempfile.TemporaryDirectory() + temp_model_path = Path(tempdir_context.__enter__()) + else: + temp_model_path = Path(subset_model_path).expanduser().resolve() + if temp_model_path.exists(): + shutil.rmtree(temp_model_path) + temp_model_path.mkdir(parents=True, exist_ok=True) + + try: filtered_weight_map, subset_shards = materialize_subset_checkpoint( model_path=model_path, temp_model_path=temp_model_path, @@ -374,6 +384,9 @@ def load_layer_subset_model( model, loading_info = kimi_cls.from_pretrained(str(temp_model_path), **model_kwargs) finally: kimi_cls.base_model_prefix = original_base_model_prefix + finally: + if tempdir_context is not None: + tempdir_context.__exit__(None, None, None) unexpected_keys = loading_info["unexpected_keys"] missing_keys = loading_info["missing_keys"] @@ -403,6 +416,7 @@ def load_kimi_k25_layer_subset_model( num_experts_per_tok: int = NUM_EXPERTS_PER_TOKEN, dtype=torch.float32, seed: int = 1234, + subset_model_path: Path | None = None, ): set_deterministic(seed) resolved_model_path = Path(model_path) if model_path is not None else resolve_model_path() @@ -418,6 +432,7 @@ def load_kimi_k25_layer_subset_model( loaded_expert_ids=loaded_expert_ids, num_experts_per_tok=num_experts_per_tok, dtype=dtype, + subset_model_path=subset_model_path, ) model.vision_tower.patch_embed.pos_emb.interpolation_mode = "bilinear" return model.eval().to("cpu"), tokenizer, processor