M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining

1 INTRODUCTION . Foundation models with self-supervised learning on big data have become an emerging paradigm of artificial intelligence systems [ 2 ] , as they mostly possess high transferability to a wide range of downstream tasks and even multiple modalities . The scale of foundation models across domains , including natural language processing , computer vision , and cross-modality representation learning , have been growing tremendously from millions to trillions of parameters [ 7 ; 23 ; 24 ; 26 ; 33 ; 37 ; 3 ; 29 ; 14 ; 11 ; 16 ; 47 ; 46 ; 42 ; 44 ; 38 ] thanks to the concurrent advancement in distributed training framework [ 20 ; 27 ; 31 ; 28 ; 1 ; 12 ; 30 ] and hardware design , and these studies have made a demonstration of the neural scaling law [ 13 ] . However , the training of these transformerbased models incurs high financial costs and even environmental damage due to the massive carbon footprint and thus training extreme-scale models under a decent amount of resources but with high efficiency should be a fundamental goal for both the research and industrial communities to achieve , which promotes the progress of greener AI [ 22 ; 34 ] . Generally there are two tracks of research in large-scale pretraining , dense models and sparse expert models respectively . A typical case of large-scale dense models is GPT-3 [ 3 ] , a 175-billionparameter transformer model trained with 10 , 000 GPUs for months , incurring striking financial and environmental costs . Researchers have been searching for methods to training large-scale models with a decent amount of costs . Solutions include effective management of memory with gradient and optimizer state partitioning [ 27 ] or more efficient model parallelism and pipeline parallelism [ 37 ; 20 ; 12 ] . A series of following studies apply those techniques to realize fast training of 10-billion-parameter transformers with hundreds of GPUs in 1 − 2 months . [ 16 ; 46 ; 44 ; 38 ] Sparse expert models with large model capacity are capable of fast training owing to the combination of data parallelism and expert parallelism [ 35 ; 14 ; 11 ; 32 ] , and it is even accessible to train a 1-trillion-parameter transformer with no more than 500 GPUs [ 42 ] . Be there as it may , a question emerges in our mind : is it possible to train an extreme-scale model with only a decent amount of resources , e.g. , training a 10-trillion-parameter model with around 500 GPUs ? Tackling the problem requires the utilization of external memory except for GPU memory , for instance , CPU memory or even NVMe storage [ 31 ; 28 ] . These methods resolve the problem of high memory footprint , but instead , their extra cost is low training efficiency caused by the frequent swap in-and-out between memories . In this paper , we provide a solution to training large models that require high memory footprint , and we demonstrate a successful practice of pretraining an unprecedented extreme-scale model with over 10 trillion parameters , an order of magnitude larger than the previous state-of-the-arts [ 11 ; 42 ] . The whole pretraining was conducted on solely 512 NVIDIA-V100 GPUs and lasted around 10 days . A simple and effective training strategy called “ Pseudo-to-Real ” enables the sharing and delinking of parameters . This training strategy is compatible with architectures built by stacking layers of an identical structure , including dense models like GPT [ 23 ; 24 ; 3 ] , BERT [ 7 ] , or sparse expert models like M6 [ 16 ; 42 ] . It is essentially a two-stage training . It first trains a relatively small model but with a computation graph of a large one with the utilization of cross-layer parameter sharing , and we name it “ Pseudo Giant ” . Then it builds a correspondingly large model and delinks the parameters of the shared layer for second-stage model initialization . In this way , we achieve fast convergence in the first stage as the training costs much less memory and speeds up with large batches . Parameter sharing that addresses the communication overhead improves training speed as well . The second-stage training is responsible for the final convergence for better performance . We unlock the secret of pretraining an unprecedented extreme-scale model with over 10 trillion parameters on limited resources of 512 GPUs . Compared with the previous M6-T on around 500 GPUs , we do not have a significant increase in computation resources but level up the model scale by an order of magnitude . Besides the application of the “ Pseudo-to-Real ” training strategy , we provide a faster offloading mechanism for both management of CPU memory for parameter storage and utility of GPUs . We successfully train the M6-10T within 10 days to reach strong performance in log perplexity evaluation and outperform the baseline M6-T . Contributions at a glance are below : • We illustrate the training difficulty of extreme-scale models on limited resources and provide a simple but effective solution called “ Pseudo-to-Real ” . Upstream and downstream evaluation demonstrates the effectiveness of the strategy . • We further demonstrate a successful practice of pretraining a 10-trillion-parameter model on 512 GPUs and reach an outstanding performance within 10 days . 2 RELATED WORK . Large-Scale Pretraining In recent years , pretrained language models with growing magnitudes of parameters have been proposed , keeping to raise the validated upper limit of scaling law for model capacity w.r.t the number of parameters [ 13 ] . Earlier milestones of extreme-large models come from GPT-2 [ 24 ] and Megatron-LM [ 37 ] , which demonstrates that scaling the transformer model up to billions of parameters can result in improvement on language modeling benchmarks [ 19 ; 21 ] . TuringNLG [ 33 ] , as a successor , implements a 17-billion-parameter transformer and achieves further lower perplexity . Similar phenomena are also observed on classification language tasks by T5 model [ 26 ] . The GPT-3 [ 3 ] pushes the boundary of model scale to 1,750 billion parameters and demonstrates its striking effectiveness on downstream tasks in even zero-shot settings . Furthermore , large-scale pretraining has recently demonstrated success in other fields , including pretraining on other languages or cross-lingual pretraining [ 41 ; 46 ; 44 ; 6 ] , cross-modal pretraining [ 25 ; 29 ; 16 ; 8 ; 48 ] and code generation [ 4 ] . Along with the benefit from increasing the model scale , the concern of unaffordable pretraining cost in time , computation resource and energy keeps emerging [ 22 ; 2 ] , resulting in the strong demand for more efficient and greener large-scale pretraining [ 34 ] . Methods to Train Larger Models Faster Our work is about designing algorithm to train extremescale models efficiently . Researchers have been demonstrating different types of methods to reach this objective . To reduce the computational cost during training , some studies introduce sparsity to the model . Mixture-of-Experts ( MoE ) [ 35 ] was proposed and was reintroduced in Mesh Tensorflow [ 36 ] . It shows that MoE with sparsity can significantly improve the model scale efficiently without increasing computation , and the models achieved state-of-the-art performance in language modeling and machine translation . GShard [ 14 ] extended it to a tremendously large scale of 600B parameters with the sophisticated collaborated design of model architecture and demonstrated its effectiveness across over 100 languages . Similarly , Switch Transformer reached 1.6 trillion parameters and showed its strong performance on NLU tasks . Those models with high sparsity are computationally efficient , and therefore though they possess large model capacity they still can be trained with high efficiency . Most other studies still focus on training dense models to validate the scaling law , and thus the emerged problem is the distributed training of large models . The most influential distributed framework should be DeepSpeed that proposed ZERO [ 27 ] that partitions optimizer states and gradients to multiple GPU devices , and ZERO-offload [ 31 ] as well as ZERO-Infinity [ 28 ] can even offload parameters to CPU memory and NVMe storage . The memory footprint management makes training extremely large models on limited resources possible . However , such offloading mechanisms still have some defects that they may fail to fully utilize the fast hardware . For example , when offloading all parameters to the CPU , the GPU memory can be idle . In this work , we tackle this issue by proposing a granular offloading mechanism that can determine which parameters to be offloaded . In addition to reducing the amount of computation in a single iteration , another route to speedup training is to reduce the needed iterations for model convergence . Child et al . [ 5 ] and Xiong et al . [ 40 ] propose to put forward the layernorm operations in transformer blocks for more stabilized and faster convergence . You et al . [ 43 ] employs a layer-wise adaptive optimizer to enable super-large pretraining batches . Zhang & He [ 45 ] progressively increases the layer dropping rate in a stochastic manner which significantly speedups pretraining . 3 APPROACH . 3.1 MODEL ARCHITECTURE . Choice in model architecture depends on several factors . First , the architecture should contain a sequence of stacking layers , as the sequential structure enables cross-layer parameter sharing . We prefer a simple encoder or decoder architecture , instead of an encoder-decoder framework where there are cross attentions that bring extra parameters and incur difficulties in activation checkpointing . Second , a model of such architecture should be compatible with different types of downstream tasks , including understanding and generation , and it is more preferable that it can be compatible with multiple modalities . Third , as we mention that dense models and sparse expert models are two main tracks of large-scale pretraining , we prefer the model that can flexibly become whether dense or sparse expert models . Therefore , we select M6 [ 16 ] as an option , and we evaluate the effects of our method on M6 of different scales and types . M6 is built with stacking transformer layers , which includes self attention and feed-forward neural nets ( FFN ) . For the transformation from dense models to sparse expert models , we should only replace FFN layers with the Mixture-of-Expert ( MoE ) layers . MoE consists of multiple experts , which are usually FFNs distributed on different devices . A gating network decides the dispatching and combining behaviors of each token and thus tokens can be processed in diverse devices . Such mechanism is a combination of data parallelism and expert parallelism , and thus it is highly efficient though with large model capacity . For the training , to realize the learning of both understanding and generation , the model is trained with text denoising and language modeling on plain text data and with image-based text denoising and image captioning on multimodal data . The model is compatible with different types of downstream tasks and can process information of multiple modalities . 3.2 PSEUDO-TO-REAL . This section demonstrates the details of “ Pseudo-to-Real ” two-stage training strategy that enables fast training of high-memory-footprint-required transformer models . The strategy consists of two stages . The first stage trains a model with many fewer parameters but with a large computation graph ( “ Pseudo Giant ” ) , and the second stage trains a corresponding large model ( “ Real Giant ” ) initialized with the delinked weights of the shared layer . Thus we name the strategy “ Pseudo-to-Real ” , and the general idea is demonstrated in Figure 1 .

Summary: The authors propose to train very large neural language models via a "Sharing-Delinking" paradigm. The proposed method first trains a model with weights shared across layers. In this way, the model appears to be smaller and it can fit into fewer GPUs. At some point, the authors delink the weights, and continue training the model in the conventional way.  The authors also propose a granular CPU offloading mechanism to save CPU memory.