Optimal parallel strategy for non-uniform heterogeneous chips, and search method and apparatus therefor

WO2026179041A1PCT designated stage Publication Date: 2026-09-03SHANGHAI ARTIFICIAL INTELLIGENCE INNOVATION CENT
View PDF 0 Cites 0 Cited by

Patent Information

Application Number
PCT/CN2025/109487
Authority / Receiving Office
WO · WO
Patent Type
Applications
Current Assignee / Owner
Priority Date
2025-02-28
Filing Date
2025-07-21
Publication Date
2026-09-03

Smart Images

  • Figure CN2025109487_03092026_PF_FP_ABST
    Figure CN2025109487_03092026_PF_FP_ABST
Patent Text Reader

Abstract

The present invention relates to an optimal parallel strategy for non-uniform heterogeneous chips, and a search method and apparatus therefor. The method comprises: on the basis of a data parallelism dimension constraint, a heterogeneous chip quantity constraint, and a tensor parallelism dimension constraint, acquiring a first candidate set, the first candidate set comprising combinations of data parallelism dimensions, tensor parallelism dimensions, and pipeline parallelism dimensions for all types of heterogeneous chips; on the basis of a mapping constraint, a search space constraint, and a memory constraint, performing allocation of the number of model layers, and determining a parallel strategy candidate set, the parallel strategy candidate set comprising combinations of the data parallelism dimensions, the tensor parallelism dimensions, the pipeline parallelism dimensions, and the number of model layers for all types of the heterogeneous chips; and for each combination in the parallel strategy candidate set, calculating a single iteration time, and selecting a combination having the shortest single iteration time as an optimal parallel strategy. Compared with the prior art, the present invention has the advantages of realizing load distribution and scheduling across thousand-card-scale massive heterogeneous chips, and achieving optimal pipeline execution efficiency.
Need to check novelty before this filing date? Find Prior Art

Description

Optimal Parallelism Strategy for Non-Uniform Heterogeneous Chips and its Search Method and Device Technical Field

[0001] This invention relates to the field of distributed training technology, and in particular to an optimal parallel strategy for non-uniform heterogeneous chips, as well as its search method and apparatus. Background Technology

[0002] Currently, the main challenges in heterogeneous cluster training stem from three aspects: bandwidth, computing power, and GPU memory. Specifically, compared to the interconnection between homogeneous chips, the heterogeneous interconnection bandwidth of various newly manufactured computing chips is lower. At the same time, there are significant differences in GPU memory and computing power among the new chips, resulting in varying processing capabilities among heterogeneous chips. Existing parallel strategies designed for homogeneous clusters are insufficient to fully utilize the computing power of heterogeneous chips in heterogeneous scenarios.

[0003] Other parallelism strategies besides pipelined parallelism, such as data parallelism and tensor parallelism, require chips with similar computing power in the same parallel dimension and high-speed interconnectivity. However, existing pipelined parallelism strategies also have certain limitations. On the one hand, in heterogeneous distributed hybrid training systems, hardware components have different characteristics. How the system can uniformly or non-uniformly divide tasks based on model characteristics and the system's real-time training state under uneven computing power distribution, ensuring maximum utilization of computing power, has a significant impact on efficient training of large-scale heterogeneous clusters. However, existing pipelined parallel strategies are designed for smaller scales, mainly concentrated at the hundred-card level. At large scales, the search space for heterogeneous pipelined parallelism is much larger, and existing pipelined parallel strategies struggle to find the optimal heterogeneous parallel strategy at scales of thousands of cards or more. On the other hand, most existing heterogeneous pipelined parallel strategies are designed for heterogeneous interconnects between various chip models from a single vendor (such as NVIDIA). Since the computing power and interconnect bandwidth of new chips from different vendors vary greatly, existing pipelined parallel strategies struggle to achieve efficient load distribution and scheduling in complex scenarios involving multiple vendors and various new computing power chips.

[0004] In addition, the use of parameter servers is another method to solve the problems of heterogeneous parallel training. It mainly divides different types of heterogeneous chips into different groups. The main problem with this method is that training is difficult to converge quickly, and there are problems of uneven computing power between different groups and large synchronization overhead between groups, resulting in poor execution efficiency. Summary of the Invention

[0005] The purpose of this invention is to overcome the shortcomings of the prior art by providing an optimal parallel strategy and its search method and apparatus for non-uniform heterogeneous chips. This strategy can search for the optimal heterogeneous parallel strategy in heterogeneous clusters of more than 1,000 kilocalories, achieve efficient and accurate chip load distribution and scheduling in complex new computing chip scenarios, and achieve optimal pipeline execution efficiency.

