Commercial server and large language model-oriented efficient heterogeneous pipeline parallel training method

Abstract

The invention belongs to the field of parallel training of large language models, and particularly relates to a commercial server and large language model-oriented efficient heterogeneous pipeline parallel training method, which comprises an offline performance data acquisition stage, an offline model segmentation stage, an offline model preparation stage, an offline model initialization stage and an online training stage. In the online training stage, a prefetch perception layer packet division strategy is firstly adopted to segment a model into layer packets, so that prefetch overhead is completely covered, and GPU memory occupation is optimized; by removing redundant parameters in the CPU memory and a memory reuse strategy, the situation that the size of the trainable model is limited by the CPU memory capacity is avoided. In addition, by asynchronously executing exchange of data among heterogeneous devices by taking a layer as a unit, asynchronous parameter prefetching and CPU parameter updating are realized. According to the method, memory occupation of the GPU and the CPU during the training period is remarkably reduced, and the training performance is improved.

Description

Technical Field

[0001] This invention belongs to the field of parallel training of large language models, specifically involving an efficient heterogeneous pipelined parallel training method for commercial servers and large language models. Background Technology

[0002] With the emergence of Large Language Models (LLMs) based on the Transformer architecture, the number of model parameters has increased significantly. Well-known large language models such as GPT-3 and LLaMA have exceeded hundreds of billions of parameters, leading to memory requirements for training models far exceeding the memory limit of a single GPU. Despite the remarkable performance of LLMs, currently only large enterprises equipped with expensive, data center-level servers can explore and develop them. For example, training GPT-3 requires millions of dollars. Clearly, for the vast majority of researchers, limited by commercial servers consisting of only a few low-cost graphics cards and limited CPU memory, training or fine-tuning LLMs is difficult.

[0003] To address this issue, some methods alleviate GPU memory pressure by offloading data from the GPU to CPU memory or solid-state drives (SSDs). However, most of these works are based on observations of convolutional neural networks (CNNs) and do not consider the characteristics of LLMs. For example, virtual deep neural networks (vDNNs) selectively offload the output of convolutional layers because activation values ​​consume the most GPU memory during CNN training. Similarly, the Superneurous technique chooses to offload convolutional layer outputs rather than recompile them, also based on the significant computational overhead of convolutional layers. More importantly, these methods are designed for single-GPU scenarios and cannot fully utilize the parallel computing power and aggregated GPU memory of multi-GPU setups.

[0004] Parallel training methods leverage the aggregated memory of multiple GPUs to distribute memory overhead during training. For example, pipelined parallelism and tensor parallelism employ inter-layer and intra-layer parameter partitioning methods, respectively, enabling larger trainable model sizes by having each GPU store only a subset of the model's parameters. However, in LLM training scenarios, model parameters and optimizer states are the primary sources of GPU memory overhead. Especially in commercial servers equipped with only a few low-cost GPUs, even integrating the memory resources of all GPUs cannot meet the massive GPU memory requirements of LLM. For instance, fine-tuning GPT-3 requires the computing power of nine DGX-2 servers. Therefore, relying solely on aggregated GPU memory for training or fine-tuning LLM remains prohibitively expensive for most data scientists.

[0005] Heterogeneous data parallelism methods, such as ZeRO-Offload and ZeRO-Infinity, combine data parallelism and external storage techniques to further expand the scale of trainable models. ZeRO-Offload offloads the optimizer state to CPU memory, but the scale of model training it can support is still limited by GPU memory, as each GPU needs to store a complete set of model parameters. ZeRO-Infinity, on the other hand, partitions the parameters of each network layer and offloads the partitioned parameters to SSD. However, because the partitioned parameters need to be restored to their complete form before use, coupled with the cumulative effect of parameter prefetching, this method does not show a significant reduction in peak GPU memory usage compared to ZeRO-Offload. Furthermore, both methods require frequent aggregation communication of parameters and gradients, resulting in poor performance on commercial servers that rely solely on PCIe (Peripheral Component Interconnect Express) for inter-device communication.

[0006] The heterogeneous tensor parallel method Stronghold also introduces external storage devices, which are unloaded into CPU memory after parameter computation. This method effectively reduces the memory pressure on the GPU by partitioning the parameters of each layer of the model, allowing each GPU to store only a portion of the parameters from each layer. However, it is worth noting that because the parameters are partitioned within each layer, the results calculated independently by each GPU still require frequent ensemble communication to sum them up in order to generate complete activations or gradients, significantly reducing training efficiency.

[0007] The heterogeneous pipeline parallel method WA-Pipe only needs to transfer activation values ​​between two adjacent GPUs without reconstructing parameters, significantly reducing communication overhead compared to data-parallel and tensor-parallel methods. However, this method has the following problems:

[0008] 1. Severe GPU Memory Consumption: WA-Pipe's suboptimal model splitting method results in significant GPU memory consumption. Specifically, WA-Pipe breaks down the training process into a series of consecutive computational tasks, each of which can be viewed as a multi-layered package. This method employs a heuristic packaging strategy, using GPU hardware capacity as a constraint to automatically split the model. Its core logic only ensures that the GPU memory consumption of the generated layer packages during runtime does not exceed the hardware's rated capacity. Clearly, this design inevitably leads to the model being split into only a very small number of layer packages, and each layer package generating substantial memory consumption during runtime.

[0009] 2. High CPU memory usage limits the size of trainable models: WA-Pipe lacks a proactive CPU memory management mechanism, directly causing parameter redundancy and resulting in huge CPU memory consumption. Specifically, multiple identical copies of the model are repeatedly stored in CPU memory. This deficiency makes the size of trainable models easily limited by the actual capacity of CPU memory.

