L ego MT 2: Selective Asynchronous Sharded Data Parallel Training for Massive Neural Machine Translation (2025)

Abstract: It is a critical challenge to learn a single model for massive languages. Prior methods focus on increasing the model size and training data size. However, large models are difficult to optimize efficiently even with distributed parallel training and translation capacity can interfere among languages. To address the challenge, we propose LegoMT2, an efficient training approach with an asymmetric multi-way model architecture for massive multilingual neural machine translation. LegoMT2 shards 435 languages into 8 language-centric groups and attributes one local encoder for each group’s languages and a mix encoder-decoder for all languages. LegoMT2 trains the model through local data parallel and asynchronous distributed updating of parameters. LegoMT2 is 16.2 \times faster than the distributed training method for M2M-100-12B (which only for 100 languages) while improving the translation performance by an average of 2.2 BLEU on Flores-101 , especially performing better for low-resource languages .

## introduction recent years have witnessed great success in multilingual neural machine translation (mnmt; @xcite @xcite that uses a single model for translating all directions. to construct an mnmt system that supports high-quality translation for massive directions, many efforts have been put into scaling up the model size and training corpus @xcite . for example, @xcite constructed a 54.5b nllb model to support translation among 200 languages. however, with the increasing model size, training a single model over massive data brings new challenges. specifically, the challenge is twofold: (1) huge training costs. existing distributed training methods still exhibit large communication overhead @xcite ; (2) parameter interference. languages compete for model capacity within an mnmt model. @xcite @xcite is a popular solution to reduce parameter interference, but it also introduces substantial memory and computational overheads. to address these challenges, we propose legomt2, an efficient approach to train massive mnmt. legomt2 consists of three key designs: a proper language-based data sharding scheme, an asymmetric multi-way model architecture, and asynchronous distributed updating. first, legomt2 shards data into 8 carefully designed language groups. under this scheme, we arrange all languages based on the size of the language-centric data (sentence pairs that are from or to a specific language) each group's data is stored on a dedicated set of gpu servers, therefore no moving of training data is needed. second, we design an asymmetric multi-way transformer to alleviate parameter interference. our key insight is the separation of the model used for training and inference and splitting language capacity into different model components. our model consists of a mix encoder-decoder for all languages and one language-specific encoder for each shard. at inference time, it only uses the mix encoder-decoder. third, we design an asynchronous distributed updating algorithm to accelerate the training. our key insight is that at training time, we no longer need to load all model parameters into all workers, thanks to our sharding scheme and model architecture. we only load and train two encoders and one decoder responsible for the shard. we only aggregate a fraction of parameters using a parameter server asynchronously. there is no need to transfer language-specific encoders, thereby reducing communication costs. parameter communication is efficient and does not block the training on local workers. we construct a large-scale mnmt translation dataset to train legomt2. the proposed dataset contains 25b parallel pairs, covering 435 languages and 22,613 translation directions. we use our method to obtain a 1.6b legomt2 model for 435 languages. our experiments on flores-101 show that legomt2 achieves 16.2× speedups over m2m-1001-12b and 2.2 bleu gains over the prior best approach nllb of the similar size. ## related work the most common approach in mnmt is using a single model to handle all translation directions @xcite @xcite @xcite , which has promising generalization abilities by transferring knowledge from high-resources and low-resources. recently, scaling up the size of mnmt models has brought significant quantitative improvements and new qualitative capabilities (m2m-100, @xcite ; inter alia). however, employing a single model for all translation directions @xcite @xcite @xcite @xcite is often subject to capacity bottlenecks and tradeoffs between translation quality and the number of languages @xcite @xcite . to accelerate the training of models on large data with large models, there are some existing acceleration strategies @xcite @xcite . distributed data parallel @xcite is a method for distributed training that creates model copies on each worker generates gradients independently, and shares them to keep the models consistent. fully shared data parallel @xcite is a high-quality solution for training large models; megatron-lm @xcite shows how combining different parallelism methods to scale up to thousands of gpus and models with trillions of parameters. nevertheless, researchers @xcite have observed that there is a trade-off between translation quality and language number when using a single model for inference. based on this, lass @xcite is designed to learn distinct sub-networks for each language direction. and @xcite propose a detachable model architecture further en-hancing the performance of machine translation by language-specific module. ## legomt2 approach we present legomt2 approach to train a single unified model to translate a massive number of languages (@xmath0. previous training methods such as distributed data parallel need to load the full model parameter to each worker's gpus. our approach only selects part of the model and trains them on each worker. these model shards are communicated from and to a parameter server asynchronously (figure 1 ). this reduces the memory consumption during training and therefore enables a larger batch size and lower communication overhead during training. ## analysis of legomt2 why does asynchronous training work? legomt2 introduces asynchronous training to reduce communication delays to accelerate training. each client pulls the latest parameters every α minutes and pushes current parameters into the federated server every β minutes. it represents that all clients do not always enjoy and latest parameters. to prove whether such delay affects final performance, we conduct experiments by using global modules from other clients for inference. figure 2 shows that delayed global parameters do not affect model training. the client can use delayed global parameters from other clients for inference without any performance drops. impact of push/pull (α/β) intervals setting test the effect of different save/load intervals on system performance, i.e. the effect of α and β in the algorithm. theoretically, when α and β are small enough, the localized training by legomt2 is approximately equal to synchronous training. here, we conduct two different settings: 1) α = 10min, β = 20min and 2) α = 20min, β = 40min. we test two settings legomt2 on the flores-101 devtest. figure 3 shows the winning rate of the first configuration over the second. the results show that more frequent updates does improve the performance at the cost of more communication. impact of language/data sharding in this work, we sort languages based on the size of languagecentric data and split languages into different equalsize groups. we adopt this split method because we find that balanced training flows in different clients help multilingual machine translation. in addition, the common strategy of language clustering is by similarity. therefore, we use similarity clustering to construct a baseline. given an mnmt model, here we use the single multilingual model to get language id embedding, then directly apply kmeans (by sklearn) on those embedding. the resulting ## conclusion the typical multilingual neural machine translation is training a single model for all directions with a centralized training schema, which faces many challenges in practice including parameter competition and efficiency problems. in this paper, we propose a new mnmt pre-training framework with federated learning, legomt2. extensive experiments verify the effectiveness of legomt2. it brings 16.2× training speedups and large performance gains. we build a translation system that supports 435 languages, the supported language number outperforming all existing open-source multilingual machine translation systems.