[0006] The objective of this invention can be achieved through the following technical solutions:

[0007] According to a first aspect of the present invention, an optimal parallel strategy search method for non-uniform heterogeneous chips is provided. The optimal parallel strategy search method is implemented based on preset mapping constraints, search space constraints, data parallelism dimension constraints, heterogeneous chip quantity constraints, tensor parallelism dimension constraints, and memory constraints, and includes the following steps: S101, obtaining a first candidate set according to the data parallelism dimension constraints, the heterogeneous chip quantity constraints, and the tensor parallelism dimension constraints, wherein the first candidate set includes combinations of data parallelism dimension, tensor parallelism dimension, and pipeline parallelism dimension for all types of heterogeneous chips; S102, based on the first candidate set, allocating model layers according to the mapping constraints, the search space constraints, and the memory constraints to determine a parallel strategy candidate set, wherein the parallel strategy candidate set includes combinations of data parallelism dimension, tensor parallelism dimension, pipeline parallelism dimension, and model layer number for all types of heterogeneous chips; S103, calculating the single iteration time for each combination in the parallel strategy candidate set, and selecting the combination with the lowest single iteration time as the optimal parallel strategy.

[0008] As a preferred technical solution, in the optimal parallel strategy search method: the mapping constraint includes that each neural network layer of the model to be trained is mapped to a heterogeneous chip for execution; the search space constraint includes that, among all heterogeneous chips, the number of model layers in each pipeline stage corresponding to the same type of chip is the same, and the number of model layers in the pipeline stages corresponding to heterogeneous chips is the same or different; the data parallel dimension constraint includes that the micro-batch size of each pipeline stage is 1, and thus each pipeline stage uses the same data parallel dimension; the heterogeneous chip quantity constraint includes that the total number of heterogeneous chips... The quantity is equal to the sum of the total number of heterogeneous chips in each pipeline stage, and the total number of heterogeneous chips in each pipeline stage is equal to the product of the tensor parallel dimension and the data parallel dimension of the corresponding pipeline stage; the tensor parallel dimension constraint includes that the tensor parallel dimension is a power of 2, and the tensor parallel dimension is less than the maximum tensor parallel dimension independently specified for each heterogeneous chip; the memory constraint includes that the chips in the pipeline are arranged in descending order of memory, and the sum of the memory occupied by the model, optimizer state and activation value in the first pipeline stage of each heterogeneous chip is less than the maximum memory occupied by the corresponding chip.

[0009] As a preferred technical solution, the optimal parallel strategy search method is also based on a preset computing card constraint, which is used to constrain the allocation of model layers; the computing card constraint includes that different pipeline stages containing the same computing card use the same tensor parallel dimension.

[0010] As a preferred technical solution, S101 specifically includes: S111, selecting a global set of candidate data parallel dimensions according to the data parallel dimension constraint, and then calculating the number of micro-batches; S112, for a given candidate data parallel dimension, selecting a candidate tensor parallel dimension for a single chip according to the tensor parallel dimension constraint, and then calculating the number of pipeline stages for a single chip according to the heterogeneous chip quantity constraint, and then combining the types of heterogeneous chips to obtain candidate combinations of multiple pipeline stage quantities and tensor parallel dimensions under a given candidate data parallel dimension; S113, based on the set of candidate data parallel dimensions and the candidate combinations of multiple pipeline stage quantities and tensor parallel dimensions, obtaining a candidate set of data parallel dimensions, tensor parallel dimensions, and pipeline parallel dimensions for all types of heterogeneous chips under all candidate data parallel dimensions, i.e., the first candidate set.

[0011] As a preferred technical solution, the number of micro-batches is the quotient of the training batch size and the data parallel dimension.

[0012] As a preferred technical solution, S102 specifically includes: obtaining a candidate set of model layer allocation schemes based on the mapping constraints and the search space constraints; pruning the model layer allocation schemes using a heuristic method; balancing the time required for each heterogeneous chip to process one layer of the model before and after backpropagation in the current tensor parallel dimension among the candidate combinations of pipeline stage number and tensor parallel dimension, and fine-tuning until the total number of model layers is satisfied to obtain a preliminary set of parallel strategy candidates; and removing parallel strategy candidates that do not meet the constraints from the preliminary set of parallel strategy candidates based on the memory constraints to obtain a final set of parallel strategy candidates.

