Heterogeneous sequence load-oriented flexible large model training method and system

Abstract

The invention discloses a heterogeneous sequence load-oriented flexible large model training method and system, and belongs to the technical field of large models. The method comprises the following steps: dividing a training data batch into a plurality of micro-batches; wherein one training data batch is composed of a plurality of sequences; for each micro-batch, distributing the sequences in the micro-batch to different SP groups by minimizing the maximum execution time of the micro-batch so as to generate an optimal sequence parallel strategy; and based on the optimal sequence parallel strategy, executing the training of the micro-batch on the large model, and accumulating gradients obtained by back propagation to update the parameters of the large model. According to the method, the optimal sequence parallel scheme can be solved according to the sequence length condition of the real-time workload during actual training, low-efficiency cross-node communication caused by high sequence parallelism is reduced, and the training efficiency of a large model in a long sequence scene is improved.

Description

Technical Field

[0001] The present invention relates to the field of large model technology, and in particular to a flexible large model training method and system for heterogeneous sequence loads. Background Art

[0002] In recent years, large language models (LLMs), represented by large Transformer models, have achieved amazing achievements in the field of artificial intelligence. Currently, there is a growing demand to extend the context length of LLMs (i.e., the maximum sequence length supported by LLMs). Therefore, the training of long-context LLMs has attracted widespread attention from academia and industry.

[0003] Training LLMs with long context lengths significantly increases the demand for device (e.g., GPU) video memory. To address this challenge, an effective strategy is to process the training input in parallel on multiple devices. In particular, sequence parallelism (SP) has become a key technology in long-context LLM training. In short, SP splits the training input in the sequence dimension and distributes different shards to multiple devices to amortize memory consumption. At the same time, to complete the training, devices also need to communicate during the forward and backward propagation processes to exchange intermediate results.

[0004] There are two existing systems that support LLM long sequence training, namely DeepSpeed ​​and Megatron-LM. They use different sequence parallel schemes to support long sequence LLM training.

[0005] DeepSpeed ​​uses the Ulysses-style sequence parallelism method to split the sequence dimension in the linear projection, dropout and normalization modules of the MLP and attention modules, and uses All-to-All operations to aggregate and distribute the sequences. In the Ulysses-style sequence parallelism, each device holds a part of the sequence And the complete model parameters Where N is the total sequence length, P is the SP parallelism, and d is the hidden layer dimension size. After completing the linear projection of the query, key, and value (Formula 1), the system collects the complete sequence on each device through three rounds of All-to-All communication (Formula 2), then performs multi-head attention operations on the complete sequence (Formula 3), and distributes the sequence to each device through another round of All-to-All communication (Formula 4).

[0006]

[0007] Megatron-LM uses context parallelism (CP). CP also splits on the sequence dimension. Compared with sequence parallelism, which splits the activation values ​​of the dropout and normalization modules but still requires the full sentence input for attention operations, CP goes a step further and distributes the attention operation to different devices. Specifically, CP distributes the sequence dimensions of queries, keys, and values ​​to multiple devices, and collects queries and values ​​through additional ring communications to complete the attention calculation. Although the additional communication volume is large, CP achieves the overlap of communication overhead and calculation by performing ring communication and attention calculation in blocks, and CP also achieves load balancing of calculations on different devices.

[0008] Although sequence parallel SP can support LLM training of long sequences, and the above-mentioned DeepSpeed ​​and Megatron-LM have also reached the most advanced level in distributed training systems, the training performance on variable-length sequence corpora is still unsatisfactory and there are serious efficiency issues. Specifically, existing sequence parallel schemes and SOTA training systems (such as Megatron-LM and DeepSpeed) are designed for homogeneous training data, assuming that the length of input training data is consistent, using a pre-set parallel degree and homogeneous parallel strategy, and only supporting training of fixed-length sequences.

