Add fix_octen_dml.py reproduction script

fix_octen_dml.py

@@ -0,0 +1,171 @@
+"""
+Fix Octen-Embedding-0.6B ONNX for DirectML compatibility.
+
+Root cause: dynamo torch.onnx.export generates `val_41 = [-1]` used in Reshape
+shapes for multi-head attention Q/K/V projections. ONNX resolves -1 differently
+per input (16 for Q with 2048-dim output, 8 for K/V with 1024-dim output).
+DirectML's execution provider needs concrete shape values at graph-capture time
+and crashes on symbolic -1 dims.
+
+Fix: Replace val_41=[-1] with three concrete head-count constants:
+  - val_41_q = [16]  (Q projection: num_attention_heads=16)
+  - val_41_k = [8]   (K projection: num_key_value_heads=8)
+  - val_41_v = [8]   (V projection: num_key_value_heads=8)
+
+Create separate Concat nodes for each shape and reconnect the Reshape consumers
+based on the Q/K/V naming pattern:
+  - node_view{4L+0} (Q)  -> val_50_q
+  - node_view{4L+1} (K)  -> val_50_k
+  - node_view{4L+2} (V)  -> val_50_v
+"""
+
+import sys
+import os
+import shutil
+from pathlib import Path
+
+import onnx
+import numpy as np
+from onnx import helper, TensorProto
+
+# octen-embedding-0.6b attention heads
+NUM_Q_HEADS = 16
+NUM_KV_HEADS = 8
+HEAD_DIM = 128
+
+
+def patch_octen_for_directml(input_path: str, output_path: str):
+    print(f"Loading: {input_path}")
+    m = onnx.load(input_path)
+    g = m.graph
+    initials = {i.name: i for i in g.initializer}
+
+    # --- 1. Create new head-count initializers ---
+    def make_int64_initializer(name: str, value: int) -> TensorProto:
+        return helper.make_tensor(name, TensorProto.INT64, [1], [value])
+
+    new_inits = [
+        make_int64_initializer("val_41_q", NUM_Q_HEADS),
+        make_int64_initializer("val_41_k", NUM_KV_HEADS),
+        make_int64_initializer("val_41_v", NUM_KV_HEADS),
+    ]
+
+    # Add them to the graph (before the original val_41 so they're available)
+    val_41_idx = next(i for i, init in enumerate(g.initializer) if init.name == "val_41")
+    for offset, init in enumerate(new_inits):
+        g.initializer.insert(val_41_idx + offset, init)
+
+    # --- 2. Create new Concat nodes for val_50_q, val_50_k, val_50_v ---
+    # Original: val_50 = Concat([val_0 (batch), val_1 (seq), val_41 (-1), val_49 (128)])
+    # New:      val_50_q = Concat([val_0, val_1, val_41_q(16), val_49])
+    #           val_50_k = Concat([val_0, val_1, val_41_k(8),  val_49])
+    #           val_50_v = Concat([val_0, val_1, val_41_v(8),  val_49])
+    #
+    # The original val_50 Concat has attribute axis=0
+    val_50_node = next(n for n in g.node if "val_50" in n.output)
+
+    for suffix, val_name in [("q", "val_41_q"), ("k", "val_41_k"), ("v", "val_41_v")]:
+        new_concat = helper.make_node(
+            "Concat",
+            inputs=[val_50_node.input[0], val_50_node.input[1], val_name, val_50_node.input[3]],
+            outputs=[f"val_50_{suffix}"],
+            name=f"node_val_50_{suffix}",
+            axis=0,
+        )
+        # Insert after the original val_50 Concat
+        val_50_pos = next(i for i, n in enumerate(g.node) if n.name == val_50_node.name)
+        g.node.insert(val_50_pos + 1 + {"q": 0, "k": 1, "v": 2}[suffix], new_concat)
+
+    # --- 3. Reconnect Reshape consumers ---
+    # Per layer L (0..27):
+    #   linear_{7L+0} (Q weight dims [1024,2048]) -> node_view_{4L+0}
+    #   linear_{7L+1} (K weight dims [1024,1024]) -> node_view_{4L+1}
+    #   linear_{7L+2} (V weight dims [1024,1024]) -> node_view_{4L+2}
+    #
+    # So: node_view_{4L+0} uses val_50_q, node_view_{4L+1} uses val_50_k,
+    #     node_view_{4L+2} uses val_50_v
+
+    consumers = [n for n in g.node if "val_50" in n.input]
+    import re
+
+    q_patches = 0
+    k_patches = 0
+    v_patches = 0
+
+    for n in consumers:
+        name = n.name
+        match = re.match(r"node_view_(\d+)$", name)
+        if match:
+            idx = int(match.group(1))
+            # Q: idx % 4 == 0, K: idx % 4 == 1, V: idx % 4 == 2
+            if idx % 4 == 0:
+                replacement = "val_50_q"
+                q_patches += 1
+            elif idx % 4 == 1:
+                replacement = "val_50_k"
+                k_patches += 1
+            elif idx % 4 == 2:
+                replacement = "val_50_v"
+                v_patches += 1
+            else:
+                print(f"  WARNING: unexpected index {idx} for {name}, skipping")
+                continue
+        elif name == "node_view":
+            # node_view (no suffix) is Q for layer 0
+            replacement = "val_50_q"
+            q_patches += 1
+        else:
+            print(f"  WARNING: unexpected consumer {name}, skipping")
+            continue
+
+        # Replace val_50 with the specific variant in the Reshape's inputs
+        new_inputs = [
+            replacement if inp == "val_50" else inp for inp in n.input
+        ]
+        del n.input[:]
+        n.input.extend(new_inputs)
+
+    print(f"  Q reshapes patched: {q_patches}")
+    print(f"  K reshapes patched: {k_patches}")
+    print(f"  V reshapes patched: {v_patches}")
+    print(f"  Total: {q_patches + k_patches + v_patches}")
+
+    # --- 4. Clean up: remove original val_41 and val_50 (optional, but cleaner) ---
+    # We keep them to avoid breaking anything else that might reference them.
+    # val_50 only has the 84 Reshape consumers; val_41 might be used elsewhere.
+
+    # --- 5. Check op checksums / validate ---
+    try:
+        onnx.checker.check_model(m)
+        print("  ONNX validation: PASSED")
+    except Exception as e:
+        print(f"  ONNX validation WARNING: {e}")
+        print("  (non-critical, DirectML may still accept it)")
+
+    # --- 6. Save ---
+    # Handle external data files
+    output_dir = os.path.dirname(output_path) or "."
+    os.makedirs(output_dir, exist_ok=True)
+
+    # If the model has external data, copy the .data file
+    input_dir = os.path.dirname(input_path)
+    data_file = input_path + ".data"
+    if os.path.exists(data_file):
+        output_data = output_path + ".data"
+        print(f"  Copying external data: {data_file} -> {output_data}")
+        shutil.copy2(data_file, output_data)
+
+    onnx.save(m, output_path)
+    size_mb = os.path.getsize(output_path) / (1024 * 1024)
+    print(f"  Saved: {output_path} ({size_mb:.1f} MB)")
+    print("Done!")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print("Usage: python fix_octen_dml.py <input.onnx> [output.onnx]")
+        sys.exit(1)
+
+    input_path = sys.argv[1]
+    output_path = sys.argv[2] if len(sys.argv) > 2 else input_path.replace(".onnx", "_dml.onnx")
+    patch_octen_for_directml(input_path, output_path)