Large-scale language model automatic parallel training method and system based on parallel monte carlo algorithm

Abstract

The application discloses a large-scale language model automatic parallel training method and system based on a parallel Monte Carlo algorithm, and the method steps comprise the following steps: in step S01, a large-scale language model and a device topology structure of a computing cluster are acquired; in step S02, the acquired large-scale language model is converted into an operator sequence and a corresponding calculation graph, and a model representation of the model is generated; in step S03, a unified parallel strategy search space containing at least three dimensions of data parallelism, model parallelism and pipeline parallelism is constructed according to the operator sequence, the calculation graph, the model representation and the device topology structure; in step S04, an optimal parallel training strategy is searched in the unified parallel strategy search space by using a parallel Monte Carlo tree search algorithm; and in step S05, the large-scale language model is trained according to the optimal parallel training strategy. The application can realize automatic parallel training of the large-scale language model, reduce the time and resources required for training, and improve the training efficiency.

Description

Technical Field

[0001] This invention relates to the field of large-scale language model training technology, and in particular to an automatic parallel training method and system for large-scale language models based on the parallel Monte Carlo algorithm. Background Technology

[0002] In the field of natural language processing, Transformer models have laid the foundation for the scalability of models, such as BERT, GPT-2, Megatron-LM, T5, and the latest large-scale model GPT4. The number of parameters has grown exponentially, and with the increase in parameters, the model's capacity and representational power have also increased accordingly. Therefore, models with larger parameter sizes and more layers have become the current research trend. To meet the needs of training large-scale neural network models with hundreds of billions or even trillions of parameters, distributed parallel training strategies have emerged. Mainstream strategies include data parallelism, model parallelism, and pipelined parallelism. The development of distributed parallel training strategies has also promoted the continuous increase in the size of neural network models. Automated parallel training aims to find the optimal parallel training strategies for large-scale language models in data parallelism, model parallelism, and pipelined parallelism, attempting to eliminate the manual design work of distributed algorithms. How to achieve efficient automated parallel training of large-scale language models is a current research focus.

[0003] Traditional distributed parallel training methods mainly include three training approaches: data parallelism, model parallelism, and pipeline parallelism. Data parallelism involves dividing the training data equally according to the number of computing nodes and replicating the model on each node. This reduces the computational overhead of each node and the corresponding training time. However, communication is required during backpropagation to ensure convergence. While data parallelism is simple to implement, communication overhead is a problem that needs to be addressed. Furthermore, replicating the model on each node can lead to significant memory redundancy. Therefore, data parallelism cannot be used alone for training large-scale models. Model parallelism divides the model into training layers into several sub-models, with each computing node storing one sub-model. This effectively utilizes the GPU memory of computing nodes, but it is difficult to implement, requiring extensive manual experience and specialized knowledge in model training. Pipeline parallelism divides the model into several pipeline stages between training layers, training the entire model through these pipeline stages. Similar to model parallelism, pipeline parallelism can effectively utilize the memory of each computing node, but the partitioning of pipeline parallelism requires sufficient research experience.

[0004] Traditional distributed parallel training methods mainly suffer from the following problems: 1) Inability to meet the increasing training demands of trillion-scale parameter models As the number of parameters in large-scale language models continues to increase, traditional automated parallel training methods, including data parallelism, model parallelism, and pipeline parallelism, can no longer meet the current training needs of trillions of parameters. While new parallel dimensions such as recomputation and ZeRO can further improve parallel efficiency, the search space for these new parallel dimensions in automated parallel training scenarios exhibits an exponential growth trend, making it difficult to find the optimal parallel strategy. Therefore, these new parallel dimensions cannot be directly adapted to automated parallel training scenarios.

[0005] 2) Solving the search algorithm takes a long time. Currently, the search algorithms commonly used in automated parallel training are time-consuming. For example, the commonly used Monte Carlo search algorithm does not rely on prior knowledge of the problem domain or manually designed heuristic functions. It explores the decision space through a large number of random simulations and can effectively discover potential optimal strategies. However, the Monte Carlo search algorithm usually requires a large number of random simulations, which consumes a lot of computing time and resources. When dealing with complex problems or when high accuracy is required, the computational cost will be very high, and its convergence speed is slow, especially when the problem size is large or the decision space is very complex, it takes a long time to reach the convergence state. Summary of the Invention

[0006] The technical problem to be solved by this invention is: in view of the technical problems existing in the prior art, this invention provides a method and system for automatic parallel training of large-scale language models based on parallel Monte Carlo algorithm, which is simple to implement, requires less training time and resources, and is highly efficient.

[0007] To solve the above-mentioned technical problems, the technical solution proposed by this invention is as follows: An automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm, comprising the following steps: Step S01. Data Acquisition: Obtain the device topology of the large-scale language model and computing cluster; Step S02. Model Representation Generation: Convert the acquired large-scale language model into a sequence of operators and the corresponding computation graph, and generate the model representation of the model; Step S03. Multidimensional search space construction: Based on the operator sequence, computation graph, model representation and device topology, construct a unified parallel strategy search space that includes at least three dimensions: data parallelism, model parallelism and pipeline parallelism. Step S04: Parallel Monte Carlo Search: Within the unified parallel strategy search space, multiple computing units use the parallel Monte Carlo tree search algorithm to search for the optimal parallel training strategy. During the search process, the search space is divided into different sub-search spaces, and the search process of each sub-search space is assigned to an independent computing unit. The results of each computing unit are summarized and the node statistics are updated. The optimal parallel training strategy includes the optimal operator parallel strategy and the optimal pipeline partitioning label. Step S05. Model Training: Train the large-scale language model according to the optimal parallel training strategy.

