[Training] [ShapeInferenceError] Dimension could not be inferred: incompatible shapes

[srijanie03]

Describe the issue

I am trying to get ONNX training graph for Lllama2_7b model. I can get the forward graph, no problem. But the issue occurs when I use generate artifacts.

I don't get this error when I run a signle layer Transformer block (attention+MLP) with similar dimensions. What is causing the issue here?

Additionally, the loss function is throwing an error too : expected 2 but got 66 arguments. Please explain. Thank you!

To reproduce

base_model = llama2_7b()


batch = torch.tensor([[    1,  7569,  7225, 16229,   366],
        [    1,  7569,  2462,  8640,   263]])

model_outputs = base_model(batch)
if isinstance(model_outputs, torch.Tensor):
    model_outputs = [model_outputs]

input_names = ["input"]
output_names = ["output"]
dynamic_axes = {"input": {0: "batch_size"}, "output": {0: "batch_size"}}


f = io.BytesIO()
torch.onnx.export(
    base_model,
    batch,
    "torchtune_llama2.onnx",
    input_names=input_names,
    output_names=output_names,
    opset_version=14,
    do_constant_folding=False,
    training=torch.onnx.TrainingMode.TRAINING,
    dynamic_axes=dynamic_axes,
    export_params=True,
    keep_initializers_as_inputs=False,
)
requires_grad = [name for name, param in base_model.named_parameters() if param.requires_grad]

frozen_params = [name for name, param in base_model.named_parameters() if not param.requires_grad]

artifacts.generate_artifacts(
    "torchtune_llama2.onnx",
    #optimizer=artifacts.OptimType.AdamW,
    #loss=artifacts.LossType.CrossEntropyLoss, 
    #loss=artifacts.LossType.MSELoss,
    requires_grad=requires_grad,
    frozen_params=frozen_params,
    artifact_directory="llama2",
    additional_output_names=["output"])

Urgency

Urgent

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.19.0

PyTorch Version

2.4.0

Execution Provider

CUDA

Execution Provider Library Version

CUDA 12.4

[carzh]

Hi, do you have the full script for exporting & generating artifacts? Or could you provide the forward graph ONNX file?

Generally, we see these errors when the generated forward graph is incorrect. Especially for the loss function, which expects a certain number of graph outputs. For LLM's especially, we see that unless the base Torch model that is being exported is in training mode, then usually it uses a key-value cache (an inference-only optimization that adds inputs and outputs to the graph).

The Torch model passed to torch.onnx.export (the base_model) must be in training mode (ie, you should be able to train with it), and the input and output names passed to the export function should correlate with the input names and output names of the Torch model.

If you have a working PyTorch training script for Llama2_7b, you can use that to determine the correct input names and output names, and what inputs you need to pass in for it to be in training mode.

[srijanie03]

Hi @carzh. Thanks a lot for the comment. Yes, I exactly did that and was able to resolve the issue. There was a mismatch with the input dimensions which was generating the error.

[carzh]

Glad to hear it! Closing as resolved -- feel free to reopen if you run into further issues.