Source-linked AI summary
Scaling Neural Machine Translation
Myle Ott, Sergey Edunov, David Grangier, Michael Auli
TL;DR
NMT training on large datasets can take several days on a single machine, motivating faster scalable methods. The paper combines reduced precision, very large batches, and distributed training, achieving state-of-the-art results including 29.3 BLEU for En–De and 43.2 BLEU for En–Fr.
Problem
Several days may be required to train state-of-the-art NMT models on large datasets using a single machine with 8 GPUs.
Method
The paper combines reduced-precision computation, very large batches with increased learning rates, and synchronous multi-machine parallelization.
Results
The approach achieves state-of-the-art results of 29.3 BLEU on En–De and 43.2 BLEU on En–Fr, with En–De reaching 29.8 BLEU using Paracrawl data.
Takeaways & Limitations
Careful implementation enables nearly 5x faster big-Transformer training on one 8-GPU machine and rapid large-scale NMT training on 128 GPUs.
Takeaways & Limitations
Multi-node parallelization still incurs significant overhead: 16-node training is only approximately 10x faster than 1-node training.
Abstract
from arXiv · showhide
Sequence to sequence learning models still require several days to reach state of the art performance on large benchmark datasets using a single machine. This paper shows that reduced precision and large batch training can speedup training by nearly 5x on a single 8-GPU machine with careful tuning and implementation. On WMT'14 English-German translation, we match the accuracy of Vaswani et al. (2017) in under 5 hours when training on 8 GPUs and we obtain a new state of the art of 29.3 BLEU after training for 85 minutes on 128 GPUs. We further improve these results to 29.8 BLEU by training on the much larger Paracrawl dataset. On the WMT'14 English-French task, we obtain a state-of-the-art BLEU of 43.2 in 8.5 hours on 128 GPUs.
1 Introduction
The paper targets computationally intensive NMT training, studying how distributed efficiency, batch size, and reduced precision affect training. It combines these approaches to substantially shorten training while maintaining or improving translation accuracy.
- Motivation: Several days of training on a machine with 8 highend GPUs motivate improving NMT training efficiency.Scaling across machines introduces stragglers, while larger batches raise questions about optimization and generalization.
- Approach: 65% lower training time is achieved with reduced floating point precision without affecting accuracy.The paper then increases batch size from 25k to over 400k tokens using gradient accumulation.
- Approach: 40% further training-time reduction on one machine comes from large batches and a higher learning rate.Large batches are less data-efficient but can be parallelized.
- Distributed scaling: 90% additional training-time reduction is obtained by parallelizing across 16 machines compared with one machine.This extends the single-machine efficiency improvements to distributed training.
- Results: 29.3 BLEU is reached on WMT’16 En-De in 85 minutes on 128 GPUs, while matching Vaswani et al. accuracy in under 5 hours on 8 GPUs.The same setup reaches 43.2 BLEU on WMT’14 En-Fr in 8.5 hours on 128 GPUs.
2 Related Work
The related work covers reduced-precision computation and distributed neural-network training. The paper uses synchronous data parallelism and differs from approaches that abandon stragglers by retaining all worker computation.
- Reduced precision: Prior work studies reduced numerical precision for neural-network training and inference.This paper relies on half-precision computation with loss scaling to avoid gradient underflow or overflow.
- Distributed training: Distributed neural-network training commonly uses model parallelism or data parallelism.The paper uses data parallelism, keeping a model copy on each worker and distributing batches.
- Distributed training: Synchronous SGD distributes gradient computation across machines before synchronously updating model weights.The paper follows this strategy, which the cited work describes as more efficient than asynchronous SGD.
- Stragglers: Coding-based approaches improve communication efficiency by abandoning stragglers, whereas this paper does not discard workers’ computation.The distinction separates the paper’s strategy from redundancy-based straggler handling.
3 Experimental Setup
Experiments evaluate large Transformer models for English–German and English–French translation using established WMT settings, with additional experiments scaling training data through filtered Paracrawl. Evaluation uses BLEU under specified decoding and reporting conventions.
- Datasets and tasks: Experiments cover En–De and En–Fr translation using WMT training data and established prior-work setups.En–De uses WMT’16 with 4.5M sentence pairs; En–Fr uses WMT’14 with 36M pairs.
- Datasets and tasks: More than 4.5B En–De and 4.2B En–Fr Paracrawl pairs are explored to scale training beyond the standard datasets.The paper filters this large, noisy corpus before training.
- Model: The models use a six-block “big” Transformer encoder and decoder implemented in PyTorch fairseq.Encoder blocks use self-attention and feed-forward layers; decoder blocks additionally use encoder–decoder attention.
- Optimization: Optimization uses Adam with the Vaswani et al. learning-rate schedule, including warmup and inverse-square-root decay.Label smoothing is applied with weight 0.1.
- Hardware: Experiments run on DGX-1 nodes containing 8 NVIDIA V100 GPUs connected by Infiniband.NCCL2 and torch.distributed provide inter-GPU communication.
4 Experiments and Results
The experiments combine reduced precision, larger batches, learning-rate scaling, memory savings, and distributed techniques to accelerate Transformer training while preserving or improving translation quality. These methods yield substantial speedups on one machine and across nodes, with strong BLEU results on WMT and Paracrawl data.
- Large-Batch Training: 36% higher tokens/sec results when gradients are accumulated across multiple batches before each weight update.Accumulation reduces communication and workload variance between workers; larger batches also support increased learning rates and faster convergence per optimization step.
- Large-Batch Training: 5% additional speedup comes from increasing the per-GPU maximum from 3.5k to 5k tokens after reduced precision lowers memory consumption.The larger sub-batches are enabled by the memory savings of reduced-precision computation.
- Distributed Training: 37 minutes falls to 32 minutes when gradient communication overlaps with backward computation during 16-node training.Buffered gradients are synchronized in a background thread while back-propagation continues through the rest of the network.
- Translation Results: 29.3 BLEU for En-De in 1h 25min and 43.2 BLEU for En-Fr in 8h 32min establish new state-of-the-art results on the reported datasets.The experiments use newstest14 test results and exclude settings with additional training data.
5 Analysis of Stragglers
Synchronous distributed training is slowed by variable sub-batch processing times, creating stragglers that force faster workers to wait. The paper analyzes these timing differences and proposes balancing sub-batch processing time across workers.
- 3.5k tokens: token-budgeted sub-batches can contain different sentence counts and maximum source or target lengths.These shape differences contribute to unequal processing times across GPU workers.
- 0.049–0.228 seconds: sub-batch processing times vary substantially across the 44K En-De sub-batches.The mean processing time is 0.11 seconds.
- Similar-shape batches show much less timing variability than batches grouped only by a maximum token budget.The cited example considers batches with 23 ≤ src len ≈ tgt len ≤ 27.
- Gradient accumulation over multiple sub-batches reduces variance between workers before model updates.This is presented as a way to mitigate stragglers without assigning every worker similarly shaped batches.
- Assigning all workers similarly shaped sub-batches increases gradient variance across batches and can prevent convergence to validation perplexity 4.32.The paper therefore considers shape-based batching a naive solution despite its timing benefit.
6 Conclusions
The paper combines reduced precision, very large batches, and larger learning rates to accelerate Transformer training on single- and multi-node GPU systems. It reports faster training alongside state-of-the-art results on English–French and English–German translation, while noting substantial multi-node overhead remains.
- Nearly 5x: the implementation speeds training of a big Transformer model on one machine with 8 GPUs.The approach investigates lower precision computation, batches up to 400k tokens, and larger learning rates.
- 43.2 BLEU vs. 41.5: WMT’14 En-Fr improves over Shaw et al. (2018) with training in less than 9 hours on 128 GPUs.
- 29.3 BLEU vs. 29.2: WMT’14 En-De reaches a higher score than Shaw et al. (2018) in 85 minutes on 128 GPUs.Scaling the training set with Paracrawl further improves BLEU to 29.8.
- 16-node training is only ∼10x faster than 1-node training, indicating significant overhead in multi-node parallelization.The paper identifies better batching and communication strategies as future work.