[ML] Add quantized model ops to pytorch_inference allowlist

Add aten::mul_ and quantized::linear_dynamic to the allowed operations
list, fixing validation failures for dynamically quantized models such
as ELSER v2 when imported via Eland with torch.quantization.quantize_dynamic.

Also update the model extraction tooling to support a "quantize" flag in
reference_models.json so that quantized variants are traced with dynamic
quantization applied before graph extraction, mirroring the Eland import
pipeline.

Made-with: Cursor

Author: edsavage

bin/pytorch_inference/CSupportedOperations.cc

@@ -39,6 +39,8 @@ const CSupportedOperations::TStringViewSet CSupportedOperations::FORBIDDEN_OPERA
 // elastic/test-elser-v2.
 // Additional ops from Elasticsearch integration test models
 // (PyTorchModelIT, TextExpansionQueryIT, TextEmbeddingQueryIT).
+// Quantized operations from dynamically quantized variants of the above
+// models (torch.quantization.quantize_dynamic on nn.Linear layers).
 const CSupportedOperations::TStringViewSet CSupportedOperations::ALLOWED_OPERATIONS = {
     // aten operations (core tensor computations)
     "aten::Int"sv,
@@ -79,6 +81,7 @@ const CSupportedOperations::TStringViewSet CSupportedOperations::ALLOWED_OPERATI
     "aten::mean"sv,
     "aten::min"sv,
     "aten::mul"sv,
+    "aten::mul_"sv,
     "aten::ne"sv,
     "aten::neg"sv,
     "aten::new_ones"sv,
@@ -124,6 +127,8 @@ const CSupportedOperations::TStringViewSet CSupportedOperations::ALLOWED_OPERATI
     "prim::dtype"sv,
     "prim::max"sv,
     "prim::min"sv,
+    // quantized operations (dynamically quantized models, e.g. ELSER v2)
+    "quantized::linear_dynamic"sv,
 };
 }
 }

bin/pytorch_inference/unittest/testfiles/reference_model_ops.json

