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: This work studies the problem of efficient pretraining of large-scale models for language and vision representations, namely the issue of significant memory requirements for models with billions to trillions of parameters. Authors propose two modifications: first, to reduce the memory load and improve convergence at the initial stage of training, they suggest to train a multilayer model with shared parameters and then unshare them. Second, to maximize GPU utilization with offloading, authors develop a method for granular CPU offloading, which keeps larger chunks of the model in GPU memory. When combined, the proposed methods allow the authors to train a 10 trillion parameter model on 512 GPUs.