[0010] 3. Low Training Efficiency: WA-Pipe's suboptimal model splitting method also prevents it from fully utilizing the parallelism of multiple GPUs. For example, when training GPT-2 XL on a 4-GPU 3090, this method can only generate 3 forward computation tasks, meaning it can only use 3 GPUs to perform forward computation. Furthermore, WA-Pipe cannot effectively utilize GPU computation for parameter update operations hidden on the CPU. Finally, WA-Pipe's layer packet computation on the GPU and the data transfer of layer packets from CPU memory to GPU memory are executed serially, further reducing training performance. Summary of the Invention

[0011] The purpose of this invention is to address the problem that existing heterogeneous parallel training methods struggle to efficiently train large language models (LLMs) on low-cost servers with limited GPU and CPU memory. This invention proposes an efficient heterogeneous pipelined parallel training method for commercial servers and large language models. This method significantly reduces reliance on GPU and CPU memory in low-cost commercial server environments with limited GPU memory, CPU memory, and communication bandwidth, overcoming the limitations on the size of trainable models while maintaining high training efficiency. This provides a feasible path for researchers to conduct LLM research more broadly.

[0012] To achieve the above objectives, the present invention adopts the following technical solution:

[0013] An efficient heterogeneous pipelined parallel training method for commercial servers and large language models includes the following steps:

[0014] S1. Offline performance data collection: The large language model is split into original layers, and the performance parameters of each original layer on the GPU are collected. The performance parameters include forward computation time, backward computation time, and data transfer time of each original layer parameter from external storage medium to GPU memory.

[0015] S2. Offline model segmentation: Based on the performance data collected in step S1, a layer packet size is determined using a prefetch-aware layer packet partitioning strategy, and consecutive original layers are combined into layer packets according to the layer packet size, which serve as the basic computational units during training.

[0016] S3. Memory Management Configuration: Choose one of the following two configurations to manage model parameter storage: Configuration 1, saves the complete model parameters in CPU memory; Configuration 2, saves the complete model parameters in Solid State Drive (SSD).

[0017] S4. Training Initialization: Allocate GPU memory according to the layer packet size determined in step S2, and register hook functions for each layer of the model to trigger parameter prefetching, memory reclamation and gradient unloading;

[0018] S5. Pipeline Execution: During training, forward computation, recomputation, and backward computation are sequentially executed on the GPU, layer by layer; specifically, during the execution of the current layer... During the calculation, asynchronous triggering will trigger the next layer of packages. The parameters are prefetched from the storage location corresponding to the configuration selected in step S3 to the GPU memory;

[0019] S6. Asynchronous Gradient Unloading: After the GPU performs backward computation to generate the gradient for the i-th original layer within the layer pack, it immediately unloads the gradient. The gradient of the layer is asynchronously unloaded to CPU memory;

[0020] S7. Asynchronous parameter update: After the CPU memory receives the gradient, it asynchronously executes the parameter update of the corresponding original layer.

[0021] Furthermore, the prefetch sensing layer packet partitioning strategy described in step S2 includes: dynamically setting the layer packet size based on the collected raw layer performance data; wherein the setting of the layer packet size satisfies: under ideal conditions, the current layer packet... The forward computation time on the GPU is no less than that of the next layer of the packet. The parameters are calculated based on the time required to transfer from external storage media to GPU memory, in order to achieve overlap between computation and prefetch data transfer.

[0022] Furthermore, step S2, which dynamically sets the layer packet size and combines multiple consecutive original layers into a layer packet based on this layer packet size, includes:

[0023] Based on the forward computation time of all raw layers collected in step S1, the average forward computation time of a single raw layer is calculated. ;

[0024] Based on the raw layer data transfer time collected in step S1, calculate the average data transfer time of a single raw layer from external storage medium to GPU memory. ;

[0025] Calculate the prefetch cost ratio ;

[0026] The layer package size is set to the larger of a preset minimum value and K, and the original layers are combined into a layer package according to this value. The minimum value is 2.

[0027] Furthermore, the large language model is a model based on the Transformer architecture, and the original layer is a Transformer layer; the Transformer layers are combined into layer packages according to the set layer package size in a uniform combination manner.

[0028] Furthermore, in step S3:

[0029] When configuration one is selected, the complete model parameters are stored in the CPU's shared memory for all training processes to access;

[0030] When configuration two is selected, the complete model parameters will be saved in the SSD;

[0031] In both configurations, a layer buffer for temporarily storing prefetch parameters and a gradient buffer for temporarily storing unloaded gradients are created in the CPU's page-locked memory.

[0032] Furthermore, the GPU memory space allocated in step S4 can simultaneously accommodate at least one layer packet that is currently being computed and one layer packet that is being prefetched.

[0033] Furthermore, the hook functions registered in step S4 include:

[0034] The forward pre-hook function is triggered before the forward computation or recomputation of the original layer and is used to prefetch the parameters of the original layer in the next layer package to be executed from the layer buffer to GPU memory.

[0035] The forward hook function is triggered after the forward computation of the original layer is completed, and is used to release the space occupied by the parameters of that layer in GPU memory;

[0036] The reverse hook function is triggered after the reverse computation of the original layer is completed. It is used to trigger the asynchronous unloading operation of the gradient of the previous original layer that has completed the reverse computation in step S6.

[0037] Furthermore, during the recomputation process, after the original layer parameters are prefetched into GPU memory, their cache space in the layer buffer is reserved for receiving the gradient data of that layer unloaded from the GPU after reverse computation.