[0009] First, long sequences require larger parallel groups, which leads to low efficiency. Megatron-LM and DeepSpeed ​​simply assume that the sequence length is uniform, so uniform parallelism is applied in each training task. To better illustrate the defects of uniform parallel design, the researchers evaluate its efficiency in the face of different sequence lengths through small tests. Specifically, for each row in Table 1, fixed-length sequences are generated until the total number of words reaches 4 million, and these sequences are trained using different parallelisms. Table 1 shows the end-to-end execution time and the proportion of all-to-all communication. Finally, it is found that due to memory limitations, long sequences require larger parallel groups. For example, a 32K sequence can be processed by a parallel group with a parallelism of 8, while a 128K sequence requires at least 32 parallelism to fit into the device memory. In addition, larger parallel groups may lead to inefficiency, mainly due to increased communication overhead. Smaller parallel groups generally perform better; for example, for 512 8K sequences, a parallel group with a parallelism of 64 results in more than 50% of communication time due to slow inter-node bandwidth. In contrast, the group with a parallelism of 8 reduces communication to only 7.8%, benefiting from the high-bandwidth intra-node connections.

[0010]

[0011] Table 1 Iteration time and All-to-All communication ratio

[0012] Second, real-world datasets exhibit skewness in the distribution of sequence lengths, exhibiting a long-tail distribution. Figure 1 The sequence length distributions of three popular large language model (LLM) training datasets are shown, including dataset 1 (GitHub dataset), dataset 2 (CommonCrawl dataset), and dataset 3 (Wikipedia dataset). Observations show that all three datasets show a clear unimodal long-tail distribution, with most sequences less than 8K in length, and only a few sequences exceeding 32K. Dataset 1 has the most overlong sequences, followed by dataset 2, and dataset 3 has the least. Therefore, in real-world datasets, sequence lengths show significant skewness and follow a long-tail distribution.

[0013] Finally, it is critical to match heterogeneous parallel groups with heterogeneous sequence lengths. Current systems mainly target uniform sequence lengths and adopt a single static parallel strategy throughout the training process. However, processing sequences with long-tailed length distributions requires larger parallel groups to accommodate overly long sequences, which will reduce the efficiency of shorter sequences. For example, when there are 128K sequences in the dataset, all 8K sequences are forced to use parallel groups of at least 32 and cannot enjoy the efficient small parallel groups. In addition, this inefficiency is particularly evident because short sequences are more common in skewed distributions.

[0014] In summary, the sequence lengths of real-world datasets vary and the loads are uneven, and there is an inherent contradiction between the isomorphic design of these systems and the heterogeneous characteristics of the datasets. This causes existing technical solutions to show efficiency defects when training on mixed-length corpora, mainly reflected in low communication efficiency. In scenarios where real-world datasets have varying sequence lengths and heterogeneous loads, how to efficiently design a flexible sequence parallel strategy that is adaptive to heterogeneous loads, thereby reducing unnecessary communication costs and achieving efficient training of heterogeneous loads, has become a technical problem to be overcome in this field. Summary of the invention

[0015] The present invention discloses a flexible large model training method and system (Flexible Sequence Parallelism, FlexSP) for heterogeneous sequence loads, which can solve the optimal sequence parallel scheme according to the sequence length of the real-time workload during actual training, reduce the inefficient cross-node communication caused by high sequence parallelism, and improve the efficiency of large model training in long sequence scenarios.

[0016] To achieve the above objectives, the technical solution of the present invention includes the following contents.

[0017] A flexible large model training method for heterogeneous sequence loads, the method comprising:

[0018] Divide a training data batch into several micro-batches; wherein a training data batch consists of several sequences;

[0019] For each micro-batch, the sequences in the micro-batch are assigned to different SP groups by minimizing the maximum execution time of the micro-batch to generate the optimal sequence parallel strategy;

[0020] Based on the optimal sequential parallel strategy, the micro-batch training is performed on the large model, and the gradients obtained by backpropagation are accumulated to update the parameters of the large model.

[0021] Furthermore, dividing a training data batch into a plurality of micro-batches includes:

[0022] sorting the sequences according to their lengths;

[0023] Based on minimizing the maximum total number of words per micro-batch, the micro-batch division of the sequence is performed on the sorted sequence queue to obtain a number of micro-batches.

[0024] Furthermore, based on minimizing the maximum total number of words in each micro-batch, the micro-batch division of the sequence is performed on the sorted sequence queue to obtain a plurality of micro-batches, including:

