File size: 4,208 Bytes
9c1d5d9 | 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 | import torch
from torch.export import Dim
from transformers import MoonshineForConditionalGeneration
from transformers.cache_utils import DynamicCache, EncoderDecoderCache
import os
MODEL = os.environ.get('MS_MODEL', 'moonshine-ai/moonshine-tiny-uk')
OUT = os.environ.get('MS_OUT', '/work/out')
m = MoonshineForConditionalGeneration.from_pretrained(MODEL, attn_implementation='eager').eval()
cfg = m.config
L = cfg.decoder_num_hidden_layers
print('layers:', L)
def flatten(cache):
"""Cache object -> flat tensor list, self layers first, then cross."""
out = []
for sub in (cache.self_attention_cache, cache.cross_attention_cache):
for layer in sub.layers:
out += [layer.keys, layer.values]
return out
def build(self_kv, cross_kv):
"""Flat tensors -> EncoderDecoderCache, with cross marked filled so the model
reuses it instead of projecting the encoder output again on every step."""
sa, ca = DynamicCache(), DynamicCache()
for i in range(L):
sa.update(self_kv[2*i], self_kv[2*i+1], i)
ca.update(cross_kv[2*i], cross_kv[2*i+1], i)
cache = EncoderDecoderCache(sa, ca)
cache.is_updated = {i: True for i in range(L)}
return cache
class DecoderInit(torch.nn.Module):
def __init__(self, m): super().__init__(); self.m = m
def forward(self, input_ids, encoder_hidden_states):
out = self.m.model.decoder(input_ids=input_ids,
encoder_hidden_states=encoder_hidden_states,
use_cache=True)
return (self.m.proj_out(out.last_hidden_state), *flatten(out.past_key_values))
class DecoderStep(torch.nn.Module):
def __init__(self, m): super().__init__(); self.m = m
def forward(self, input_ids, cache_position, encoder_hidden_states, *flat):
self_kv, cross_kv = list(flat[:2*L]), list(flat[2*L:])
cache = build(self_kv, cross_kv)
# The cross cache alone is not enough: without the encoder states the decoder
# skips cross-attention entirely and quietly answers from nothing. Passed here and
# marked already-updated, so it is read rather than recomputed.
out = self.m.model.decoder(input_ids=input_ids,
past_key_values=cache,
cache_position=cache_position,
encoder_hidden_states=encoder_hidden_states,
use_cache=True)
new = flatten(out.past_key_values)[:2*L]
return (self.m.proj_out(out.last_hidden_state), *new)
hs = torch.zeros(1, 123, 288)
start = torch.tensor([[cfg.decoder_start_token_id]])
init = DecoderInit(m).eval()
with torch.no_grad():
init_out = init(start, hs)
print('init outputs:', len(init_out), '| logits', tuple(init_out[0].shape),
'| self k', tuple(init_out[1].shape), '| cross k', tuple(init_out[1+2*L].shape))
self_names = [f'{p}_self_{i}' for i in range(L) for p in ('key','value')]
cross_names = [f'{p}_cross_{i}' for i in range(L) for p in ('key','value')]
frames = Dim('frames', min=4, max=2000)
past = Dim('past', min=1, max=448)
torch.onnx.export(
init, (start, hs), f'{OUT}/decoder_init.onnx', dynamo=True,
input_names=['input_ids', 'encoder_hidden_states'],
output_names=['logits'] + [f'present_{n}' for n in self_names + cross_names],
dynamic_shapes={'input_ids': {}, 'encoder_hidden_states': {1: frames}},
)
print('decoder_init exported')
flat = list(init_out[1:])
step = DecoderStep(m).eval()
pos = torch.tensor([1], dtype=torch.long)
nxt = torch.tensor([[100]])
with torch.no_grad():
step_out = step(nxt, pos, hs, *flat)
print('step outputs:', len(step_out), '| logits', tuple(step_out[0].shape))
torch.onnx.export(
step, (nxt, pos, hs, *flat), f'{OUT}/decoder_step.onnx', dynamo=True,
input_names=['input_ids', 'cache_position', 'encoder_hidden_states']
+ [f'past_{n}' for n in self_names + cross_names],
output_names=['logits'] + [f'present_{n}' for n in self_names],
dynamic_shapes=({}, {}, {1: frames},
tuple([{2: past}] * len(self_names) + [{2: frames}] * len(cross_names))),
)
print('decoder_step exported')
|