[0038] Furthermore, in step S7, the asynchronous parameter update is performed in parallel with the recalculation and reverse computation of other layer packages on the GPU; and step S7 also includes: after the parameter update is completed, if the current configuration is two, the updated parameters are written back to the SSD, and the cache space occupied by the layer parameters in the layer buffer is reclaimed.

[0039] By adopting the above technical solution, the present invention has the following beneficial effects:

[0040] (1) This invention divides the model into multiple layer packets by constructing a layer packet partitioning strategy for prefetching perception, so that the number of layers contained in each layer packet is significantly lower than that of existing layer packet partitioning strategies. During the partitioning process, a prefetching overhead constraint is introduced to achieve full coverage of prefetching overhead and optimize memory usage. This avoids exhausting GPU memory in a single layer packet during runtime, while making fuller use of the parallelism of multiple GPUs. It also enables better use of GPU reverse computation to hide most of the CPU parameter update operations, thereby improving training performance.

[0041] (2) In terms of memory management, this invention effectively avoids limiting the size of trainable models by removing redundant parameters in CPU memory and employing memory reuse strategies. Furthermore, it provides two training configurations for memory management: CPU offload and SSD offload. CPU offload uses only CPU memory to alleviate GPU memory pressure, while SSD offload further utilizes SSDs to distribute memory overhead during training, reducing the CPU memory usage of large language models during training.

[0042] (3) During operation, the present invention performs data exchange asynchronously on a layer-by-layer basis, realizing asynchronous parameter prefetching, gradient unloading, and CPU parameter updating. The asynchronous nature means that the reading of parameters from external storage to GPU memory, the unloading of gradients from GPU memory to CPU memory, and the updating of parameters on the CPU and GPU computation are performed simultaneously, without blocking GPU computation, thus significantly improving training performance. Attached Figure Description

[0043] Figure 1 This is a schematic diagram of the parallel training method of the present invention;

[0044] Figure 2 This is a flowchart illustrating the prefetch operation triggered during the forward computation of the layer packet in an embodiment.

[0045] Figure 3 This is a flowchart of the prefetch operation triggered during the recalculation of the layer package;

[0046] Figure 4 This involves a temporary inverse computation process within the layer package, asynchronous gradient unloading, and asynchronous CPU parameter update process.

[0047] Figure 5 This involves the reverse computation process within the layer package when configuration two is adopted, as well as the asynchronous gradient unloading and CPU asynchronous parameter update process.

[0048] Figure 6This is a comparison of the peak GPU memory usage of each method in the examples when training a 1.5 billion parameter LLaMA model;

[0049] Figure 7 This is a comparison of the peak GPU memory usage of each method in the examples when training a 5 billion parameter LLaMA model;

[0050] Figure 8 A comparison of CPU memory usage for training an LLaMA model with 1.5 billion parameters using different methods in the examples;

[0051] Figure 9 A comparison of CPU memory usage for training an LLaMA model with 2.5 billion parameters using different methods in the examples;

[0052] Figure 10 This example compares the throughput of each method in training an LLaMA model with 1.5 billion parameters.

[0053] Figure 11 This is a comparison of the maximum trainable model size for each method in the embodiments. Detailed Implementation

[0054] The technical solution of the present invention will be described in detail below with reference to the accompanying drawings and embodiments.

[0055] This embodiment provides an efficient heterogeneous pipelined parallel training method for commercial servers and large language models, which includes three stages: offline information collection and planning, resource configuration and initialization, and online training execution. Figure 1 This is a schematic diagram of the parallel training method of the present invention, where rectangles represent sub-layers of the model, and dashed rectangles indicate that the layer has not yet been read from the device. Figure 1 As shown, each stage is explained in detail below:

[0056] Phase 1: Offline Information Collection and Planning

[0057] Existing layer-packet partitioning schemes, represented by the WA-Pipe algorithm, are constrained by the principle that the number of parameters in a single layer-packet should not exceed the maximum memory capacity of a single GPU. Based on this constraint, large language models are evenly split. Using GPU memory capacity alone as the partitioning criterion leads to three technical problems: First, excessively large parameters in a single layer-packet result in a surge in activation data during computation, exhausting the GPU's memory resources and preventing the completion of the normal training process. Second, an insufficient number of layer-packets prevents the pipelined parallel mode from fully utilizing the parallel computing capabilities of multiple GPUs. For example, when training a GPT-2 XL model on a commercial server with four GPUs each with 24GB of memory, this scheme can only partition three layer-packets for forward computation, directly resulting in one GPU being idle during the forward computation phase, wasting hardware resources. Third, after the GPU completes the backward computation of all layer-packets, it must wait for a long time for the CPU to perform parameter update operations on the last layer-packet that has completed backward computation, during which time the GPU is completely idle.

[0058] To address this, this embodiment proposes a prefetch perception layer packet partitioning strategy. Under the premise of satisfying GPU capacity constraints, and using pre-substitution as a constraint, the model is divided into layer packets composed of a small number of original layers. "Small number" refers to a layer packet containing significantly fewer layers than those obtained by existing layer packet partitioning strategies. Large Language Models (LLMs) consist of a large number of identical Transformer layers. This embodiment first uniformly partitions the Transformer part of the model according to the layer packet size, and then distributes the resulting layer packets to various computing nodes using a round-robin scheduling algorithm. The round-robin scheduling algorithm involves performing a modulo operation on the layer packet index using the number of GPUs. Let the number of GPUs be N, then the... The layer package is assigned to the () ( ) GPUs, thereby ensuring load balancing across all GPUs.