[0013] As a preferred technical solution, the memory constraint is specifically: number of model layers × (memory of the parameters and optimizer state of model layer 1 under the current data parallel dimension and tensor parallel dimension + number of preheating micro-batches in the current pipeline stage × memory of the activation value of 1 sample processed by model layer under the current tensor parallel dimension) < maximum memory of the chip.

[0014] As a preferred technical solution, in S103, the time for a single iteration is the sum of the global forward and backward computation time, the global update optimizer parameter time, the pipeline bubble time, and the pipeline communication time.

[0015] According to a second aspect of the present invention, an optimal parallel strategy search apparatus for non-uniform heterogeneous chips is provided. The apparatus includes preset mapping constraints, search space constraints, data parallelism dimension constraints, heterogeneous chip quantity constraints, tensor parallelism dimension constraints, and memory constraints. It comprises a first candidate set acquisition module, a parallel strategy candidate set determination module, and an optimal parallel strategy selection module. The first candidate set acquisition module is used to acquire a first candidate set based on the data parallelism dimension constraints, the heterogeneous chip quantity constraints, and the tensor parallelism dimension constraints. The first candidate set includes combinations of data parallelism dimensions, tensor parallelism dimensions, and pipeline parallelism dimensions for all types of heterogeneous chips. The parallel strategy candidate set determination module is used to determine a parallel strategy candidate set based on the first candidate set, according to the mapping constraints, the search space constraints, and the memory constraints, by allocating model layers. The parallel strategy candidate set includes combinations of data parallelism dimensions, tensor parallelism dimensions, pipeline parallelism dimensions, and model layer numbers for all types of heterogeneous chips. The optimal parallel strategy selection module is used to calculate the single iteration time for each combination in the parallel strategy candidate set and select the combination with the lowest single iteration time as the optimal parallel strategy.

[0016] According to a third aspect of the present invention, an optimal parallel strategy for non-uniform heterogeneous chips is provided, the optimal parallel strategy being obtained based on the method described above, wherein the optimal parallel strategy is a combination of data parallel dimension, tensor parallel dimension, pipeline parallel dimension and model layer number, wherein the data parallel dimension is an integer, the tensor parallel dimension, the pipeline parallel dimension and the model layer number are all arrays, and the dimension of each array is limited by the type of heterogeneous chip.

[0017] Compared with the prior art, the present invention has the following beneficial effects:

[0018] 1. The optimal parallel strategy search method provided by the present invention is used for non-uniform heterogeneous chips. The method is based on a variety of preset constraints. These constraints ensure that the optimal parallel strategy includes the optimal model partitioning scheme, the optimal parallel dimension, and the optimal resource mapping, so that different types of heterogeneous chips on the pipeline can correspond to different parallel strategies, thereby searching for the optimal heterogeneous parallel strategy in heterogeneous clusters of more than 1000 kilocalories.

[0019] 2. This invention fully considers the basic information of heterogeneous clusters, such as the information of the model to be trained, the specifications of heterogeneous chips, cluster configuration, network configuration, heterogeneous interconnection configuration, etc., to ensure that the heterogeneous chips and the corresponding pipeline stage loads are matched to each other, so as to achieve the optimal pipeline execution efficiency.

[0020] 3. This invention obtains a candidate set of model layer allocation schemes based on mapping constraints and search space constraints, and uses a heuristic method to prune the model layer allocation schemes to achieve optimal model splitting, ensuring that the splitting method and chip computing power are matched. At the same time, the communication volume between different model slices and the interconnection bandwidth between chips are matched, which can effectively simplify the search space and reduce the search complexity. In addition, based on the candidate combinations of pipeline stage number and tensor parallel dimension, the time required for each heterogeneous chip to process the forward and backward computation of one layer of the model under the current tensor parallel dimension is used to obtain a preliminary set of parallel strategy candidates. The set is then filtered according to memory constraints to achieve fine-grained scheduling based on hardware characteristics, optimize the forward and backward computation during training, ensure that the computing resources of heterogeneous chips are fully utilized, and improve the computing utilization rate of heterogeneous chips.

[0021] 4. This invention eliminates the need to synchronize parameters or gradients of different types of heterogeneous chips through a parameter server, thereby reducing synchronization overhead. Furthermore, through fine-grained pipeline partitioning and the selection of parallel strategies, it effectively solves the problems of unbalanced load in traditional parallel strategies and inefficient convergence in asynchronous parameter servers.