[0008] Further, in step S02, the acquired large-scale language model is converted into an operator sequence Q, which includes N operators. A computation graph G is constructed based on the dependencies between the operators in the operator sequence Q. The two vertices of each edge in graph G are the two operators in the operator sequence Q.

[0009] Further, in step S02, the model representation includes JaxPr representation and HLO representation. The JaxPr representation is used to obtain the communication overhead and computation overhead of each operator, and the HLO representation is used to obtain the point-to-point communication overhead between pipelined label candidates.

[0010] Furthermore, in step S03, the unified parallel strategy search space also includes a tensor recomputation dimension and / or a zero-redundancy optimizer dimension. The tensor recomputation dimension is configured with a flag to indicate whether recomputation is activated during training to optimize memory usage. The zero-redundancy optimizer dimension is configured with a flag to indicate whether the zero-redundancy optimizer is used during training and to evaluate the additional cost of the zero-redundancy optimizer dimension during training by calculating the additional time cost of the zero-redundancy optimizer.

[0011] Further, step S04 includes: Step S401. Divide the unified parallel strategy search space into multiple sub-search spaces according to the number of computing units, and allocate an independent computing unit to each sub-search space; Step S402. Each computing unit independently executes a Monte Carlo tree search loop within its allocated subspace to search for the optimal candidate strategy and evaluates the cost corresponding to the optimal candidate strategy found. Step S403. Summarize the optimal candidate strategies and their evaluated costs found by each computing unit in the subspace, and select the strategy with the lowest cost as the search result for the current round. The cost includes computational cost and / or communication cost. Step S404. Redivide the search space based on the search results of the current round, return to step S402, and continue until the iteration stopping condition is met to obtain the final optimal parallel training strategy output.

[0012] Further, step S402 includes: Step S421. The current computing unit selects one policy from all candidate policies as the current policy; Step S422. The current computing unit randomly generates a new policy from all candidate policies. Calculate the current policy and the new policy respectively. The corresponding costs depend on the current strategy and the new strategy. The corresponding cost determines whether to accept the new strategy. If so, then the new strategy will be implemented. As the current strategy, a new strategy is adopted. probability for:

[0013] in, , These are the current strategy and the new strategy. The corresponding cost, This represents a hyperparameter used to control the weight of the difference between the two parallel strategies before and after the update; Step S423. The current computing unit determines whether the search stopping condition has been met. If so, the final optimal candidate strategy is output; otherwise, return to step S422.

[0014] Furthermore, in step S04, the solution objective when searching the search space in the three dimensions of data parallelism, model parallelism, and pipeline parallelism is:

[0015] in, For operators The computational overhead, For operators Communication overhead, For operators in the operator sequence Q and The overhead of re-partitioning between them Pipeline partitioning occurs in the operator Subsequent point-to-point communication overhead, N This represents the number of operators in an operator sequence Q. This represents the total cost of searching the search space across the three dimensions of data parallelism, model parallelism, and pipeline parallelism. S is the set of all operator parallelism strategies, P is the set of all pipeline partitioning tags, and R is the set of repartitioning operator pairs.

[0016] Furthermore, when the unified parallel strategy search space also includes the tensor recomputation dimension to form a 4D parallel search space, the solution objective for searching in the 4D parallel search space is:

[0017]

[0018] in, Let represent the objective function for solving the 4D parallel search space, i.e., minimizing the total cost of the 4D parallel search space. Let represent the objective function for solving the 3D parallel search space, i.e., minimizing the total cost of the 3D parallel search space. This represents the time cost of memory swapping in 4D parallel training mode. It is the PCIe bandwidth of the training cluster devices. It is the unloading operator space The Middle The memory overhead of each JaxPr operator.

[0019] Furthermore, when the unified parallel strategy search space also includes the tensor recomputation dimension and the zero-redundancy optimizer dimension to form a 5D parallel search space, the solution objective for searching in the 5D parallel search space is:

[0020]

[0021] in, Let represent the objective function for solving the 5D parallel search space, i.e., minimizing the total cost of the 5D parallel search space. Let represent the objective function for solving the 4D parallel search space, i.e., minimizing the total cost of the 4D parallel search space. This indicates the additional cost of zero-redundancy optimization parallelism during the training process in 5D parallel training mode. It refers to the number of model parameters. It refers to the bandwidth between GPUs. This indicates the number of optimizer states.

[0022] Furthermore, in step S04, during the process of multiple computing units searching for the optimal parallel training strategy using the parallel Monte Carlo tree search algorithm, the method also includes using GPU for data parallel acceleration and using JAX sampling before the search begins to complete the selection of the search space. Specifically, the operator-level policy space is first mapped to a multidimensional dense tensor in GPU memory. In the search preprocessing stage, the cost model is transformed into a GPU kernel function using JAX. During the search process, multiple Markov chains are driven in parallel on the GPU using multiple independent random seeds. Each chain performs policy sampling through tensor indexing operations and calls the cost model function in the GPU for performance evaluation.

[0023] A computer system includes a processor and a memory, the memory being used to store a computer program, and the processor being used to execute the computer program to perform the method described above.

[0024] Compared with the prior art, the advantages of the present invention are as follows: 1. This invention addresses the automatic parallel training of large-scale language models by constructing a unified parallel strategy search space encompassing at least three dimensions: data parallelism, model parallelism, and pipeline parallelism. This avoids the suboptimal selection defects of traditional hierarchical solution methods. Furthermore, by implementing a parallelized Monte Carlo search algorithm, the search tree is divided into different subtrees, with the search process of each subtree assigned to an independent computational unit. Simulation tasks are performed simultaneously in different computational units, and the results are finally aggregated through shared memory or message passing mechanisms to update node statistics. This reduces the search time and resource consumption of the search algorithm in high-dimensional search spaces, thereby improving training efficiency.