[0059] To accurately quantify the impact of layer size on GPU performance, this embodiment first performs offline training information collection on the original split layers. For each original layer, multiple complete forward and backward training processes are performed under multiple micro-batches, and the performance parameters of each original layer on the GPU are collected. The performance parameters include the forward computation time of each original layer on the GPU. Reverse calculation time And the communication time of this layer parameter between heterogeneous devices. The communication time covers the copying time from the SSD to the CPU memory for each original layer. Copy time from CPU memory to GPU memory The micro-batch refers to a small batch of data units further subdivided in each training iteration. During performance testing, the above process is repeated multiple times, and the average value of the collected time indicators is calculated to ensure the validity and statistical accuracy of the obtained information.

[0060] The prefetch-aware layer packet partitioning method in this embodiment is based on two important observations: (1) once the layer packet size is determined, the time for data transfer between external storage and GPU memory remains constant; (2) in heterogeneous pipeline parallelism, each layer packet must complete the computation of all micro-batches before starting the computation of the next layer packet. Figure 2 This is a flowchart illustrating the prefetching operation triggered during the forward computation of an example layer; in the diagram, dashed arrows indicate prefetching of layers that will be used subsequently; see reference. Figure 2 It is known that the GPU only begins computation on layer 2 after layer 0 has completed the processing of micro-batch 0 and micro-batch 1. This characteristic ensures that even if the number of layers in the current layer 0 is much less than that generated by existing methods, sufficient computation time can still mask the overhead of prefetching the next layer 0. Based on this, this embodiment introduces a prefetch overhead constraint to determine the layer 0 size. The core principle is that when the computation time of a layer 0 is equal to or exceeds the communication time for prefetching the next layer 0, the layer 0 size is no longer increased. At this point, the computation time of the current layer 0 is sufficient to completely mask the communication overhead of the prefetch operation. Continuing to increase the layer 0 size will not improve computational efficiency but will instead lead to problems such as excessive memory consumption and insufficient parallelism.

[0061] The mathematical expression for the prefetch cost constraint is: when At that time, stop increasing the size of the layer package; among which, Indicates the calculation time. , Indicates the sequence number of the Transformer layer. Indicates the batch size. Indicates the size of the micro-batch; Indicates the layer packet to be processed. Indicates data transmission time. This indicates the next layer package to be executed.

[0062] The data transmission time of the next layer packet Determined based on the type of external storage medium: When only the CPU is used as external storage. ,in, This indicates the time it takes for data to be transferred from CPU memory to GPU memory; when both SSD and CPU are used as external storage, ,in, This indicates the time it takes for data to be transferred from the SSD storage medium to the CPU memory.

[0063] Since the computation time of forward propagation for a single layer of the model is significantly lower than that of back propagation, this embodiment selects the computation time of forward propagation as the overhead constraint parameter for the prefetching operation, thereby improving the rationality and execution efficiency of the data prefetching strategy. In subsequent calculations, based on the performance parameters of offline acquisition, the average forward computation time of a single Transformer layer is calculated as follows: ;

[0064] The average transmission time is: ;in, This indicates the total number of Transformer layers. Indicates the sequence number of the Transformer layer. Indicates batch size, Indicates the micro-batch size.

[0065] The prefetch overhead constraint is configured to determine whether the computation time of a single Transformer layer is sufficient to cover the data transfer time of a single Transformer layer. The mathematical relationship is as follows:

[0066] ;

[0067] in, This represents the average forward computation time of a primitive layer on the GPU. This represents the average transfer time for a single raw layer to be transferred from CPU memory to the GPU.

[0068] against This embodiment dynamically sets the layer packet size based on different values:

[0069] when The computation time of a single Transformer layer is sufficient to cover the transmission time of a single Transformer layer. Combined with the recomputation technique used in this implementation to increase the size of the trainable model during the training phase, the input activations of each layer packet are unloaded into CPU memory. Setting the layer packet size to 1 would result in a large number of activations being unloaded into CPU memory during training; therefore, this embodiment sets the layer packet size to 2. This setting completely hides the prefetching overhead while minimizing the layer packet size and significantly reduces CPU memory usage during training.

[0070] when The size of the layer will be set to the same as the size of the package. Similarly, further increasing the layer packet size at this point will not improve performance; instead, it will lead to more prefetching operations on layers in the next layer packet. Therefore, this embodiment significantly reduces GPU memory usage while maintaining minimal prefetch overhead.

[0071] The offline information collection and planning phase of this embodiment specifically performs the following steps:

[0072] S1. Offline performance data collection: The large language model is split into multiple raw layers, and the performance parameters of each raw layer on the GPU are collected. The performance parameters include forward computation time, backward computation time, and data transfer time of each raw layer parameter from external storage medium to GPU memory.

[0073] S2. Offline Model Segmentation: Based on the performance data collected in step S1, the size of a layer packet is determined using a prefetch-aware layer packet partitioning strategy. For large language models with repetitive basic structures, the repetitive structural parts are uniformly segmented according to the dynamically set optimal layer packet size, and the segmented layer packets are allocated to each computing node according to the above-mentioned cyclic scheduling algorithm to ensure balanced GPU load on each node.

[0074] The prefetching sensing layer packet partitioning strategy adopted in step S2 has two core performance advantages because it divides the model into layer packets with a significantly lower number of layers than existing strategies:

[0075] First, the number of layers in the layer package is less, which can better utilize the parallelism of multiple GPUs. Under any circumstances, all GPUs can be used to perform forward and backward calculations, and the computing tasks are distributed more evenly among the GPUs, avoiding the waste of hardware resources caused by GPU idleness.