[0022] 5. In the optimal parallel strategy provided by this invention, the data parallel dimension is an integer, while the tensor parallel dimension, pipeline parallel dimension, and model layer number are all arrays. Furthermore, the dimension of each array is limited by the type of heterogeneous chip. When this parallel strategy is used for training on a large-scale cluster of non-uniform heterogeneous chips, it can ensure the training efficiency and load balance of distributed training. Attached Figure Description

[0023] Figure 1 is a flowchart illustrating the method provided by the present invention;

[0024] Figure 2 is a schematic diagram of the optimal parallel strategy in an embodiment of the present invention;

[0025] Figure 3 is a schematic diagram of the search device structure in an embodiment of the present invention. Detailed Implementation

[0026] In the context of this invention, a training batch (Mini-Batch) refers to a set of samples processed by the model in a single training step, and is the most basic batch unit in deep learning; a micro-batch is a subset of a mini-batch, typically used for finer-grained control and optimization. In some distributed training frameworks, micro-batches are used to further decompose mini-batches to improve the efficiency of pipelined parallel processing. In heterogeneous pipelined parallel systems, different stages of different micro-batches of the same training batch (Mini-Batch) can be pipelined in parallel.

[0027] In the context of this invention, pipeline parallelism is abbreviated as PP; tensor parallelism is abbreviated as TP; and data parallelism is abbreviated as DP. Further, DP Size represents the data parallelism dimension, TP Size represents the tensor parallelism dimension, and PP Size represents the number of pipeline stages, i.e., the pipeline parallelism dimension.

[0028] The present invention will now be described in detail with reference to the accompanying drawings and specific embodiments. These embodiments are based on the technical solution of the present invention and provide detailed implementation methods and specific operating procedures. However, the scope of protection of the present invention is not limited to the following embodiments.

[0029] Example

[0030] As shown in Figure 1, this embodiment provides an optimal parallel strategy search method for non-uniform heterogeneous chips. This method is based on preset mapping constraints, search space constraints, data parallelism dimension constraints, heterogeneous chip quantity constraints, tensor parallelism dimension constraints, and memory constraints. It can obtain the optimal parallel strategy, including the optimal partitioning scheme, optimal parallelism dimension, and optimal resource mapping, based on the basic information of the heterogeneous cluster, such as model information, heterogeneous chip specifications, cluster configuration, network configuration, and heterogeneous interconnect configuration. Furthermore, this search method fully considers the specifications of the heterogeneous chips, such as memory size and computational information, thereby ensuring that the load of the heterogeneous chips and the corresponding pipeline stages are matched to achieve optimal pipeline execution efficiency. It can be efficiently implemented on various new heterogeneous chips at the kilocalorie level.

[0031] The specific details of each constraint are as follows:

[0032] Constraint 1: Mapping Constraint

[0033] This constraint requires that each neural network layer of the model to be trained be mapped to a heterogeneous chip for execution. The relationship between the total number of layers in the model and the number of neural network layers in each stage is as follows:

[0034] Total number of neural network layers in the model = sum(number of neural network layers in each pipeline stage)

[0035] Constraint 2: Search Space Constraint

[0036] This constraint requires that, among all heterogeneous chips, the number of pipeline stage model layers must be the same for chips of the same type, while the number of pipeline stage model layers for heterogeneous chips can be different or the same. This constraint can effectively reduce the search space.

[0037] Constraint 3: Data Parallelism Dimension Constraint

[0038] This constraint requires that the micro-batch size be 1 in each pipeline stage, thus ensuring that each pipeline stage uses the same data parallelism dimension (DP Size). This constraint is primarily set considering the search space size, the memory constraints of activation values, and the cavitation ratio in pipeline parallelism. From this constraint, we can derive:

[0039] Data parallelism dimension × number of micro-batches = training batch size

[0040] Constraint 4: Heterogeneous Chip Quantity Constraint

[0041] In other words, the parallel strategy is constrained by the number of heterogeneous chips. This constraint requires that the total number of heterogeneous chips equals the sum of the total number of heterogeneous chips in each pipeline stage, and the total number of heterogeneous chips in each pipeline stage equals the product of the tensor parallelism dimension and the data parallelism dimension of the corresponding pipeline stage.

[0042] Total number of heterogeneous chips = sum(total number of heterogeneous chips in each pipeline stage) = sum(tensor parallelism dimension of each pipeline stage × data parallelism dimension)

[0043] Constraint 5: Calculation Card Constraint

