The Sockeye 2 Neural Machine Translation Toolkit at AMTA 2020

Introduction

Sockeye (Hieber et al.,, 2017) is a versatile toolkit for research in the fast-moving field of NMT. Since the initial release, it has been used in at least 25 scientific publications, including winning submissions to WMT evaluations (Schamper et al.,, 2018). Sockeye also powers Amazon Translate, showing industrial-strength performance in addition to the flexibility needed in academic environments. Moreover, we are excited to see that hardware manufacturers are contributing to optimizing MXNet (Chen et al.,, 2015) and Sockeye for speed. Intel has demonstrated large performance gains for Sockeye inference on Intel Skylake processors.https://www.intel.ai/amazing-inference-performance-with-intel-xeon-scalable-processors/#gs.wrgsji NVIDIA is working on significant performance improvements for Sockeye’s Transformer (Vaswani et al.,, 2017) implementation through fused operators and an optimized beam search. This paper discusses Sockeye 2’s streamlined Gluon implementation (§2), support for state of the art architectures and efficient decoding (§3), and improved model training (§4).

Gluon Implementation

Sockeye 2 adopts Gluon, the latest and preferred API of MXNet. Gluon simplifies the code while improving overall performance. Developers can define building blocks of neural network architectures as Python classes and seamlessly switch between eager execution for step-by-step debugging and cached computation graphs for maximum performance. Migration to Gluon significantly simplifies training and inference code in Sockeye 2, reducing the overall number of lines of Python code by 25%. The hybridized Gluon transformer implementation in Sockeye 2 improves training speed by 14%, compared to Sockeye.

Focus on State of the Art Models

Due to the success of self-attentional models, we concentrate development of Sockeye 2 on the Transformer architecture (Vaswani et al.,, 2017). Our starting point is the “base” transformer with 6 encoder and decoder layers, model dimensionality of 512, and feed-forward layer size of 2048. An exploration of different encoder and decoder depths shows that deep encoders with shallow decoders are competitive in BLEU and significantly faster for decoding. Table 1 shows results with different numbers of encoder and decoder layers, denoted by $x$:$y$ where $x$ is the number of encoder layers and $y$ the number of decoder layers. For FI-EN and EN-FI, the 20:2 model outperforms both the 6:6 model and the 10:10 model in terms of BLEU. The 20:2 model also has roughly half the decoding latency of the 6:6 model and roughly one third the latency of the 10:10 model. The relative efficiency of encoder versus decoder layers can be attributed to (1) the ability to parallelize across input tokens, (2) attention to only input tokens, and (3) not needing to run beam search on the source side.

Sockeye supports source factors in the spirit of Sennrich and Haddow, (2016), additional representations that are combined with word embeddings prior to the first encoder layer. In Sockeye 2, we improve source factor support by allowing different types of embedding combinations (concatenation, summation, or average), as well as weight sharing between source factor and word embeddings.

As an example application, we use source factors to represent input case. Variations in case pose a challenge for machine translation systems as different orthographic variations are considered to be independent by the translation model (e.g., “case” is different from “Case” and both are different from “CASE”). We address these variations by lowercasing the input and encoding the original case information as a source factor (“lowercase”, “capitalized”, “all uppercase” or “mixed”). We refer to this method as “SF-case”. An alternative is to lowercase and include the original cased word itself as a source factor, which we refer to as “SF-word”. As the original and lowercased versions of many words will be the same, it is useful to share the embeddings, a variant we refer to as “SF-word-share”.

To evaluate the robustness of these strategies, we modify test sets by either entirely lowercasing, entirely uppercasing, or capitalizing the first character of each word. We compare a baseline model that was trained on cased input (no source factors) against all “SF-*” methods. The factored models also use BPE type factors as introduced by Sennrich and Haddow, (2016). Models use the 20:2 transformer architecture and training settings described in §3. Shown in Table 2, encoding case information with source factors is an effective way to improve robustness against case variation with the two versions of “SF-case” performing best.