[0025] Calculate the optimal value when splitting the first k sequences into i mini-batches;

[0026] According to the micro-batch division scheme corresponding to the optimal value, micro-batch division of the sequence is performed on the sorted sequence queue.

[0027] Furthermore, by minimizing the maximum execution time of the micro-batch, the sequences in the micro-batch are assigned to different SP groups to generate an optimal sequence parallel strategy, including:

[0028] Based on the number of GPUs, determine the total number P of SP groups;

[0029] The optimization problem between the allocation of build sequences to different SP groups and the maximum execution time of micro-batches is: Where C represents the maximum execution time of a micro-batch, K represents the total number of sequences in the micro-batch, and the group selection vector m represents the number of selected SP groups and the parallelism of the SP groups. <m1,m2,…,m P >,m P =1 indicates the Pth SP group G P Selected, m P =0 indicates the Pth SP group G P Not selected, element A in the sequence allocation matrix Ak,p =1 indicates sequence S k Assigned to the Pth SP group G p , the element A in the sequence allocation matrix A k,p =0 means sequence S k Not assigned to the Pth SP group G p ;

[0030] Generate the constraints of the optimization problem; wherein the constraints include: SP group G p The execution time of each device on the micro-batch is not greater than the maximum execution time of the SP group G p The activation memory consumption of each device is not greater than the memory budget, the total parallelism of the selected SP group does not exceed the number of cluster devices, no sequence is assigned to an unselected SP group, and each sequence must and can only be assigned to one SP group;

[0031] Estimated SP Group G p Calculate the execution time and activation memory consumption of each device, and determine the constraints of the optimization problem based on the estimation results;

[0032] The optimization problem is solved under certain constraints to obtain an optimal sequential parallel strategy.

[0033] Furthermore, the estimated SP group G p Activation memory consumption per device on the , including:

[0034] Get the model state memory consumption M fixed on all devices ms and the activation memory cost M for each token token ;

[0035] Based on the sequence allocation matrix A, SP group parallelism d p 、Model state memory consumption M ms and activation memory cost M token , calculate SP group G p Activation memory consumption per device on .

[0036] Furthermore, the estimated SP group G p The execution time of each device on the network, including:

[0037] Extending the α-β model by taking the sequence length as an independent variable;

[0038] Estimation of G in SP group based on the extended α-β model p The computing time and communication time of each device on the network;

[0039] Comprehensive SP Group G p The calculation time and communication time of each device are obtained byp The execution time of each device on the

[0040] Furthermore, solving the optimization problem under the determined constraints to obtain the optimal sequence parallel strategy includes:

[0041] The sequences of similar length in the micro-batch are packaged by a bucketing algorithm and represented by a uniform length. When the bucketing algorithm is executed, the sequences in the micro-batch are sorted in ascending order of length, and the sequences are transferred based on the minimum error of dividing the first k sequences into q buckets to obtain the bucketing result.

[0042] Based on the bucketing results, a multi-process mixed integer linear programming problem is solved for the optimization problem to obtain an optimal sequential parallel strategy.

[0043] Furthermore, the performing of the micro-batch training on the large model based on the optimal sequence parallel strategy includes:

[0044] The optimal sequence parallel strategy is collected based on a distributed storage system, and an NCCL group pool is constructed to manage the SP group; wherein the SP group can be reused;

[0045] After the GPU iteratively and sequentially reads at least one optimal sequence parallel strategy from the distributed storage system, it calls the corresponding SP group in the NCCL group pool to train the large model; wherein the time for generating the optimal sequence parallel strategy for each micro-batch in a training data batch is equivalent to the time for the GPU to train the large model based on the training data batch.

[0046] A flexible large model training system for heterogeneous sequence loads, the system comprising:

[0047] A solver module is used to divide a training data batch into a number of micro-batches; wherein a training data batch consists of a number of sequences; for each micro-batch, by minimizing the maximum execution time of the micro-batch, the sequences in the micro-batch are assigned to different SP groups to generate an optimal sequence parallel strategy;