[0044] This constraint requires that different pipeline stages containing the same compute cards use the same tensor parallelism dimension. Specifically, the optimal tensor parallelism dimension is the same for the same compute card; however, the optimal parallelism strategy may differ for different stages containing different compute cards, and therefore different tensor parallelism dimensions can be used.

[0045] Constraint 6: Tensor Parallelism Dimension Constraint

[0046] This constraint requires that the tensor parallelism dimension be a power of 2, and that the tensor parallelism dimension be less than the maximum tensor parallelism dimension (TP_MAX) independently specified for each heterogeneous chip. Since the number of attention heads in a large language model is generally a power of 2, and the TP needs efficient communication, the TP should, based on hardware topology and other factors, be less than the number of compute cards in a single machine / single NUMA / single PCIE switch configuration. Here, NUMA stands for Non-Uniform Memory Access, and PCIE refers to the high-speed serial computer expansion bus standard.

[0047] Constraint 7: Memory Constraint

[0048] This constraint requires that chips be arranged in descending order of memory size in the pipeline, and that the sum of the memory usage of the model, optimizer state, and activation values ​​in the first pipeline stage for each heterogeneous chip is less than the maximum available memory of the corresponding chip. This constraint is primarily designed to account for the larger number of preheating micro-batches in earlier pipeline stages, resulting in higher memory usage for activation values. This can be expressed as:

