1 Introduction↩︎

Neural sequence-to-sequence models have been successful on a variety of text generation tasks, including machine translation, abstractive document summarization, and language modeling. Accordingly, both researchers and industry professionals can benefit from a fast and easily extensible sequence modeling toolkit.

There are several toolkits with similar basic functionality, but they differ in focus area and intended audiences. For example, OpenNMT [1] is a community-built toolkit written in multiple languages with an emphasis on extensibility. MarianNMT [2] focuses on performance and the backend is written in C++ for fast automatic differentiation. OpenSeq2Seq [3] provides reference implementations for fast distributed and mixed precision training. Tensor2tensor [4] and Sockeye [5] focus on production-readiness.

In this paper, we present fairseq, a sequence modeling toolkit written in PyTorch that is fast, extensible, and useful for both research and production. fairseq features: (i) a common interface across models and tasks that can be extended with user-supplied plug-ins (§2); (ii) efficient distributed and mixed precision training, enabling training over datasets with hundreds of millions of sentences on current hardware (§3); (iii) state-of-the-art implementations and pretrained models for machine translation, summarization, and language modeling (§4); and (iv) optimized inference with multiple supported search algorithms, including beam search, diverse beam search [6], and top-k sampling. fairseq is distributed with a BSD license and is available on GitHub at https://github.com/pytorch/fairseq.

2 Design↩︎

2.0.0.1 Extensibility.

fairseq can be extended through five types of user-supplied plug-ins, which enable experimenting with new ideas while reusing existing components as much as possible.

2.0.0.2 Models

define the neural network architecture and encapsulate all learnable parameters. Models extend the BaseFairseqModel class, which in turn extends torch.nn.Module. Thus any fairseq model can be used as a stand-alone module in other PyTorch code. Models can additionally predefine named architectures with common network configurations (e.g., embedding dimension, number of layers, etc.). We also abstracted the methods through which the model interacts with the generation algorithm, e.g., beam search, through step-wise prediction. This isolates model implementation from the generation algorithm.

2.0.0.3 Criterions

compute the loss given the model and a batch of data, roughly: loss = criterion(model, batch). This formulation makes criterions very expressive, since they have complete access to the model. For example, a criterion may perform on-the-fly generation to support sequence-level training [7] or online backtranslation [8], [9]. Alternatively, in a mixture-of-experts model, a criterion may implement EM-style training and backpropagate only through the expert that produces the lowest loss [10].

2.0.0.4 Tasks

store dictionaries, provide helpers for loading and batching data and define the training loop. They are intended to be immutable and primarily interface between the various components. We provide tasks for translation, language modeling, and classification.

2.0.0.5 Optimizers

update the model parameters based on the gradients. We provide wrappers around most PyTorch optimizers and an implementation of Adafactor [11], which is a memory-efficient variant of Adam.

2.0.0.6 Learning Rate Schedulers

update the learning rate over the course of training. We provide several popular schedulers, e.g., the inverse square-root scheduler from [12] and cyclical schedulers based on warm restarts [13].

2.0.0.7 Reproducibility and forward compatibility.

fairseq includes features designed to improve reproducibility and forward compatibility. For example, checkpoints contain the full state of the model, optimizer and dataloader, so that results are reproducible if training is interrupted and resumed. fairseq also provides forward compatibility, i.e., models trained using old versions of the toolkit will continue to run on the latest version through automatic checkpoint upgrading.

3 Implementation↩︎

fairseq is implemented in PyTorch and it provides efficient batching, mixed precision training, multi-GPU as well as multi-machine training.

3.0.0.1 Batching.

There are multiple strategies to batch input and output sequence pairs [14]. fairseq minimizes padding within a mini-batch by grouping source and target sequences of similar length. The content of each mini-batch stays the same throughout training, however mini-batches themselves are shuffled randomly every epoch. When training on more than one GPU or machine, then the mini-batches for each worker are likely to differ in the average sentence length which results in more representative updates.

3.0.0.2 Multi-GPU training.

fairseq uses the NCCL2 library and torch.distributed for inter-GPU communication. Models are trained in a synchronous optimization setup where each GPU has a copy of the model to process a sub-batch of data after which gradients are synchronized between GPUs; all sub-batches constitute a mini-batch. Even though sub-batches contain a similar number of tokens, we still observe a high variance in processing times. In multi-GPU or multi-machine setups, this results in idle time for most GPUs while slower workers are finishing their work (Figure 1 (a)). fairseq mitigates the effect of stragglers by overlapping gradient synchronization between workers with the backward pass and by accumulating gradients over multiple mini-batches for each GPU [15].

