Source-linked AI summary
fairseq: A Fast, Extensible Toolkit for Sequence Modeling
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, Michael Auli
TL;DR
Researchers need a fast, extensible sequence-modeling toolkit spanning text-generation tasks. FAIRSEQ provides a PyTorch-based toolkit with scalable training and optimized inference, including 54% faster FP16 decoding than FP32 without accuracy loss.
Problem
Researchers and industry professionals need a fast, easily extensible toolkit for diverse neural text-generation tasks.
Method
FAIRSEQ provides a PyTorch toolkit with extensible plug-ins, distributed and mixed-precision training, reusable model interfaces, and multiple generation algorithms.
Results
54% faster FP16 decoding than FP32 was achieved with no loss in accuracy, while FAIRSEQ also reported improved BLEU scores over Vaswani et al. (2017).
Takeaways & Limitations
FAIRSEQ is a scalable, fast, extensible toolkit suitable for research and production across many sequence-modeling applications.
Abstract
from arXiv · showhide
fairseq is an open-source sequence modeling toolkit that allows researchers and developers to train custom models for translation, summarization, language modeling, and other text generation tasks. The toolkit is based on PyTorch and supports distributed training across multiple GPUs and machines. We also support fast mixed-precision training and inference on modern GPUs. A demo video can be found at https://www.youtube.com/watch?v=OtgDdWtHvto
1 Introduction
FAIRSEQ is a fast, extensible PyTorch sequence modeling toolkit intended for research and production across text generation tasks. It combines a common extensible interface with efficient training, pretrained models, and optimized inference.
- Motivation: Neural sequence-to-sequence models support machine translation, abstractive document summarization, and language modeling, motivating fast and extensible tooling.The toolkit is positioned as useful to both researchers and industry professionals.
- Related work: Existing toolkits differ in focus and audience, spanning extensibility, C++ performance, and distributed or mixed-precision reference implementations.The introduction contrasts OpenNMT, MarianNMT, and OpenSeq2Seq.
- Contributions: FAIRSEQ provides a common interface across models and tasks that can be extended with user-supplied plug-ins.This interface is part of the toolkit’s extensibility design.
- Contributions: FAIRSEQ enables efficient distributed and mixed-precision training over datasets with hundreds of millions of sentences on current hardware.The training capabilities are described as a core toolkit feature.
- Contributions: FAIRSEQ includes state-of-the-art implementations and pretrained models for machine translation, summarization, and language modeling.These models and implementations target the toolkit’s supported generation tasks.
- Contributions: FAIRSEQ optimizes inference with beam search, diverse beam search, and top-k sampling.Multiple search algorithms are supported for inference.
2 Design
FAIRSEQ is designed around extensible plug-ins that reuse shared components while separating models, generation, training objectives, and data handling. It also supports reproducibility and forward compatibility through complete checkpoints and automatic checkpoint upgrading.
- Extensibility: FAIRSEQ supports five types of user-supplied plug-ins, enabling experimentation with new ideas while reusing existing components.This plug-in design is the central extensibility mechanism described in the section.
- Extensibility: Models define neural architectures and parameters, support named configurations, and isolate model implementation from generation algorithms through step-wise prediction.Models extend BaseFairseqModel and can function as stand-alone modules in other PyTorch code.
- Extensibility: Criterions compute losses with complete access to the model, supporting sequence-level training, online backtranslation, and mixture-of-experts EM-style updates.A criterion can generate data on the fly or backpropagate only through the expert producing the lowest loss.
- Extensibility: Tasks store dictionaries, load and batch data, define the training loop, and interface between components for translation, language modeling, and classification.Tasks are intended to be immutable.
- Reproducibility and forward compatibility: Complete model, optimizer, and dataloader checkpoints support reproducible interruption and resumption, while automatic checkpoint upgrading preserves forward compatibility across toolkit versions.Models trained with older toolkit versions continue to run on the latest version.
3 Implementation
FAIRSEQ is implemented in PyTorch with efficient batching, mixed-precision computation, distributed multi-GPU and multi-machine training, and fast inference. Its implementation reduces padding, overlaps gradient synchronization with back-propagation, preserves FP32 parameter updates during FP16 computation, and accelerates FP16 decoding.
- Batching: FAIRSEQ minimizes padding by grouping source and target sequences of similar length, while shuffling mini-batches randomly each epoch.The content of each mini-batch remains fixed throughout training.
- Multi-GPU training: FAIRSEQ synchronizes gradients across GPUs using NCCL2 and torch.distributed in a synchronous optimization setup.Each GPU processes a sub-batch, then synchronizes gradients; all sub-batches form one mini-batch.
- Multi-GPU training: FAIRSEQ overlaps gradient synchronization with back-propagation and accumulates gradients across multiple sub-batches to reduce processing-time variance and communication time.Gradients are buffered and synchronized in a background thread once the buffer reaches a predefined threshold.
- Mixed precision: FAIRSEQ supports FP32 and FP16 training and inference, performing forward-backward computation and gradient all-reduce in FP16 while retaining FP32 parameter updates.Dynamic loss scaling prevents underflows caused by FP16’s limited precision.
- Inference: 54% faster decoding is achieved with FP16 inference than FP32, with no loss in accuracy.Inference batches contain a variable number of examples up to a user-specified token limit.
4 Applications
FAIRSEQ supports a broad range of sequence-modeling applications through reference implementations of major architectures, including machine translation, language modeling, and abstractive summarization. The toolkit also provides evaluation setups, tutorials, and pretrained models for these tasks.
- FAIRSEQ has been used for machine translation, language modeling, abstractive document summarization, story generation, error correction, multilingual sentence embeddings, and dialogue.
- Machine translation: Reference implementations include LSTM, convolutional, and Transformer sequence-to-sequence models for machine translation.
- Machine translation: FAIRSEQ evaluates a “big” Transformer on WMT English–German and English–French using beam search, and reports improved BLEU scores over Vaswani et al. (2017).The reported decoding uses beam width 4 and length penalty 0.6; improvements came from a bigger batch size and increased learning rate.
- Language modeling: FAIRSEQ supports language modeling with gated convolutional and Transformer models, adaptive softmax, adaptive inputs, and multiple input and output representations.Tutorials and pretrained models replicate results on WikiText-103 and the One Billion Word datasets.
- Abstractive summarization: FAIRSEQ applies a base Transformer encoder-decoder to abstractive summarization on CNN-DailyMail and evaluates generated summaries with ROUGE-1, ROUGE-2, and ROUGE-L.Articles are truncated to 400 tokens, use 30K-operation BPE, and decoding disallows repeated trigrams; a pretrained language-model representation configuration is also evaluated.
5 Conclusion
The paper concludes that FAIRSEQ is a fast, extensible, scalable toolkit for sequence modeling with broad applicability, and that its development will continue to support further research advances.
- FAIRSEQ is presented as a fast and extensible toolkit for sequence modeling.
- The toolkit is described as scalable and suitable for many applications.
- Future development will aim to enable further research advances.