[0048] The executor module is used to execute the training of the micro-batch on the large model based on the optimal sequence parallel strategy, and accumulate the gradients obtained by back propagation to update the parameters of the large model.

[0049] An electronic device, characterized in that the electronic device comprises: a processor and a memory storing computer program instructions; when the processor executes the computer program instructions, it implements any of the above-mentioned flexible large model training methods for heterogeneous sequence loads.

[0050] Compared with the prior art, the present invention has at least the following beneficial effects.

[0051] Compared with the previous SOTA systems DeepSpeed ​​and Megatron-LM, the FlexSP system achieved a speedup of up to 72% and 98% in end-to-end training time, respectively.

[0052] The heterogeneous load adaptive flexible sequential parallel method of the FlexSP system can be applied to a variety of existing sequential parallel modes and has better adaptability and scalability.

[0053] The FlexSP system has better scalability. When the cluster size is expanded from 16 cards to 32 cards, FlexSP achieves a 1.91x speedup ratio, while the previous SOTA systems DeepSpeed ​​and Megatron-LM can only achieve 1.65x and 1.71x speedups respectively; when the cluster is expanded from 32 cards to 64 cards, FlexSP also has a 1.82x near-linear speedup ratio.

[0054] The heterogeneous load-adaptive flexible serial parallel technology proposed in the present invention can adaptively adjust the parallel strategy according to the actual workload. It has been verified that it has the best communication efficiency compared with the SOTA system under multiple data sets and multiple groups of longest sequence lengths, which is something that previous large model training systems could not do. BRIEF DESCRIPTION OF THE DRAWINGS

[0055] Figure 1 Distribution of sequence lengths of the training dataset.

[0056] Figure 2 Flexible sequential parallel example graph.

[0057] Figure 3 Architectural diagram of a flexible large-model training system for heterogeneous sequence workloads. DETAILED DESCRIPTION

[0058] The technical solutions in the embodiments of the present invention will be fully described below in conjunction with the accompanying drawings in the implementation of the present invention.

[0059] The flexible sequence parallel method for accelerating large model training of the present invention can adaptively adjust the parallel strategy and match the heterogeneous workload caused by sequences of different lengths through heterogeneous parallel groups. Specifically, the present invention forms small groups for short sequences to improve efficiency, while retaining large groups for long sequences to avoid memory overflow errors. In addition, the distribution of sequences is reasonably controlled to balance the execution time between all parallel groups. This flexible and heterogeneous strategy will significantly improve communication efficiency and overall system performance.

[0060] Figure 2 An example is shown in which a total of 64 devices are included and the context length of the training task is 192K. If the existing technology is used, a configuration with an SP degree of 32 is required, thus forming two SP=32 groups. Assume that the current training step needs to process five sequences of lengths 100K, 48K, 48K, 48K, and 48K, respectively. Due to the homogeneous parallel design adopted in the existing work, the shorter sequence (48K) must be processed with an SP degree of 32 no matter how the sequences are packaged. However, if heterogeneous SP groups are allowed (i.e., groups are allowed to have different SP degrees), a SP=32 group can be formed to process the long sequence (100K), and four SP=8 groups can be formed to process the short sequences (48K) respectively. By doing so, the processing of short sequences can be accelerated due to the reduced communication overhead. The common large language model (LLM) training corpus follows a long-tail distribution, i.e., there are few long sequences, while short sequences dominate. Therefore, by adjusting the SP group, the training of most sequences can be accelerated, thereby improving the overall performance.

[0061] Figure 3 The architecture diagram of a flexible sequential parallel system for accelerating large model training. The core of the FlexSP system consists of two modules: the solver module and the executor module.

[0062] 1. Solver module.

[0063] The solver mainly consists of two parts: Sequence Blaster (sequence splitter) and Parallelism Planner (parallel strategy solver). Given a batch of sequences with different lengths, Sequence Blaster is responsible for dividing different sequences in a batch into multiple micro-batches for gradient accumulation, while Parallelism Planner is responsible for solving the optimal sequence parallel strategy for each micro-batch. The executor executes a single iteration of training tasks based on the optimal sequence parallel strategy obtained by the solver.

