Fix Onnx Export for Composer HuggingFaceModels

[nik-mosaic]

Due to how torch.onnx.export processes arguments that are dictionaries (and since our HuggingFaceModel inputs are generally dictionaries), Onnx export for HuggingFaceModels is currently broken. This PR should fix this.

We fix this by writing
export_for_inference(..., sample_input=(input_batch, {})
instead of
export_for_inference(..., sample_input=(input_batch, ).

This works for both tensor inputs (ResNets, etc.) and dictionary inputs.

In addition to the above change, we move the model and sample_input to CPU, since ONNX export models and inputs cannot be on the GPU. We add a test to verify this works.

We also manually specify the ONNX export opset_version to 13. This is currently the default version for PyTorch 1.12.1, but not for older versions of PyTorch (opset_version 9 was the previous default), which was causing an error. Ideally we would not specify this manually, but I cannot find a way to get the latest supported opset_version. Once we stop supporting PyTorch < 1.12.1, we should stop manually specifying this.

Finally, we also add an optional dynamic_axes input, which is required for older (1.10.2) PyTorch CPU versions when exporting HuggingFaceModels. If users in this scenario do not provide dynamic_axes for their inputs, they may see the following UserWarning and their exported model may not be correct.
UserWarning: Type cannot be inferred, which might cause exported graph to produce incorrect results.

[dskhudia]

@nik-mosaic Thank you for hunting this down. I think a better solution would be if the sample_input is a dict we extract the names and pass those to input_names argument at export time.

See this: https://github.com/pytorch/pytorch/blob/108b25db25ef180bfbfaed2347c9b99103aa68ef/torch/onnx/utils.py#L288

[dskhudia]

Also while we are at this, we may want to assign the output names as well. HF models, depending on a config value, return a dict of outputs. https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L1373-L1378

[dskhudia]

See the implementation of this function as well. https://github.com/huggingface/transformers/blob/02b176c4ce14340d26d42825523f406959c6c202/src/transformers/onnx/convert.py#L84

[dskhudia]

Could you also check if the output type (dict or tuple of tensor) is the same between eager model output vs onnx exported model?

[nik-mosaic]

Also while we are at this, we may want to assign the output names as well. HF models, depending on a config value, return a dict of outputs. https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L1373-L1378

Wouldn't this require writing Hugging Face-specific code for the output names? Can we do this and remain agnostic to the model type?

I think the model/onnx_config config has output_names. Not sure if this is model or onnx specific config.

[nik-mosaic]

Could you also check if the output type (dict or tuple of tensor) is the same between eager model output vs onnx exported model?

They are not, the onnx output is just a list of arrays, the eager output is a dict of tensors.

[dskhudia]

Could you also check if the output type (dict or tuple of tensor) is the same between eager model output vs onnx exported model?

They are not, the onnx output is just a list of arrays, the eager output is a dict of tensors.

I am ok keeping this as is (i.e., without output names) for now. However, someone directly using onnx output downstream where they were using eager may run into issues if the output of eager is a dict. Could you create a task for this problem?

[dskhudia]

Also side comment for a future PR: Are we better off using export_pytorch function from HF for HF models? They might have handled a few more edge cases that we might be missing.

[nik-mosaic]

Also side comment for a future PR: Are we better off using export_pytorch function from HF for HF models? They might have handled a few more edge cases that we might be missing.

Yes, for HuggingFace models, I think we are better off calling this function. They do a better job with their OnnxConfig class, and have utilities to convert a HuggingFace config to an OnnxConfig so the inputs and outputs are named well for many classes of models.