Auto-detect framework for large models at ONNX export

[WangYizhang01]

System Info

• transformers version: 4.26.1
• Platform: Linux-4.4.0-142-generic-x86_64-with-glibc2.23
• Python version: 3.9.15
• Huggingface_hub version: 0.11.1
• PyTorch version (GPU?): 1.13.0 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: no
• Using distributed or parallel set-up in script?: no

Who can help?

@sgugger @muellerzr

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

import torch
import torch.nn as nn
from transformers import GPT2Config, GPT2Tokenizer, GPT2Model

num_attention_heads = 40
num_layers = 40
hidden_size = 5120

configuration = GPT2Config( 
                    n_embd=hidden_size,
                    n_layer=num_layers,
                    n_head=num_attention_heads
                )

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2Model(configuration)

tokenizer.save_pretrained('gpt2_checkpoint')
model.save_pretrained('gpt2_checkpoint')

python -m transformers.onnx --model=gpt2_checkpoint onnx/

Expected behavior

I created a GPT2 with a parameter volume of 13B. Just for testing, refer to https://huggingface.co/docs/transformers/serialization, I save it to gpt2_checkpoint. Then convert it to onnx using transformers.onnx. Due to the large amount of parameters, save_pretrained saves the model as *-0001.bin, *-0002.bin and so on. Later, when running ‘python -m transformers.onnx --model=gpt2_checkpoint onnx/’, an error FileNotFoundError: Cannot determine framework from given checkpoint location. There should be a pytorch_model.bin for PyTorch or tf_model.h5 for TensorFlow. So, I would like to ask how to convert a model with a large number of parameters into onnx for inference.

[sgugger]

cc @michaelbenayoun as well. This might be a feature necessary to implement in optimum.

[michaelbenayoun]

@fxmarty do you know if we support that in optimum?

[fxmarty]

Hi @WangYizhang01 ,

Using Optimum export, optimum-cli export onnx --model gpt2_checkpoint --task causal-lm --framework pt gpt2_onnx/ works fine. Not sure for transformers export if you can pass this --framework pt argument or not.

We could auto-detect the framework for split bins, right.