File size: 2,274 Bytes
47d8ad6 | 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 | """Repro of prior spike's crasher: 2 cross-dependent states + attention.
State A (mem_bank) is READ -> attention -> result WRITTEN to state B (ptr_bank).
Prior result on torch 2.13 + coremltools 9.0: KeyError in optimize_state.py:93.
Now testing under torch 2.7.0 (coremltools' tested ceiling).
"""
import numpy as np
import torch
import torch.nn as nn
import coremltools as ct
C, T, S = 256, 7, 64 # hidden dim, num_maskmem, spatial tokens
class CrossStateModel(nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("mem_bank", torch.zeros(T, S, C)) # state A
self.register_buffer("ptr_bank", torch.zeros(T, C)) # state B
self.attn = nn.MultiheadAttention(C, 4, batch_first=True)
self.enc = nn.Linear(C, C)
def forward(self, x): # x: (1, S, C)
# read BOTH states
mem = self.mem_bank.reshape(1, T * S, C)
ptrs = self.ptr_bank.reshape(1, T, C)
kv = torch.cat([mem, ptrs], dim=1)
y, _ = self.attn(x, kv, kv)
# write state B from a value derived from reading state A (cross-dependent)
new_ptr = y.mean(dim=1) # (1, C)
ptr_updated = torch.cat([self.ptr_bank[1:], new_ptr], dim=0)
self.ptr_bank.copy_(ptr_updated)
# write state A from current input (shift + cat, static shape)
new_mem = self.enc(y) # (1, S, C)
mem_updated = torch.cat([self.mem_bank[1:], new_mem], dim=0)
self.mem_bank.copy_(mem_updated)
return y
def main():
torch.manual_seed(0)
m = CrossStateModel().eval()
m.requires_grad_(False)
x = torch.randn(1, S, C)
with torch.no_grad():
ep = torch.export.export(m, (x,))
ep = ep.run_decompositions({})
print("torch.export OK")
mlmodel = ct.convert(
ep,
minimum_deployment_target=ct.target.iOS18,
compute_units=ct.ComputeUnit.CPU_ONLY,
)
print("coremltools convert OK")
state = mlmodel.make_state()
out1 = mlmodel.predict({"x": x.numpy()}, state=state)
out2 = mlmodel.predict({"x": x.numpy()}, state=state)
k = list(out1)[0]
print("predict OK; outputs differ across calls (state evolving):",
not np.allclose(out1[k], out2[k]))
if __name__ == "__main__":
main()
|