File size: 3,592 Bytes
0415692 | 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 | """Export FP32 AMOS medium one component at a time to bound conversion memory."""
import gc,json,sys
from pathlib import Path
import torch
from convert_decoder import save,Heads
from qwen_export_helpers import DecisionVision
class Embedding(torch.nn.Module):
def __init__(self,decoder):super().__init__();self.embed_tokens=decoder.embed_tokens;self.rotary=decoder.rotary_emb
def forward(self,input_ids,position_ids,visual_embeddings,visual_mask):
h=torch.where(visual_mask,visual_embeddings,self.embed_tokens(input_ids))
cos,sin=self.rotary(h,position_ids[1:]);return h,cos,sin
class Layer(torch.nn.Module):
def __init__(self,layer):super().__init__();self.layer=layer
def forward(self,hidden_states,cos,sin,position_ids,attention_mask):
return self.layer(hidden_states,position_embeddings=(cos,sin),position_ids=position_ids,attention_mask=attention_mask if self.layer.block_type=='full_attention' else None,use_cache=False)
class Finish(torch.nn.Module):
def __init__(self,model,letters):
super().__init__();self.norm=model.model.language_model.norm
self.register_buffer('readout',model.get_output_embeddings().weight[letters].detach().float().clone())
def forward(self,hidden_states):
h=self.norm(hidden_states);return h,h@self.readout.T
def main():
source=Path('upstream/amos-1-medium');root=Path('bundles/amos-1-medium-CoreAI');root.mkdir(exist_ok=True,parents=True)
sys.path.insert(0,str((source/'serve').resolve()))
from amos.decoder import load_model,load_heads
torch.set_num_threads(4)
m,codec=load_model(str(source/'model'),'cpu',dtype=torch.bfloat16);m.eval();heads=load_heads(source/'model',m,codec)
from stable_delta import install
install()
decoder=m.model.language_model;L=512;d=decoder.config.hidden_size
manifest={'sequence_length':L,'hidden_size':d,'precision':'float32','vision_precision':'float32','echo':codec.echo,'kinds':heads.kinds,'keep':heads.keep,'has_box':heads.box is not None,'image_side':224,'noncausal':False,'text_layout':'layer-sharded','layers':[]}
path=root/'embedding.aimodel'
decoder.embed_tokens.float();emb=Embedding(decoder).eval()
args=(torch.ones(1,L,dtype=torch.int32),torch.arange(L,dtype=torch.int32).view(1,1,-1).expand(4,1,-1),torch.zeros(1,L,d),torch.zeros(1,L,1,dtype=torch.bool))
with torch.no_grad():_,cos,sin=emb(*args)
print('Embedding',flush=True);manifest['embedding_outputs']=save(emb,args,path)
decoder.embed_tokens.to('meta');del emb;gc.collect()
layer_args=(torch.zeros(1,L,d),cos,sin,torch.arange(L,dtype=torch.int32)[None],torch.ones(1,1,L,L,dtype=torch.bool).tril())
for i,layer in enumerate(decoder.layers):
name=f'layer-{i:02}.aimodel';print(name,flush=True)
layer.float().eval();outputs=save(Layer(layer).eval(),layer_args,root/name)
manifest['layers'].append({'asset':name,'output':outputs[0],'kind':layer.block_type})
layer.to('meta');gc.collect()
print('Finish',flush=True);decoder.norm.float();manifest['outputs']=save(Finish(m,codec.letters).eval(),(torch.zeros(1,L,d),),root/'finish.aimodel')
m.get_output_embeddings().to('meta');decoder.norm.to('meta');gc.collect()
print('Vision',flush=True);m.model.visual.float();manifest['vision_output']=save(DecisionVision(m.model.visual).eval(),(torch.zeros(196,1536),),root/'vision.aimodel')[0]
m.model.visual.to('meta');gc.collect()
print('Heads',flush=True);manifest['head_outputs']=save(Heads(heads).eval(),(torch.zeros(d),torch.zeros(L,d),torch.zeros(49,d)),root/'heads.aimodel')
(root/'manifest.json').write_text(json.dumps(manifest,indent=2));print('Saved FP32 sharded model',flush=True)
if __name__=='__main__':main()
|