[0025] 2. This invention employs the Monte Carlo search algorithm to search for the optimal distributed parallel strategy for large-scale language models within a unified 3D parallel search space, thereby eliminating the dependence on mathematical models and improving the feasibility of practical applications.

[0026] 3. This invention further extends high-dimensional distributed parallel technology by using a unified three-dimensional parallel strategy search space. It couples recomputation technology and zero-redundancy optimization with the 3D search space, and combines Monte Carlo search algorithm to search for the optimal parallel training strategy in the multi-dimensional search space. This enables 4D / 5D parallel training, further improves the search space, and meets the training needs of large-scale language models. Attached Figure Description

[0027] Figure 1 This is a schematic diagram illustrating the implementation process of the automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm in this embodiment.

[0028] Figure 2This is a schematic diagram illustrating the principle of constructing a 3D parallel strategy search space in this embodiment.

[0029] Figure 3 This is a schematic diagram illustrating the principle of constructing a 4D parallel strategy search space in this embodiment.

[0030] Figure 4 This is a schematic diagram illustrating the principle of the large-scale language model automatic parallel training system architecture constructed in this embodiment. Detailed Implementation

[0031] The present invention will be further described below with reference to the accompanying drawings and specific preferred embodiments, but this does not limit the scope of protection of the present invention.

[0032] To facilitate understanding, the relevant technical background of the present invention will first be introduced by way of example.

[0033] For training large-scale models with hundreds of millions or even trillions of parameters, current GPUs cannot meet the memory requirements of training high-volume models using data parallelism. To train such large-scale language models, model parallelism and pipelined parallelism have been proposed. These three dimensions of distributed parallel training—data parallelism, model parallelism, and pipelined parallelism—together constitute 3D parallel training technology. Training large-scale language models requires expensive and time-consuming expert-designed strategies, while automatic parallel training requires automatically generating the optimal training strategy for the input model based on the input model and cluster topology. Existing parallel strategies include algorithms such as OptCNN, FlexFlow, Alpa, and Unity. Among them, OptCNN uses a graph search-based dynamic programming method to find the optimal parallel strategy for convolutional neural networks. FlexFlow constructs a search space combining data parallelism and model parallelism based on samples, operators, attributes, and parameters, and uses the Monte Carlo search algorithm to find the optimal parallel strategy in this search space. However, these two types of algorithms can only explore parallel strategies in a two-dimensional search space and do not consider pipelined parallelism, which cannot meet the needs of large-scale model training. The ALPA algorithm categorizes operator-level parallelism methods into two types: inter-operator parallelism and intra-operator parallelism. Intra-operator parallelism consists of data parallelism and model parallelism, while inter-operator parallelism specifies pipeline parallelism. It employs a hierarchical approach, using Integer Linear Programming (ILP) to solve the dynamic programming problems of intra-operator and inter-operator parallelism to seek the optimal parallel training strategy. However, this approach suffers from problems such as suboptimal solutions due to hierarchical approaches and reliance on complex mathematical models like ILP / DP. Furthermore, its search space is limited to the traditional 3D parallel strategy search space. The Unity algorithm combines operator fusion and operator parallelism methods, employing a Monte Carlo search algorithm to solve for the optimal parallel training strategy for large-scale language models within the entire 3D parallel strategy search space. However, the search space dimension of this algorithm is still limited to 3D space, and the Monte Carlo search algorithm suffers from slow speed and difficulty in convergence during the search process. Furthermore, given the model and training cluster topology, traditional algorithms for finding the optimal parallel strategy often rely heavily on mathematical formulas, and the traditional 3D automatic parallel search process uses a non-uniform hierarchical space (such as the Alpa algorithm), which cannot achieve a unified search space for 3D parallelism, thus posing serious difficulties for exploring more parallel training techniques.

[0034] Automated parallel training of large-scale language models aims to find the optimal partitioning strategy within a 3D parallel search space and attempts to move away from the manual design of distributed training algorithms. The partitioning strategy in distributed parallel training refers to the planning method for rationally dividing and distributing the model and data when using multiple computing devices (such as multiple GPUs, multiple servers, etc.) for model training. Existing methods for finding the optimal partitioning strategy for large-scale models lack a unified 3D parallel model and heavily rely on mathematical models, making them unsuitable for practical applications. Furthermore, the limited search space results in high training costs and restricted application scope for large-scale language models.

[0035] Monte Carlo Tree Search (MCTS) is a heuristic search algorithm for decision-making processes. It constructs a search tree where each node represents an intermediate state and edges represent actions between states. Starting from a given state, it performs random simulations, simulating results according to a random strategy without relying on a search algorithm, until the final state. Each node records statistical information such as the number of visits and cumulative rewards to evaluate the merits of playing the game from that node. The algorithm flow includes: 1) Selection: Traversing the search tree from the root node to select a leaf node based on a strategy balancing exploration and exploitation; 2) Expansion: Expanding non-termined leaf nodes with unexplored actions to create new child nodes; 3) Simulation: Randomly simulating from the expanded nodes until the game terminates, with the results used to update node statistics; 4) Backpropagation: Updating the statistics of all nodes from the leaf nodes back to the root node based on the simulation results.

