aws-rhsoln
commented
1 year ago As a note: We are looking to add better support for these encoder-decoder models using trace. We are working on performance for these kind of models.
For the issue: We tried reproducing the issue, but we didn't have access to your custom model. Hence, we used a t5-small to reproduce (changed the MODEL_PATH="t5-small"). We noticed that you are moving the model and inputs to xla device, that is not necessary and hence, we recommend removing that from this script. With this change, we went pass the compilation pass the compilation phase but ran into an assertion error related to the inputs. The error looks something like this:
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch/jit/_trace.py", line 983, in trace_module
argument_names,
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch/nn/modules/module.py", line 1182, in _slow_forward
result = self.forward(*input, **kwargs)
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch_neuronx/xla_impl/trace.py", line 85, in forward
tensors = self.flattener(tensors)
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch_neuronx/xla_impl/structure.py", line 206, in __call__
assert self.layout == layout
AssertionErrorThis error happens because we are passing same inputs twice (this case is not supported at this moment during trace and we are working on fixing this). We recommend making the following changes:
model_neuron = torch_neuronx.trace(model,(tpu_input_ids,
tpu_mask,
tpu_input_ids.clone().detach()
))After making these changes, we were able to compile but ran into a final error related to jit.trace. The error looks like this:
Traceback (most recent call last):
File "t5.py", line 40, in <module>
tpu_input_ids.clone().detach()
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch_neuronx/xla_impl/trace.py", line 323, in trace
return torch.jit.trace(result, example_inputs, strict=False)
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch/jit/_trace.py", line 768, in trace
_module_class,
File "/home/ec2-user/release-211/lib64/python3.7/site-packages/torch/jit/_trace.py", line 983, in trace_module
argument_names,
RuntimeError: Tracer cannot infer type of {'logits': tensor([[[0., 0., 0., ..., 0., 0., 0.],This happens because HF T5 model outputs a mixture of tensor and dictionary of tensors https://github.com/huggingface/transformers/blob/v4.26.0/src/transformers/models/t5/modeling_t5.py#L1465 . This doesn't go well with torch.jit.trace as it expects the outputs to be of single type. If we can make a change to just output a tuple, something along this:
class Model(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, input_ids, attention_mask, decoder_input_ids):
outputs = self.model(input_ids, attention_mask, decoder_input_ids)
return outputs.logits # Note: you can return any number here as output.logits, outputs.pass_key_values ..
model_wrapper = Model(model)After making the above changes the model should compile and we should be able to get the results.
Here is the script that reaches this point:
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch
import torch_neuronx
# MODEL_PATH = "/home/ubuntu/ipt_t5"
NEURON_PATH = "model_neuron_t5_v1.pt"
tokenizer = AutoTokenizer.from_pretrained("t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
def encode(tokenizer, text, max_length=128, batch_size=1):
tokens = tokenizer.encode_plus(text,
max_length=max_length,
padding=True,
truncation=True,
pad_to_multiple_of = 32,
return_tensors = 'pt'
)
return tokens
device = 'xla'
sequence = "What a wonderful day!"
tokenized_sequence = encode(tokenizer,sequence)
print(tokenized_sequence)
tpu_input_ids = tokenized_sequence['input_ids']
tpu_mask = tokenized_sequence['attention_mask']
class Model(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, input_ids, attention_mask, decoder_input_ids):
outputs = self.model(input_ids, attention_mask, decoder_input_ids)
return outputs.logits # Note: you can return any number here
model_wrapper = Model(model)
model_neuron = torch_neuronx.trace(model_wrapper,(tpu_input_ids,
tpu_mask,
tpu_input_ids.clone().detach()
))
#save as torchscript for inference
torch.jit.save(model_neuron,NEURON_PATH)
preds_cpu = model(input_ids=tokenized_sequence['input_ids'], attention_mask=tokenized_sequence['attention_mask'],
decoder_input_ids=tokenized_sequence['input_ids'])
preds_neuron = model_neuron(tokenized_sequence['input_ids'], tokenized_sequence['attention_mask'],
tokenized_sequence['input_ids'])
print("="*20)
print("PREDS CPU")
print(preds_cpu.logits)
print("="*20)
print("PREDS NEURON")
print(preds_neuron)
As a recommendation, since these encoder-decoder are autoregressive and since the encoder is run once and the decoder is run multiple times, encoder and decoder can be traced separately. We would run the encoder once and then run the decoder with padded inputs and run it till the stopping criteria is reached.
seanbenhur
ghost
jluntamazon
I am trying to compile a T5 model with torch-neuronx on inf2 instance, when I compile the model, I am getting the segmentation fault error
Here is my code
Here is a full error trace