2 Quantization for Inference

Sockeye 2 now supports 8-bit quantized matrix multiplication (Quinn and Ballesteros,, 2018) on CPUs based on the intgemm library.https://github.com/kpu/intgemm By scaling values such that 127 corresponds to the maximum absolute value found in a tensor, matrix multiplication can be conducted with 8-bit integer representations in place of the default 32-bit floating-point representations without significant degradation of overall model accuracy. Parameters can either be quantized offline and stored in a smaller model file or quantized on the fly at loading time. Activations are quantized on the fly while other operators that consume far less runtime remain as 32-bit floats.

Latency-sensitive applications typically run with batch size 1 and small beam sizes, leaving little opportunity for batch parallelism. Instead, matrix multiplication parallelizes over outputs of a layer. To reduce latency, matrix multiplication and quantization are both parallelized with OpenMP.https://www.openmprtl.org Layer outputs can be computed independently and the layer size is typically much larger than the batch size. Parallelizing over layer inputs would require summing across threads.

Shown in Table 3, quantization significantly reduces non-batched decoding times with minimal effect on BLEU scores. Improvement is most pronounced when running on a single CPU core while models using up to 4 cores still see a significant benefit.For 1 and 2 cores, we set the number of OpenMP threads to 1 and 2 respectively. For 4 cores, we set the number of OpenMP threads to 3 for best interaction with MXNet’s own parallelization over operators.

Training Improvements

Sockeye 2 significantly accelerates training with Horovodhttps://github.com/horovod/horovod integration (Sergeev and Balso,, 2018) and MXNet’s automatic mixed precision (AMP). Horovod extends synchronous training to any number of GPUs (including across nodes) while AMP automatically detects and converts parts of the model that can run in FP16 mode without loss of quality. These methods also require additional computation per update (synchronizing data across distributed GPUs and checking reduced precision operations for overflow). This overhead can be amortized by significantly increasing the effective batch size; gradients are aggregated per-GPU for several batches, then combined and checked for overflow for a single parameter update. In practice, scaling the effective batch size by $N$, the learning rate by $\sqrt{N}$ (Krizhevsky,, 2014), and leaving other hyper parameters unchanged works well for batches of up to 260K tokens.

Sockeye also provides a data-driven alternative to the popular “inverse square root” learning schedule used by Vaswani et al., (2017) and Ott et al., (2018). Termed “plateau-reduce”, this scheduler keeps the same learning rate until validation perplexity does not increase for several checkpoints, at which time it reduces the learning rate and rewinds all model and optimizer parameters to the best previous point. Training concludes when validation perplexity reaches an extended plateau. In a WMT19 benchmark (Barrault et al.,, 2019), plateau-reduce training produces stronger models in slightly less time than the setup described by Ott et al., (2018). The results are presented in Table 4 where all values are averages over 3 independent training runs with different random initializations and all models train until validation perplexity reaches a plateau.

The relevant hyper parameters for Sockeye 2’s large batch training are an effective batch size of 262,144 tokens, a learning rate of 0.00113 with 2000 warmup steps and a reduce rate of 0.9, a checkpoint interval of 125 steps, and learning rate reduction after 8 checkpoints without improvement. After an extended plateau of 60 checkpoints, the 8 checkpoints with the lowest validation perplexity are averaged to produce the final model parameters. While Horovod enables scaling to any number of GPUs, we find that training on 8 GPUs on a single node still delivers the best value when considering both speed and cost.

Licensing and availability

Sockeye 2 is availablehttps://github.com/awslabs/sockeye under the Apache 2.0 license. It includes Docker builds to easily run training or inference with all of the latest features on any supported platform.

Conclusion

Sockeye 2 provides out-of-the-box support for quickly training strong Transformer models for research or production. Extensive configuration options and the simplified Gluon code base enable rapid development and experimentation. As an open source project, we invite the community to contribute their ideas to Sockeye 2 and hope that the new programming model and various performance improvements enable others to conduct effective and successful research.

References