mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-31 02:02:21 +06:00
Add Switch transformers (#19323)
* first commit * add more comments * add router v1 * clean up - remove `tf` modeling files * clean up - remove `tf` modeling files * clean up * v0 routers * added more router - Implemented `ExpertsChooseMaskedRouter` - added tests - 2 more routers to implement * last router * improved docstring - completed the docstring in `router.py` - added more args in the config * v0 sparse mlp * replace wrong naming * forward pass run * update MOE layer * small router update * fixup * consistency * remove scatter router * remove abstract layer * update test and model for integration testing * v1 conversion * update * hardcode hack * all keys match * add gin conversion, without additional libraries * update conversion sctipy * delete router file * update tests wrt router deletion * fix router issues * update expert code * update, logits match, code needsREFACTORING * Refactor code Co-authored-by: Younes Belkada <younesbelkada@users.noreply.github.com> * add generate tests Co-authored-by: younesbelkada <younesbelkada@gmail.com> * add support for router loss Co-authored-by: Younes Belkada <younesbelkada@users.noreply.github.com> * fix forward error * refactor a bit * remove `FlaxSwitchTransformers` modules * more tests pass * Update code Co-authored-by: Younes Belkada <younesbelkada@users.noreply.github.com> * fixup * fix tests * fix doc * fix doc + tokenization * fix tokenizer test * fix test * fix loss output * update code for backward pass * add loss support * update documentation * fix documentation, clean tokenizer * more doc fix, cleanup example_switch * fix failing test * fix test * fix test * fix loss issue * move layer * update doc and fix router capacity usage * fixup * add sparse mlp index for documentation on hub * fixup * test sparse mix architecture * Apply suggestions from code review * Update docs/source/en/model_doc/switch_transformers.mdx * fixup on update * fix tests * fix another test * attempt fix * Update src/transformers/models/switch_transformers/configuration_switch_transformers.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Update src/transformers/models/switch_transformers/convert_switch_transformers_original_flax_checkpoint_to_pytorch.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * try * all tests pass * fix jitter noise * Apply suggestions from code review * doc tests pass * Update src/transformers/models/switch_transformers/modeling_switch_transformers.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Update src/transformers/models/switch_transformers/modeling_switch_transformers.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * remove assert * change config order * fix readme japanese * Apply suggestions from code review Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> * remove parallelizable tests + add one liners * remove ONNX config * fix nits - add `T5Tokenizer` in auto mapping - remove `Switch Transformers` from ONNX supported models * remove `_get_router` * remove asserts * add check in test for `router_dtype` * add `SwitchTransformersConfig` in `run_pipeline_test` * Update tests/pipelines/test_pipelines_summarization.py * add huge model conversion script * fix slow tests - add better casting for `Linear8bitLt` - remove `torchscript` tests * add make dir * style on new script * fix nits - doctest - remove `_keys_to_ignore_on_load_unexpected` * Update src/transformers/models/switch_transformers/configuration_switch_transformers.py * add google as authors * fix year * remove last `assert` statements * standardize vertical spaces * fix failing import * fix another failing test * Remove strange àuthorized_keys` * removing todo and padding that is never used Co-authored-by: Arthur Zucker <arthur.zucker@gmail.com> Co-authored-by: ybelkada <younes@huggingface.co> Co-authored-by: Younes Belkada <younesbelkada@users.noreply.github.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Co-authored-by: Arthur Zucker <arthur@huggingface.co>
This commit is contained in:
Younes Belkada
GitHub
parent
55ba31908a
commit
163ac3d3ee
@ -379,6 +379,7 @@ Current number of checkpoints: ** (from Berkeley) released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
|
||||
1. **[Swin Transformer](https://huggingface.co/docs/transformers/model_doc/swin)** (from Microsoft) released with the paper [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo.
|
||||
1. **[Swin Transformer V2](https://huggingface.co/docs/transformers/model_doc/swinv2)** (from Microsoft) released with the paper [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) by Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo.
|
||||
1. **[SwitchTransformers](https://huggingface.co/docs/transformers/main/model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](https://huggingface.co/docs/transformers/model_doc/t5)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1)** (from Google AI) released in the repository [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[Table Transformer](https://huggingface.co/docs/transformers/model_doc/table-transformer)** (from Microsoft Research) released with the paper [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) by Brandon Smock, Rohith Pesala, Robin Abraham.
|
||||
|
||||
@ -379,6 +379,7 @@ Número actual de puntos de control: ** (from Berkeley) released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
|
||||
1. **[Swin Transformer](https://huggingface.co/docs/transformers/model_doc/swin)** (from Microsoft) released with the paper [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo.
|
||||
1. **[Swin Transformer V2](https://huggingface.co/docs/transformers/model_doc/swinv2)** (from Microsoft) released with the paper [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) by Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo.
|
||||
1. **[SwitchTransformers](https://huggingface.co/docs/transformers/main/model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](https://huggingface.co/docs/transformers/model_doc/t5)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1)** (from Google AI) released in the repository [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[Table Transformer](https://huggingface.co/docs/transformers/model_doc/table-transformer)** (from Microsoft Research) released with the paper [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) by Brandon Smock, Rohith Pesala, Robin Abraham.
|
||||
|
||||
@ -414,6 +414,7 @@ Flax、PyTorch、TensorFlowをcondaでインストールする方法は、それ
|
||||
1. **[SqueezeBERT](https://huggingface.co/docs/transformers/model_doc/squeezebert)** (from Berkeley) released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
|
||||
1. **[Swin Transformer](https://huggingface.co/docs/transformers/model_doc/swin)** (from Microsoft) released with the paper [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo.
|
||||
1. **[Swin Transformer V2](https://huggingface.co/docs/transformers/model_doc/swinv2)** (from Microsoft) released with the paper [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) by Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo.
|
||||
1. **[SwitchTransformers](https://huggingface.co/docs/transformers/main/model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](https://huggingface.co/docs/transformers/model_doc/t5)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1)** (from Google AI) released in the repository [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[Table Transformer](https://huggingface.co/docs/transformers/model_doc/table-transformer)** (from Microsoft Research) released with the paper [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) by Brandon Smock, Rohith Pesala, Robin Abraham.
|
||||
|
||||
@ -329,6 +329,7 @@ Flax, PyTorch, TensorFlow 설치 페이지에서 이들을 conda로 설치하는
|
||||
1. **[SqueezeBERT](https://huggingface.co/docs/transformers/model_doc/squeezebert)** (from Berkeley) released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
|
||||
1. **[Swin Transformer](https://huggingface.co/docs/transformers/model_doc/swin)** (from Microsoft) released with the paper [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo.
|
||||
1. **[Swin Transformer V2](https://huggingface.co/docs/transformers/model_doc/swinv2)** (from Microsoft) released with the paper [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) by Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo.
|
||||
1. **[SwitchTransformers](https://huggingface.co/docs/transformers/main/model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](https://huggingface.co/docs/transformers/model_doc/t5)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1)** (from Google AI) released in the repository [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[Table Transformer](https://huggingface.co/docs/transformers/model_doc/table-transformer)** (from Microsoft Research) released with the paper [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) by Brandon Smock, Rohith Pesala, Robin Abraham.
|
||||
|
||||
@ -353,6 +353,7 @@ conda install -c huggingface transformers
|
||||
1. **[SqueezeBERT](https://huggingface.co/docs/transformers/model_doc/squeezebert)** (来自 Berkeley) 伴随论文 [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) 由 Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer 发布。
|
||||
1. **[Swin Transformer](https://huggingface.co/docs/transformers/model_doc/swin)** (来自 Microsoft) 伴随论文 [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) 由 Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo 发布。
|
||||
1. **[Swin Transformer V2](https://huggingface.co/docs/transformers/model_doc/swinv2)** (来自 Microsoft) 伴随论文 [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) 由 Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo 发布。
|
||||
1. **[SwitchTransformers](https://huggingface.co/docs/transformers/main/model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](https://huggingface.co/docs/transformers/model_doc/t5)** (来自 Google AI) 伴随论文 [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) 由 Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu 发布。
|
||||
1. **[T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1)** (来自 Google AI) 伴随论文 [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) 由 Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu 发布。
|
||||
1. **[Table Transformer](https://huggingface.co/docs/transformers/model_doc/table-transformer)** (来自 Microsoft Research) 伴随论文 [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) 由 Brandon Smock, Rohith Pesala, Robin Abraham 发布。
|
||||
|
||||
@ -365,6 +365,7 @@ conda install -c huggingface transformers
|
||||
1. **[SqueezeBERT](https://huggingface.co/docs/transformers/model_doc/squeezebert)** (from Berkeley) released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
|
||||
1. **[Swin Transformer](https://huggingface.co/docs/transformers/model_doc/swin)** (from Microsoft) released with the paper [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo.
|
||||
1. **[Swin Transformer V2](https://huggingface.co/docs/transformers/model_doc/swinv2)** (from Microsoft) released with the paper [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) by Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo.
|
||||
1. **[SwitchTransformers](https://huggingface.co/docs/transformers/main/model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](https://huggingface.co/docs/transformers/model_doc/t5)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1)** (from Google AI) released with the paper [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[Table Transformer](https://huggingface.co/docs/transformers/model_doc/table-transformer)** (from Microsoft Research) released with the paper [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) by Brandon Smock, Rohith Pesala, Robin Abraham.
|
||||
|
||||
@ -351,6 +351,8 @@
|
||||
title: Splinter
|
||||
- local: model_doc/squeezebert
|
||||
title: SqueezeBERT
|
||||
- local: model_doc/switch_transformers
|
||||
title: SwitchTransformers
|
||||
- local: model_doc/t5
|
||||
title: T5
|
||||
- local: model_doc/t5v1.1
|
||||
|
||||
@ -167,6 +167,7 @@ The documentation is organized into five sections:
|
||||
1. **[SqueezeBERT](model_doc/squeezebert)** (from Berkeley) released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
|
||||
1. **[Swin Transformer](model_doc/swin)** (from Microsoft) released with the paper [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo.
|
||||
1. **[Swin Transformer V2](model_doc/swinv2)** (from Microsoft) released with the paper [Swin Transformer V2: Scaling Up Capacity and Resolution](https://arxiv.org/abs/2111.09883) by Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo.
|
||||
1. **[SwitchTransformers](model_doc/switch_transformers)** (from Google) released with the paper [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
1. **[T5](model_doc/t5)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[T5v1.1](model_doc/t5v1.1)** (from Google AI) released in the repository [google-research/text-to-text-transfer-transformer](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
|
||||
1. **[Table Transformer](model_doc/table-transformer)** (from Microsoft Research) released with the paper [PubTables-1M: Towards Comprehensive Table Extraction From Unstructured Documents](https://arxiv.org/abs/2110.00061) by Brandon Smock, Rohith Pesala, Robin Abraham.
|
||||
@ -322,6 +323,7 @@ Flax), PyTorch, and/or TensorFlow.
|
||||
| SqueezeBERT | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
| Swin Transformer | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| Swin Transformer V2 | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| SwitchTransformers | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| T5 | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| Table Transformer | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| TAPAS | ✅ | ❌ | ✅ | ✅ | ❌ |
|
||||
|
||||
64
docs/source/en/model_doc/switch_transformers.mdx
Normal file
64
docs/source/en/model_doc/switch_transformers.mdx
Normal file
@ -0,0 +1,64 @@
|
||||
<!--Copyright 2022 The HuggingFace Team. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||
the License. You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
|
||||
specific language governing permissions and limitations under the License.
|
||||
-->
|
||||
|
||||
# SwitchTransformers
|
||||
|
||||
## Overview
|
||||
|
||||
The SwitchTransformers model was proposed in [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.
|
||||
|
||||
The Switch Transformer model uses a sparse T5 encoder-decoder architecure, where the MLP are replaced by a Mixture of Experts (MoE). A routing mechanism (top 1 in this case) associates each token to one of the expert, where each expert is a dense MLP. While switch transformers have a lot more weights than their equivalent dense models, the sparsity allows better scaling and better finetuning performance at scale.
|
||||
During a forward pass, only a fraction of the weights are used. The routing mecanism allows the model to select relevant weights on the fly which increases the model capacity without increasing the number of operations.
|
||||
|
||||
|
||||
The abstract from the paper is the following:
|
||||
|
||||
*In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model -- with outrageous numbers of parameters -- but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs and training instability -- we address these with the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle the instabilities and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the "Colossal Clean Crawled Corpus" and achieve a 4x speedup over the T5-XXL model.*
|
||||
|
||||
Tips:
|
||||
|
||||
- SwitchTransformers uses the [`T5Tokenizer`], which can be loaded directly from each model's repository.
|
||||
- The released weights are pretrained on English [Masked Language Modeling](https://moon-ci-docs.huggingface.co/docs/transformers/pr_19323/en/glossary#general-terms) task, and should be finetuned.
|
||||
|
||||
This model was contributed by [Younes Belkada](https://huggingface.co/ybelkada) and [Arthur Zucker](https://huggingface.co/ArtZucker) .
|
||||
The original code can be found [here](https://github.com/google/flaxformer/tree/main/flaxformer/architectures/moe).
|
||||
|
||||
|
||||
## SwitchTransformersConfig
|
||||
|
||||
[[autodoc]] SwitchTransformersConfig
|
||||
|
||||
## SwitchTransformersTop1Router
|
||||
|
||||
[[autodoc]] SwitchTransformersTop1Router
|
||||
- _compute_router_probabilities
|
||||
- forward
|
||||
|
||||
## SwitchTransformersSparseMLP
|
||||
|
||||
[[autodoc]] SwitchTransformersSparseMLP
|
||||
- forward
|
||||
|
||||
## SwitchTransformersModel
|
||||
|
||||
[[autodoc]] SwitchTransformersModel
|
||||
- forward
|
||||
|
||||
## SwitchTransformersForConditionalGeneration
|
||||
|
||||
[[autodoc]] SwitchTransformersForConditionalGeneration
|
||||
- forward
|
||||
|
||||
## SwitchTransformersEncoderModel
|
||||
|
||||
[[autodoc]] SwitchTransformersEncoderModel
|
||||
- forward
|
||||
@ -354,6 +354,7 @@ _import_structure = {
|
||||
"models.squeezebert": ["SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "SqueezeBertConfig", "SqueezeBertTokenizer"],
|
||||
"models.swin": ["SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "SwinConfig"],
|
||||
"models.swinv2": ["SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Swinv2Config"],
|
||||
"models.switch_transformers": ["SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP", "SwitchTransformersConfig"],
|
||||
"models.t5": ["T5_PRETRAINED_CONFIG_ARCHIVE_MAP", "T5Config"],
|
||||
"models.table_transformer": ["TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "TableTransformerConfig"],
|
||||
"models.tapas": ["TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP", "TapasConfig", "TapasTokenizer"],
|
||||
@ -2025,6 +2026,17 @@ else:
|
||||
"load_tf_weights_in_t5",
|
||||
]
|
||||
)
|
||||
_import_structure["models.switch_transformers"].extend(
|
||||
[
|
||||
"SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST",
|
||||
"SwitchTransformersEncoderModel",
|
||||
"SwitchTransformersForConditionalGeneration",
|
||||
"SwitchTransformersModel",
|
||||
"SwitchTransformersPreTrainedModel",
|
||||
"SwitchTransformersTop1Router",
|
||||
"SwitchTransformersSparseMLP",
|
||||
]
|
||||
)
|
||||
_import_structure["models.trajectory_transformer"].extend(
|
||||
[
|
||||
"TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
|
||||
@ -3465,6 +3477,7 @@ if TYPE_CHECKING:
|
||||
from .models.squeezebert import SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, SqueezeBertConfig, SqueezeBertTokenizer
|
||||
from .models.swin import SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinConfig
|
||||
from .models.swinv2 import SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP, Swinv2Config
|
||||
from .models.switch_transformers import SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP, SwitchTransformersConfig
|
||||
from .models.t5 import T5_PRETRAINED_CONFIG_ARCHIVE_MAP, T5Config
|
||||
from .models.table_transformer import TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, TableTransformerConfig
|
||||
from .models.tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig, TapasTokenizer
|
||||
@ -4818,6 +4831,15 @@ if TYPE_CHECKING:
|
||||
Swinv2Model,
|
||||
Swinv2PreTrainedModel,
|
||||
)
|
||||
from .models.switch_transformers import (
|
||||
SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
|
||||
SwitchTransformersEncoderModel,
|
||||
SwitchTransformersForConditionalGeneration,
|
||||
SwitchTransformersModel,
|
||||
SwitchTransformersPreTrainedModel,
|
||||
SwitchTransformersSparseMLP,
|
||||
SwitchTransformersTop1Router,
|
||||
)
|
||||
from .models.t5 import (
|
||||
T5_PRETRAINED_MODEL_ARCHIVE_LIST,
|
||||
T5EncoderModel,
|
||||
|
||||
@ -286,6 +286,91 @@ class BaseModelOutputWithPastAndCrossAttentions(ModelOutput):
|
||||
cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoEModelOutput(ModelOutput):
|
||||
"""
|
||||
Base class for model's outputs, with potential hidden states and attentions.
|
||||
|
||||
Args:
|
||||
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
|
||||
Sequence of hidden-states at the output of the last layer of the model.
|
||||
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
||||
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
||||
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
|
||||
|
||||
Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
|
||||
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
|
||||
heads.
|
||||
router_probs (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
|
||||
|
||||
Raw router probabilities that are computed by MoE routers, these terms are used to compute the auxiliary
|
||||
loss and the z_loss for Mixture of Experts models.
|
||||
"""
|
||||
|
||||
last_hidden_state: torch.FloatTensor = None
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
||||
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
router_probs: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoEModelOutputWithPastAndCrossAttentions(ModelOutput):
|
||||
"""
|
||||
Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding) as well as
|
||||
Mixture of Expert's router hidden states terms, to train a MoE model.
|
||||
|
||||
Args:
|
||||
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
|
||||
Sequence of hidden-states at the output of the last layer of the model.
|
||||
|
||||
If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,
|
||||
hidden_size)` is output.
|
||||
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
||||
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
|
||||
`(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if
|
||||
`config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads,
|
||||
encoder_sequence_length, embed_size_per_head)`.
|
||||
|
||||
Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
|
||||
`config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values`
|
||||
input) to speed up sequential decoding.
|
||||
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
||||
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
||||
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
|
||||
|
||||
Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
|
||||
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
|
||||
heads.
|
||||
cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
|
||||
weighted average in the cross-attention heads.
|
||||
router_probs (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
|
||||
|
||||
Raw router probabilities that are computed by MoE routers, these terms are used to compute the auxiliary
|
||||
loss and the z_loss for Mixture of Experts models.
|
||||
"""
|
||||
|
||||
last_hidden_state: torch.FloatTensor = None
|
||||
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
||||
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
router_probs: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Seq2SeqModelOutput(ModelOutput):
|
||||
"""
|
||||
@ -347,6 +432,78 @@ class Seq2SeqModelOutput(ModelOutput):
|
||||
encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Seq2SeqMoEModelOutput(ModelOutput):
|
||||
"""
|
||||
Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential
|
||||
decoding.
|
||||
|
||||
Args:
|
||||
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
|
||||
Sequence of hidden-states at the output of the last layer of the decoder of the model.
|
||||
|
||||
If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,
|
||||
hidden_size)` is output.
|
||||
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
||||
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
|
||||
`(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
|
||||
`(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
|
||||
|
||||
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
|
||||
blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
|
||||
decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
||||
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
||||
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
|
||||
|
||||
Hidden-states of the decoder at the output of each layer plus the optional initial embedding outputs.
|
||||
decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
|
||||
self-attention heads.
|
||||
decoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
|
||||
|
||||
Router logits of the decoder model, useful to compute the auxiliary loss for Mixture of Experts models.
|
||||
cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
|
||||
weighted average in the cross-attention heads.
|
||||
encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
||||
Sequence of hidden-states at the output of the last layer of the encoder of the model.
|
||||
encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
||||
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
||||
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
|
||||
|
||||
Hidden-states of the encoder at the output of each layer plus the optional initial embedding outputs.
|
||||
encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
|
||||
self-attention heads.
|
||||
encoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
|
||||
|
||||
Router logits of the encoder model, useful to compute the auxiliary loss and the z_loss for the sparse
|
||||
modules.
|
||||
"""
|
||||
|
||||
last_hidden_state: torch.FloatTensor = None
|
||||
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
||||
decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
||||
decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
decoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None
|
||||
cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
encoder_last_hidden_state: Optional[torch.FloatTensor] = None
|
||||
encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
||||
encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
encoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CausalLMOutput(ModelOutput):
|
||||
"""
|
||||
@ -581,6 +738,81 @@ class Seq2SeqLMOutput(ModelOutput):
|
||||
encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Seq2SeqMoEOutput(ModelOutput):
|
||||
"""
|
||||
Base class for sequence-to-sequence language models outputs.
|
||||
|
||||
Args:
|
||||
loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
|
||||
Language modeling loss.
|
||||
logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
|
||||
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
|
||||
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
|
||||
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
|
||||
`(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
|
||||
`(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
|
||||
|
||||
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
|
||||
blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
|
||||
decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
||||
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
||||
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
|
||||
|
||||
Hidden-states of the decoder at the output of each layer plus the initial embedding outputs.
|
||||
decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
|
||||
self-attention heads.
|
||||
decoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
|
||||
|
||||
Router logits of the decoder model, useful to compute the auxiliary loss for Mixture of Experts models.
|
||||
cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
|
||||
weighted average in the cross-attention heads.
|
||||
encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
|
||||
Sequence of hidden-states at the output of the last layer of the encoder of the model.
|
||||
encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
|
||||
Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
|
||||
one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
|
||||
|
||||
Hidden-states of the encoder at the output of each layer plus the initial embedding outputs.
|
||||
encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
|
||||
sequence_length)`.
|
||||
|
||||
Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
|
||||
self-attention heads.
|
||||
encoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`):
|
||||
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
|
||||
|
||||
Router logits of the encoder model, useful to compute the auxiliary loss and z_loss for Mixture of Experts
|
||||
models.
|
||||
"""
|
||||
|
||||
loss: Optional[torch.FloatTensor] = None
|
||||
logits: torch.FloatTensor = None
|
||||
encoder_z_loss: torch.FloatTensor = None
|
||||
decoder_z_loss: torch.FloatTensor = None
|
||||
encoder_aux_loss: torch.FloatTensor = None
|
||||
decoder_aux_loss: torch.FloatTensor = None
|
||||
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
||||
decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
||||
decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
decoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None
|
||||
cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
encoder_last_hidden_state: Optional[torch.FloatTensor] = None
|
||||
encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
||||
encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
||||
encoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class NextSentencePredictorOutput(ModelOutput):
|
||||
"""
|
||||
|
||||
@ -141,6 +141,7 @@ from . import (
|
||||
squeezebert,
|
||||
swin,
|
||||
swinv2,
|
||||
switch_transformers,
|
||||
t5,
|
||||
table_transformer,
|
||||
tapas,
|
||||
|
||||
@ -137,6 +137,7 @@ CONFIG_MAPPING_NAMES = OrderedDict(
|
||||
("squeezebert", "SqueezeBertConfig"),
|
||||
("swin", "SwinConfig"),
|
||||
("swinv2", "Swinv2Config"),
|
||||
("switch_transformers", "SwitchTransformersConfig"),
|
||||
("t5", "T5Config"),
|
||||
("table-transformer", "TableTransformerConfig"),
|
||||
("tapas", "TapasConfig"),
|
||||
@ -273,6 +274,7 @@ CONFIG_ARCHIVE_MAP_MAPPING_NAMES = OrderedDict(
|
||||
("squeezebert", "SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
("swin", "SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
("swinv2", "SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
("switch_transformers", "SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
("t5", "T5_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
("table-transformer", "TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
("tapas", "TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
|
||||
@ -429,6 +431,7 @@ MODEL_NAMES_MAPPING = OrderedDict(
|
||||
("squeezebert", "SqueezeBERT"),
|
||||
("swin", "Swin Transformer"),
|
||||
("swinv2", "Swin Transformer V2"),
|
||||
("switch_transformers", "SwitchTransformers"),
|
||||
("t5", "T5"),
|
||||
("t5v1.1", "T5v1.1"),
|
||||
("table-transformer", "Table Transformer"),
|
||||
|
||||
@ -133,6 +133,7 @@ MODEL_MAPPING_NAMES = OrderedDict(
|
||||
("squeezebert", "SqueezeBertModel"),
|
||||
("swin", "SwinModel"),
|
||||
("swinv2", "Swinv2Model"),
|
||||
("switch_transformers", "SwitchTransformersModel"),
|
||||
("t5", "T5Model"),
|
||||
("table-transformer", "TableTransformerModel"),
|
||||
("tapas", "TapasModel"),
|
||||
@ -203,6 +204,7 @@ MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
|
||||
("roc_bert", "RoCBertForPreTraining"),
|
||||
("splinter", "SplinterForPreTraining"),
|
||||
("squeezebert", "SqueezeBertForMaskedLM"),
|
||||
("switch_transformers", "SwitchTransformersForConditionalGeneration"),
|
||||
("t5", "T5ForConditionalGeneration"),
|
||||
("tapas", "TapasForMaskedLM"),
|
||||
("transfo-xl", "TransfoXLLMHeadModel"),
|
||||
@ -277,6 +279,7 @@ MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict(
|
||||
("roformer", "RoFormerForMaskedLM"),
|
||||
("speech_to_text", "Speech2TextForConditionalGeneration"),
|
||||
("squeezebert", "SqueezeBertForMaskedLM"),
|
||||
("switch_transformers", "SwitchTransformersForConditionalGeneration"),
|
||||
("t5", "T5ForConditionalGeneration"),
|
||||
("tapas", "TapasForMaskedLM"),
|
||||
("transfo-xl", "TransfoXLLMHeadModel"),
|
||||
@ -519,6 +522,7 @@ MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
|
||||
("pegasus_x", "PegasusXForConditionalGeneration"),
|
||||
("plbart", "PLBartForConditionalGeneration"),
|
||||
("prophetnet", "ProphetNetForConditionalGeneration"),
|
||||
("switch_transformers", "SwitchTransformersForConditionalGeneration"),
|
||||
("t5", "T5ForConditionalGeneration"),
|
||||
("xlm-prophetnet", "XLMProphetNetForConditionalGeneration"),
|
||||
]
|
||||
|
||||
@ -258,6 +258,13 @@ else:
|
||||
"squeezebert",
|
||||
("SqueezeBertTokenizer", "SqueezeBertTokenizerFast" if is_tokenizers_available() else None),
|
||||
),
|
||||
(
|
||||
"switch_transformers",
|
||||
(
|
||||
"T5Tokenizer" if is_sentencepiece_available() else None,
|
||||
"T5TokenizerFast" if is_tokenizers_available() else None,
|
||||
),
|
||||
),
|
||||
(
|
||||
"t5",
|
||||
(
|
||||
|
||||
84
src/transformers/models/switch_transformers/__init__.py
Normal file
84
src/transformers/models/switch_transformers/__init__.py
Normal file
@ -0,0 +1,84 @@
|
||||
# flake8: noqa
|
||||
# There's no way to ignore "F401 '...' imported but unused" warnings in this
|
||||
# module, but to preserve other warnings. So, don't check this module at all.
|
||||
|
||||
# Copyright 2022 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ...utils import (
|
||||
OptionalDependencyNotAvailable,
|
||||
_LazyModule,
|
||||
is_flax_available,
|
||||
is_sentencepiece_available,
|
||||
is_tf_available,
|
||||
is_tokenizers_available,
|
||||
is_torch_available,
|
||||
)
|
||||
|
||||
|
||||
_import_structure = {
|
||||
"configuration_switch_transformers": [
|
||||
"SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP",
|
||||
"SwitchTransformersConfig",
|
||||
"SwitchTransformersOnnxConfig",
|
||||
]
|
||||
}
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
_import_structure["modeling_switch_transformers"] = [
|
||||
"SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST",
|
||||
"SwitchTransformersEncoderModel",
|
||||
"SwitchTransformersForConditionalGeneration",
|
||||
"SwitchTransformersModel",
|
||||
"SwitchTransformersPreTrainedModel",
|
||||
"SwitchTransformersTop1Router",
|
||||
"SwitchTransformersSparseMLP",
|
||||
]
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .configuration_switch_transformers import (
|
||||
SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP,
|
||||
SwitchTransformersConfig,
|
||||
SwitchTransformersOnnxConfig,
|
||||
)
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
pass
|
||||
else:
|
||||
from .modeling_switch_transformers import (
|
||||
SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
|
||||
SwitchTransformersEncoderModel,
|
||||
SwitchTransformersForConditionalGeneration,
|
||||
SwitchTransformersModel,
|
||||
SwitchTransformersPreTrainedModel,
|
||||
SwitchTransformersSparseMLP,
|
||||
SwitchTransformersTop1Router,
|
||||
)
|
||||
|
||||
|
||||
else:
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
|
||||
@ -0,0 +1,203 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2022, Google and HuggingFace Inc.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" Switch Transformers model configuration"""
|
||||
from ...configuration_utils import PretrainedConfig
|
||||
from ...utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
|
||||
"google/switch-base-8": "https://huggingface.co/google/switch-base-8/blob/main/config.json",
|
||||
}
|
||||
|
||||
|
||||
class SwitchTransformersConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`SwitchTransformersModel`]. It is used to
|
||||
instantiate a SwitchTransformers model according to the specified arguments, defining the model architecture.
|
||||
Instantiating a configuration with the defaults will yield a similar configuration to that of the
|
||||
SwitchTransformers [google/switch-base-8](https://huggingface.co/google/switch-base-8) architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
Arguments:
|
||||
vocab_size (`int`, *optional*, defaults to 32128):
|
||||
Vocabulary size of the SwitchTransformers model. Defines the number of different tokens that can be
|
||||
represented by the `inputs_ids` passed when calling [`SwitchTransformersModel`].
|
||||
d_model (`int`, *optional*, defaults to 512):
|
||||
Size of the encoder layers and the pooler layer.
|
||||
d_kv (`int`, *optional*, defaults to 64):
|
||||
Size of the key, query, value projections per attention head. `d_kv` has to be equal to `d_model //
|
||||
num_heads`.
|
||||
d_ff (`int`, *optional*, defaults to 2048):
|
||||
Size of the intermediate feed forward layer in each `SwitchTransformersBlock`.
|
||||
expert_capacity (`int`, *optional*, defaults to 64):
|
||||
Number of tokens that can be stored in each expert. If set to 1, the model will behave like a regular
|
||||
Transformer.
|
||||
num_layers (`int`, *optional*, defaults to 12):
|
||||
Number of dense hidden layers in the Transformer encoder layer.
|
||||
num_sparse_encoder_layers (`int`, *optional*, defaults to 6):
|
||||
Number of sparse (MoE) dense hidden layers in the Transformer encoder layer.
|
||||
num_decoder_layers (`int`, *optional*, defaults to 12):
|
||||
Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set.
|
||||
num_sparse_decoder_layers (`int`, *optional*, defaults to 12):
|
||||
Number of sparse (MoE) dense hidden layers in the Transformer decoder layer.
|
||||
num_heads (`int`, *optional*, defaults to 8):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
num_experts (`int`, *optional*, defaults to 8):
|
||||
Number of experts for each SwitchTransformer layer.
|
||||
router_type (`str`, *optional*, defaults to `"tokens_masked"`):
|
||||
Router type - choose between `"tokens_masked", `"tokens_scatter"` and `"experts_masked"`.
|
||||
router_bias (`bool`, *optional*, defaults to `True`):
|
||||
Whether to add a bias to the router.
|
||||
router_jitter_noise (`float`, *optional*, defaults to 0.1):
|
||||
Amount of noise to add to the router.
|
||||
router_dtype (`str`, *optional*, default to `"float32"`):
|
||||
The `dtype` used for the routers. It is preferable to keep the `dtype` to `"float32"` as specified in the
|
||||
*selective precision* discussion in [the paper](https://arxiv.org/abs/2101.03961).
|
||||
router_ignore_padding_tokens (`bool`, *optional*, defaults to `False`):
|
||||
Whether to ignore padding tokens when routing.
|
||||
relative_attention_num_buckets (`int`, *optional*, defaults to 32):
|
||||
The number of buckets to use for each attention layer.
|
||||
relative_attention_max_distance (`int`, *optional*, defaults to 128):
|
||||
The maximum distance of the longer sequences for the bucket separation.
|
||||
dropout_rate (`float`, *optional*, defaults to 0.1):
|
||||
The ratio for all dropout layers.
|
||||
layer_norm_eps (`float`, *optional*, defaults to 1e-6):
|
||||
The epsilon used by the layer normalization layers.
|
||||
router_z_loss_coef (`float`, *optional*, defaults to 0.001):
|
||||
The z loss factor for the total loss.
|
||||
router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
|
||||
The aux loss factor for the total loss.
|
||||
initializer_factor (`float`, *optional*, defaults to 1):
|
||||
A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
|
||||
testing).
|
||||
feed_forward_proj (`string`, *optional*, defaults to `"relu"`):
|
||||
Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. SwitchTransformersv1.1
|
||||
uses the `"gated-gelu"` feed forward projection. Original SwitchTransformers uses `"relu"`.
|
||||
add_router_probs (`bool`, *optional*, defaults to `False`):
|
||||
Whether to output router probabilities to compute router auxiliary loss.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models).
|
||||
"""
|
||||
model_type = "switch_transformers"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
attribute_map = {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=32128,
|
||||
d_model=768,
|
||||
d_kv=64,
|
||||
d_ff=2048,
|
||||
expert_capacity=64,
|
||||
num_layers=12,
|
||||
num_sparse_encoder_layers=3,
|
||||
num_decoder_layers=12,
|
||||
num_sparse_decoder_layers=3,
|
||||
num_heads=12,
|
||||
num_experts=8,
|
||||
router_type="tokens_masked",
|
||||
router_bias=False,
|
||||
router_jitter_noise=0.01,
|
||||
router_dtype="float32",
|
||||
router_ignore_padding_tokens=False,
|
||||
relative_attention_num_buckets=32,
|
||||
relative_attention_max_distance=128,
|
||||
dropout_rate=0.1,
|
||||
layer_norm_epsilon=1e-6,
|
||||
router_z_loss_coef=0.001,
|
||||
router_aux_loss_coef=0.001,
|
||||
initializer_factor=1.0,
|
||||
feed_forward_proj="relu",
|
||||
is_encoder_decoder=True,
|
||||
add_router_probs=False,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
eos_token_id=1,
|
||||
**kwargs
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.d_model = d_model
|
||||
self.d_kv = d_kv
|
||||
self.d_ff = d_ff
|
||||
|
||||
self.num_sparse_encoder_layers = num_sparse_encoder_layers
|
||||
|
||||
self.num_layers = num_layers
|
||||
self.num_decoder_layers = (
|
||||
num_decoder_layers if num_decoder_layers is not None else self.num_layers
|
||||
) # default = symmetry
|
||||
self.num_sparse_decoder_layers = num_sparse_decoder_layers
|
||||
|
||||
# This tells us, each how many encoder layer we'll have to set a sparse layer.
|
||||
if self.num_sparse_encoder_layers > 0:
|
||||
self.encoder_sparse_step = self.num_layers // self.num_sparse_encoder_layers
|
||||
else:
|
||||
self.encoder_sparse_step = self.num_layers # HACK: this will create 0 sparse layers
|
||||
|
||||
# This tells us, each how many encoder layer we'll have to set a sparse layer.
|
||||
if self.num_sparse_decoder_layers > 0:
|
||||
self.decoder_sparse_step = self.num_decoder_layers // self.num_sparse_decoder_layers
|
||||
else:
|
||||
self.decoder_sparse_step = self.num_decoder_layers # HACK: this will create 0 sparse layers
|
||||
|
||||
self.num_heads = num_heads
|
||||
self.router_type = router_type
|
||||
self.num_experts = num_experts
|
||||
self.expert_capacity = expert_capacity
|
||||
self.router_bias = router_bias
|
||||
self.router_jitter_noise = router_jitter_noise
|
||||
if router_dtype not in ["float32", "float16", "bfloat16"]:
|
||||
raise ValueError(f"`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}")
|
||||
self.router_dtype = router_dtype
|
||||
|
||||
self.router_ignore_padding_tokens = router_ignore_padding_tokens
|
||||
self.relative_attention_num_buckets = relative_attention_num_buckets
|
||||
self.relative_attention_max_distance = relative_attention_max_distance
|
||||
|
||||
self.dropout_rate = dropout_rate
|
||||
self.layer_norm_epsilon = layer_norm_epsilon
|
||||
self.initializer_factor = initializer_factor
|
||||
self.feed_forward_proj = feed_forward_proj
|
||||
self.use_cache = use_cache
|
||||
self.add_router_probs = add_router_probs
|
||||
|
||||
self.router_z_loss_coef = router_z_loss_coef
|
||||
self.router_aux_loss_coef = router_aux_loss_coef
|
||||
|
||||
act_info = self.feed_forward_proj.split("-")
|
||||
self.dense_act_fn = act_info[-1]
|
||||
self.is_gated_act = act_info[0] == "gated"
|
||||
|
||||
if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
|
||||
raise ValueError(
|
||||
f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
|
||||
"Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
|
||||
"'gated-gelu' or 'relu'"
|
||||
)
|
||||
|
||||
# for backwards compatibility
|
||||
if feed_forward_proj == "gated-gelu":
|
||||
self.dense_act_fn = "gelu_new"
|
||||
|
||||
super().__init__(
|
||||
pad_token_id=pad_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
is_encoder_decoder=is_encoder_decoder,
|
||||
**kwargs,
|
||||
)
|
||||
@ -0,0 +1,193 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import torch
|
||||
from tensorflow.io import gfile
|
||||
|
||||
import tensorstore as ts
|
||||
from flax import serialization
|
||||
from flax.traverse_util import flatten_dict, unflatten_dict
|
||||
from transformers.modeling_utils import dtype_byte_size
|
||||
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
|
||||
rename_keys,
|
||||
)
|
||||
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
|
||||
from transformers.utils.hub import convert_file_size_to_int
|
||||
|
||||
|
||||
def rename_base_flax_keys(flax_key_tuple, flax_tensor):
|
||||
"""
|
||||
Post renaming of basic JAX keys to pytorch.
|
||||
"""
|
||||
if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
|
||||
# expert layer
|
||||
flax_key_tuple = flax_key_tuple[:-1] + ("weight",)
|
||||
flax_tensor = torch.permute(flax_tensor, (0, 2, 1))
|
||||
elif flax_key_tuple[-1] == "kernel" and ".".join(flax_key_tuple):
|
||||
# linear layer
|
||||
flax_key_tuple = flax_key_tuple[:-1] + ("weight",)
|
||||
flax_tensor = flax_tensor.T
|
||||
elif flax_key_tuple[-1] in ["scale", "embedding"]:
|
||||
flax_key_tuple = flax_key_tuple[:-1] + ("weight",)
|
||||
|
||||
return flax_key_tuple, flax_tensor
|
||||
|
||||
|
||||
def get_key_and_tensorstore_dict(layer, checkpoint_info, switch_checkpoint_path):
|
||||
if "metadata" in layer:
|
||||
split_layer = layer.split("metadata")
|
||||
curr_real_layer_name = "".join(split_layer[0])[:-1]
|
||||
split_layer = [tuple(("metadata" + split_layer[1]).split("/"))]
|
||||
elif "kvstore" in layer:
|
||||
split_layer = layer.split("kvstore")
|
||||
curr_real_layer_name = "".join(split_layer[0])[:-1]
|
||||
split_layer = [tuple(("kvstore" + split_layer[1]).split("/"))]
|
||||
|
||||
else:
|
||||
split_layer = layer.split("/")
|
||||
curr_real_layer_name = "/".join(split_layer[:-1])
|
||||
split_layer[-1] = (split_layer[-1],)
|
||||
|
||||
if "kvstore/path" in layer:
|
||||
content = f"{switch_checkpoint_path}/{checkpoint_info[layer]}"
|
||||
elif "kvstore/driver" in layer:
|
||||
content = "file"
|
||||
else:
|
||||
content = checkpoint_info[layer]
|
||||
|
||||
return curr_real_layer_name, split_layer, content
|
||||
|
||||
|
||||
def rename_and_save_block(current_block, save_path):
|
||||
current_block = rename_keys(current_block)
|
||||
new_current_block = {}
|
||||
for k, v in current_block.items():
|
||||
new_current_block[k.replace("/", ".")] = v
|
||||
current_block = new_current_block
|
||||
torch.save(current_block, save_path)
|
||||
|
||||
|
||||
def shard_on_the_fly(switch_checkpoint_path, dump_path, max_shard_size, dtype, weights_name: str = WEIGHTS_NAME):
|
||||
max_shard_size = convert_file_size_to_int(max_shard_size)
|
||||
sharded_state_dicts = []
|
||||
current_block = {}
|
||||
current_block_size = 0
|
||||
total_size = 0
|
||||
|
||||
os.makedirs(dump_path, exist_ok=True)
|
||||
with gfile.GFile(switch_checkpoint_path + "/checkpoint", "rb") as fp:
|
||||
checkpoint_info = serialization.msgpack_restore(fp.read())["optimizer"]["target"]
|
||||
checkpoint_info = flatten_dict(checkpoint_info, sep="/")
|
||||
|
||||
all_layers = {}
|
||||
for layer in checkpoint_info.keys():
|
||||
curr_real_layer_name, split_layer, content = get_key_and_tensorstore_dict(
|
||||
layer, checkpoint_info, switch_checkpoint_path
|
||||
)
|
||||
if curr_real_layer_name in all_layers:
|
||||
all_layers[curr_real_layer_name][split_layer[-1]] = content
|
||||
else:
|
||||
all_layers[curr_real_layer_name] = {split_layer[-1]: content}
|
||||
|
||||
for key in all_layers.keys():
|
||||
# open tensorstore file
|
||||
raw_weights = ts.open(unflatten_dict(all_layers[key])).result().read().result()
|
||||
raw_weights = torch.tensor(raw_weights)
|
||||
weight_size = raw_weights.numel() * dtype_byte_size(raw_weights.dtype)
|
||||
|
||||
# use the renaming pattern from the small conversion scripts
|
||||
key, raw_weights = rename_base_flax_keys(tuple(key.split("/")), raw_weights)
|
||||
key = "/".join(key)
|
||||
|
||||
# If this weight is going to tip up over the maximal size, we split.
|
||||
if current_block_size + weight_size > max_shard_size:
|
||||
save_path = os.path.join(
|
||||
dump_path, weights_name.replace(".bin", f"-{len(sharded_state_dicts)+1:05d}-of-???.bin")
|
||||
)
|
||||
rename_and_save_block(current_block, save_path)
|
||||
sharded_state_dicts.append(current_block.keys())
|
||||
del current_block
|
||||
current_block = {}
|
||||
current_block_size = 0
|
||||
|
||||
current_block[key] = raw_weights.to(getattr(torch, dtype))
|
||||
current_block_size += weight_size
|
||||
total_size += weight_size
|
||||
|
||||
# Add the last block
|
||||
save_path = os.path.join(dump_path, weights_name.replace(".bin", f"-{len(sharded_state_dicts)+1:05d}-of-???.bin"))
|
||||
rename_and_save_block(current_block, save_path)
|
||||
sharded_state_dicts.append(current_block.keys())
|
||||
|
||||
# If we only have one shard, we return it
|
||||
if len(sharded_state_dicts) == 1:
|
||||
return {weights_name: sharded_state_dicts[0]}, None
|
||||
|
||||
# Otherwise, let's build the index
|
||||
weight_map = {}
|
||||
shards = {}
|
||||
for idx, shard in enumerate(sharded_state_dicts):
|
||||
shard_file = weights_name.replace(
|
||||
".bin", f"-{idx+1:05d}-of-{len(sharded_state_dicts):05d}.bin"
|
||||
) # len(sharded_state_dicts):05d}
|
||||
temp_filename = os.path.join(dump_path, weights_name.replace(".bin", f"-{idx+1:05d}-of-???.bin"))
|
||||
os.rename(temp_filename, os.path.join(dump_path, shard_file))
|
||||
shards[shard_file] = shard
|
||||
for key in shard:
|
||||
weight_map[key] = shard_file
|
||||
|
||||
# Add the metadata
|
||||
metadata = {"total_size": total_size}
|
||||
index = {"metadata": metadata, "weight_map": weight_map}
|
||||
|
||||
with open(os.path.join(dump_path, WEIGHTS_INDEX_NAME), "w", encoding="utf-8") as f:
|
||||
content = json.dumps(index, indent=2, sort_keys=True) + "\n"
|
||||
f.write(content)
|
||||
|
||||
return metadata, index
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
# Required parameters
|
||||
parser.add_argument(
|
||||
"--switch_t5x_checkpoint_path",
|
||||
default="/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600",
|
||||
type=str,
|
||||
required=False,
|
||||
help="Path to a directory containing a folder per layer. Follows the original Google format.",
|
||||
)
|
||||
parser.add_argument("--max_shard_size", default="10GB", required=False, help="Max shard size")
|
||||
parser.add_argument("--dtype", default="bfloat16", type=str, required=False, help="dtype of the saved model")
|
||||
parser.add_argument(
|
||||
"--pytorch_dump_folder_path",
|
||||
default="/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted",
|
||||
type=str,
|
||||
required=False,
|
||||
help="Path to the output pytorch model.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
shard_on_the_fly(
|
||||
args.switch_t5x_checkpoint_path,
|
||||
args.pytorch_dump_folder_path,
|
||||
args.max_shard_size,
|
||||
args.dtype,
|
||||
)
|
||||
|
||||
|
||||
def sanity_check():
|
||||
from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, T5Tokenizer
|
||||
|
||||
config = SwitchTransformersConfig.from_pretrained("google/switch-base-8")
|
||||
config.save_pretrained("/home/arthur_huggingface_co/transformers/switch_converted")
|
||||
model = SwitchTransformersForConditionalGeneration.from_pretrained(
|
||||
"/home/arthur_huggingface_co/transformers/switch_converted", device_map="auto"
|
||||
)
|
||||
|
||||
tokenizer = T5Tokenizer.from_pretrained("t5-small")
|
||||
text = "A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>."
|
||||
|
||||
input_ids = tokenizer(text, return_tensors="pt").input_ids
|
||||
out = model.generate(input_ids, decoder_start_token_id=0)
|
||||
print(tokenizer.decode(out[0]))
|
||||
@ -0,0 +1,203 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2022 The HuggingFace Inc. team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Convert SwitchTransformersX checkpoints from the original repository to JAX/FLAX model."""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
|
||||
from flax.traverse_util import flatten_dict, unflatten_dict
|
||||
from t5x import checkpoints
|
||||
from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
|
||||
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logging.set_verbosity_info()
|
||||
|
||||
|
||||
# should not include what is already done by the `from_pt` argument
|
||||
MOE_LAYER_NAME_MAPPING = {
|
||||
"/attention/": "/0/SelfAttention/",
|
||||
"/self_attention/": "/0/SelfAttention/",
|
||||
"/encoder_decoder_attention/": "/1/EncDecAttention/",
|
||||
"value": "v",
|
||||
"query": "q",
|
||||
"key": "k",
|
||||
"out": "o",
|
||||
"pre_self_attention_layer_norm": "0/layer_norm",
|
||||
"pre_cross_attention_layer_norm": "1/layer_norm",
|
||||
"pre_attention_layer_norm": "0/layer_norm", # previously 1, but seems wrong
|
||||
"token_embedder": "shared",
|
||||
"encoder_norm": "final_layer_norm",
|
||||
"decoder_norm": "final_layer_norm",
|
||||
"relpos_bias/rel_embedding": "block/0/layer/0/SelfAttention/relative_attention_bias/weight",
|
||||
"router/router_weights/w/": "router/classifier/",
|
||||
"roer/roer_weights/w/": "router/classifier/",
|
||||
"logits_dense": "lm_head",
|
||||
}
|
||||
|
||||
|
||||
def rename_keys(s_dict):
|
||||
# 1. in HF T5, we have block.{x}.layer.{y}. which corresponds to layer.{x} in
|
||||
# the original model
|
||||
keys = list(s_dict.keys())
|
||||
for key in keys:
|
||||
layer_to_block_of_layer = r".*/layers_(\d+)"
|
||||
new_key = key
|
||||
if re.match(layer_to_block_of_layer, key):
|
||||
new_key = re.sub(r"layers_(\d+)", r"block/\1/layer", new_key)
|
||||
|
||||
layer_to_block_of_layer = r"(encoder|decoder)\/"
|
||||
|
||||
if re.match(layer_to_block_of_layer, key):
|
||||
groups = re.match(layer_to_block_of_layer, new_key).groups()
|
||||
if groups[0] == "encoder":
|
||||
new_key = re.sub(r"/mlp/", r"/1/mlp/", new_key)
|
||||
new_key = re.sub(r"/pre_mlp_layer_norm/", r"/1/layer_norm/", new_key)
|
||||
|
||||
elif groups[0] == "decoder":
|
||||
new_key = re.sub(r"/mlp/", r"/2/mlp/", new_key)
|
||||
new_key = re.sub(r"/pre_mlp_layer_norm/", r"/2/layer_norm/", new_key)
|
||||
|
||||
# 2. Convert other classic mappings
|
||||
for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
|
||||
if old_key in new_key:
|
||||
new_key = new_key.replace(old_key, temp_key)
|
||||
|
||||
print(f"{key} -> {new_key}")
|
||||
s_dict[new_key] = s_dict.pop(key)
|
||||
|
||||
if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
|
||||
s_dict["encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"] = s_dict[
|
||||
"encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
|
||||
].T
|
||||
if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
|
||||
s_dict["decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"] = s_dict[
|
||||
"decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
|
||||
].T
|
||||
|
||||
# 3. Take extra care of the EXPERTS layer
|
||||
for key in list(s_dict.keys()):
|
||||
if "expert" in key:
|
||||
|
||||
num_experts = s_dict[key].shape[0]
|
||||
expert_weihts = s_dict[key]
|
||||
for idx in range(num_experts):
|
||||
s_dict[key.replace("expert/", f"experts/expert_{idx}/")] = expert_weihts[idx]
|
||||
print(f"{key} -> {key.replace('expert/', f'experts/expert_{idx}/')}")
|
||||
|
||||
s_dict.pop(key)
|
||||
|
||||
return s_dict
|
||||
|
||||
|
||||
GIN_TO_CONFIG_MAPPING = {
|
||||
"NUM_ENCODER_LAYERS": "num_layers",
|
||||
"NUM_DECODER_LAYERS": "num_decoder_layers",
|
||||
"NUM_HEADS": "num_heads",
|
||||
"HEAD_DIM": "d_kv",
|
||||
"EMBED_DIM": "d_model",
|
||||
"MLP_DIM": "d_ff",
|
||||
"NUM_SELECTED_EXPERTS": "num_selected_experts",
|
||||
"NUM_ENCODER_SPARSE_LAYERS": "num_sparse_encoder_layers",
|
||||
"NUM_DECODER_SPARSE_LAYERS": "num_sparse_decoder_layers",
|
||||
"dense.MlpBlock.activations": "feed_forward_proj",
|
||||
}
|
||||
|
||||
|
||||
def convert_gin_to_config(gin_file, num_experts):
|
||||
# Convert a google style config to the hugging face fromat
|
||||
import regex as re
|
||||
|
||||
with open(gin_file, "r") as f:
|
||||
raw_gin = f.read()
|
||||
|
||||
regex_match = re.findall(r"(.*) = ([0-9.]*)", raw_gin)
|
||||
args = {}
|
||||
for param, value in regex_match:
|
||||
if param in GIN_TO_CONFIG_MAPPING and value != "":
|
||||
args[GIN_TO_CONFIG_MAPPING[param]] = float(value) if "." in value else int(value)
|
||||
|
||||
activation = re.findall(r"(.*activations) = \(\'(.*)\',\)", raw_gin)[0]
|
||||
args[GIN_TO_CONFIG_MAPPING[activation[0]]] = str(activation[1])
|
||||
|
||||
args["num_experts"] = num_experts
|
||||
config = SwitchTransformersConfig(**args)
|
||||
return config
|
||||
|
||||
|
||||
def convert_flax_checkpoint_to_pytorch(
|
||||
flax_checkpoint_path, config_file, gin_file=None, pytorch_dump_path="./", num_experts=8
|
||||
):
|
||||
# Initialise PyTorch model
|
||||
|
||||
print(f"Loading flax weights from : {flax_checkpoint_path}")
|
||||
flax_params = checkpoints.load_t5x_checkpoint(flax_checkpoint_path)
|
||||
|
||||
if gin_file is not None:
|
||||
config = convert_gin_to_config(gin_file, num_experts)
|
||||
else:
|
||||
config = SwitchTransformersConfig.from_pretrained(config_file)
|
||||
|
||||
pt_model = SwitchTransformersForConditionalGeneration(config)
|
||||
|
||||
flax_params = flax_params["target"]
|
||||
flax_params = flatten_dict(flax_params, sep="/")
|
||||
flax_params = rename_keys(flax_params)
|
||||
flax_params = unflatten_dict(flax_params, sep="/")
|
||||
|
||||
# Load the flax params in the PT model
|
||||
load_flax_weights_in_pytorch_model(pt_model, flax_params)
|
||||
|
||||
print(f"Save PyTorch model to {pytorch_dump_path}")
|
||||
pt_model.save_pretrained(pytorch_dump_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
# Required parameters
|
||||
parser.add_argument(
|
||||
"--switch_t5x_checkpoint_path",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
help=(
|
||||
"The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the"
|
||||
" model architecture. If not provided, a `gin_file` has to be provided."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gin_file",
|
||||
default=None,
|
||||
type=str,
|
||||
required=False,
|
||||
help="Path to the gin config file. If not provided, a `config_file` has to be passed ",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--config_name", default=None, type=str, required=False, help="Config name of SwitchTransformers model."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output pytorch model."
|
||||
)
|
||||
parser.add_argument("--num_experts", default=8, type=int, required=False, help="Number of experts")
|
||||
args = parser.parse_args()
|
||||
convert_flax_checkpoint_to_pytorch(
|
||||
args.switch_t5x_checkpoint_path,
|
||||
args.config_name,
|
||||
args.gin_file,
|
||||
args.pytorch_dump_folder_path,
|
||||
args.num_experts,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@ -5064,6 +5064,51 @@ class Swinv2PreTrainedModel(metaclass=DummyObject):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST = None
|
||||
|
||||
|
||||
class SwitchTransformersEncoderModel(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class SwitchTransformersForConditionalGeneration(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class SwitchTransformersModel(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class SwitchTransformersPreTrainedModel(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class SwitchTransformersSparseMLP(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
class SwitchTransformersTop1Router(metaclass=DummyObject):
|
||||
_backends = ["torch"]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
requires_backends(self, ["torch"])
|
||||
|
||||
|
||||
T5_PRETRAINED_MODEL_ARCHIVE_LIST = None
|
||||
|
||||
|
||||
|
||||
0
tests/models/switch_transformers/__init__.py
Normal file
0
tests/models/switch_transformers/__init__.py
Normal file
File diff suppressed because it is too large
Load Diff
@ -20,6 +20,7 @@ from transformers import (
|
||||
LEDConfig,
|
||||
LongT5Config,
|
||||
SummarizationPipeline,
|
||||
SwitchTransformersConfig,
|
||||
T5Config,
|
||||
pipeline,
|
||||
)
|
||||
@ -54,8 +55,8 @@ class SummarizationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseMe
|
||||
)
|
||||
self.assertEqual(outputs, [{"summary_text": ANY(str)}])
|
||||
|
||||
if not isinstance(model.config, (T5Config, LongT5Config, LEDConfig)):
|
||||
# LED, T5, LongT5 can handle it.
|
||||
if not isinstance(model.config, (SwitchTransformersConfig, T5Config, LongT5Config, LEDConfig)):
|
||||
# Switch Transformers, LED, T5, LongT5 can handle it.
|
||||
# Too long.
|
||||
with self.assertRaises(Exception):
|
||||
outputs = summarizer("This " * 1000)
|
||||
|
||||
@ -39,6 +39,7 @@ PRIVATE_MODELS = [
|
||||
"LongT5Stack",
|
||||
"RealmBertModel",
|
||||
"T5Stack",
|
||||
"SwitchTransformersStack",
|
||||
"TFDPRSpanPredictor",
|
||||
]
|
||||
|
||||
|
||||
@ -5,6 +5,7 @@ docs/source/en/autoclass_tutorial.mdx
|
||||
docs/source/en/task_summary.mdx
|
||||
docs/source/en/model_doc/markuplm.mdx
|
||||
docs/source/en/model_doc/speech_to_text.mdx
|
||||
docs/source/en/model_doc/switch_transformers.mdx
|
||||
docs/source/en/model_doc/t5.mdx
|
||||
docs/source/en/model_doc/t5v1.1.mdx
|
||||
docs/source/en/model_doc/byt5.mdx
|
||||
|
||||
Loading…
Reference in New Issue
Block a user