Skip to content

Errors when converting LLaMA to ONNX using dynamo export #104903

Description

@kunal-vaishnavi

🐛 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

Metadata

Metadata

Labels

module: onnxRelated to torch.onnxtriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

Type

No type

Projects

Status
Done

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions