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?

Sign up or log in to comment