[0064] 1. Sequence segmenter module.

[0065] During the training process of a large model, when the input batch contains too many sequences, these sequences may not be processed simultaneously due to limited memory capacity, so the optimization problem will have no feasible solution due to memory constraints. Gradient accumulation is a common technique for dealing with such situations. It splits the global data batch into multiple micro-batches, executes each micro-batch in turn, and accumulates model gradients for parameter updates. For training systems designed to handle uniform sequence lengths, micro-batch splitting is relatively simple - the present invention can simply fix the number of sequences in each micro-batch. However, in the scenario of the present invention where the input sequence lengths are uneven, micro-batch splitting is not an easy task.

[0066] To this end, the present invention designs a sequence cutter to divide the sequence into micro batches so that the parallelism planner can determine the optimal sequence parallel strategy. Given an input data batch B = {S k} has K sequences S k , the sequence cutter splits the sequence into M disjoint micro-batches {M i},satisfy The core idea of ​​FlexSP's Sequence Blaster is to minimize the number of micro-batches, reduce the variance of sequence lengths in a micro-batch, and ensure that the number of words in each micro-batch is as uniform as possible. This design of the splitter can make the video memory more stable during model training while ensuring training efficiency.

[0067] The detailed algorithm is as follows:

[0068] (1) The sequences are sorted in ascending order of length;

[0069] (2) Dynamic programming algorithm solution;

[0070] In order to balance the number of words in each micro-batch, the goal of the present invention is to minimize the maximum total number of words in each micro-batch, that is:

[0071]

[0072] Among them, M represents the number of batches to be divided, and j1,j2,…,j M Indicates the boundary of each batch.

[0073] DP[k][i] represents the optimal value when the first k sequences are divided into i micro-batches. The dynamic programming state transition equation is as follows:

[0074]

[0075] Among them, ∑ l∈[j+1,k] s l Indicates that in S j The total number of tokens in the i-th micro-batch when splitting the micro-batch. DP[K][M] represents the optimal solution. The optimized j i Split the global batch data into M micro-batches with balanced memory consumption.

[0076] 2. Parallel Strategy Solver Module

[0077] As mentioned above, given a batch of variable-length sequences, adopting a homogeneous sequence parallel strategy often cannot adapt to the heterogeneity of the workload, which will limit the training to inefficient operations. Choosing a suitable heterogeneous parallel strategy will help improve the efficiency of training. The current SOTA systems all use static homogeneous parallel strategies, and there is no mature algorithm to solve the optimal parallel strategy based on the workload.

[0078] To this end, the present invention aims to allocate sequence parallel groups of different sizes to sequences of different lengths according to the heterogeneity of sequence lengths, so as to automatically adapt to the heterogeneity of workloads during training and achieve higher acceleration ratios when training large models. The core idea of ​​the parallel strategy solver algorithm is to formalize the solution of the problem as a mixed integer linear programming problem, use the α-β model to model the computation and communication time overhead, and use the bucketing algorithm and multi-process solving to improve the solving speed.

[0079] (1) Formalization of optimization problem.

[0080] Given a sequence {S k}, and N devices with a memory budget E, the present invention needs to determine the following factors: (1) the number of SP groups, (2) the parallelism of each SP group, and (3) to which SP group each sequence should be assigned. At the same time, since sequences in different SP groups are processed in parallel, the optimization goal is to minimize the maximum execution time C of the micro-batch.

[0081] Since the candidate set of SP parallelism is very small and each sequence can only be assigned to one SP group, the present invention converts all decision variables into 0-1 integer variables. Specifically, the present invention assumes that there are P virtual SP groups, where the pth SP group G p The SP parallelism is d p For example, if there are 2 GPUs, the present invention has three virtual SP groups, whose SP parallelisms are 1, 1 and 2 respectively. Then, the present invention defines the group selection vector m= <m1,m2,…,m P >∈{0,1} P , where m p =1 indicates that SP group G is selected p , and m p =0 means that the SP group G is not selected p In this way, the number of SP groups and the parallelism of each SP group can be concisely represented by m. Next, the present invention further defines a sequence allocation matrix A∈{0,1} K×P , where A k,p =1 indicates sequence S k Assigned to G pBased on this, the present invention can construct a joint optimization problem to perform SP group selection and sequence allocation.