[0076] Second, it allows for a more thorough overlap between GPU computation and CPU parameter update operations, effectively hiding a large number of CPU parameter update operations. Fewer layers mean faster backward computation, allowing gradients to be unloaded to CPU memory and CPU parameter updates to begin earlier. Since CPU parameter updates are asynchronous operations that can be executed in parallel with GPU computation, most layer parameter update operations can be executed in parallel with the recomputation and backward computation of subsequent layers.

[0077] In contrast, existing layer-packet partitioning strategies result in fewer computational tasks on GPUs and CPUs, and longer execution times for individual computational tasks. For the last layer-packet on the GPU, while the CPU performs numerous parameter update operations, the GPU remains completely idle, making it impossible to hide this overhead through parallel computing. For example, WA-Pipe's model partitioning method can only generate one layer-packet for reverse computation on each GPU for the GPT-2 2.6B model. No subsequent GPU computations can be executed in parallel with the parameter update operations of this layer-packet on the CPU, causing each process to spend approximately 25% of its time in a single iteration entirely on parameter updates.

[0078] Phase Two: Resource Configuration and Initialization

[0079] S3. Memory Management Configuration: This embodiment selects one of the following two configurations to manage model parameter storage, eliminating redundant parameters in CPU memory through proactive CPU memory management methods, and balancing CPU memory usage and training performance:

[0080] Configuration 1: CPU Offloading. This configuration stores complete model parameters in the CPU's shared memory, accessible to all training processes. It is suitable for scenarios with ample CPU memory or small models, relying solely on CPU memory for parameter storage without accessing SSDs, thus avoiding the throughput loss introduced by SSDs. Addressing the issue of redundant parameter usage in CPU memory in existing heterogeneous pipeline and training methods, such as WA-Pipe, this embodiment optimizes memory usage through fine-grained page-locked memory management. The specific design is as follows:

[0081] Existing heterogeneous pipelined parallel methods suffer from two aspects of parameter consumption: firstly, each process maintains complete parameters in CPU page-locked memory, resulting in memory consumption of [missing information]. ,in, For the number of GPUs, This represents the total memory usage of the model. This represents the total number of floors. The first is the memory usage of a single original layer; the second is that the complete model is retained in the CPU shared memory for parameter updates and to ensure the consistency of parameters across processes.

[0082] This implementation addresses the redundancy issue of parameter usage in CPU page-locked memory by creating a layer buffer for each process within that memory. This buffer consists of several cache layers, fewer than the total number of layers a process needs to process during training, and only stores parameters for layers about to be used. The layer buffer is used for data transfer between the CPU and GPU during parameter prefetching. After use, the cache layers are released for reuse in subsequent prefetching operations, achieving efficient memory reuse and significantly reducing page-locked memory usage. Simultaneously, the complete model is maintained in CPU shared memory for parameter sharing and updates between processes. Furthermore, a complete gradient buffer is configured on the CPU specifically to receive gradients unloaded from the GPU; its memory usage can be represented as… .

[0083] To avoid cache exhaustion, the layer buffer size is set according to the following rule: Assume the number of independent Transformer layers the GPU will run is... The retention ratio is The size of the layer buffer is At this point, the memory usage of page-locked parameters drops to [a certain value]. .

[0084] Configuration 2: SSD Offloading. This configuration stores the complete model parameters on an SSD, suitable for scenarios with insufficient CPU memory. By introducing an SSD, it alleviates the CPU memory pressure, storing only a portion of the core data in CPU memory, completely eliminating parameter redundancy in shared memory. The core difference between Configuration 2 and Configuration 1 lies in the parameter storage medium and gradient buffer size design:

[0085] Configuration 2 cancels the full model storage in CPU shared memory and migrates all parameters to SSD storage, thus reducing the shared memory parameter usage from... When the shared memory is reduced to 0, it retains only the complete optimizer state, supporting CPU-side parameter update capabilities. Configuration 2's design rules are consistent with Configuration 1; the CPU page-locked memory still retains layer buffers and gradient buffers, where the gradient buffer size is the same as the layer buffer size. It does not require maintaining a gradient space of the same size as the full model.

[0086] During parameter prefetching, parameters are read directly from the SSD into the layer buffer. After parameter updates are completed on the cache layer, they are asynchronously unloaded to the SSD for storage. Although this configuration introduces a small amount of additional overhead due to SSD communication, resulting in slightly lower training efficiency than configuration one, this overhead still allows sufficient time for CPU-side parameter updates, ensuring that the gradient buffer always has free space to receive new gradient data.

[0087] S4. Training Initialization: This step aims to complete GPU memory allocation and hook function registration, laying the foundation for automated triggering of the subsequent training process. The specific operation process is as follows:

[0088] GPU memory is allocated based on the layer pack size determined in step S2. Each GPU only reserves memory space of twice the layer pack size, eliminating the need to store all layers of the complete model and significantly reducing GPU memory usage.

[0089] Register forward pre-hook functions, forward hook functions, and backward hook functions for each layer of the large language model. Hook functions are callback functions registered at the model layers and can be automatically executed when specific training events are triggered, thereby ensuring the automated execution of subsequent operations such as parameter prefetching, GPU memory reclamation, and asynchronous gradient unloading. The main functions of each type of hook function are as follows:

[0090] Forward pre-hook function: Triggered before the forward computation and recomputation of the corresponding layer, to realize asynchronous prefetching operation of the layers that will be used later;

[0091] Forward hook function: Triggered after the forward computation of the corresponding layer is completed, realizing the reclamation of GPU memory occupied by the parameters of that layer;

