diff --git a/QEfficient/exporter/weight_free/checkpoint_key_resolver.py b/QEfficient/exporter/weight_free/checkpoint_key_resolver.py index 1a064f682f..e4d8c554ac 100644 --- a/QEfficient/exporter/weight_free/checkpoint_key_resolver.py +++ b/QEfficient/exporter/weight_free/checkpoint_key_resolver.py @@ -39,6 +39,8 @@ "original_inv_freq", "embed_positions", "embed_scale", + "position_ids", + "token_type_ids", } diff --git a/QEfficient/transformers/models/bert/modeling_bert.py b/QEfficient/transformers/models/bert/modeling_bert.py index 05a83c2a15..13670bc908 100644 --- a/QEfficient/transformers/models/bert/modeling_bert.py +++ b/QEfficient/transformers/models/bert/modeling_bert.py @@ -16,12 +16,24 @@ Fix: override `_create_attention_masks` to use `_prepare_4d_attention_mask` (standard tensor ops, fully ONNX-traceable) for the encoder (non-decoder) path. + +Separately, `RobertaEmbeddings` / `XLMRobertaEmbeddings` / `NomicBertEmbeddings` +`.forward` build `token_type_ids` from a non-persistent `token_type_ids` buffer +(created on a concrete device) gathered against `position_ids`. Under +FakeTensor/meta-device tracing (e.g. dynamo + use_onnx_subfunctions, +weight-free export) `position_ids` can be a fake tensor on the `meta` device +while the buffer stays on a concrete device, so `torch.gather` fails with +"found at least two devices". The `QEff*Embeddings` classes below override +`forward` to move the buffer to `position_ids.device` before the gather; +otherwise identical to the upstream implementation. """ +import torch from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask from transformers.models.bert.modeling_bert import BertModel -from transformers.models.roberta.modeling_roberta import RobertaModel -from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaModel +from transformers.models.nomic_bert.modeling_nomic_bert import NomicBertEmbeddings +from transformers.models.roberta.modeling_roberta import RobertaEmbeddings, RobertaModel +from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaEmbeddings, XLMRobertaModel class _QEffBertFamilyMixin: @@ -76,3 +88,103 @@ class QEffRobertaModel(_QEffBertFamilyMixin, RobertaModel): class QEffXLMRobertaModel(_QEffBertFamilyMixin, XLMRobertaModel): pass + + +class _QEffRobertaFamilyEmbeddingsMixin: + """ + Shared fixed `forward` for `RobertaEmbeddings` / `XLMRobertaEmbeddings` (identical + upstream implementations): moves the buffered `token_type_ids` to `position_ids.device` + before the `torch.gather` call. See module docstring for the FakeTensor/meta-device + rationale. + """ + + def forward( + self, + input_ids=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + past_key_values_length=0, + ): + if position_ids is None: + if input_ids is not None: + position_ids = self.create_position_ids_from_input_ids( + input_ids, self.padding_idx, past_key_values_length + ) + else: + position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds, self.padding_idx) + + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + batch_size, seq_length = input_shape + + if token_type_ids is None: + if hasattr(self, "token_type_ids"): + buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1) + buffered_token_type_ids = buffered_token_type_ids.to(position_ids.device) + buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids) + token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length) + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + embeddings = inputs_embeds + token_type_embeddings + + position_embeddings = self.position_embeddings(position_ids) + embeddings = embeddings + position_embeddings + + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + +class QEffRobertaEmbeddings(_QEffRobertaFamilyEmbeddingsMixin, RobertaEmbeddings): + pass + + +class QEffXLMRobertaEmbeddings(_QEffRobertaFamilyEmbeddingsMixin, XLMRobertaEmbeddings): + pass + + +class QEffNomicBertEmbeddings(NomicBertEmbeddings): + """ + Fixes the same buffered-`token_type_ids` device mismatch as `QEffRobertaEmbeddings`, + for `NomicBertEmbeddings` (native in Transformers v5.5, generated from + `modular_nomic_bert.py`). Otherwise identical to the upstream implementation. + """ + + def forward( + self, + input_ids=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + ): + embeddings = inputs_embeds + if inputs_embeds is None: + embeddings = self.word_embeddings(input_ids) + + input_shape = embeddings.shape[:-1] + device = embeddings.device + + if token_type_ids is None: + if hasattr(self, "token_type_ids"): + buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1) + buffered_token_type_ids = buffered_token_type_ids.to(position_ids.device) + buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids) + token_type_ids = buffered_token_type_ids.expand(*input_shape) + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = embeddings + token_type_embeddings + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + + return embeddings diff --git a/QEfficient/transformers/models/modeling_auto.py b/QEfficient/transformers/models/modeling_auto.py index fa874b5570..cd2d8ab689 100755 --- a/QEfficient/transformers/models/modeling_auto.py +++ b/QEfficient/transformers/models/modeling_auto.py @@ -173,7 +173,7 @@ def _disable_unsupported_weight_free(kwargs: dict, qeff_auto_class_name: str) -> return logger.warning( - "weight_free=True is only supported for QEFFAutoModelForCausalLM; disabling it for %s.", + "weight_free=True is only supported for QEFFAutoModelForCausalLM, QEffAutoModel; disabling it for %s.", qeff_auto_class_name, ) @@ -394,7 +394,8 @@ class QEFFTransformersBase(QEFFBaseModel): def __init__(self, model: nn.Module, **kwargs) -> None: _configure_proxy_for_model(self, kwargs.pop("enable_proxy", False)) - _disable_unsupported_weight_free(kwargs, self.__class__.__name__) + if self.__class__ is not QEFFAutoModel: + _disable_unsupported_weight_free(kwargs, self.__class__.__name__) if ( hasattr(model, "config") @@ -538,6 +539,7 @@ class QEFFAutoModel(QEFFTransformersBase): _pytorch_transforms = [CustomOpsTransform, AwqToMatmulNbitsTransform, GPTQToMatmulNbitsTransform] # FP16Clip inlines external weights; without Split the saved protobuf exceeds 2GB for large embedders. _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform] + _checkpoint_transforms = [DtypeConversionCheckpointTransform] def __init__(self, model: nn.Module, pooling=None, **kwargs): """ @@ -570,7 +572,7 @@ def __init__(self, model: nn.Module, pooling=None, **kwargs): @classmethod @with_replaced_quantizers - def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **kwargs): + def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, weight_free=False, *args, **kwargs): """ Load a QEfficient transformer model from a pretrained HuggingFace model or local path. @@ -589,6 +591,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k - "avg": Average pooling - Callable: A custom pooling function - None: No pooling applied. Default is None. + weight_free : bool, optional + If True, the model will be loaded in weight-free mode, which avoids materializing checkpoint weights. Default is False. *args : Positional arguments passed directly to `cls._hf_auto_class.from_pretrained`. **kwargs : @@ -602,9 +606,11 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k QEFFAutoModel An instance initialized with the pretrained weights. """ - _disable_unsupported_weight_free(kwargs, cls.__name__) enable_proxy = kwargs.pop("enable_proxy", False) + if weight_free: + validate_dynamo_export_requirements("weight_free=True") + if kwargs.get("attn_implementation", None) not in {None, "eager"}: logger.warning('Updating attn_implementation="eager"') @@ -614,7 +620,13 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False}) _resolve_torch_dtype(kwargs) - model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + if weight_free: + # Weight-free mode: build the model on the meta device so no + # checkpoint weights are ever materialized here. The real weights + # are supplied later at export time via pretrained_model_name_or_path. + model = _build_meta_model(cls._hf_auto_class, pretrained_model_name_or_path, kwargs) + else: + model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) # This is support models that should be classified to in a different auto class but transformers load them via this class kv_offload = kwargs.pop("kv_offload", None) @@ -626,7 +638,13 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k model, kv_offload=kv_offload, **kwargs ) - return cls(model, pretrained_model_name_or_path=pretrained_model_name_or_path, pooling=pooling, **kwargs) + return cls( + model, + pretrained_model_name_or_path=pretrained_model_name_or_path, + pooling=pooling, + weight_free=weight_free, + **kwargs, + ) @property def get_model_config(self) -> dict: @@ -640,7 +658,7 @@ def get_model_config(self) -> dict: """ return self.model.config.__dict__ - def export(self, export_dir: str | None = None, **kwargs) -> str: + def export(self, export_dir: str | None = None, dynamo: bool = False, **kwargs) -> str: """ Export the model to ONNX format using ``torch.onnx.export``. @@ -663,6 +681,11 @@ def export(self, export_dir: str | None = None, **kwargs) -> str: bs = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE seq_len = constants.ONNX_EXPORT_EXAMPLE_SEQ_LEN + dynamo = kwargs.get("dynamo", False) or self._weight_free + if dynamo: + # torch.export requires example inputs to satisfy dynamic_shapes min=2; gpt_oss non-CB keeps bs=1. + bs = max(2, bs) + example_inputs = { "input_ids": torch.zeros((bs, seq_len), dtype=torch.int64), "attention_mask": torch.ones((bs, seq_len), dtype=torch.int64), @@ -677,6 +700,7 @@ def export(self, export_dir: str | None = None, **kwargs) -> str: output_names=output_names, dynamic_axes=dynamic_axes, export_dir=export_dir, + dynamo=dynamo, use_onnx_subfunctions=kwargs.get("use_onnx_subfunctions", False), ) diff --git a/QEfficient/transformers/models/pytorch_transforms.py b/QEfficient/transformers/models/pytorch_transforms.py index abac3c9e1a..6457964d7e 100644 --- a/QEfficient/transformers/models/pytorch_transforms.py +++ b/QEfficient/transformers/models/pytorch_transforms.py @@ -216,6 +216,7 @@ except ImportError: from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLRMSNorm as Qwen2_5RMSNorm from transformers.models.bert.modeling_bert import BertModel +from transformers.models.nomic_bert.modeling_nomic_bert import NomicBertEmbeddings from transformers.models.qwen3.modeling_qwen3 import ( Qwen3Attention, Qwen3DecoderLayer, @@ -288,7 +289,7 @@ Qwen3VLMoeVisionAttention, Qwen3VLMoeVisionModel, ) -from transformers.models.roberta.modeling_roberta import RobertaModel +from transformers.models.roberta.modeling_roberta import RobertaEmbeddings, RobertaModel from transformers.models.starcoder2.modeling_starcoder2 import ( Starcoder2Attention, Starcoder2DecoderLayer, @@ -313,7 +314,7 @@ WhisperModel, WhisperPositionalEmbedding, ) -from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaModel +from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaEmbeddings, XLMRobertaModel from QEfficient.base.pytorch_transforms import ( ExternalModuleMapperTransform, @@ -325,7 +326,10 @@ from QEfficient.transformers.embeddings.embedding_utils import POOLING_MAP, PooledModel, validate_user_pooling_function from QEfficient.transformers.models.bert.modeling_bert import ( QEffBertModel, + QEffNomicBertEmbeddings, + QEffRobertaEmbeddings, QEffRobertaModel, + QEffXLMRobertaEmbeddings, QEffXLMRobertaModel, ) from QEfficient.transformers.models.codegen.modeling_codegen import ( @@ -724,6 +728,11 @@ class CustomOpsTransform(ModuleMappingTransform): BertModel: QEffBertModel, RobertaModel: QEffRobertaModel, XLMRobertaModel: QEffXLMRobertaModel, + # *Embeddings: fix a FakeTensor/meta-device mismatch in the buffered + # token_type_ids gather (see QEff*Embeddings docstrings in modeling_bert.py). + RobertaEmbeddings: QEffRobertaEmbeddings, + XLMRobertaEmbeddings: QEffXLMRobertaEmbeddings, + NomicBertEmbeddings: QEffNomicBertEmbeddings, Qwen3_5RMSNorm: GemmaCustomRMSNormAIC, Qwen3_5MoeRMSNorm: GemmaCustomRMSNormAIC, Qwen3_5RMSNormGated: QEffQwen3_5GatedDeltaNetCustomRMSNormAIC, diff --git a/examples/embeddings/text_embeddings.py b/examples/embeddings/text_embeddings.py index e69e6f1afb..d3ace2cb44 100644 --- a/examples/embeddings/text_embeddings.py +++ b/examples/embeddings/text_embeddings.py @@ -48,6 +48,16 @@ def main(): default="32,64", help="Sequence length(s) - single int (e.g., '32') or comma-separated list (e.g., '32,64')", ) + parser.add_argument( + "--weight-free", + action="store_true", + help="Build the model on meta tensors and load weights at compile time", + ) + parser.add_argument( + "--use-onnx-subfunctions", + action="store_true", + help="Use subfunctions while exporting", + ) args = parser.parse_args() # Parse seq_len argument @@ -67,15 +77,20 @@ def main(): # You can specify the pooling strategy either as a string (e.g., "max") or by passing a custom pooling function. # If no pooling is specified, the model will return its default output (typically token embeddings). if args.pooling == "max": - qeff_model = AutoModel.from_pretrained(args.model_name, pooling=max_pooling) + qeff_model = AutoModel.from_pretrained(args.model_name, pooling=max_pooling, weight_free=args.weight_free) elif args.pooling == "mean": - qeff_model = AutoModel.from_pretrained(args.model_name, pooling="mean") + qeff_model = AutoModel.from_pretrained(args.model_name, pooling="mean", weight_free=args.weight_free) else: - qeff_model = AutoModel.from_pretrained(args.model_name) + qeff_model = AutoModel.from_pretrained(args.model_name, weight_free=args.weight_free) # Compile the model # seq_len can be a list of seq_len or single int - qeff_model.compile(num_cores=args.num_cores, seq_len=seq_len) + qeff_model.compile( + num_cores=args.num_cores, + seq_len=seq_len, + dynamo=args.weight_free, + use_onnx_subfunctions=args.use_onnx_subfunctions, + ) # Tokenize sentences encoded_input = tokenizer(args.sentences, return_tensors="pt") diff --git a/tests/unit_test/models/test_modeling_auto_cpu.py b/tests/unit_test/models/test_modeling_auto_cpu.py index 2d6061aa63..ec5b12aa5f 100644 --- a/tests/unit_test/models/test_modeling_auto_cpu.py +++ b/tests/unit_test/models/test_modeling_auto_cpu.py @@ -230,7 +230,6 @@ def test_init_passes_without_quantization_config(self): @pytest.mark.parametrize( ("wrapper_cls", "model_factory"), [ - pytest.param(QEFFAutoModel, make_tiny_bert, id="automodel"), pytest.param(QEFFAutoModelForSequenceClassification, make_tiny_bert_seq_cls, id="sequence-classification"), pytest.param(QEFFAutoModelForSpeechSeq2Seq, make_tiny_whisper, id="speech-seq2seq"), pytest.param(QEFFAutoModelForCTC, make_tiny_wav2vec2, id="ctc"), @@ -257,14 +256,13 @@ def fake_from_pretrained(_model_id, *args, **kwargs): @pytest.mark.parametrize( ("wrapper_cls", "model_factory"), [ - pytest.param(QEFFAutoModel, make_tiny_bert, id="automodel"), pytest.param(QEFFAutoModelForSequenceClassification, make_tiny_bert_seq_cls, id="sequence-classification"), pytest.param(QEFFAutoModelForSpeechSeq2Seq, make_tiny_whisper, id="speech-seq2seq"), pytest.param(QEFFAutoModelForCTC, make_tiny_wav2vec2, id="ctc"), ], ) def test_direct_init_disables_unsupported_weight_free(self, wrapper_cls, model_factory, caplog): - """Non-CausalLM direct construction must not enable the weight-free export path.""" + """Direct construction disables weight-free mode for unsupported wrappers.""" caplog.set_level(logging.WARNING, logger="QEfficient") qeff_model = wrapper_cls(model_factory()[0], weight_free=True) @@ -282,6 +280,32 @@ def test_causal_lm_direct_init_preserves_weight_free(self, caplog): assert qeff_model._weight_free is True assert UNSUPPORTED_WEIGHT_FREE_WARNING not in caplog.text + def test_embedding_direct_init_preserves_weight_free(self): + """Embedding models retain weight-free mode for Dynamo export.""" + qeff_model = QEFFAutoModel(make_tiny_bert()[0], weight_free=True) + + assert qeff_model._weight_free is True + + def test_embedding_export_defaults_to_dynamo_in_weight_free_mode(self, monkeypatch): + """Weight-free embedding export supplies Dynamo-compatible dynamic inputs.""" + qeff_model = QEFFAutoModel(make_tiny_bert()[0], weight_free=True) + captured = {} + + def fake_export(example_inputs, **kwargs): + captured["example_inputs"] = example_inputs + captured.update(kwargs) + return "embedding.onnx" + + monkeypatch.setattr(qeff_model, "_export", fake_export) + qeff_model.export(dynamo=False) + + assert captured["dynamo"] is True + assert captured["dynamic_axes"] == { + "input_ids": {0: "batch_size", 1: "seq_len"}, + "attention_mask": {0: "batch_size", 1: "seq_len"}, + } + assert captured["example_inputs"]["input_ids"].shape[0] >= 2 + # --------------------------------------------------------------------------- # Stage 2: QEFFAutoModelForCausalLM — logic methods diff --git a/tests/weight_free/_helpers.py b/tests/weight_free/_helpers.py index 64858d041c..bd1aed9888 100644 --- a/tests/weight_free/_helpers.py +++ b/tests/weight_free/_helpers.py @@ -16,14 +16,13 @@ import copy from pathlib import Path -from typing import Dict, Tuple import numpy as np import onnx import onnxruntime import pytest import torch -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer from QEfficient.exporter.weight_free import load_weight_free_ort_inputs from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM @@ -31,7 +30,7 @@ # --------------------------------------------------------------------------- # Worker-level model cache # --------------------------------------------------------------------------- -_HF_MODEL_CACHE: Dict[str, Tuple[AutoModelForCausalLM, AutoTokenizer]] = {} +_HF_MODEL_CACHE: dict[str, tuple[AutoModelForCausalLM, AutoTokenizer]] = {} # --------------------------------------------------------------------------- # Model registry — same tiny-random models as tests/dynamo/ @@ -61,6 +60,26 @@ "starcoder2": "hf-tiny-v2/tiny-random-Starcoder2ForCausalLM", } +WEIGHT_FREE_EMBEDDING_MODELS = { + "bge_reranker": ("BAAI/bge-reranker-v2-m3", "mean"), + "bge_small": ("BAAI/bge-small-en-v1.5", "mean"), + "granite_embedding_30m": ("ibm-granite/granite-embedding-30m-english", "mean"), + "granite_multilingual_107m": ("ibm-granite/granite-embedding-107m-multilingual", "mean"), + "jina": ("jinaai/jina-embeddings-v2-base-code", "mean"), + "mpnet": ("sentence-transformers/multi-qa-mpnet-base-cos-v1", "mean"), + "nomic": ("nomic-ai/nomic-embed-text-v1.5", "mean"), +} + +_WEIGHT_FREE_EMBEDDING_XFAILS = {"jina", "nomic"} + + +def _embedding_test_param(model_type, model_id, pooling): + marks = [pytest.mark.xdist_group(name=f"qaic-runtime-{model_type}")] + if model_type in _WEIGHT_FREE_EMBEDDING_XFAILS: + marks.append(pytest.mark.xfail(reason="Weight-free is not enabled for this embedding model yet", strict=False)) + return pytest.param(model_type, model_id, pooling, marks=marks, id=model_type) + + WEIGHT_FREE_QAIC_MODEL_PARAMS = [ pytest.param( model_type, @@ -71,6 +90,11 @@ for model_type, model_id in sorted(WEIGHT_FREE_CAUSAL_LM_MODEL_IDS.items()) ] +WEIGHT_FREE_EMBEDDING_MODEL_PARAMS = [ + _embedding_test_param(model_type, model_id, pooling) + for model_type, (model_id, pooling) in sorted(WEIGHT_FREE_EMBEDDING_MODELS.items()) +] + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -107,6 +131,21 @@ def load_hf_model(model_id: str) -> AutoModelForCausalLM: return copy.deepcopy(model) +def load_hf_embedding_model(model_id: str) -> AutoModel: + """Load a real encoder model for HF reference inference.""" + if model_id not in _HF_MODEL_CACHE: + model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + **MODEL_KWARGS, + ) + model.eval() + tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + _HF_MODEL_CACHE[model_id] = (model, tokenizer) + model, _ = _HF_MODEL_CACHE[model_id] + return copy.deepcopy(model) + + def load_tokenizer(model_id: str) -> AutoTokenizer: if model_id not in _HF_MODEL_CACHE: load_hf_model(model_id) diff --git a/tests/weight_free/test_export.py b/tests/weight_free/test_export.py index 92bb475895..b70c682085 100644 --- a/tests/weight_free/test_export.py +++ b/tests/weight_free/test_export.py @@ -19,11 +19,16 @@ from __future__ import annotations +import numpy as np +import onnx +import onnxruntime import pytest -from transformers import AutoConfig +import torch +from transformers import AutoConfig, AutoTokenizer -from QEfficient.exporter.weight_free import resolve_weight_spec_path -from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM +from QEfficient.exporter.weight_free import load_weight_free_ort_inputs, resolve_weight_spec_path +from QEfficient.transformers.embeddings.embedding_utils import POOLING_MAP +from QEfficient.transformers.models.modeling_auto import QEFFAutoModel, QEFFAutoModelForCausalLM from QEfficient.utils import get_num_layers_from_config from QEfficient.utils.run_utils import ApiRunner @@ -32,12 +37,14 @@ CTX_LEN, PROMPT_LEN, WEIGHT_FREE_CAUSAL_LM_MODEL_IDS, + WEIGHT_FREE_EMBEDDING_MODEL_PARAMS, assert_has_subfunctions, assert_no_int64_kv_cache_inputs, assert_retained_state_outputs, assert_subfunction_names_match_decoder_class, assert_unique_graph_input_names, exported_onnx_path, + load_hf_embedding_model, load_hf_model, load_tokenizer, run_weight_free_ort, @@ -45,6 +52,51 @@ ) +@pytest.mark.weight_free +@pytest.mark.weight_free_export +@pytest.mark.parametrize("model_type,model_id,pooling", WEIGHT_FREE_EMBEDDING_MODEL_PARAMS) +def test_weight_free_embedding_export_ort_parity(model_type, model_id, pooling, tmp_export_dir): + """Export an embedding model without weights and validate HF-to-ORT parity.""" + try: + model_hf = load_hf_embedding_model(model_id) + tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + except (ImportError, OSError, RuntimeError, ValueError) as exc: + if model_type in {"jina", "nomic"}: + raise + skip_on_model_fetch_error(exc, model_id) + + inputs = tokenizer("hello world", return_tensors="pt") + with torch.no_grad(): + hf_hidden_states = model_hf(**inputs).last_hidden_state + hf_output = POOLING_MAP[pooling](hf_hidden_states, inputs["attention_mask"]) + + qeff_model = QEFFAutoModel.from_pretrained(model_id, pooling=pooling, weight_free=True) + onnx_path = exported_onnx_path(qeff_model.export(tmp_export_dir, dynamo=True, offload_pt_weights=False)) + assert onnx_path.is_file(), f"Expected ONNX file at {onnx_path}" + + weight_spec_path = resolve_weight_spec_path(onnx_path) + assert weight_spec_path.is_file(), f"Expected weight spec at {weight_spec_path}" + + onnx_model = onnx.load(str(onnx_path), load_external_data=False) + graph_input_names = [value.name for value in onnx_model.graph.input] + assert len(graph_input_names) == len(set(graph_input_names)) + assert {"input_ids", "attention_mask"}.issubset(graph_input_names) + + session = onnxruntime.InferenceSession(str(onnx_path)) + session_input_names = {value.name for value in session.get_inputs()} + ort_inputs = {name: value.detach().cpu().numpy() for name, value in inputs.items() if name in session_input_names} + ort_inputs = load_weight_free_ort_inputs(weight_spec_path, ort_inputs) + ort_output = session.run(None, ort_inputs)[0] + + np.testing.assert_allclose( + hf_output.detach().cpu().numpy(), + ort_output, + rtol=1e-4, + atol=1e-4, + err_msg=f"HF PT vs weight-free ORT parity failed for {model_type}", + ) + + @pytest.mark.weight_free @pytest.mark.weight_free_export @pytest.mark.parametrize( diff --git a/tests/weight_free/test_on_qaic.py b/tests/weight_free/test_on_qaic.py index 1134a7a31a..4e216f60c7 100644 --- a/tests/weight_free/test_on_qaic.py +++ b/tests/weight_free/test_on_qaic.py @@ -24,15 +24,17 @@ import numpy as np import pytest import torch -from transformers import AutoConfig +from transformers import AutoConfig, AutoModel, AutoTokenizer -from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM +from QEfficient.transformers.embeddings.embedding_utils import POOLING_MAP +from QEfficient.transformers.models.modeling_auto import QEFFAutoModel, QEFFAutoModelForCausalLM from QEfficient.utils import get_num_layers_from_config from ._helpers import ( BATCH_SIZE, CTX_LEN, PROMPT_LEN, + WEIGHT_FREE_EMBEDDING_MODEL_PARAMS, WEIGHT_FREE_QAIC_MODEL_PARAMS, exported_onnx_path, load_hf_model, @@ -41,6 +43,57 @@ ) +@pytest.mark.weight_free +@pytest.mark.on_qaic +@pytest.mark.embedding_audio_model +@pytest.mark.parametrize("model_type,model_id,pooling", WEIGHT_FREE_EMBEDDING_MODEL_PARAMS) +def test_weight_free_embedding_hw_hf_parity(model_type, model_id, pooling, tmp_export_dir): + """HF embedding output matches weight-free embedding output on QAIC.""" + try: + tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + model_hf = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="eager", + low_cpu_mem_usage=False, + torch_dtype=torch.float32, + ) + model_hf.eval() + except (ImportError, OSError, RuntimeError, ValueError) as exc: + if model_type in {"jina", "nomic"}: + raise + skip_on_model_fetch_error(exc, model_id) + + inputs = tokenizer("hello world", return_tensors="pt") + with torch.no_grad(): + hf_hidden_states = model_hf(**inputs).last_hidden_state + hf_output = POOLING_MAP[pooling](hf_hidden_states, inputs["attention_mask"]).cpu().numpy() + + qeff_model = QEFFAutoModel.from_pretrained(model_id, pooling=pooling, weight_free=True) + onnx_path = exported_onnx_path( + qeff_model.export( + tmp_export_dir / f"{model_type}_export", + dynamo=True, + offload_pt_weights=False, + ) + ) + qeff_model.compile( + onnx_path=str(onnx_path), + compile_dir=str(tmp_export_dir / f"{model_type}_compile"), + seq_len=32, + batch_size=1, + num_cores=16, + ) + qaic_output = qeff_model.generate(inputs=inputs) + + assert qaic_output is not None, f"Weight-free QAIC generate returned None for {model_type}" + qaic_embeddings = qaic_output["output"] + if qaic_embeddings.ndim == 3: + qaic_embeddings = qaic_embeddings[:, 0, :] + mad = np.mean(np.abs(hf_output - qaic_embeddings)) + assert mad <= 1e-2, f"HF PT vs weight-free QAIC parity failed for {model_type}: MAD={mad}" + + @pytest.mark.weight_free @pytest.mark.on_qaic @pytest.mark.llm_model