TorchScale is a PyTorch library that allows researchers and developers to scale up Transformers efficiently and effectively.
It has the implementation of fundamental research to improve modeling generality and capability, as well as training stability and efficiency of scaling Transformers.
It takes only several lines of code to create a model with the above fundamental research features enabled. Here is how to quickly obtain a BERT-like encoder:
```python
>>> from torchscale.architecture.config import EncoderConfig
>>> from torchscale.architecture.encoder import Encoder
>>> config = EncoderConfig(vocab_size=64000)
>>> model = Encoder(config)
>>> print(model)
```
We also support the `Decoder` architecture and the `EncoderDecoder` architecture:
```python
# Creating a decoder model
>>> from torchscale.architecture.config import DecoderConfig
>>> from torchscale.architecture.decoder import Decoder
>>> config = DecoderConfig(vocab_size=64000)
>>> decoder = Decoder(config)
>>> print(decoder)
# Creating a encoder-decoder model
>>> from torchscale.architecture.config import EncoderDecoderConfig
>>> from torchscale.architecture.encoder_decoder import EncoderDecoder
We plan to provide more examples regarding different tasks (e.g. vision pretraining and speech recognition) and various deep learning toolkits (e.g. [DeepSpeed](https://github.com/microsoft/DeepSpeed) and [Megatron-LM](https://github.com/NVIDIA/Megatron-LM)). Any comments or PRs are welcome!
## Results
### Stability Evaluation
<palign="center">
<imgsrc="./assets/convergence.png"width="800"/>
</p>
The training curve is smooth by using TorchScale, while the baseline Transformer cannot converge.
### Scaling-up Experiments
<palign="center">
<imgsrc="./assets/scaling_curve.png"width="800"/>
</p>
TorchScale supports arbitrary depths and widths, successfully scaling-up the models without pain.
## Acknowledgments
Some implementations in TorchScale are either adapted from or inspired by the [FairSeq](https://github.com/facebookresearch/fairseq) repository and the [UniLM](https://github.com/microsoft/unilm) repository.
## Citations
If you find this repository useful, please consider citing our work:
```
@article{deepnet,
author = {Hongyu Wang and
Shuming Ma and
Li Dong and
Shaohan Huang and
Dongdong Zhang and
Furu Wei},
title = {{DeepNet}: Scaling Transformers to 1,000 Layers},
journal = {CoRR},
volume = {abs/2203.00555},
year = {2022},
}
```
```
@article{magneto,
author = {Hongyu Wang and
Shuming Ma and
Shaohan Huang and
Li Dong and
Wenhui Wang and
Zhiliang Peng and
Yu Wu and
Payal Bajaj and
Saksham Singhal and
Alon Benhaim and
Barun Patra and
Zhun Liu and
Vishrav Chaudhary and
Xia Song and
Furu Wei},
title = {Foundation Transformers},
journal = {CoRR},
volume = {abs/2210.06423},
year = {2022}
}
```
```
@article{xmoe,
author = {Zewen Chi and
Li Dong and
Shaohan Huang and
Damai Dai and
Shuming Ma and
Barun Patra and
Saksham Singhal and
Payal Bajaj and
Xia Song and
Furu Wei},
title = {On the Representation Collapse of Sparse Mixture of Experts},