[0082]

[0083] ∑ p d p ×m p ≤N (8)

[0084]

[0085] Here, Time({s k ,A k,p};d p ) and Memory({s k ,A k,p};d p ) respectively represent the SP group G p The optimization goal is to minimize the maximum execution time of all SP groups (condition (6)). Condition (7) represents the memory constraint of each device in each SP group. Condition (8) indicates that the total parallelism of all selected SP groups should not exceed the number of cluster devices N. Condition (9) ensures that no sequence will be assigned to an unselected virtual SP group. Condition (10) requires that each sequence must and can only be assigned to one group.

[0086] (2) Cost Estimation

[0087] In order to solve the problem of optimizing the above formula (5), we need to accurately estimate Time ({s k ,A k,p};d p ) and Memory({s k ,A k,p};d p ). Given a model, when adopting Ulysses-style sequential parallelism, the memory overhead of the model state depends only on the ZeRO stage of the application and the current number of devices, and is not affected by the sequential parallel communication group and sequence allocation. Therefore, for the sequential parallel group G p For each device, its video memory cost can be estimated as:

[0088]

[0089] Among them, M ms represents the model state memory consumption fixed on all devices, M token represents the activation memory overhead of each word, A k,p Represents the sequence S k Is it assigned to G p .

[0090] Previous work has proposed effective distributed training cost estimation models for fixed-length sequences. They usually use the α-β model T = αW + β to estimate the time cost of computation and communication. However, existing work assumes that the sequence lengths are homogeneous and cannot accurately estimate the time cost of variable-length sequences. FlexSP extends the α-β model and uses sequence length as an independent variable to process real variable-length sequence data sets. The present invention estimates the computation and communication time as follows:

[0091]

[0092]

[0093] The α and β parameters represent the parameters of the α-β model, which are obtained through performance analysis. The present invention combines the two parts of the overhead to estimate the overall parallel execution time of sequences of different lengths, as shown below:

[0094] Time({s k , A k，p};d p )=T comp +T comm . (14)

[0095] (3) Problem Solving

[0096] According to the problem formalization and cost estimation, the problem of optimizing the above formula (5) is a mixed integer linear programming (MILP) problem. Although existing MILP solvers such as SCIP can solve MILP problems, there are too many decision variables involved in formula (5), which makes it difficult to obtain a feasible solution within the specified time. The present invention adopts a bucketing algorithm to reduce the complexity of the solution, and uses multi-process solution to improve the solution efficiency of linear programming, thereby obtaining the optimal sequential parallel strategy for each micro-batch.

[0097] The core idea of ​​bucketing is to pack sequences of similar length and represent them with a uniform length (usually the length of the longest sequence). The optimization goal of bucketing is to minimize the estimation error, that is:

[0098]

[0099] FlexSP uses a dynamic programming algorithm to perform efficient bucketing. The present invention first sorts the sequences in ascending order of length, and uses err[k][q] to represent the minimum error of dividing the first k sequences into q buckets. The transfer equation is as follows:

[0100]

[0101] here, Indicates that you are selecting As bucket B (q-1) q-1 The bucketing error of the qth bucket when the upper limit is .

[0102] 3. Actuator module

[0103] There are two stages for each batch of training data, namely, the solver derives the optimal plan (i.e., a combination of parallel strategy and sequence allocation), and the executor performs training. Since problem solving is performed on the CPU and training is performed on the GPU, the present invention decouples these two stages to facilitate overlapping processing. Specifically, on each node (machine), the present invention establishes a solver service that takes the length of a data batch as input and derives the optimal sequence parallel strategy for each micro-batch on the CPU. Subsequently, the present invention manages a distributed storage system to collect these optimal sequence parallel strategies, and the executor sequentially reads an optimal sequence parallel strategy for each iteration and performs training on the GPU. By doing so, FlexSP solves problems for multiple batches of data simultaneously, and the problem solving and training processes overlap.

