Large language model training acceleration method and device based on Fast Parallelism framework

Abstract

The invention provides a large language model training acceleration method and a large language model training acceleration device based on a Fast Parallelism framework. The method comprises the following steps: training a large language model on a plurality of GPUs (Graphics Processing Unit) by adopting a Fast Parallelism framework, wherein the Fast Parallelism framework comprises a tensor parallel module, a data parallel module and a hyper-parameter adjustment module; the method comprises the following steps of: when the large language model is trained, setting a sample number batch-size and a gradient accumulation which are trained in each iteration in a hyper-parameter adjustment module; a Deepspeed ZeRO data parallel processing strategy is adopted in the data parallel module; in the tensor parallel module, tensor segmentation is carried out in two dimensions and four dimensions, and large language model accelerated training is carried out in combination with the data parallel module and the hyper-parameter adjustment module. According to the method, a FastParallelism framework is innovatively designed, the framework adopts modular design, data parallelism, tensor parallelism, hyper-parameter adjustment and other strategies are applied, the internal performance of a computer system can be improved, and under the condition of limited resources, the large model training speed is increased.

Description

Technical Field

[0001] The present invention relates to the field of large model technology, and in particular to a large language model training acceleration method and device based on a FastParallelism framework. Background Art

[0002] With the application and development of the Transformer architecture, deep learning models have expanded to hundreds of billions of parameters. As the scale of parameters increases, a large amount of data, computing resources, and training time are often required. How to accelerate the training of large models with limited resources is the focus of current research.

[0003] The current status of the development of large model training acceleration technology is mainly reflected in distributed training, mixed precision training, hardware accelerators and automatic parallelization. By distributing training tasks to multiple computing nodes for parallel processing, using mixed precision training technology, adopting various hardware accelerators specifically used to accelerate deep learning calculations, and the development of automatic parallelization technology, the training time can be significantly reduced and the large model training process can be accelerated. Representative technical solutions include Megatron-LM and Deepspeed. Megatron-LM uses distributed training and mixed precision training to accelerate model training, which uses pipeline parallelism and tensor parallelism strategies. DeepSpeed ​​mainly accelerates model training by reducing memory usage and uses data parallelism strategies.

[0004] Although large model training acceleration technology has achieved unprecedented development, it still has some shortcomings. Existing acceleration technologies may have a certain impact on the effect of training models. For example, mixed precision training may lead to a decrease in model accuracy, while tensor parallelization may increase communication overhead. Therefore, while accelerating model training, it is necessary to balance the trade-off between training speed and model effect.

[0005] Large-scale model training usually requires the use of multiple computing nodes and high-performance hardware accelerators such as GPUs or TPUs. These hardware devices are expensive and may not be easy to obtain for some research institutions or individuals. In addition, the use of a large number of hardware devices will also increase energy consumption and maintenance costs.

[0006] In distributed training, data transmission and communication overhead become bottlenecks. As the model size increases, communication overhead may increase accordingly, resulting in reduced training efficiency. Summary of the invention

[0007] The technical problem to be solved by the present invention is how to accelerate model training while ensuring model performance under limited resource conditions, and to provide a large language model training acceleration method and device based on the Fast Parallelism framework.

[0008] A large language model training acceleration method based on a FastParallelism framework according to an embodiment of the present invention includes:

[0009] The Fast Parallelism framework is used to train a large language model on multiple GPUs, wherein the Fast Parallelism framework includes a tensor parallel module, a data parallel module, and a hyperparameter adjustment module;

[0010] When training a large language model, the batch-size and gradient accumulation of samples trained in each iteration are set in the hyperparameter adjustment module; the Deepspeed ZeRO data parallel processing strategy is adopted in the data parallel module; in the tensor parallel module, the tensor is split in 2D and 4D, and the data parallel module and the hyperparameter adjustment module are combined to accelerate the training of the large language model.

[0011] According to some embodiments of the present invention, when training a large language model, the batch-size of the number of samples trained in each iteration is set to 512 in the hyperparameter adjustment module, the gradient accumulation is set to 8, and the gradients of 8 batches are accumulated in the model training and then back-propagated.

[0012] In some embodiments of the present invention, when training a large language model, each GPU only saves part of the tensors of the large language model and part of the model parameters and optimizer states that need to be trained, and inputs data in a distributed manner.

[0013] According to some embodiments of the present invention, the data parallel module reduces the memory consumption of each GPU by splitting the large language model training state among the various GPUs through Deepspeed ZeRO.

[0014] In some embodiments of the present invention, Deepspeed ZeRO splits large language model training states including weights, gradients, and optimizer states among various GPUs.

[0015] According to some embodiments of the invention, Deepspeed ZeRO is implemented as a step-by-step optimization phase, including:

[0016] S1, split the optimizer state into various processes, and each process only updates the corresponding partition;

[0017] S2, splits the gradients used to update the model weights, and each process only retains the gradients corresponding to the optimizer state part;