[0092] Reverse hook function: Triggers an asynchronous unloading operation on the gradient of the previous layer that has completed reverse calculation after the reverse calculation of the corresponding layer is completed.

[0093] The third stage, online training execution: In this stage, forward computation, recomputation, and backward computation are executed sequentially on the GPU at the layer-pack level. Simultaneously, an asynchronous prefetching strategy masks the communication overhead between devices, addressing the efficiency bottlenecks and scenario adaptability issues of existing methods such as WA-Pipe and Stronghold. This includes steps S5 to S7.

[0094] Existing methods have significant drawbacks: WA-Pipe only prefetches layer packet input activations, not parameters, leading to serial GPU computation and parameter copying; while Stronghold supports parameter prefetching, it suffers from high communication redundancy and fails to fully utilize the optimization opportunities provided by recomputation steps. Furthermore, it cannot adapt to the non-contiguous and repetitive layer execution requirements of a single GPU in a multi-GPU heterogeneous pipeline. Therefore, a new pipeline execution method is designed in step S5 of this embodiment to specifically address the aforementioned problems. This pipeline execution method has the following characteristics:

[0095] Prefetching only model parameters without carrying gradients effectively reduces communication by half;

[0096] Trigger a prefetch operation before recalculation to fully utilize the recalculation time to mask communication overhead;

[0097] For multi-GPU heterogeneous pipelined parallel training scenarios, each process triggers the prefetching of the layer corresponding to the next layer only before the forward computation and recomputation of the first micro-batch. This prefetching operation can mask the communication overhead of prefetching by leveraging the computation time of the current and subsequent layers on the first micro-batch and the overall computation time of the layer pack on subsequent micro-batches.

[0098] Figure 4 The process of reverse computation within a configured time-layer package, asynchronous gradient unloading, and asynchronous CPU parameter update is demonstrated. Figure 5 This involves the reverse computation process using a two-time layer package configuration, asynchronous gradient unloading, and asynchronous CPU parameter update process. Figure 3 This is a flowchart of the prefetch operation triggered during the recalculation of the layer package. Combining the hook function and the layer buffer, the specific pipeline execution flow of step S5 is as follows:

[0099] Forward computation: The first layer package loaded by each GPU during the initialization phase. Begin by executing the layer-by-layer sequence defined in step S2. The process is as follows:

[0100] Layered packaging Before the inner target layer performs calculations, the forward pre-hook function is triggered to load parameters according to the layer buffer logic defined in step S3, namely, configuration one is taken from CPU shared memory, configuration two is taken from SSD, and then transferred to GPU memory;

[0101] The target layer performs forward computation. This process is only used to calculate the final loss value of the model and does not perform the GPU to CPU offload operation.

[0102] After the forward computation is completed, the forward hook function is triggered to reclaim the GPU memory occupied by the parameters of that layer.

[0103] Recalculation: Recalculation is a layered package Before performing the reverse calculation, the forward calculation is re-executed to restore the required activation value. The specific process is as follows:

[0104] Before the target layer performs computation, the forward pre-hook function is triggered to load parameters into the GPU using the layer buffer prefetch logic mentioned above.

[0105] The target layer performs recomputation, and after completion, it does not reclaim GPU memory but retains the parameters for subsequent reverse computation.

[0106] S6. Gradient Asynchronous Unloading: In this step, the GPU performs gradient asynchronous unloading on the first layer within the package. After the target layer performs backpropagation to generate gradient data, the backpropagation hook function registered for that layer is automatically triggered to execute the backpropagation hook function for the target layer. The asynchronous unloading operation of the layer's gradient follows the following process:

[0107] No. Each target layer completes the reverse computation to generate gradient data;

[0108] The reverse hook function is triggered, which will... The gradients of the layer are asynchronously unloaded to the gradient buffer in CPU page-locked memory.

[0109] S7. Asynchronous parameter update:

[0110] After receiving the gradient data, the CPU asynchronously executes the parameter update operation for the corresponding layer. This process is executed in parallel with the reverse computation currently underway on the GPU and does not cause blocking. At the same time, through proper design, it is ensured that the gradient buffer always has sufficient free space. The specific process is as follows:

[0111] First, wait for the gradient of the target layer to be asynchronously unloaded into the gradient buffer; after the gradient unloading is completed, the asynchronous update thread immediately reclaims the GPU memory space occupied by that layer;

[0112] Subsequently, the asynchronous update thread performs parameter updates based on the gradient buffer data, a process that proceeds in parallel with the recomputation and reverse computation in the GPU;

[0113] If configuration two in step S3 is enabled, after the parameter update is completed, the updated layer parameters need to be asynchronously unloaded to SSD for storage.

[0114] Finally, after all the above operations are completed, that is, after the parameter update and optional SSD unloading are completed, the cache layer in the CPU layer buffer occupied by this layer is reclaimed for reuse in subsequent prefetch operations.

[0115] The aforementioned asynchronous mechanism and resource reuse design not only completely avoid the blocking of GPU computation by CPU parameter updates, but also ensure the continuous availability of layer buffers and gradient buffers.

[0116] Experimental verification:

[0117] The experimental environment consisted of a typical commercial server, comprising four commercial graphics cards with 24GB of RAM each, two CPUs, one 2TB SSD, and 256GB of RAM. Each pair of GPUs was connected to one CPU via PCIe 3.0 x16. The training accuracy was FP32.

[0118] The model employs LLaMA, which is widely used to evaluate the effectiveness of training methods for large language models. To fully verify the effectiveness of this invention, experiments were conducted using four variants of LLaMA with different sizes, ranging from 1.5 billion to 7 billion parameters, as shown in the table below. Training was performed using the well-known WikiText dataset with a sequence length of 1024, and tests were conducted at different batch sizes.

