Tensor storage management method for large model training

Abstract

The invention discloses a tensor storage management method for large model training, and relates to the technical field of computers, in particular to a tensor storage management method for large model training, which comprises memory management software for large model heterogeneous training, and the memory management software comprises a memory manager and a memory information counter. The memory management software enables tensors to be dynamically distributed in a storage space of a CPU-GPU in a training process, so that model training breaks through a memory wall of the GPU, the memory manager is responsible for model data tensors and is used for marking state information of the tensors, and large model training comprises a preheating stage and a formal stage. According to the method, the memory use condition of a CPU or a GPU of a system is collected regularly, non-model data numerical values are accurately counted according to the statistical moment of each group, large model training breaks through a memory wall by managing the storage mode of tensors in the CPU and the GPU, and training of a larger model can be completed under the same storage hardware configuration.

Description

technical field [0001] The invention relates to the technical field of computers, in particular to a tensor storage management method oriented to large model training. Background technique [0002] The emergence of pre-trained models (PTMs) represented by BERT and GPT today is a milestone in the field of natural language processing (NLP). NLP is entering the era of pre-training. PTM uses a neural network of stacked multi-Transformer structures to pre-train common language feature representations on a large amount of text, and then transfer the learned knowledge to different downstream tasks through fine-tuning. Using massive text data from the Internet, PTM can capture subtle features of natural language and make downstream tasks more impressive. Therefore, the consensus in the AI ​​community is to adopt PTMs as the backbone for specific NLP tasks, rather than training models from scratch on task-related datasets. [0003] The source of power of PTM is its parameter scale...

AI Technical Summary

Problems solved by technology

[0006] The purpose of the present invention is to provide a tensor storage management method for large-scale model training, to solve the problem that DeepSpeed, the current best heterogeneous training solution proposed in the above-mentioned background technology, still has a lot of room for optimization. In case, DeepSpeed ​​statically divides model data between CPU and GPU memory, and their memory layout is constant for different training configurations, when GPU memory or CPU memory is insufficient to meet its corresponding model data requirements, even if other There is still free memory on the device, and the system will crash

Images