🐛 Describe the bug
While exporting LLaMA from PyTorch to ONNX using the dynamo exporter, the following error occurs.
While executing %full : [num_users=2] = call_function[target=torch.full](args = ((8, 8), %tensor), kwargs = {device: cpu})
Original traceback:
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 688, in forward
outputs = self.model(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1514, in _call_impl
return forward_call(*args, **kwargs)
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 537, in forward
attention_mask = self._prepare_decoder_attention_mask(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 465, in _prepare_decoder_attention_mask
combined_attention_mask = _make_causal_mask(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 49, in _make_causal_mask
mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 590, in dynamo_export
return Exporter(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 455, in export
graph_module = pre_export_passes(self.options, graph_module, updated_model_args)
File "<@beartype(torch.onnx._internal.exporter.pre_export_passes) at 0x7f096de17af0>", line 72, in pre_export_passes
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 622, in pre_export_passes
module = passes.Decompose(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/diagnostics/infra/decorator.py", line 142, in wrapper
ctx.log_and_raise_if_error(diag)
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 265, in log_and_raise_if_error
raise RuntimeErrorWithDiagnostic(
torch.onnx._internal.diagnostics.infra.context.RuntimeErrorWithDiagnostic: Running Decompose pass. Raised from:
DataDependentOutputException: aten._local_scalar_dense.default
Here is the code used to export LLaMA.
import torch
from transformers import LlamaForCausalLM
llama = LlamaForCausalLM.from_pretrained("decapoda-research/llama-7b-hf")
batch_size, seq_len = 2, 8
input_ids = torch.randint(low=0, high=6, size=(batch_size, seq_len), dtype=torch.int64)
attn_mask = torch.randint(low=0, high=2, size=(batch_size, seq_len), dtype=torch.int64)
torch.onnx.dynamo_export(
llama,
input_ids,
attn_mask,
).save("llama-7b-dynamo.onnx")
This error appears to arise from how torch.full is used. The documentation says that fill_value should be a scalar value. Hugging Face's implementation defines fill_value as torch.tensor(...) though.
# Copied from transformers.models.bart.modeling_bart._make_causal_mask
def _make_causal_mask(
input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
):
"""
Make causal mask used for bi-directional self-attention.
"""
bsz, tgt_len = input_ids_shape
mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
...
After trying a quick workaround to change mask to torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device), another error occurs.
[2023-07-10 19:36:59,057] torch.onnx: [ERROR] Cannot find symbolic function for aten::index.Tensor, which should be registered under aten.index.Tensor.
[2023-07-10 19:36:59,058] torch.onnx: [ERROR] None
[2023-07-10 19:36:59,059] torch.onnx: [ERROR] Cannot find symbolic function for aten::index.Tensor, which should be registered under aten.index.Tensor.
[2023-07-10 19:36:59,060] torch.onnx: [ERROR] None
...
[2023-07-10 19:36:59,100] torch.onnx: [ERROR] Cannot find symbolic function for aten::index.Tensor, which should be registered under aten.index.Tensor.
[2023-07-10 19:36:59,100] torch.onnx: [ERROR] None
[2023-07-10 19:36:59,100] torch.onnx: [ERROR] Cannot find symbolic function for aten::index.Tensor, which should be registered under aten.index.Tensor.
[2023-07-10 19:36:59,100] torch.onnx: [ERROR] None
[2023-07-10 19:36:59,101] torch.onnx: [ERROR] Unsupported FX nodes: {'call_function': ['aten.index.Tensor']}.
[2023-07-10 19:36:59,101] torch.onnx: [ERROR] None
Traceback (most recent call last):
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 590, in dynamo_export
return Exporter(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 455, in export
graph_module = pre_export_passes(self.options, graph_module, updated_model_args)
File "<@beartype(torch.onnx._internal.exporter.pre_export_passes) at 0x7f23bf50baf0>", line 72, in pre_export_passes
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/exporter.py", line 646, in pre_export_passes
analysis.UnsupportedFxNodesAnalysis(
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/fx/analysis/unsupported_nodes.py", line 83, in analyze
self._lint(analysis_result, diagnostic_level)
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/fx/analysis/unsupported_nodes.py", line 37, in _lint
self.diagnostic_context.log_and_raise_if_error(diagnostic)
File "/home/kvaishnavi/.local/lib/python3.8/site-packages/torch/onnx/_internal/diagnostics/infra/context.py", line 265, in log_and_raise_if_error
raise RuntimeErrorWithDiagnostic(
torch.onnx._internal.diagnostics.infra.context.RuntimeErrorWithDiagnostic: Unsupported FX nodes: {'call_function': ['aten.index.Tensor']}.
Note: There was an initial error with tabulate that I resolved with this PR. The PR change is in my version of torch.
Versions
Torch: v2.1.0.dev20230630+cu118
Transformers: v4.30.0
🐛 Describe the bug
While exporting LLaMA from PyTorch to ONNX using the dynamo exporter, the following error occurs.
Here is the code used to export LLaMA.
This error appears to arise from how
torch.fullis used. The documentation says thatfill_valueshould be a scalar value. Hugging Face's implementation definesfill_valueastorch.tensor(...)though.After trying a quick workaround to change
masktotorch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device), another error occurs.Note: There was an initial error with
tabulatethat I resolved with this PR. The PR change is in my version of torch.Versions
Torch: v2.1.0.dev20230630+cu118
Transformers: v4.30.0