[0119] This embodiment introduces four heterogeneous parallel training methods for comparison: Method 1, Method 2, Method 3, and Method 4, which correspond to ZeRO-Offload, ZeRO-Infinity, Stronghold, and WA-Pipe, respectively.

[0120] Table 1: Specific configurations of four different LLaMA models

[0121]

[0122] As shown in the experimental results and Table 1, the heterogeneous pipeline parallel training method in this embodiment achieves the following effects:

[0123] Lower GPU memory footprint: This embodiment introduces prefetching overhead as a constraint on model splitting through a prefetch-aware model sublayer packaging method, dividing the model into multiple fixed-size, low-GPU-memory-footprint layer packets. Therefore, this invention only requires maintaining two layer packet sizes in GPU memory for the execution of the target layer packet and the prefetching of future layer packets. This constraint avoids meaningless increases in layer packet size: While increasing the layer packet size can extend computation time when the computation time of a layer packet cannot cover the prefetching communication time of the next layer packet, the number of prefetching operations is positively correlated with the layer packet size, which simultaneously increases communication overhead. Therefore, increasing the layer packet size cannot effectively mask the communication overhead.

[0124] Existing heterogeneous parallel training methods all suffer from excessive GPU memory consumption: the heterogeneous pipelined parallel method WA-Pipe only uses GPU memory capacity as a constraint for splitting, ensuring that memory does not exceed the limit during layer-by-layer runtime, resulting in huge GPU memory consumption without any performance gain; the heterogeneous data parallel method ZeRO-Offload requires storing the complete model on each GPU, resulting in extremely high space consumption; although ZeRO-Infinity achieves intra-layer splitting, the splitting parameters need to be restored to their complete form through ensemble communication, and a large amount of the restored parameters are stored in GPU memory, thus still resulting in significant GPU memory consumption on commercial servers; the heterogeneous tensor parallel method Stronghold uses intra-layer parameter splitting, but this method requires temporarily storing a large number of activation values ​​in GPU memory, also leading to high memory consumption. These activation values ​​are the results generated by forward computation of each layer.

[0125] In this embodiment, only the current computational layer packet and adjacent prefetched layer packets need to be saved at any given time (a single layer packet has only two layers in this embodiment), significantly reducing GPU memory usage. Figure 6 As shown, when training a 1.5 billion parameter LLaMA model with a batch size of 16, the peak GPU memory usage in both configurations of this embodiment is the lowest. When using CPU offload, it reduces memory usage by 56.4%, 45.9%, 39.7%, and 34.4% compared to ZeRO-Offload, ZeRO-Infinity, Stronghold, and WA-Pipe, respectively. Figure 7 As shown, when training a 5 billion parameter model with a batch size of 32, other heterogeneous parallel training methods all fail due to insufficient GPU or CPU memory. In this embodiment, both configurations can successfully execute training.

[0126] Lower CPU memory usage: This embodiment completely eliminates parameter redundancy in WA-Pipe through proactive CPU memory management. Simultaneously, by constructing buffers for the model's original layers in CPU memory, it achieves efficient reuse of CPU memory space during training. Compared to existing heterogeneous data parallelism, tensor parallelism, and pipeline parallelism methods, this embodiment exhibits the lowest CPU memory usage after the initialization phase, thus providing more ample CPU memory space for the training phase.

[0127] like Figure 8 and Figure 9 As shown, this embodiment maintains the lowest CPU memory usage for both model scales. Taking training a model with 1.5 billion parameters as an example, the CPU memory usage in this embodiment is only 26.6GB when using SSD offload, which is 42.9%, 61.1%, 38.9%, and 55.1% lower than ZeRO-Offload, ZeRO-Infinity, Stronghold, and WA-Pipe, respectively.

[0128] Superior Training Performance: The performance advantage of this embodiment stems from four design aspects. First, the prefetch-aware model sub-layer packaging method achieves fine-grained model splitting, generating a far greater number of layer packets than GPUs, fully leveraging the parallel potential of multi-GPUs. In contrast, WA-Pipe can only generate a small number of layer packets containing a large number of sub-layers, resulting in low parallel utilization. For example, when training an LLaMA model with 1.5 billion parameters on a commercial server with four GPUs each with 24GB of memory, only two forward computation layer packets can be generated, leaving two GPUs idle. Second, because the packaging method of this invention generates layer packets containing fewer layers, a layer packet can complete backward computation faster, thus allowing gradients to be unloaded to CPU memory and parameter updates to begin on the CPU earlier. Furthermore, since parameter updates on the CPU are asynchronous operations that can be executed in parallel with GPU computation, this means that the parameter update operations of most layer packets can be executed in parallel with the recomputation and backward computation of subsequent layer packets. Third, the online training execution phase of this embodiment can synchronously prefetch the next layer packet to be used during the forward / recomputation of layer packets, minimizing the parameter transfer overhead between the CPU and GPU. This allows the GPU to immediately begin computing the next layer packet after completing the computation of one layer packet, without needing to serially read the next layer packet. Fourth, this invention asynchronously unloads gradients on a layer-by-layer basis, enabling the update operation of a layer on the CPU to be executed in parallel with the reverse computation of other layers within the same layer packet. In contrast, WA-Pipe requires sequential execution of the reverse computation of layer packets and gradient unloading, and gradient unloading and CPU update operations are also executed sequentially, resulting in low efficiency.

