add ONNX support for BLOOM (#17961)

* add onnx support for BLOOM

* use TYPE_CHECKING for type annotations

* fix past_shape for bloom (different from gpt2)

* use logical_or instead of `+` for onnx support

* bigger `atol_for_validation` for larger bloom models

* copied -> taken because it's no longer an exact copy

* remove "copied from" comment

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
This commit is contained in:
Nouamane Tazi
2022-07-01 16:44:42 +02:00
committed by GitHub
parent 462b7f3a94
commit b68d408f1b
6 changed files with 116 additions and 19 deletions

View File

@@ -205,6 +205,7 @@ PYTORCH_EXPORT_MODELS = {
}
PYTORCH_EXPORT_WITH_PAST_MODELS = {
("bloom", "bigscience/bloom-350m"),
("gpt2", "gpt2"),
("gpt-neo", "EleutherAI/gpt-neo-125M"),
}