[0036] Monte Carlo search algorithms do not rely on prior knowledge of the problem domain or manually designed heuristics, making them relatively easy to apply to various problem types, including complex problems where manually designing heuristics is difficult. By exploring the decision space through numerous random simulations, they effectively discover potential optimal policy paths, especially effective when dealing with complex, high-dimensional decision spaces, avoiding getting trapped in local optima. Furthermore, the performance of Monte Carlo search algorithms improves with increased computational resources; by increasing the number of simulations, the decision space can be explored more deeply, allowing the algorithm to approach the optimal solution as close as possible within a limited time. However, to obtain accurate results, Monte Carlo search algorithms typically require a large number of random simulations, consuming significant computational time and resources, especially when the problem size is large or the decision space is highly complex, taking a long time to reach convergence.

[0037] This invention addresses the automatic parallel training of large-scale language models by constructing a unified parallel strategy search space encompassing at least three dimensions: data parallelism, model parallelism, and pipeline parallelism. This avoids the suboptimal selection defects of traditional hierarchical solution methods. Furthermore, by implementing a parallelized Monte Carlo search algorithm, the search tree is divided into different subtrees, with the search process of each subtree assigned to an independent computing unit. The simulation task is performed simultaneously on different computing units (such as different CPU cores or GPU threads). The results are then aggregated through shared memory or message passing mechanisms to update node statistics. This reduces the search time and resource consumption of the search algorithm in the high-dimensional search space, improving training efficiency. Moreover, by employing the Monte Carlo search algorithm to search for the optimal distributed parallel strategy for large-scale language models within a unified 3D parallel search space, it also eliminates the dependence on mathematical models, enhancing its feasibility in practical applications.

[0038] like Figure 1 As shown, the detailed steps of the automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm in this embodiment are as follows: Step S01. Data Acquisition: Acquire the device topology of the large-scale language model and computing cluster.

[0039] In this embodiment, the model structure of the large language model and the device topology of the computing cluster are obtained simultaneously. This device topology is a cascaded device structure (X, Y), representing the connections between computing graphics cards, where X represents the number of nodes in the computing cluster and Y represents the number of computing devices in each node. The model structure of the large language model and the cascaded device structure (X, Y) of the computing cluster are obtained, where X is the number of nodes in the computing cluster and Y is the number of computing devices in each node.

[0040] Step S02. Model Representation Generation: Convert the acquired large-scale language model into a sequence of operators and the corresponding computation graph, and generate the model representation of the model.

[0041] In this embodiment, the acquired large-scale language model is converted into an operator sequence Q, which includes N operators. A computation graph G is constructed based on the dependencies between operators in the operator sequence Q. The two vertices of each edge in graph G are the two operators in the operator sequence Q. The model representation includes JaxPr representation and HLO representation. JaxPr representation is used to obtain the communication overhead and computation overhead of each operator, and HLO representation is used to obtain the point-to-point communication overhead between pipeline label candidates.

[0042] Specifically, firstly, based on the model form of the input large-scale language model M, its mathematical expression under the input data and corresponding labels (X, Y) is formalized. For the input model M, it can be transformed into a sequence of operators Q, consisting of N operators. Composition, input data is Therefore, there is To meet the constraints and tensor shape requirements of deep neural networks, this embodiment also establishes a computational graph G based on the M model. Assume that the set of edges in G is E, and each edge consists of two vertices, i.e. and In fact, and These are also two items in Q. Assume... and There are respectively and Segmentation strategy. and yes and Partitioning strategy in S. and After the division, the shapes should match.

[0043] Step S03. Multidimensional search space construction: Based on the operator sequence, computation graph, model representation and the device topology, construct a unified parallel strategy search space that includes at least three dimensions: data parallelism, model parallelism and pipeline parallelism.

[0044] In the process of automated parallel training, key training elements such as data and model operators are partitioned to realize all possible solutions for distributed training, thus forming the search space of the automated parallel system. This embodiment adds pipeline parallelism to the search space, constructing a unified 3D parallel strategy search space encompassing data parallelism, model parallelism, and pipeline parallelism. This allows for strategy search based on a unified target model, decision variables, and search algorithm during the strategy search process.

[0045] Considering the device topology and assuming an e-dimensional tensor to be computed in a j×k device topology, it has e×j×k+1 candidate partitioning strategies, where j is the number of hosts and k is the number of devices on each host. Then, N operators have a total of (e×j×k+1)N candidate partitioning strategies. In this embodiment, when pipeline parallelism is added to the search space, the number of candidate partitioning strategies increases to (e×j×k+1)PN, where P is the number of pipeline stage candidates. All partitioning strategy candidates constitute the search space. This embodiment assumes There are n search dimensions. If a grid search algorithm is used, the time complexity is O((e×j×k +1)MN). Clearly, the search candidate partitioning strategies are finite, and the search can be performed completely enumerated in finite time. This applies when the Markov chain is irreducible, aperiodic, and follows a stationary distribution. At that time, the Monte Carlo search algorithm converges. Therefore, for each constant... There exists a and ,in . The objective of this embodiment is to find the optimal parallel training strategy {S, P} based on the input model M and the device topology T, where S is the set of parallel strategies within all operators, and P is the set of pipeline partitioning labels. For each operator... It has computational overhead and communication overhead In each parallel strategy, the operator and There is a repartitioning overhead, R(i,j). For pipeline partitioning, there is point-to-point communication between each pipeline stage. If this point-to-point communication occurs in the operator... Its expenses thereafter are .