[0129] Furthermore, existing heterogeneous data parallelism and tensor parallelism methods require frequent communication of parameters and gradients between multiple GPUs, severely limiting training efficiency. This embodiment only needs to transfer activation values ​​between adjacent GPUs, significantly reducing communication overhead and resulting in superior performance. Figure 10 As shown, this embodiment achieves optimal training performance when CPU is offloaded. Taking the training of a 1.5 billion parameter LLaMA model with a batch size of 16 as an example, the throughput is improved by 87.3%, 670.5%, 67.8%, and 53.4% ​​respectively compared to ZeRO-Offload, ZeRO-Infinity, Stronghold, and WA-Pipe.

[0130] Support for training larger-scale language models: This embodiment effectively avoids memory overflow errors through joint optimization of GPU and CPU memory, enabling the training of larger-scale language models (LLMs) with limited memory resources. Figure 11 As shown, compared with existing heterogeneous parallel training methods, the memory optimization capability of this embodiment can cover the training needs of models with larger parameter scales, thus expanding the resource adaptation boundary of LLM training.

Claims

1. A highly efficient heterogeneous pipelined parallel training method for commercial servers and large language models, characterized in that, Includes the following steps: S1. Offline performance data collection: The large language model is split into original layers, and the performance parameters of each original layer on the GPU are collected. The performance parameters include forward computation time, backward computation time, and data transfer time of each original layer parameter from external storage medium to GPU memory. S2. Offline model segmentation: Based on the performance data collected in step S1, a prefetch-aware layer packet partitioning strategy is used to determine the size of a layer packet, and multiple consecutive original layers are combined into a layer packet according to the layer packet size, which serves as the basic computational unit during training. S3. Memory Management Configuration: Choose one of the following two configurations to manage model parameter storage: Configuration 1, saves the complete model parameters in CPU memory; Configuration 2, saves the complete model parameters in solid-state drive; S4. Training Initialization: Allocate GPU memory according to the layer packet size determined in step S2, and register hook functions for each layer of the model to trigger parameter prefetching, memory reclamation and gradient unloading; S5. Pipeline Execution: During training, forward computation, recomputation, and backward computation are sequentially executed on the GPU, layer by layer; specifically, during the execution of the current layer... During the calculation, asynchronous triggering will trigger the next layer of packages. The parameters are prefetched from the storage location corresponding to the configuration selected in step S3 to the GPU memory; S6. Asynchronous Gradient Unloading: After the GPU performs backward computation to generate the gradient for the i-th original layer within the layer pack, it immediately unloads the gradient. The gradient of the layer is asynchronously unloaded to CPU memory; S7. Asynchronous parameter update: After the CPU memory receives the gradient, it asynchronously executes the parameter update of the corresponding original layer.

2. The method according to claim 1, characterized in that, The prefetch-sensing layer packet partitioning strategy described in step S2 includes: dynamically setting the layer packet size based on the collected raw layer performance data; wherein the setting of the layer packet size satisfies: under ideal conditions, the current layer packet... The forward computation time on the GPU is no less than that of the next layer of the packet. The parameters are calculated based on the time required to transfer from external storage media to GPU memory, in order to achieve overlap between computation and prefetch data transfer.

3. The method according to claim 2, characterized in that, The process of dynamically setting the layer packet size in step S2 and combining multiple consecutive original layers into a layer packet according to the layer packet size includes: Based on the forward computation time of all raw layers collected in step S1, the average forward computation time of a single raw layer is calculated. ; Based on the raw layer data transfer time collected in step S1, calculate the average data transfer time of a single raw layer from external storage medium to GPU memory. ; Calculate the prefetch cost ratio ; The layer package size is set to the larger of the preset minimum value and K, and the original layers are combined into a layer package according to this value; the minimum value is 2.

4. The method according to claim 3, characterized in that, The large language model is a model based on the Transformer architecture, and the original layer is a Transformer layer; the Transformer layers are combined into layer packages according to the set layer package size in a uniform combination manner.

5. The method according to claim 1, characterized in that, In step S3: When configuration one is selected, the complete model parameters are stored in the CPU's shared memory for all training processes to access; When configuration two is selected, the complete model parameters will be saved to the solid-state drive; In both configurations, a layer buffer for temporarily storing prefetch parameters and a gradient buffer for temporarily storing unloaded gradients are created in the CPU's page-locked memory.

6. The method according to claim 1 or 5, characterized in that, The GPU memory space allocated in step S4 can simultaneously accommodate at least one layer packet that is currently being computed and one layer packet that is being prefetched.

7. The method according to claim 5, characterized in that, The hook functions registered in step S4 include: The forward pre-hook function is triggered before the forward computation or recomputation of the original layer and is used to prefetch the parameters of the original layer in the next layer package to be executed from the layer buffer to GPU memory. The forward hook function is triggered after the forward computation of the original layer is completed, and is used to release the space occupied by the parameters of that layer in GPU memory; The reverse hook function is triggered after the reverse computation of the original layer is completed. It is used to trigger the asynchronous unloading operation of the gradient of the previous original layer that has completed the reverse computation in step S6.

8. The method according to claim 7, characterized in that, During the recomputation process, after the original layer parameters are prefetched into GPU memory, their cache space in the layer buffer is reserved to receive the gradient data of that layer unloaded from the GPU after reverse computation.

9. The method according to claim 1, characterized in that, In step S7, the asynchronous parameter update is performed in parallel with the recalculation and reverse calculation of other layers on the GPU; and step S7 also includes: after the parameter update is completed, if the current configuration is two, the updated parameters are written back to the solid-state drive, and the cache space occupied by the layer parameters in the layer buffer is reclaimed.

Images