@@ -267,6 +267,7 @@
     },
     "elastic-eis-elser-v2": {
       "model_id": "elastic/eis-elser-v2",
+      "quantized": false,
       "ops": [
         "aten::Int",
         "aten::ScalarImplicit",
@@ -303,6 +304,7 @@
     },
     "elastic-elser-v2": {
       "model_id": "elastic/elser-v2",
+      "quantized": false,
       "ops": [
         "aten::Int",
         "aten::ScalarImplicit",
@@ -337,6 +339,44 @@
         "prim::NumToTensor"
       ]
     },
+    "elastic-elser-v2-quantized": {
+      "model_id": "elastic/elser-v2",
+      "quantized": true,
+      "ops": [
+        "aten::Int",
+        "aten::ScalarImplicit",
+        "aten::__and__",
+        "aten::add",
+        "aten::arange",
+        "aten::contiguous",
+        "aten::dropout",
+        "aten::embedding",
+        "aten::expand",
+        "aten::gather",
+        "aten::ge",
+        "aten::gelu",
+        "aten::index",
+        "aten::layer_norm",
+        "aten::mul_",
+        "aten::new_ones",
+        "aten::reshape",
+        "aten::scaled_dot_product_attention",
+        "aten::select",
+        "aten::size",
+        "aten::slice",
+        "aten::tanh",
+        "aten::to",
+        "aten::transpose",
+        "aten::unsqueeze",
+        "aten::view",
+        "prim::Constant",
+        "prim::DictConstruct",
+        "prim::GetAttr",
+        "prim::ListConstruct",
+        "prim::NumToTensor",
+        "quantized::linear_dynamic"
+      ]
+    },
     "elastic-hugging-face-elser": {
       "model_id": "elastic/hugging-face-elser",
       "ops": [

dev-tools/extract_model_ops/extract_model_ops.py

@@ -41,19 +41,40 @@
 DEFAULT_CONFIG = SCRIPT_DIR / "reference_models.json"
 
 
-def load_reference_models(config_path: Path) -> dict[str, str]:
-    """Load the architecture-to-model mapping from a JSON config file."""
-    with open(config_path) as f:
-        return json.load(f)
+def load_reference_models(config_path: Path) -> dict[str, dict]:
+    """Load the architecture-to-model mapping from a JSON config file.
 
+    Each entry is either a plain model name string or a dict with
+    ``model_id`` and optional ``quantize`` flag.  All entries are
+    normalised to ``{"model_id": ..., "quantize": ...}`` dicts.
+    Keys starting with ``_comment`` are ignored.
+    """
+    with open(config_path) as f:
+        raw = json.load(f)
 
-def extract_ops_for_model(model_name: str) -> set[str] | None:
+    models = {}
+    for key, value in raw.items():
+        if key.startswith("_comment"):
+            continue
+        if isinstance(value, str):
+            models[key] = {"model_id": value, "quantize": False}
+        else:
+            models[key] = {
+                "model_id": value["model_id"],
+                "quantize": value.get("quantize", False),
+            }
+    return models
+
+
+def extract_ops_for_model(model_name: str,
+                          quantize: bool = False) -> set[str] | None:
     """Trace a HuggingFace model and return its TorchScript op set.
 
     Returns None if the model could not be loaded or traced.
     """
-    print(f"  Loading {model_name}...", file=sys.stderr)
-    traced = load_and_trace_hf_model(model_name)
+    label = f"{model_name} (quantized)" if quantize else model_name
+    print(f"  Loading {label}...", file=sys.stderr)
+    traced = load_and_trace_hf_model(model_name, quantize=quantize)
     if traced is None:
         return None
     return collect_inlined_ops(traced)
@@ -90,8 +111,9 @@ def main():
           file=sys.stderr)
 
     failed = []
-    for arch, model_name in reference_models.items():
-        ops = extract_ops_for_model(model_name)
+    for arch, spec in reference_models.items():
+        ops = extract_ops_for_model(spec["model_id"],
+                                    quantize=spec["quantize"])
         if ops is None:
             failed.append(arch)
             print(f"  {arch}: FAILED", file=sys.stderr)
@@ -109,7 +131,8 @@ def main():
             "pytorch_version": torch.__version__,
             "models": {
                 arch: {
-                    "model_id": reference_models[arch],
+                    "model_id": reference_models[arch]["model_id"],
+                    "quantized": reference_models[arch]["quantize"],
                     "ops": sorted(ops),
                 }
                 for arch, ops in sorted(per_model_ops.items())
@@ -125,7 +148,11 @@ def main():
 
     if args.per_model:
         for arch, ops in sorted(per_model_ops.items()):
-            print(f"\n=== {arch} ({reference_models[arch]}) ===")
+            spec = reference_models[arch]
+            label = spec["model_id"]
+            if spec["quantize"]:
+                label += " (quantized)"
+            print(f"\n=== {arch} ({label}) ===")
             for op in sorted(ops):
                 print(f"  {op}")
 

dev-tools/extract_model_ops/reference_models.json

@@ -16,5 +16,10 @@
     "elastic-hugging-face-elser": "elastic/hugging-face-elser",
     "elastic-multilingual-e5-small-optimized": "elastic/multilingual-e5-small-optimized",
     "elastic-splade-v3": "elastic/splade-v3",
-    "elastic-test-elser-v2": "elastic/test-elser-v2"
+    "elastic-test-elser-v2": "elastic/test-elser-v2",
+
+    "_comment:quantized": "Quantized variants: Eland applies torch.quantization.quantize_dynamic on nn.Linear layers when importing models. These produce quantized::* ops not present in the standard traced graphs above.",
+    "elastic-elser-v2-quantized": {"model_id": "elastic/elser-v2", "quantize": true},
+    "elastic-eis-elser-v2-quantized": {"model_id": "elastic/eis-elser-v2", "quantize": true},
+    "elastic-test-elser-v2-quantized": {"model_id": "elastic/test-elser-v2", "quantize": true}
 }

dev-tools/extract_model_ops/torchscript_utils.py

@@ -35,9 +35,13 @@ def collect_inlined_ops(module) -> set[str]:
     return collect_graph_ops(graph)
 
 
-def load_and_trace_hf_model(model_name: str):
+def load_and_trace_hf_model(model_name: str, quantize: bool = False):
     """Load a HuggingFace model, tokenize sample input, and trace to TorchScript.
 
+    When *quantize* is True the model is dynamically quantized (nn.Linear
+    layers converted to quantized::linear_dynamic) before tracing.  This
+    mirrors what Eland does when importing models for Elasticsearch.
+
     Returns the traced module, or None if the model could not be loaded or traced.
     """
     token = os.environ.get("HF_TOKEN")
@@ -53,6 +57,16 @@ def load_and_trace_hf_model(model_name: str):
         print(f"    LOAD ERROR: {exc}", file=sys.stderr)
         return None
 
+    if quantize:
+        try:
+            model = torch.quantization.quantize_dynamic(
+                model, {torch.nn.Linear}, dtype=torch.qint8)
+            print("    Applied dynamic quantization (nn.Linear -> qint8)",
+                  file=sys.stderr)
+        except Exception as exc:
+            print(f"    QUANTIZE ERROR: {exc}", file=sys.stderr)
+            return None
+
     inputs = tokenizer(
         "This is a sample input for graph extraction.",
         return_tensors="pt", padding="max_length",

dev-tools/extract_model_ops/validate_allowlist.py

@@ -103,10 +103,12 @@ def check_ops(ops: set[str],
 def validate_model(model_name: str,
                    allowed: set[str],
                    forbidden: set[str],
-                   verbose: bool) -> bool:
+                   verbose: bool,
+                   quantize: bool = False) -> bool:
     """Validate one HuggingFace model. Returns True if all ops pass."""
-    print(f"  {model_name}...", file=sys.stderr)
-    traced = load_and_trace_hf_model(model_name)
+    label = f"{model_name} (quantized)" if quantize else model_name
+    print(f"  {label}...", file=sys.stderr)
+    traced = load_and_trace_hf_model(model_name, quantize=quantize)
     if traced is None:
         print(f"    FAILED (could not load/trace)", file=sys.stderr)
         return False
@@ -152,13 +154,27 @@ def main():
     results: dict[str, bool] = {}
 
     with open(args.config) as f:
-        models = json.load(f)
+        raw_models = json.load(f)
+
+    models = {}
+    for key, value in raw_models.items():
+        if key.startswith("_comment"):
+            continue
+        if isinstance(value, str):
+            models[key] = {"model_id": value, "quantize": False}
+        else:
+            models[key] = {
+                "model_id": value["model_id"],
+                "quantize": value.get("quantize", False),
+            }
+
     print(f"Validating {len(models)} HuggingFace models from "
           f"{args.config.name}...", file=sys.stderr)
 
-    for arch, model_id in models.items():
+    for arch, spec in models.items():
         results[arch] = validate_model(
-            model_id, allowed, forbidden, args.verbose)
+            spec["model_id"], allowed, forbidden, args.verbose,
+            quantize=spec["quantize"])
 
     if args.pt_dir and args.pt_dir.is_dir():
         pt_files = sorted(args.pt_dir.glob("*.pt"))
@@ -178,7 +194,11 @@ def main():
         if key.startswith("pt:"):
             print(f"  {key}: {status}", file=sys.stderr)
         else:
-            print(f"  {key} ({models[key]}): {status}", file=sys.stderr)
+            spec = models[key]
+            label = spec["model_id"]
+            if spec["quantize"]:
+                label += " (quantized)"
+            print(f"  {key} ({label}): {status}", file=sys.stderr)
 
     print("=" * 60, file=sys.stderr)
     if all_pass:

dev-tools/extract_model_ops/validation_models.json

@@ -19,6 +19,10 @@
     "elastic-splade-v3": "elastic/splade-v3",
     "elastic-test-elser-v2": "elastic/test-elser-v2",
 
+    "elastic-elser-v2-quantized": {"model_id": "elastic/elser-v2", "quantize": true},
+    "elastic-eis-elser-v2-quantized": {"model_id": "elastic/eis-elser-v2", "quantize": true},
+    "elastic-test-elser-v2-quantized": {"model_id": "elastic/test-elser-v2", "quantize": true},
+
     "ner-dslim-bert-base": "dslim/bert-base-NER",
     "sentiment-distilbert-sst2": "distilbert-base-uncased-finetuned-sst-2-english",