[0046] As a preferred implementation, in the 3D parallel strategy search space of data parallelism, model parallelism, and pipeline parallelism (e.g. Figure 2 As shown in the figure, for each strategy s S, the objective of the search is: (1) in, For operators The computational overhead, For operators Communication overhead, For operators in the operator sequence Q and The overhead of re-partitioning between them Pipeline partitioning occurs in the operator Subsequent point-to-point communication overhead, N This represents the number of operators in an operator sequence Q. S represents the total cost of searching the search space across the three dimensions of data parallelism, model parallelism, and pipeline parallelism, where S is the parallelism strategy for all operators. s The set, It is a set of all pipeline partition markers. It is a set of repartitioning operator pairs.

[0047] Furthermore, the unified parallel strategy search space also includes tensor recomputation dimensions and / or zero-redundancy optimizer (ZeRO) dimensions to construct high-dimensional distributed parallel strategies. The tensor recomputation dimension uses a flag to indicate whether recomputation is activated during training to optimize memory usage. The zero-redundancy optimizer dimension uses a flag to indicate whether the zero-redundancy optimizer is used during training, and evaluates the additional cost of the zero-redundancy optimizer dimension during training by calculating its additional time cost. The search method for the recomputation dimension involves layer-by-layer or extruding tensors during the forward pass of training the deep neural network and recomputing them during the backward pass, i.e., the checkpointing technique. The zero-redundancy optimizer is implemented by partitioning the optimizer state, gradients, and model parameters. With a data parallelism of N, the optimizer state, gradients, and parameters are divided into N groups, and each GPU or process is responsible for storing and updating 1 / N of these values.

[0048] Specifically, based on the 3D parallel policy search space, a search for tensor recalculation dimensions can be added to the search space to form a 4D parallel policy search space, that is, a checkpoint-parallel policy search is added to the search space, such as... Figure 3 As shown. For the fourth dimension of checkpoint parallelism, a flag is set to determine whether memory optimization is used during training. Assume a sharding strategy. The memory overhead is That is, the parallel strategy searched in the 3D parallel search space. The video memory overhead is The fourth checkpoint of parallel training begins when parallelism starts. The corresponding strategy search will be initiated at that time. This refers to the available memory in each GPU device. For the fourth dimension of checkpoint parallelism, a flag CkptPara is set to determine whether checkpoint parallelism is enabled during training. When memory consumption in each GPU exceeds the memory budget, the flag CkptPara is set to True; otherwise, it is set to False. The policy search process attempts to find the optimal partitioning policy in the new search space. That is, when... When this happens, the flag will be true, and a reduction is needed ( Memory is allocated to accommodate the training of large-scale language models. This embodiment enables memory release by using a recomputation method as a checkpoint parallel selection in the training system.

[0049] The 4D parallel training mode needs to consider the additional time overhead caused by operator swapping. The time cost of memory swapping can be calculated as follows: (2) in It is the PCIe bandwidth of the training cluster devices. It is the unloading operator space The Middle The memory overhead of each JaxPr operator. The goal of Parallel Scale in the 4D parallel search space is: (3) in, Let represent the objective function for solving the 4D parallel search space, i.e., minimizing the total cost of the 4D parallel search space. Let represent the objective function for solving the 3D parallel search space, i.e., minimizing the total cost of the 3D parallel search space. This represents the time cost of memory swapping in 4D parallel training mode.

[0050] Extending 4D parallel automation based on traditional methods requires reconstructing the parameter partitioning logic. In this embodiment, the ParaScale parallel strategy only requires adding a checkpoint cost term (i.e., formula (2)) to formula (1) to extend to the 4D parallel solution objective (formula (3)). Subsequent addition of corresponding cost terms can also support 5D parallelism (such as adding expert parallelism) without modifying the search framework. This provides a scalable foundation for the automation of the latest hybrid parallel technologies (such as MoE + recomputation).

[0051] Furthermore, based on the 4D parallel policy search space, a zero-redundancy optimizer dimension is added to the search space, forming a 5D parallel policy search space. For the fifth dimension, a flag is also set to determine whether a zero-redundancy optimizer is used during training. Similar to recomputation, this embodiment defines... Evaluate the additional cost of zero-redundancy optimization parallelism during training. This is because ZeRO-1 (the stage in the ZeRO method that implements partitioning of the optimizer state) partitions the optimizer state... Partitioning between them results in additional The communication traffic is used to collect the optimizer states, where This refers to the number of model parameters. Therefore, the additional time cost is as follows: (4) in It refers to the bandwidth between GPUs. This indicates the number of optimizer states.

[0052] Specifically, 5D parallelism adds support for ZeRO technology to 4D parallelism. ZeRO is also a runtime technology during training. In this embodiment, a runtime flag is added to the search space to indicate whether ZeRO parallelism is used during training. Therefore, the objective of Parallel Scale in the 5D parallel search space can be expressed as: (5) in, Let represent the objective function for solving the 5D parallel search space, i.e., minimizing the total cost of the 5D parallel search space. Let represent the objective function for solving the 4D parallel search space, i.e., minimizing the total cost of the 4D parallel search space. This indicates the additional cost of zero-redundancy optimization parallelism during the training process in 5D parallel training mode.

[0053] This embodiment forms a multi-dimensional parallel algorithm by coupling recomputation technology and zero-redundancy optimization on the basis of a unified 3D parallel strategy search space. It also uses the Monte Carlo search algorithm to search for the optimal parallel training strategy in the 4D or even 5D search space, which can realize 4D parallel training and 5D parallel training, and make up for the shortcomings of traditional methods in solving 4D parallel training and 5D parallel training.

[0054] This invention constructs a unified multidimensional parallel strategy search space and extends high-dimensional distributed parallel technology based on a unified three-dimensional parallel strategy search space, thus overcoming the shortcomings of current methods for automatically solving 4D and even 5D parallel training. It not only provides the cost function in the 4D and 5D parallel training search space, but also designs a search algorithm to solve the optimal training strategy for high-dimensional distributed parallel training in the 4D and 5D parallel training search space.

[0055] Step S04. Parallel Monte Carlo Search: Within the unified parallel strategy search space, multiple computing units use the parallel Monte Carlo tree search algorithm to search for the optimal parallel training strategy. During the search process, the search space is divided into different sub-search spaces, and the search process of each sub-search space is assigned to an independent computing unit. The results of each computing unit are summarized and the node statistics information is updated. The optimal parallel training strategy includes the optimal operator parallel strategy and the optimal pipeline partitioning label.

[0056] Monte Carlo search, as a greedy search algorithm, always moves towards easily attainable targets at a low cost. Furthermore, it avoids getting trapped in local minima and obtains the optimal solution. The detailed steps of implementing the policy search using the parallel Monte Carlo search method in this embodiment are as follows: Step S401. Divide the unified parallel strategy search space into multiple sub-search spaces according to the number of computing units, and allocate an independent computing unit to each sub-search space; Step S402. Each computing unit independently executes a Monte Carlo tree search loop within its allocated subspace to search for the optimal candidate strategy and evaluates the cost corresponding to the optimal candidate strategy found. Step S403. Summarize the optimal candidate strategies and their evaluated costs found by each computing unit in the subspace, and select the strategy with the lowest cost as the search result for the current round. The cost includes computational cost and / or communication cost. Step S404. Redivide the search space based on the search results of the current round, return to step S402, and continue until the iteration stopping condition is met to obtain the final optimal parallel training strategy output.

[0057] In this embodiment, the detailed steps of each computing unit performing the search, i.e., step S402, include: Step S421. The current computing unit selects one policy from all candidate policies as the current policy; Step S422. The current computing unit randomly generates a new policy from all candidate policies. Calculate the current policy and the new policy respectively. The corresponding costs depend on the current strategy and the new strategy. The corresponding cost determines whether to accept the new strategy. If so, then the new strategy will be implemented. As the current strategy, a new strategy is adopted. probability for: (6) in, , These are the current strategy and the new strategy. The corresponding cost, This represents a hyperparameter used to control the weight of the difference between the two parallel strategies before and after the update; it can specifically be set to 1. Step S423. The current computing unit determines whether the search stopping condition has been met. If so, the final optimal candidate strategy is output; otherwise, return to step S422.

[0058] Specifically, when applying the parallel Monte Carlo search algorithm to the search space, the search space is first divided into I parts, where I is the number of computational units (search workers). Each computational unit returns a partitioning strategy from the sub-search space; then, the optimal partitioning strategy is found from these I partitioning strategies, i.e., the Metropolis-Hastings algorithm for finding the optimal strategy is adopted. This algorithm consists of two steps: the first step is to randomly generate the next strategy from all candidate strategies. The second step is to calculate the corresponding computational and communication costs to determine the new strategy. Whether it is accepted depends on the adoption of new strategies. The probability is shown in equation (6).

[0059] Iterative updates of new strategies and hyperparameters during the search process Closely related, this hyperparameter That is, the time when the current search strategy is positively correlated with the current search space length. Preferably, the search termination condition is met and the search is terminated when one of the following rules is encountered: (1) The Monte Carlo search algorithm attempts up to 100 times. (2) in succession The candidate strategies found in the search within the second iteration cannot be changed further; that is, the candidate strategies will no longer be changed.

[0060] The Monte Carlo search algorithm's process of finding a strategy and calculating its cost based on a cost model is independent, indicating that multiple search processes can be executed separately, and the corresponding costs can be shared. and Integrating the result in the expression yields a function. To find the optimal search result: (7) This embodiment of the parallel Monte Carlo search algorithm includes two aspects of parallelism: firstly, each strategy and its corresponding cost can be computed and executed in parallel by multiple processors, which is called inter-parallel Monte Carlo search; secondly, each search is part of a multi-strategy search, which is called intra-parallel Monte Carlo search. During the search process in this embodiment, considering a single Monte Carlo sample, two aspects need to be considered: firstly, even if the MCMC (Markov Chain Monte Carlo algorithm) does not traverse all possibilities of the candidate strategies, it can be proven to be stable and convergent; secondly, the search results are related to the sampled candidates. Therefore, to describe the sampling rate that maps the full search space to the sampled candidates, this embodiment further defines the coefficients between the full search space and the sampled candidates. This is used to control sampling coverage and balance search efficiency with strategy optimality.

[0061] This embodiment implements a parallel Monte Carlo search algorithm in the training system to find the optimal parallel training strategy for the input model. Starting from formula (1), it begins with the case of partitioning each tensor along the batch dimension. Then, the search algorithm is used to find the next partitioning strategy and its cost is calculated according to the cost model. The above process is repeated until the best partitioning strategy and the best pipeline partitioning mark are found.

[0062] The architecture and workflow of the large-scale language model automatic parallel training system constructed in this embodiment are as follows: Figure 4 As shown, the user inputs a trainable large-scale language model and device topology into the ParaScale training system. The system then generates a search space based on these two inputs. First, the constructor in the training system transforms the entire DNN into a series of operators and a graph. During this process, the model is simultaneously represented as JaxPr and HLO modules. The HLO module represents the communication and computational overhead for obtaining each operator, while JaxPr represents the point-to-point communication overhead for obtaining pipelined label candidates. A multi-dimensional unified parallel policy search space is constructed, incorporating data parallelism, model parallelism, pipeline parallelism, checkpoint parallelism, and the ZeRo method. Then, the searcher in the training system uses a parallel Monte Carlo search algorithm within the unified parallel policy search space to find the optimal segmentation strategy for each operator and, when necessary, the optimal pipelined label. If pipeline parallelism is required, the operation sequence is divided into several pipeline stages. The model is simultaneously compiled using Ray Workers, i.e., the distributed training of the large-scale language model is automated using the Ray framework. This embodiment can improve the throughput of the training system through all-reduce (complete all reductions in one step and broadcast) communication, instead of using the reduce-scatter (reduce and distribute to each device first, then collect to each device) optimization mechanism in 2D parallelism.

[0063] Furthermore, during the process of multiple computing units using the parallel Monte Carlo tree search algorithm to search for the optimal parallel training strategy, data parallelism acceleration using GPUs is also included. That is, for the parallel Monte Carlo search algorithm, GPUs can be used to accelerate the sampling process. And since JAX can perform the random selection process—that is, using JAX sampling before the search begins to select the search space—speed can be significantly improved.

[0064] Specifically, to address the time bottleneck caused by the high-dimensional search space, this embodiment employs a fully parallel MCMC search method based on JAX. First, the operator-level policy space (ParaMesh) is mapped to a multidimensional dense tensor in GPU memory. During the search preprocessing stage, JAX's JIT compilation technology is used to transform the cost model into a GPU kernel function. During the search process, tens of thousands of independent random seeds are generated to drive multiple Markov chains in parallel on the GPU. Each chain performs policy sampling through tensor indexing operations and immediately calls the cost model function within the GPU for performance evaluation. This closed-loop 'sampling-evaluation' process is completed within GPU memory, avoiding the frequent CPU-GPU copy overhead in traditional algorithms. Compared to traditional CPU serial search, the convergence time of the optimal policy can be reduced from hours to seconds.

[0065] Large-scale language models are typically pre-trained models. Taking BERT as an example, BERT's training process involves first training on the English Wikipedia and BooksCorpus datasets to obtain the model's pre-trained weights. Then, the pre-trained model weights are fine-tuned on downstream tasks to obtain the BERT model's training results for these downstream tasks. This embodiment implements automatic parallel training of a large-scale language model based on a parallel Monte Carlo algorithm, which optimizes the training strategy for large-scale language models and is more conducive to the further development of large-scale language models.

[0066] This embodiment further provides a computer system, including a processor and a memory, wherein the memory is used to store a computer program, and the processor is used to execute the computer program to perform the method as described above.

[0067] It is understood that the method described in this embodiment can be executed by a single device, such as a computer or server, or it can be applied to a distributed scenario where multiple devices cooperate to complete the task. In a distributed scenario, one of the multiple devices may execute only one or more steps of the method described in this embodiment, and the multiple devices interact to complete the method. The processor can be implemented using a general-purpose CPU, microprocessor, application-specific integrated circuit, or one or more integrated circuits, and is used to execute relevant programs to implement the method described in this embodiment. The memory can be implemented using read-only memory (ROM), random access memory (RAM), static storage devices, and dynamic storage devices. The memory can store the operating system and other applications. When the method described in this embodiment is implemented through software or firmware, the relevant program code is stored in the memory and called and executed by the processor.

[0068] Those skilled in the art will understand that the above embodiments of this application can be provided as methods, systems, or computer program products. Therefore, this application can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, this application can take the form of a computer program product embodied on one or more computer-readable storage media (including, but not limited to, disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code. This application is described with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of this application. It should be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded processor, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, create a machine for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to operate in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1 One or more processes and / or boxes Figure 1 The functions specified in one or more boxes. These computer program instructions may also be loaded onto a computer or other programmable data processing apparatus to cause a series of operational steps to be performed on the computer or other programmable apparatus to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable apparatus for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0069] The above description is merely a preferred embodiment of the present invention and is not intended to limit the invention in any way. Although the present invention has been disclosed above with reference to preferred embodiments, it is not intended to limit the invention. Therefore, any simple modifications, equivalent changes, and alterations made to the above embodiments based on the technical essence of the present invention without departing from the scope of the present invention should fall within the protection scope of the present invention.

Claims

1. A method for automatic parallel training of large-scale language models based on parallel Monte Carlo algorithm, characterized by the following steps: include: Step S01. Data Acquisition: Obtain the device topology of the large-scale language model and computing cluster; Step S02. Model Representation Generation: Convert the acquired large-scale language model into a sequence of operators and the corresponding computation graph, and generate the model representation of the model; Step S03. Multidimensional search space construction: Based on the operator sequence, computation graph, model representation and device topology, construct a unified parallel strategy search space that includes at least three dimensions: data parallelism, model parallelism and pipeline parallelism. Step S04: Parallel Monte Carlo Search: Within the unified parallel strategy search space, multiple computing units use the parallel Monte Carlo tree search algorithm to search for the optimal parallel training strategy. During the search process, the search space is divided into different sub-search spaces, and the search process of each sub-search space is assigned to an independent computing unit. The results of each computing unit are summarized and the node statistics are updated. The optimal parallel training strategy includes the optimal operator parallel strategy and the optimal pipeline partitioning label. Step S05. Model Training: Train the large-scale language model according to the optimal parallel training strategy.

2. The automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm according to claim 1, characterized in that, In step S02, the acquired large-scale language model is converted into an operator sequence Q, which includes N operators. ; A computation graph G is constructed based on the dependencies between operators in the operator sequence Q. The two vertices of each edge in graph G are the two operators in the operator sequence Q. The model representation includes JaxPr representation and HLO representation. JaxPr representation is used to obtain the communication overhead and computation overhead of each operator, and HLO representation is used to obtain the point-to-point communication overhead between pipeline marker candidates.

3. The method for automatic parallel training of large-scale language models based on the parallel Monte Carlo algorithm according to claim 1, characterized in that, In step S03, the unified parallel strategy search space further includes a tensor recomputation dimension and / or a zero-redundancy optimizer dimension. The tensor recomputation dimension is configured with a flag to indicate whether recomputation is activated during training to optimize memory usage. The zero-redundancy optimizer dimension is configured with a flag to indicate whether the zero-redundancy optimizer is used during training and to evaluate the additional cost of the zero-redundancy optimizer dimension during training by calculating the additional time cost of the zero-redundancy optimizer.

4. The method for automatic parallel training of large-scale language models based on the parallel Monte Carlo algorithm according to claim 1, characterized in that, Step S04 includes: Step S401. Divide the unified parallel strategy search space into multiple sub-search spaces according to the number of computing units, and allocate an independent computing unit to each sub-search space; Step S402. Each computing unit independently executes a Monte Carlo tree search loop within its allocated subspace to search for the optimal candidate strategy and evaluates the cost corresponding to the optimal candidate strategy found. Step S403. Summarize the optimal candidate strategies and their evaluated costs found by each computing unit in the subspace, and select the strategy with the lowest cost as the search result for the current round. The cost includes computational cost and / or communication cost. Step S404. Redivide the search space based on the search results of the current round, return to step S402, and continue until the iteration stopping condition is met to obtain the final optimal parallel training strategy output.

5. The method for automatic parallel training of large-scale language models based on the parallel Monte Carlo algorithm according to claim 4, characterized in that, Step S402 includes: Step S421. The current computing unit selects one policy from all candidate policies as the current policy; Step S422. The current computing unit randomly generates a new policy from all candidate policies. Calculate the current policy and the new policy respectively. The corresponding costs depend on the current strategy and the new strategy. The corresponding cost determines whether to accept the new strategy. If so, then the new strategy will be implemented. As the current strategy, a new strategy is adopted. probability for: in, , The current strategy and the new strategy are respectively. The corresponding cost, This represents a hyperparameter used to control the weight of the difference between the two parallel strategies before and after the update; Step S423. The current computing unit determines whether the search stopping condition has been met. If so, the final optimal candidate strategy is output; otherwise, return to step S422.

6. The method for automatic parallel training of large-scale language models based on parallel Monte Carlo algorithm according to any one of claims 1 to 5, characterized in that, In step S04, the objective of the search in the three dimensions of data parallelism, model parallelism, and pipeline parallelism is: in, For operators The computational overhead, For operators Communication overhead, For operators in the operator sequence Q and The overhead of re-partitioning between them When partitioning the pipeline, it occurs in the operator. Subsequent point-to-point communication overhead, N This represents the number of operators in an operator sequence Q. S represents the total cost of searching the search space across the three dimensions of data parallelism, model parallelism, and pipeline parallelism, where S is the parallelism strategy for all operators. s The set, It is the set of all pipeline partition markers. It is a set of repartitioning operator pairs.

7. The automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm according to any one of claims 1 to 5, characterized in that, When the unified parallel strategy search space also includes the tensor recomputation dimension to form a 4D parallel search space, the solution objective in the 4D parallel search space is: in, Let represent the objective function for solving the 4D parallel search space, i.e., minimizing the total cost of the 4D parallel search space. Let represent the objective function for solving the 3D parallel search space, i.e., minimizing the total cost of the 3D parallel search space. This represents the time cost of memory swapping in 4D parallel training mode. It is the PCIe bandwidth of the training cluster devices. It is the unloading operator space The Middle The memory overhead of each JaxPr operator.

8. The automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm according to any one of claims 1 to 5, characterized in that, When the unified parallel strategy search space also includes the tensor recomputation dimension and the zero-redundancy optimizer dimension, forming a 5D parallel search space, the solution objective for searching in the 5D parallel search space is: in, Let represent the objective function for solving the 5D parallel search space, i.e., minimizing the total cost of the 5D parallel search space. Let represent the objective function for solving the 4D parallel search space, i.e., minimizing the total cost of the 4D parallel search space. This indicates the additional cost of zero-redundancy optimization parallelism during the training process in 5D parallel training mode. It refers to the number of model parameters. It refers to the bandwidth between GPUs. This indicates the number of optimizer states.

9. The automatic parallel training method for large-scale language models based on the parallel Monte Carlo algorithm according to any one of claims 1 to 5, characterized in that, In step S04, during the process of multiple computing units searching for the optimal parallel training strategy using the parallel Monte Carlo tree search algorithm, data parallelism acceleration using GPUs is also included. Before the search begins, JAX sampling is used to complete the selection of the search space. First, the operator-level policy space is mapped to a multidimensional dense tensor in GPU memory. In the search preprocessing stage, the cost model is transformed into a GPU kernel function using JAX. During the search process, multiple Markov chains are driven in parallel on the GPU using multiple independent random seeds. Each chain performs policy sampling through tensor indexing operations and calls the cost model function in the GPU for performance evaluation.

10. A computer system comprising a processor and a memory, the memory being used to store computer programs, characterized in that, The processor is used to execute the computer program to perform the method as described in any one of claims 1 to 9.

Images