GPT-Neo ONNX export

[michaelbenayoun]

What does this PR do?

This PR enables the export of GPT-Neo to ONNX by extending the new module transformers.onnx.

It also provides a possible way of implementing the export for specific tasks: the task can be specified when instantiating an OnnxConfig. It is a nice approach because it makes factoring most of the code for the inputs / outputs very easy, but it is less aligned with transformers DNA than having subclasses (such as OnnxConfigForSequenceClassification, etc) taking care of that.

The issue with having many subclasses is that it would have to be done everytime one wants to add the support for a model.
What do you think?

[mfuntowicz]

@sgugger @LysandreJik What do you think would be the best way to approach this exporting features for downstream task? I think we have the two possible ways:

• One config per task XOnnxConfigForY => Follow the general "duplication" pattern in transformers
• One config with task as parameter encapsulating the logic for I/O for each possible task => Potentially reduce the LoC

[sgugger]

I think using a task argument is a nice way of avoiding too many new classes which would crowd the main init of transformers.

[mfuntowicz]

@michaelbenayoun is the PR ready for review? 🥰

[michaelbenayoun]

@michaelbenayoun is the PR ready for review?
Yes, it is!

I also implemented a "factory" called FeaturesManager located at onnx/features.py from what was done before by @mfuntowicz in onnx/__main__.py which manages the mapping between features and models / onnx configs.

From what @sgugger said, I went with the "task argument" approach. Basically, a feature is the combination of a task and the potential use of past keys and values, for instance:

• sequence-classification
• sequence-classification-with-past

Any feature containing "-with-past" will be mapped by the factory to an OnnxConfig instantiated using the with_past method.

@mfuntowicz any comments on the changes I have made?

[mfuntowicz]

I think you can use torch.take_along_dim(logits, sequence_lengths, dim=1). You might need to match shape of logits and sequence_lengths.

=> It will remove the need to gather from the shape object.

[michaelbenayoun]

I applied the following changes:

          if isinstance(sequence_lengths, torch.Tensor):
              pooled_logits = torch.take_along_dim(
                  logits, indices=sequence_lengths.unsqueeze(1).unsqueeze(1), dim=1
              ).squeeze()
          else:
              pooled_logits = logits[torch.arange(batch_size), sequence_lengths]

The forward pass works and give the same output as original version, but ONNX conversion fails with:
RuntimeError: Exporting the operator take_along_dim to ONNX opset version 11 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub..

[sgugger]

LGTM, thanks a lot for the PR!

[LysandreJik]

Nice! It requires a bit more reading to understand what to add when manually adding a new configuration. It would be nice to add the ONNX part to the existing templates so that when adding new models, users would automatically make them compatible with ONNX (not in this PR though).

LGTM!