[0018] S3, splits the model parameters into various processes.

[0019] In some embodiments of the present invention, the large language model includes Baidu Wenxin large model, BERT, and ChatGPT.

[0020] According to an embodiment of the present invention, a large language model training acceleration device based on the Fast Parallelism framework includes: a memory, a processor, and a computer program stored in the memory and executable on the processor. When the computer program is executed by the processor, the steps of the large language model training acceleration method based on the Fast Parallelism framework as described above are implemented.

[0021] According to the computer storage medium of an embodiment of the present invention, a computer program is stored on the computer storage medium, and when the computer program is executed by a processor, the steps of the large language model training acceleration method based on the Fast Parallelism framework as described above are implemented.

[0022] The present invention has the following beneficial effects:

[0023] The present invention innovatively designs a FastParallelism framework, which adopts a modular design and uses strategies such as data parallelism, tensor parallelism, and hyperparameter adjustment. It can improve the internal performance of the computer system and increase the training speed of large models under limited resource conditions. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] Figure 1 A schematic diagram of the architecture of a large language model training acceleration device based on a Fast Parallelism framework according to an embodiment of the present invention;

[0025] Figure 2 Schematic diagram of existing data parallel training methods;

[0026] Figure 3 It is a schematic diagram of 1D tensor parallel processing strategy;

[0027] Figure 4 Schematic diagram of a multi-dimensional tensor parallel processing strategy according to an embodiment of the present invention. DETAILED DESCRIPTION

[0028] In order to further explain the technical means and effects adopted by the present invention to achieve the predetermined purpose, the present invention is described in detail below in conjunction with the accompanying drawings and preferred embodiments.

[0029] The description of the method flow in the specification of the present invention and the steps of the flowchart in the drawings of the specification of the present invention do not have to be strictly executed according to the step numbers, and the method steps can be executed in a different order. Moreover, some steps can be omitted, multiple steps can be combined into one step, and / or one step can be decomposed into multiple steps.

[0030] Based on the existing acceleration technology, the present invention is committed to accelerating model training while ensuring model performance under limited resource conditions.

[0031] The present invention innovatively designs the Fast Parallelism framework, which adopts modular design and uses strategies such as data parallelism, tensor parallelism, and hyperparameter adjustment to improve the training speed of large models under limited resource conditions. The FastParallelism framework consists of three modules, namely, data parallelism module, tensor parallelism module, and hyperparameter adjustment module. Figure 1 shown.

[0032] The experiment of the present invention is carried out on 8 A800 GPUs, each GPU contains 80G video memory, and the model adopted is Llama2-70B. The present invention applies the Fast Parallelism framework in the training of the Llama2-70B model, which increases the model training speed by 3 times compared with the Baseline.

[0033] The following is a detailed introduction to each module in the framework:

[0034] 1. Hyperparameter adjustment module;

[0035] In the experiment of the present invention, the batch-size is set to 512. Batch size is the number of samples trained in each iteration. The present invention reduces the video memory as much as possible through data parallelism and tensor parallelism, thereby increasing the size of the batch size. Increasing the batch size can improve the GPU utilization and data throughput, thereby shortening the training time. A larger batch size can make fuller use of the parallel computing power of the GPU, thereby reducing the GPU idle time and data transmission time, and improving the training efficiency.

[0036] Gradient accumulation is set to 8, and the gradients of 8 batches are accumulated in model training before backpropagation. Gradient accumulation can increase the batch size while reducing GPU memory usage, thereby accelerating model training. In one training iteration, gradient accumulation can accumulate the gradients of multiple batches, which is equivalent to using a larger batch size for training. This can increase the batch size and training speed without increasing GPU memory usage.

[0037] 2. Data parallel module;

[0038] Data parallelism is the most common parallelism technique. In data parallel training, the model is replicated on each device and the dataset is split into multiple segments. Figure 2 As shown, each dataset segment is fed into the model on one device.

[0039] After back-propagation, collective communication is required to synchronize parameter gradients. Data parallelism makes it easy to train models on multiple devices, and it scales sub-linearly with the number of devices.

[0040] One problem with data parallelism is that each device holds a copy of model parameters, optimizer state, and gradients, resulting in storage redundancy. When using a stateful optimizer like Adam, optimizer state (such as momentum and variance) takes up three times more memory space than model parameters. To eliminate this redundancy, a zero-redundancy optimizer is proposed in DeepSpeed ​​to partition these redundant model data onto different devices during data parallel training. Since each device only holds a partition of gradients, optimizer state, and parameters, it will only update the partition parameters instead of updating the entire model parameters on one device.

