Memory-efficient NLLB-200: Language-specific Expert Pruning of a Massively Multilingual Machine Translation Model

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 3567–3585 July 9-14, 2023 ©2023 Association for Computational Linguistics Memory-efficient NLLB-200: Language-specific Expert Pruning of a Massively Multilingual Machine Translation Model Yeskendir Koishekenov∗1,2 Alexandre Berard1 Vassilina Nikoulina1 1NAVER LABS Europe 2University of Amsterdam {first.last}@naverlabs.com yeskendir.koishekenov@student.uva.nl Abstract The recently released NLLB-200 is a set of multilingual Neural Machine Translation mod- els that cover 202 languages. The largest model is based on a Mixture of Experts architecture and achieves SoTA results across many lan- guage pairs. It contains 54.5B parameters and requires at least four 32GB GPUs just for in- ference. In this work, we propose a pruning method that enables the removal of up to 80% of experts without further finetuning and with a negligible loss in translation quality, which makes it feasible to run the model on a single 32GB GPU. Further analysis suggests that our pruning metrics can identify language-specific experts. 1 Introduction The Transformer (Vaswani et al., 2017) has be- come the dominant modeling paradigm in Natural Language Processing tasks. Many subsequent ad- vances in the field came from increasing the com- putational budget, training data, and model size. Neural Machine Translation was not an exception, where massively multilingual NMT (Aharoni et al., 2019; Fan et al., 2021; Tang et al., 2020; Zhang et al., 2020) demonstrated promising results, while attempting to overcome the curse of multilinguality (Conneau et al., 2019) by scaling up model size. However, increasing the parameter size exacer- bates the cost of training (Yang et al., 2019; Strubell et al., 2019; Patterson et al., 2021) and hurts the memory footprint and inference latency (Dai et al., 2019; Fan et al., 2021; Wang et al., 2022). Sparsely- gated Mixture-of-Experts (MoE) models are an ef- ficient alternative to dense models (Lepikhin et al., 2020; Fedus et al., 2021; Riquelme et al., 2021). For example, Du et al. (2022) demonstrates that an MoE language model results in a 7x larger model compared to GPT-3, but requires only 30% of its *Work done during an internship at NAVER LABS Europe energy for training and half of its FLOPs at infer- ence. Mixture-of-Experts models are neural networks whose set of parameters is partitioned into experts. Contrary to dense models, where all network pa- rameters are used for every input, an MoE model activates different parts of the network, the experts, depending on the input, which is typically done by a gating mechanism at the token level. MoE models are computationally efficient due to expert parallelism (Fedus et al., 2021) across a large num- ber of GPUs, by having each GPU hold a subset of all experts and communicate with the other GPUs when it needs expert outputs for its local batch. In NLLB-2001(Costa-jussà et al., 2022), a load balancing regularizer in the objective function (Shazeer et al., 2017) promotes equal distribution of the tokens across experts. This encourages the model to use all the experts and ensures that all GPUs are used …

Content truncated to 500 words