#!/usr/bin/env python3 # ============================================================================== # # Copyright (C) 2023 Sophgo Technologies Inc. All rights reserved. # # TPU-MLIR is licensed under the 2-Clause BSD License except for the # third-party components. # # ============================================================================== import os import torch import argparse from tqdm import tqdm from transformers import LlamaTokenizer, LlamaForCausalLM torch.set_grad_enabled(False) parser = argparse.ArgumentParser(description='export onnx.') parser.add_argument('--model_path', type=str, help='path to the torch model.') parser.add_argument('--seq_length', type=int, default=512, help="sequence length") parser.add_argument('--lmhead_with_topk', type=int, default=0, help="only trace the LmHeadWithTopK") args = parser.parse_args() model_path = args.model_path folder = f"./tmp/onnx" origin_model = LlamaForCausalLM.from_pretrained( model_path, trust_remote_code=True).eval() for param in origin_model.parameters(): param.requires_grad = False config = origin_model.config transformer = origin_model.model layers = transformer.layers SEQ_LENGTH = args.seq_length NUM_LAYERS = config.num_hidden_layers HIDDEN_SIZE = config.hidden_size NUM_ATTENTION_HEADS = config.num_attention_heads HEAD_DIM = HIDDEN_SIZE // NUM_ATTENTION_HEADS VOCAB_SIZE = config.vocab_size print(f'Layers: {NUM_LAYERS}\nHidden size: {HIDDEN_SIZE}\n') tokenizer = LlamaTokenizer.from_pretrained(model_path, trust_remote_code=True) class Embedding(torch.nn.Module): def __init__(self): super().__init__() def forward(self, input_ids): return transformer.embed_tokens(input_ids) class Block(torch.nn.Module): def __init__(self, layer_id): super().__init__() self.layer_id = layer_id self.layer = layers[layer_id] def forward(self, hidden_states, position_ids, attention_mask): hidden_states, past_kv = self.layer(hidden_states, attention_mask, position_ids, use_cache=True) present_k, present_v = past_kv return hidden_states, present_k, present_v class BlockCache(torch.nn.Module): def __init__(self, layer_id): super().__init__() self.layer_id = layer_id self.layer = layers[layer_id] def forward(self, hidden_states, position_ids, attention_mask, past_k, past_v): hidden_states, past_kv = self.layer(hidden_states, attention_mask, position_ids=position_ids, past_key_value=(past_k, past_v), use_cache=True) present_k, present_v = past_kv return hidden_states, present_k, present_v class LmHeadWithTopK(torch.nn.Module): def __init__(self): super().__init__() def forward(self, hidden_states): hidden_states = transformer.norm(hidden_states) m_logits = origin_model.lm_head(hidden_states) _, token = torch.topk(m_logits, 1) return token class LmHead(torch.nn.Module): def __init__(self): super().__init__() def forward(self, hidden_states): hidden_states = transformer.norm(hidden_states) m_logits = origin_model.lm_head(hidden_states) return m_logits class GreedyHead(torch.nn.Module): def __init__(self): super().__init__() def forward(self, m_logits): _, token = torch.topk(m_logits.float(), 1) return token # refs:https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py class PenaltySampleHead(torch.nn.Module): def __init__(self, top_k = 50, min_tokens_to_keep = 5): super().__init__() self.top_k = top_k self.min_tokens_to_keep = min_tokens_to_keep self.keep_matrix = torch.zeros((1, self.top_k), dtype=torch.bool) self.keep_matrix[0, :self.min_tokens_to_keep] = True def forward(self, m_logits, input_ids, top_p, temperature, penalty): # repeat penalty logits = torch.gather(m_logits, 1, input_ids) logits = torch.where(logits < 0, logits * penalty, logits / penalty) m_logits.scatter_(1, input_ids, logits) # top_k logits, token = torch.topk(m_logits.float(), self.top_k) # temperature logits = logits / temperature # top_p cumulative_probs = logits.softmax(dim=1).cumsum(dim=1) mask = cumulative_probs < top_p mask = mask + self.keep_matrix filtered_logits = torch.where(mask, logits, torch.FloatTensor([-1000.])) probs = filtered_logits.softmax(dim=1) return probs, token def convert_block(layer_id): model = Block(layer_id) hidden_states = torch.randn((1, SEQ_LENGTH, HIDDEN_SIZE)) position_ids = torch.tensor([range(SEQ_LENGTH)], dtype=torch.long) attention_mask = -1000 * torch.ones((1, 1, SEQ_LENGTH, SEQ_LENGTH), dtype=torch.float32).triu(diagonal=1) torch.onnx.export( model, (hidden_states, position_ids, attention_mask), f'{folder}/block_{layer_id}.onnx', verbose=False, input_names=['input_states', 'position_ids', 'attention_mask'], output_names=['hidden_states', 'past_k', 'past_v'], do_constant_folding=True, opset_version=15) def convert_block_cache(layer_id): model = BlockCache(layer_id) hidden_states = torch.randn((1, 1, HIDDEN_SIZE)) position_ids = torch.tensor([range(1)], dtype=torch.long) attention_mask = -1000 * torch.ones((1, 1, 1, SEQ_LENGTH + 1), dtype=torch.float32).triu(diagonal=1) past_k = torch.randn((1, SEQ_LENGTH, NUM_ATTENTION_HEADS, HEAD_DIM)) past_v = torch.randn((1, SEQ_LENGTH, NUM_ATTENTION_HEADS, HEAD_DIM)) torch.onnx.export( model, (hidden_states, position_ids, attention_mask, past_k, past_v), f'{folder}/block_cache_{layer_id}.onnx', verbose=False, input_names=[ 'input_states', 'position_ids', 'attention_mask', 'history_k', 'history_v' ], output_names=['hidden_states', 'past_k', 'past_v'], do_constant_folding=True, opset_version=15) def convert_embedding(): model = Embedding() input_ids = torch.tensor([range(SEQ_LENGTH)]) module = torch.jit.trace(model.forward, input_ids) torch.jit.save(module, f'{folder}/embedding.pt') def convert_lm_head_with_topk(): model = LmHeadWithTopK() input = torch.randn(1, 1, HIDDEN_SIZE) module = torch.jit.trace(model.forward, input) torch.jit.save(module, f'{folder}/lm_head_with_topk.pt') def convert_lm_head(): model = LmHead() input = torch.randn(1, HIDDEN_SIZE) torch.onnx.export(model, (input), f'{folder}/lm_head.onnx', verbose=False, input_names=['hidden_states'], output_names=['m_logits'], do_constant_folding=True, opset_version=15) def convert_greedy_head(): model = GreedyHead() m_logits = torch.randn(1, VOCAB_SIZE) torch.onnx.export( model, (m_logits), f'{folder}/greedy_head.onnx', verbose=False, input_names=['m_logits'], output_names=['token'], do_constant_folding=True, opset_version=15) def convert_penalty_sample_head(): model = PenaltySampleHead() m_logits = torch.randn(1, VOCAB_SIZE) input_ids = torch.tensor([range(SEQ_LENGTH)]) top_p = torch.tensor([0.8]) temperature = torch.tensor([0.98]) penalty = torch.tensor([0.98]) torch.onnx.export( model, (m_logits, input_ids, top_p, temperature, penalty), f'{folder}/penalty_sample_head.onnx', verbose=False, input_names=[ 'm_logits', 'input_ids', 'top_p', 'temperature', 'penalty' ], output_names=['probs', 'token'], do_constant_folding=True, opset_version=15) # create folder to store onnx if not os.path.exists(folder): os.makedirs(folder) # export models print(f'Convert block & block_cache') for i in tqdm(range(NUM_LAYERS)): convert_block_cache(i) convert_block(i) print(f'Convert embedding') convert_embedding() print(f'Convert lm_head') if args.lmhead_with_topk != 0: convert_lm_head_with_topk() else: convert_lm_head() convert_greedy_head() convert_penalty_sample_head() print("Done")