Overlapping gradient synchronization starts to synchronize gradients of parts of the network when they are computed. In particular, when the gradient computation for a layer finishes, fairseq adds the result to a buffer. When the size of the buffer reaches a predefined threshold, the gradients are synchronized in a background thread while back-propagation continues as usual (Figure 1 (b)). Next, we accumulate gradients for multiple sub-batches on each GPU which reduces the variance in processing time between workers since there is no need to wait for stragglers after each sub-batch (Figure 1 (c)). This also increases the effective batch size but we found that models can still be trained effectively [15].

3.0.0.3 Mixed precision.

Recent GPUs enable efficient half precision floating point (FP16) computation. fairseq provides support for both full precision (FP32) and FP16 at training and inference. We perform all forward-backward computations as well as the all-reduce for gradient synchronization between workers in FP16. However, the parameter updates remain in FP32 to preserve accuracy. fairseq implements dynamic loss scaling [16] in order to avoid underflows for activations and gradients because of the limited precision offered by FP16. This scales the loss right after the forward pass to fit into the FP16 range while the backward pass is left unchanged. After the FP16 gradients are synchronized between workers, we convert them to FP32, restore the original scale, and update the weights.

3.0.0.4 Inference.

fairseq provides fast inference for non-recurrent models [12], [17]–[19] through incremental decoding, where the model states of previously generated tokens are cached in each active beam and re-used. This can speed up a naïve implementation without caching by up to an order of magnitude, since only new states are computed for each token. For some models, this requires a component-specific caching implementation, e.g., multi-head attention in the Transformer architecture.

During inference we build batches with a variable number of examples up to a user-specified number of tokens, similar to training. fairseq also supports inference in FP16 which increases decoding speed by 54% compared to FP32 with no loss in accuracy (Table 1).

4 Applications↩︎

fairseq has been used in many applications, such as machine translation [7], [8], [17], [19]–[22], language modeling [23], [24], abstractive document summarization [25]–[27], story generation [18], [28], error correction [29], multilingual sentence embeddings [30], and dialogue [31], [32].

4.1 Machine translation↩︎

We provide reference implementations of several popular sequence-to-sequence models which can be used for machine translation, including LSTM [33], convolutional models [17], [19] and Transformer [12].

We evaluate a “big" Transformer encoder-decoder model on two language pairs, WMT English to German (En–De) and WMT English to French (En–Fr). For En–De we replicate the setup of [12] which relies on WMT’16 for training with 4.5M sentence pairs, we validate on newstest13 and test on newstest14. The 32K vocabulary is based on a joint source and target byte pair encoding (BPE; [34]). For En–Fr, we train on WMT’14 and borrow the setup of [17] with 36M training sentence pairs. We use newstest12+13 for validation and newstest14 for test. The 40K vocabulary is based on a joint source and target BPE.

We measure case-sensitive tokenized BLEU with multi-bleu [35] and de-tokenized BLEU with SacreBLEU¹ [36]. All results use beam search with a beam width of 4 and length penalty of 0.6, following [12]. fairseq results are summarized in Table 2. We reported improved BLEU scores over [12] by training with a bigger batch size and an increased learning rate [15].

4.2 Language modeling↩︎

fairseq supports language modeling with gated convolutional models [23] and Transformer models [12]. Models can be trained using a variety of input and output representations, such as standard token embeddings, convolutional character embeddings [44], adaptive softmax [45], and adaptive inputs [24]. We also provide tutorials and pre-trained models that replicate the results of [23] and [24] on WikiText-103 and the One Billion Word datasets.

We evaluate two Transformer language models, which use only a decoder network and adaptive input embeddings, following [24]. The first model has 16 blocks, inner dimension 4K and embedding dimension 1K; results on WikiText-103 are in Table 3. The second model has 24 blocks, inner dimension 8K and embedding dimension 1.5K; results on the One Billion Word benchmark are in Table 4.

4.3 Abstractive document summarization↩︎

Next, we experiment with abstractive document summarization where we use a base Transformer to encode the input document and then generate a summary with a decoder network. We use the CNN-Dailymail dataset [46], [47] of news articles paired with multi-sentence summaries. We evaluate on the full-text version with no entity anonymization [48]; we truncate articles to 400 tokens [48]. We use BPE with 30K operations to form our vocabulary following [25]. To evaluate, we use the standard rouge metric [49] and report rouge-1, rouge-2, and rouge-l. To generate summaries, we follow standard practice in tuning the minimum output length and disallow repeating the same trigram [50]. Table 5 shows results of fairseq. We also consider a configuration where we input pre-trained language model representations to the encoder network and this language model was trained on newscrawl and CNN-Dailymail, totalling 193M sentences.

5 Conclusion↩︎

We presented fairseq, a fast, extensible toolkit for sequence modeling that is scalable and suitable for many applications. In the future, we will continue the development of the toolkit to enable further research advances.

Acknowledgements↩︎

We thank Jonas Gehring for writing the original Lua/Torch version of fairseq.

References↩︎