Does anyone try to translate Llama model to MLIR ?
#23
by
hmsjwzb
- opened
I am trying to translate Llama model to MLIR. But I encounter some issues.
from transformers import AutoTokenizer, LlamaForCausalLM
from torch_mlir import fx
import torch
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
prompt = "What is the capital of France?"
inputs = tokenizer(prompt, return_tensors="pt")
print(f"Input shape: {inputs['input_ids'].shape}")
#generate_ids = model.generate(inputs.input_ids, max_length=40)
#print(tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0])
randomLongTensor = torch.randint(low=0, high=100, size=(1, 8), dtype=torch.long)
m = fx.export_and_import(model, randomLongTensor, enable_ir_printing=True,
enable_graph_printing=True)
This code will produce the following error message.
Input shape: torch.Size([1, 8])
Traceback (most recent call last):
File "/home/hmsjwzb/work/tinyLlama_0927/./lamaExample.py", line 17, in <module>
m = fx.export_and_import(model, randomLongTensor, enable_ir_printing=True,
File "/home/hmsjwzb/work/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/fx.py", line 76, in export_and_import
prog = torch.export.export(f, args, kwargs, dynamic_shapes=dynamic_shapes)
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/export/__init__.py", line 174, in export
return _export(
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/export/_trace.py", line 946, in wrapper
raise e
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/export/_trace.py", line 929, in wrapper
ep = fn(*args, **kwargs)
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/export/exported_program.py", line 88, in wrapper
return fn(*args, **kwargs)
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/export/_trace.py", line 1533, in _export
exported_program = ExportedProgram(
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/export/exported_program.py", line 245, in __init__
self.verifier().check(self)
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/_export/verifier.py", line 153, in check
self._check_graph_module(ep.graph_module)
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/_export/verifier.py", line 218, in _check_graph_module
_check_val(node)
File "/home/hmsjwzb/work/selfPython/ai/lib/python3.10/site-packages/torch/_export/verifier.py", line 61, in _check_val
raise SpecViolationError(f"Node.meta {node.name} is missing val field.")
torch._export.verifier.SpecViolationError: Node.meta _enter_autocast is missing val field.
Does anyone know how to fix it?