""" 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 [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)