File size: 6,444 Bytes
c449029 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """
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)
|