Errors when converting LLaMA to ONNX using dynamo export
Labels: module: onnx, triaged

### 🐛 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](https://pytorch.org/docs/stable/generated/torch.full.html) says that `fill_value` should be a scalar value. Hugging Face's [implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L49) 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](https://github.com/pytorch/pytorch/pull/104468). The PR change is in my version of torch.

### Versions

Torch: v2.1.0.dev20230630+cu118
Transformers: v4.30.0