[0041] The data parallel module of the present invention uses Deepspeed ZeRO technology. ZeRO uses data parallel strategies to reduce the memory and computing requirements of each device (GPU) used for model training. ZeRO reduces the memory consumption of each GPU by partitioning various model training states (weights, gradients, and optimizer states) between available devices (GPUs and CPUs) of distributed training hardware. ZeRO is implemented as a step-by-step optimization phase. Phase 1: The optimizer state (for example, for the Adam optimizer, 32-bit weights and estimates of the first and second moments) is partitioned into individual processes so that each process only updates its partition. Phase 2: The reduced 32-bit gradients used to update the model weights are also partitioned so that each process only retains the gradients corresponding to its optimizer state portion. Phase 3: The 16-bit model parameters are partitioned into individual processes. ZeRO-3 will automatically collect and partition them during forward and backward propagation.

[0042] 3. Tensor parallel module;

[0043] Tensor parallelism shards tensors onto a set of devices and requires distributed matrix-matrix multiplication algorithms for arithmetic calculations. The conventional tensor parallel method is one-dimensional tensor parallelism, which splits the linear layers in the Transformer structure in row or column dimensions. 1D tensor parallelism does not use data parallel strategies, such as Figure 3 shown.

[0044] The present invention adopts a more advanced tensor slicing mechanism to split tensors in 2D and 4D to accelerate model training. The data parallel module and hyperparameter adjustment module are combined to accelerate model training. The adjusted parallel strategy is as follows: Figure 4 As shown in the figure, each GPU only saves part of the model's tensors and part of the model parameters and optimizer states that need to be trained, and uses a distributed approach to input data to implement a low-memory, high-concurrency training strategy. The model training speed is further improved by increasing the Batch size and Gradient accumulation.

[0045] In summary, the present invention innovatively designs a Fast Parallelism training framework, innovatively proposes a training strategy that integrates tensor parallelism and data parallelism, and innovatively proposes 2D tensor parallelism and 4D tensor parallel training strategies.

[0046] In addition, it is worth emphasizing that the existing data parallel strategy of Deepspeed ZeRO and the 3D parallel strategy of Megatron-LM have not achieved the fusion training of data parallelism and tensor parallelism as a strategy for training acceleration, and the training speeds are different.

[0047] Through the description of the specific implementation methods, a deeper and more specific understanding of the technical means and effects adopted by the present invention to achieve the predetermined purpose should be obtained. However, the accompanying drawings are only for reference and illustration purposes and are not intended to limit the present invention.

Claims

1. A large language model training acceleration method based on the Fast Parallelism framework, characterized in that: include: The FastParallelism framework is used to train a large language model on multiple GPUs, wherein the FastParallelism framework includes a tensor parallel module, a data parallel module, and a hyperparameter adjustment module; When training a large language model, the batch-size and gradient accumulation of samples trained in each iteration are set in the hyperparameter adjustment module; the DeepspeedZeRO data parallel processing strategy is adopted in the data parallel module; in the tensor parallel module, the tensor is split in 2D and 4D, and the data parallel module and the hyperparameter adjustment module are combined to accelerate the training of the large language model.

2. The large language model training acceleration method based on the Fast Parallelism framework according to claim 1, characterized in that: When training a large language model, the batch-size of the number of samples trained in each iteration is set to 512 in the hyperparameter adjustment module, and the gradient accumulation is set to 8. In the model training, the gradients of 8 batches are accumulated and then back-propagated.

3. The large language model training acceleration method based on the Fast Parallelism framework according to claim 1, characterized in that: When training a large language model, each GPU only saves part of the tensors of the large language model and part of the model parameters and optimizer states that need to be trained, and inputs data in a distributed manner.

4. The large language model training acceleration method based on the Fast Parallelism framework according to claim 1, characterized in that: The data parallel module reduces the memory consumption of each GPU by splitting the large language model training state among each GPU through Deepspeed ZeRO.

5. The large language model training acceleration method based on the Fast Parallelism framework according to claim 4, characterized in that: DeepspeedZeRO splits large language model training states including weights, gradients, and optimizer states among various GPUs.

6. The large language model training acceleration method based on the Fast Parallelism framework according to claim 5, characterized in that: Deepspeed ZeRO is implemented as a gradual optimization phase consisting of: S1, split the optimizer state into various processes, and each process only updates the corresponding partition; S2, splits the gradients used to update the model weights, and each process only retains the gradients corresponding to the optimizer state part; S3, splits the model parameters into various processes.

7. The large language model training acceleration method based on the Fast Parallelism framework according to any one of claims 1-6, characterized in that: The large language models include Baidu Wenxin large model, BERT, and ChatGPT.

8. A large language model training acceleration device based on the Fast Parallelism framework, characterized in that: include: A memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the computer program is executed by the processor, the steps of the large language model training acceleration method based on the FastParallelism framework as described in any one of claims 1 to 7 are implemented.

9. A computer storage medium having a computer program stored thereon, wherein the computer program, when executed by a processor, implements the steps of the large language model training acceleration method based on the Fast Parallelism framework as described in any one of claims 1 to 7.

Images