[0104] The actuator module of the present invention mainly includes the initialization, multiplexing and switching of SP communication groups to support changing parallel strategies. FlexSP implements sequence parallelism in a hot-switching manner to handle diversified parallel strategies for different input data. For each micro-batch and its corresponding parallel strategy, FlexSP dynamically generates SP communication groups and disperses the data into corresponding groups. In order to avoid the overhead caused by redundant creation and storage of communication groups, FlexSP maintains an NCCL group pool to manage complex SP groups. Specifically, FlexSP generates communication groups in a timely manner, and only creates new groups when necessary, while existing groups are reused to optimize resource usage. FlexSP packs the sequences of a micro-batch at the same time, and uses the interface of the flash-attn library to implement attention calculations for variable-length sequence packing.

[0105] In summary, the present invention is a large-model distributed training system that can adaptively adjust the parallel strategy according to the actual workload. It has dynamic and heterogeneous characteristics, supports multiple sequence parallelisms to exist in one training step and dynamic adjustment of parallel strategies. Compared with the previous static homogeneous parallel system, FlexSP adopts dynamic heterogeneous sequence parallelism, better scalability, and better long sequence training performance.

[0106] Specifically, the technical solution proposed in the present invention can support LLM to select the optimal sequence parallel strategy according to the workload when training on a variable-length text corpus, and support the flexible formulation of heterogeneous sequence parallel strategies that are adaptive to heterogeneous loads. Existing distributed training systems that support large model long sequence training are all designed for homogeneous corpora, and do not support heterogeneous sequence lengths and heterogeneous adaptive strategies, which makes the existing systems inefficient in training on real data sets.

[0107] The functional performance of the product side can include serving as a distributed training system, providing large model training services in a variable-length sequence training environment, supporting heterogeneous and dynamically changing sequence parallelism; effectively reducing overhead during dynamic sequence parallel switching and reusing communication resources; effectively improving the problem of low communication efficiency during variable-length sequence training, making full use of high-bandwidth communication within the node, and improving the training speed.

[0108] Currently, there is no open source distributed training system on the market that supports flexible sequence parallelism that is adaptive to heterogeneous loads. Megatron-LM and DeepSpeed ​​use a homogeneous single parallel mode, which is inefficient and wastes resources when training on a variable-length sequence corpus. The FlexSP system developed by the present invention is the first large-model training system that can adaptively adjust the parallel strategy when facing heterogeneous workloads of variable-length sequences. Extensive experimental results have verified that FlexSP is always superior to the SOTA system in terms of training efficiency, especially in long-context LLM training, achieving up to 1.98 times performance improvement compared to existing SOTA systems.

[0109] It should be clear that the specific implementation and drawings of the present invention disclosed above are intended to help understand the content of the present invention and implement it accordingly. For those of ordinary skill in the art, it can be understood that various substitutions, changes and modifications are within the scope of the present invention without departing from the principles and spirit of the present invention, and the scope of the present invention is defined by the attached claims and their equivalents.

Claims

1. A flexible large model training method for heterogeneous sequence loads, characterized in that: The method comprises: Divide a training data batch into several micro-batches; wherein a training data batch consists of several sequences; For each micro-batch, the sequences in the micro-batch are assigned to different SP groups by minimizing the maximum execution time of the micro-batch to generate the optimal sequence parallel strategy; Based on the optimal sequential parallel strategy, the micro-batch training is performed on the large model, and the gradients obtained by backpropagation are accumulated to update the parameters of the large model.

2. The method according to claim 1, characterized in that The method of dividing a training data batch into several micro-batches includes: sorting the sequences according to their lengths; Based on minimizing the maximum total number of words per micro-batch, the micro-batch division of the sequence is performed on the sorted sequence queue to obtain a number of micro-batches.

3. The method according to claim 2, characterized in that Based on minimizing the maximum total number of words in each micro-batch, the micro-batch division of the sequence is performed on the sorted sequence queue to obtain a number of micro-batches, including: Calculate the optimal value when splitting the first k sequences into i mini-batches; According to the micro-batch division scheme corresponding to the optimal value, micro-batch division of the sequence is performed on the sorted sequence queue.