[0049] Number of layers × (Model layer 1 parameters and optimizer state in the current DP Size and TP Size GPU memory + number of preheating micro-batches in this pipeline stage × 1 sample of the model layer's active value in the current TP Size GPU memory) < Maximum available GPU memory of the chip

[0050] For example, based on the foregoing constraints, the given parameters required for the search method provided in this embodiment include:

[0051] (1) The quantity of each type of chip;

[0052] (2) The size of the training batch in each iteration (Mini-Batch Size);

[0053] (3) The total number of layers in the model;

[0054] (4) The time schedule required for forward and backward computation of one layer of the model when each chip processes one sample at TP=1 / 2 / 4 / 8 (some chips have 16 cards per machine, or 16 is also possible) and DP=1;

[0055] (5) The timeline required for the first-level optimizer step of the model when each chip is given a specific DP Size and TP Size;

[0056] (6) Communication and other lost time caused by each PP pipeline;

[0057] (7) Maximum available video memory for each type of chip;

[0058] (8) When the DP Size and TP Size are given in a single-layer model, the optimizer and parameters occupy the memory table. For example, if the state of the split optimizer is used to optimize the ZeRO1, the DP Size here can be regarded as the ZeRO1 Size in ZeRO1.

[0059] (9) When the TP Size is given, the single-layer model processes one sample to occupy the activation value and the video memory table.

[0060] The optimal parallel strategy search method provided in this embodiment is a staged heuristic search method. Its core idea is to first perform discrete enumeration on chip resource allocation and pipeline configuration, then balance the allocation of model layers, and finally select the optimal configuration based on the predicted iteration time. Based on the aforementioned constraints and given parameters, the specific implementation process of this method is as follows:

[0061] Step S101: Based on the data parallelism dimension constraint, heterogeneous chip quantity constraint, and tensor parallelism dimension constraint, obtain a first candidate set. The first candidate set includes combinations of data parallelism dimension, tensor parallelism dimension, and pipeline parallelism dimension for all types of heterogeneous chips. Specifically:

[0062] Step S111: Based on constraint three, namely the data parallel dimension constraint, select a global set of candidate data parallel dimensions (DP Size) such that the training batch size divided by the data parallel dimension is an integer, and then calculate the number of micro-batches (Micro Num):

[0063] Micro Num=Mini-Batch Size÷DP Size.

[0064] Step S112: Given the data parallelism dimension, determine the number of pipeline stages and candidate tensor parallelism dimension for a single chip.

[0065] For each type of chip (assuming the number of chip i is N) (i) The requirement is to divide the resources on the chip into several pipeline stages and configure them with tensor parallelism. Given a list of candidate data parallelism dimensions, according to constraint six, i.e., the tensor parallelism dimension constraint, select the candidate tensor parallelism dimension for a single chip:

[0066] The parallel dimension of the candidate tensor is ∈{1,2,4,...,TP_MAX};

[0067] Based on constraint four, namely the heterogeneous chip quantity constraint, the number of pipeline stages for a single chip is calculated:

[0068] Number of pipeline stages for a single chip = Total number of chips for a single chip ÷ (Data parallelism dimension × Tensor parallelism dimension);

[0069] The candidate combination of the number of pipeline stages and tensor parallelism dimension of a single chip is denoted as: cj=[(p1,t1),...,(pi,ti),...,(pn,tn)], where n is the type of heterogeneous chip and j is the sampling number.

[0070] Based on the types of heterogeneous chips, we obtain multiple candidate combinations of pipeline stage numbers and tensor parallel dimensions under a given candidate data parallel dimension di, denoted as:

[0071] Di = [(di,c1),...,(di,cj),...,(di,cJ)], where J is the number of samples.

[0072] Step S113: Based on the candidate set of data parallel dimensions and the candidate combinations of multiple pipeline stage numbers and tensor parallel dimensions, obtain the candidate set of data parallel dimensions, tensor parallel dimensions and pipeline parallel dimensions for all types of heterogeneous chips under all candidate data parallel dimensions, i.e., the first candidate set.

[0073] Based on steps S111 and S112, the first candidate set can be obtained, denoted as:

[0074] D = D1∪D2∪...∪DK, where K is the number of parallel dimensions for all candidate data.

[0075] Step S102: Allocate model layers to each pipeline stage. Based on the first candidate set, allocate model layers according to mapping constraints, search space constraints, and memory constraints, and determine the parallel strategy candidate set. The parallel strategy candidate set includes combinations of data parallelism, tensor parallelism, pipeline parallelism, and model layer count for all types of heterogeneous chips. Specifically:

[0076] First, select all model layer allocation schemes that satisfy constraint one (mapping constraint) and constraint two (search space constraint) to obtain a candidate set of model layer allocation schemes;

[0077] Secondly, a heuristic approach is used to prune the model layer allocation scheme;

[0078] Next, based on the time required for each heterogeneous chip to process the forward and backward computation of one layer of the model under the current ti in the (pi,ti) combination, water level balancing is performed to make the computation time of each pipeline stage of each chip as close as possible. Then, through fine-tuning until the condition of the total number of model layers is met, a preliminary set of parallel strategy candidates is obtained.

[0079] Finally, based on constraint seven (memory constraint), parallel strategy candidates that do not meet constraint seven are removed from the initial parallel strategy candidate set, resulting in the final parallel strategy candidate set. A valid parallel strategy can be denoted as (data parallelism dimension, tensor parallelism dimension, pipeline parallelism dimension, number of model layers), and the parallel strategy candidate set D is represented as:

[0080] D=[(d1,t1,p1,l1),...,(di,ti,pi,li),...,(dn,tn1,pn,ln)],

[0081] Where di is an integer, ti, pi, and li are arrays, and the size of the arrays represents the types of heterogeneous chips.

[0082] Step S103: For each combination in the candidate set of parallel strategies, that is, each candidate parallel strategy, calculate the single iteration time, and select the combination with the lowest single iteration time as the optimal parallel strategy.

[0083] Specifically, the time for each predicted iteration is expressed as:

[0084] Total prediction time = global forward and backward computation time + global optimizer parameter update time + pipeline bubble time + pipeline communication time.

[0085] Based on the aforementioned calculation formula, the parallel strategy (di,ti,pi,li) with the lowest total prediction time is selected.

[0086] In some other embodiments, the predicted iteration time may also include other existing loss time.

[0087] In heterogeneous 1000-calorie scenarios, the training throughput is doubled compared to industry-leading homogeneous distributed training frameworks (such as Megatron-LM) without any optimization.

[0088] In summary, the optimal parallel strategy search method for non-uniform heterogeneous chips provided in this embodiment has at least the following advantages:

[0089] Heterogeneity: Compared to homogeneous distributed pre-training frameworks, this method supports the heterogeneity of newly manufactured chips, especially domestically produced chips.

[0090] Large scale: Compared to other heterogeneous pipeline parallel strategies, this solution supports the search and online scheduling of efficient parallel strategies for large scale (kilocalorie level) of multiple (two or more) heterogeneous chips.

[0091] High efficiency: Compared to the parameter server approach, this solution eliminates the need to synchronize parameters or gradients of different heterogeneous chips via a parameter server, reducing synchronization overhead. Furthermore, through fine-grained pipeline partitioning and the selection of parallel strategies, it effectively solves the problems of load imbalance in traditional parallel strategies and inefficient convergence in asynchronous parameter servers.

[0092] Furthermore, by comprehensively considering fault tolerance costs and implementation costs, this method can be used in conjunction with a parameter server solution to solve the aforementioned problems in parameter servers while further increasing the training scale and achieving efficient distributed training of new heterogeneous chips at a higher scale (tens of thousands to hundreds of thousands of cards).

[0093] Furthermore, as shown in Figure 2, the optimal parallel strategy for non-uniform heterogeneous chips can be obtained using the aforementioned search method. The optimal parallel strategy includes the optimal model partitioning scheme, the optimal parallel dimension, and the optimal resource mapping. It is a combination of data parallel dimension, tensor parallel dimension, pipeline parallel dimension, and number of model layers. Among them, the data parallel dimension is an integer, and the tensor parallel dimension, pipeline parallel dimension, and number of model layers are all arrays. The dimension of each array is restricted by the type of heterogeneous chip.

[0094] Further, as shown in Figure 3, this embodiment also provides an optimal parallel strategy search device for non-uniform heterogeneous chips. This device presets mapping constraints, search space constraints, data parallelism dimension constraints, heterogeneous chip quantity constraints, tensor parallelism dimension constraints, and memory constraints. It includes a first candidate set acquisition module 301, a parallel strategy candidate set determination module 302, and an optimal parallel strategy selection module 303. The first candidate set acquisition module 301 is used to acquire a first candidate set based on the data parallelism dimension constraints, heterogeneous chip quantity constraints, and tensor parallelism dimension constraints. The first candidate set includes all types of heterogeneous chips. The system combines the data parallelism, tensor parallelism, and pipeline parallelism dimensions of the chip. The parallel strategy candidate set determination module 302, based on the first candidate set, allocates model layers according to mapping constraints, search space constraints, and memory constraints to determine the parallel strategy candidate set. The parallel strategy candidate set includes combinations of data parallelism, tensor parallelism, pipeline parallelism, and model layer numbers for all types of heterogeneous chips. The optimal parallel strategy selection module 303 calculates the single-iteration time for each combination in the parallel strategy candidate set and selects the combination with the lowest single-iteration time as the optimal parallel strategy. The specific operation flow of each module is basically the same as the implementation process of the aforementioned method steps, and will not be repeated here.

[0095] The preferred embodiments of the present invention have been described in detail above. It should be understood that those skilled in the art can make numerous modifications and variations based on the concept of the present invention without creative effort. Therefore, all technical solutions that can be obtained by those skilled in the art based on the concept of the present invention through logical analysis, reasoning, or limited experimentation on the basis of existing technology should be within the scope of protection defined by the claims.

Claims

1. A method for searching optimal parallel strategies for non-uniform heterogeneous chips, characterized in that, The optimal parallel strategy search method is implemented based on preset mapping constraints, search space constraints, data parallelism dimension constraints, heterogeneous chip quantity constraints, tensor parallelism dimension constraints, and memory constraints, and includes the following steps: S101, based on the data parallel dimension constraint, the heterogeneous chip quantity constraint and the tensor parallel dimension constraint, obtain a first candidate set, the first candidate set including all combinations of data parallel dimension, tensor parallel dimension and pipeline parallel dimension of all types of heterogeneous chips; S102, based on the first candidate set, the number of model layers is allocated according to the mapping constraints, the search space constraints and the memory constraints, and a parallel strategy candidate set is determined. The parallel strategy candidate set includes the data parallel dimension, tensor parallel dimension, pipeline parallel dimension and model layer combination of all types of heterogeneous chips. S103, For each combination in the candidate set of parallel strategies, calculate the single iteration time, and select the combination with the lowest single iteration time as the optimal parallel strategy.

2. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 1, characterized in that, In the optimal parallel strategy search method: The mapping constraint includes that each neural network layer of the model to be trained is mapped to a heterogeneous chip for execution. The search space constraints include that, among all heterogeneous chips, the number of pipeline stage model layers corresponding to the same type of chip is the same, and the number of pipeline stage model layers corresponding to heterogeneous chips is the same or different. The data parallel dimension constraint includes that the micro-batch size of each pipeline stage is 1, and thus each pipeline stage uses the same data parallel dimension. The heterogeneous chip number constraint includes that the total number of heterogeneous chips is equal to the sum of the total number of heterogeneous chips in each pipeline stage, and the total number of heterogeneous chips in each pipeline stage is equal to the product of the tensor parallel dimension and the data parallel dimension of the corresponding pipeline stage. The tensor parallel dimension constraint includes that the tensor parallel dimension is a power of 2, and that the tensor parallel dimension is less than the maximum tensor parallel dimension independently specified for each heterogeneous chip. The memory constraints include that the chips in the pipeline are arranged in descending order of memory size, and the sum of the memory occupied by the model, optimizer state, and activation value in the first pipeline stage of each heterogeneous chip is less than the maximum memory of the corresponding chip.

3. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 1, characterized in that, The optimal parallel strategy search method is also based on a preset computing card constraint, which is used to constrain the allocation of model layers. The computation card constraint includes that different pipeline stages containing the same computation card use the same tensor parallel dimension.

4. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 1, characterized in that, S101 specifically includes: S111, Based on the data parallel dimension constraint, select a global set of candidate data parallel dimensions, and then calculate the number of micro-batches; S112, For a given candidate data parallel dimension, select a candidate tensor parallel dimension for a single chip according to the tensor parallel dimension constraint, and then calculate the number of pipeline stages for a single chip according to the heterogeneous chip quantity constraint. Then, combine the types of heterogeneous chips to obtain candidate combinations of multiple pipeline stage quantities and tensor parallel dimensions under a given candidate data parallel dimension. S113. Based on the candidate set of candidate data parallel dimensions and the candidate combinations of the number of pipeline stages and tensor parallel dimensions, obtain the candidate set of data parallel dimensions, tensor parallel dimensions and pipeline parallel dimensions of all types of heterogeneous chips under all candidate data parallel dimensions, i.e., the first candidate set.

5. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 4, characterized in that, The number of micro-batches is the quotient of the training batch size and the data parallelism dimension.

6. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 4, characterized in that, S102 specifically includes: Based on the mapping constraints and the search space constraints, obtain a candidate set of model layer allocation schemes; The model layer allocation scheme is pruned using a heuristic method; Based on the candidate combinations of the number of pipeline stages and tensor parallel dimensions, the time required for each heterogeneous chip to process one layer of the model in the current tensor parallel dimension is used to balance the water level, and fine-tuning is performed until the total number of model layers is met, thus obtaining a preliminary set of parallel strategy candidates. Based on the memory constraints, parallel strategy candidates that do not meet the constraints are removed from the preliminary parallel strategy candidate set to obtain the final parallel strategy candidate set.

7. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 1, characterized in that, The specific memory constraints are as follows: Model layer number × (the memory of the parameters and optimizer state of model layer 1 under the current data parallel dimension and tensor parallel dimension + the number of preheating micro-batches in the current pipeline stage × the memory of the activation value of 1 sample processed by model layer under the current tensor parallel dimension) < chip maximum memory.

8. The optimal parallel strategy search method for non-uniform heterogeneous chips according to claim 1, characterized in that, In S103, the time for a single iteration is the sum of the global forward and backward computation time, the global update optimizer parameter time, the pipeline bubble time, and the pipeline communication time.

9. An optimal parallel strategy search device for non-uniform heterogeneous chips, characterized in that, The device includes preset mapping constraints, search space constraints, data parallel dimension constraints, heterogeneous chip quantity constraints, tensor parallel dimension constraints, and memory constraints, and includes a first candidate set acquisition module, a parallel strategy candidate set determination module, and an optimal parallel strategy selection module. The first candidate set acquisition module is used to acquire a first candidate set based on the data parallel dimension constraint, the heterogeneous chip quantity constraint and the tensor parallel dimension constraint. The first candidate set includes combinations of data parallel dimension, tensor parallel dimension and pipeline parallel dimension of all types of heterogeneous chips. The parallel strategy candidate set determination module is used to determine the parallel strategy candidate set based on the first candidate set, according to the mapping constraint, the search space constraint and the memory constraint, the model layer allocation, and the parallel strategy candidate set includes the data parallel dimension, tensor parallel dimension, pipeline parallel dimension and model layer combination of all types of heterogeneous chips. The optimal parallel strategy selection module is used to calculate the single iteration time for each combination in the candidate set of parallel strategies, and select the combination with the lowest single iteration time as the optimal parallel strategy.

10. An optimal parallel strategy for non-uniform heterogeneous chips, characterized in that, The optimal parallel strategy is obtained based on the method described in any one of claims 1-8. The optimal parallel strategy is a combination of data parallel dimension, tensor parallel dimension, pipeline parallel dimension and model layer number, wherein the data parallel dimension is an integer, the tensor parallel dimension, the pipeline parallel dimension and the model layer number are all arrays, and the dimension of each array is limited by the heterogeneous chip type.