4. The method according to claim 1, characterized in that: By minimizing the maximum execution time of the micro-batch, the sequences in the micro-batch are assigned to different SP groups to generate the optimal sequence parallel strategy, including: Based on the number of GPUs, determine the total number P of SP groups; The optimization problem between the allocation of build sequences to different SP groups and the maximum execution time of micro-batches is: Where C represents the maximum execution time of a micro-batch, K represents the total number of sequences in the micro-batch, and the group selection vector m represents the number of selected SP groups and the parallelism of the SP groups. <m1,m2,…,m P >,m P =1 indicates the Pth SP group G P Selected, m P =0 indicates the Pth SP group G P Not selected, element A in the sequence allocation matrix A k,p =1 indicates sequence S k Assigned to the Pth SP group G p , the element A in the sequence allocation matrix A k,p =0 means sequence S k Not assigned to the pth SP group G p ; Generate the constraints of the optimization problem; wherein the constraints include: SP group G p The execution time of each device on the micro-batch is not greater than the maximum execution time of the SP group G p The activation memory consumption of each device is not greater than the memory budget, the total parallelism of the selected SP group does not exceed the number of cluster devices, no sequence is assigned to an unselected SP group, and each sequence must and can only be assigned to one SP group; Estimated SP Group G p Calculate the execution time and activation memory consumption of each device, and determine the constraints of the optimization problem based on the estimation results; The optimization problem is solved under certain constraints to obtain an optimal sequential parallel strategy.

5. The method according to claim 4, characterized in that The estimated SP group G p Activation memory consumption per device on the , including: Get the model state memory consumption M fixed on all devices ms and the activation memory cost M for each token token ; Based on the sequence allocation matrix A, SP group parallelism d p 、Model state memory consumption M ms and activation memory cost M token , calculate SP group G p Activation memory consumption per device on .

6. The method according to claim 4, characterized in that The estimated SP group G p The execution time of each device on the network, including: Extending the α-β model by taking the sequence length as an independent variable; Estimation of G in SP group based on the extended α-β model p The computing time and communication time of each device on the network; Comprehensive SP Group G p The calculation time and communication time of each device are obtained by p The execution time of each device on the 7. The method according to claim 4, characterized in that Solving the optimization problem under the determined constraints to obtain the optimal sequence parallel strategy includes: The sequences of similar length in the micro-batch are packaged by a bucketing algorithm and represented by a uniform length. When the bucketing algorithm is executed, the sequences in the micro-batch are sorted in ascending order of length, and the sequences are transferred based on the minimum error of dividing the first k sequences into q buckets to obtain the bucketing result. Based on the bucketing results, a multi-process mixed integer linear programming problem is solved for the optimization problem to obtain an optimal sequential parallel strategy.

8. The method according to claim 1, characterized in that The method of executing the micro-batch training on the large model based on the optimal sequence parallel strategy includes: The optimal sequence parallel strategy is collected based on a distributed storage system, and an NCCL group pool is constructed to manage the SP group; wherein the SP group can be reused; After the GPU iteratively and sequentially reads at least one optimal sequence parallel strategy from the distributed storage system, it calls the corresponding SP group in the NCCL group pool to train the large model; wherein the time for generating the optimal sequence parallel strategy for each micro-batch in a training data batch is equivalent to the time for the GPU to train the large model based on the training data batch.

9. A flexible large model training system for heterogeneous sequence loads, characterized in that: The system comprises: A solver module is used to divide a training data batch into a number of micro-batches; wherein a training data batch consists of a number of sequences; for each micro-batch, by minimizing the maximum execution time of the micro-batch, the sequences in the micro-batch are assigned to different SP groups to generate an optimal sequence parallel strategy; The executor module is used to execute the training of the micro-batch on the large model based on the optimal sequence parallel strategy, and accumulate the gradients obtained by back propagation to update the parameters of the large model.

10. An electronic device, characterized in that: The electronic device comprises: a processor and a memory storing computer program instructions; when the processor executes the computer program instructions, the flexible large model training method for heterogeneous sequence loads as described in any one of claims 1 to 8 is implemented.

Images