Large model post-training-oriented multi-model structure perception parallel strategy search method
By constructing a multi-model structure-aware parallel strategy search method, analyzing model calculation diagrams and performance data, and optimizing resource allocation, the parallel strategy design problem of multi-model interaction in post-training of large-scale models is solved, and distributed training efficiency and performance are improved.
Patent Information
- Application Number
- CN202510462034.0
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-04-14
- Publication Date
- 2025-08-01
AI Technical Summary
In the process of post-training of large models, in multi-model interaction scenarios, it is difficult for the existing technology to efficiently design parallel strategies with excellent performance, especially when there are many models and complex structures, distributed training equipment is difficult to withstand, and the existing automated analysis methods fail to effectively consider model structure differences and training inference tasks differences.
The parallel strategy search method for multi-model structure-aware training based on large models is adopted. By analyzing the model calculation diagram and performance data, a four-layer tree search space is built, load balancing inside and outside the model, unreasonable pruning strategies, optimizing communication and computing resource allocation between models, and using prefill&decode separation deployment to realize load balancing between models and parallel strategy search.
It improves the efficiency of parallel strategy search during post-training of large models, reduces training costs, optimizes load balancing between models, and improves distributed training performance.
Smart Images

Figure CN120409631A_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the technical field of neural networks, and particularly relates to a parallel policy search method for multi-model structure perception for post-training of large models. Background Art
[0002] In recent years, with the development of industrial informatization and the Internet, large language models (LLMs) have made breakthrough progress in the field of artificial intelligence, driving the rapid development of technologies such as natural language processing, dialogue systems, and automatic code generation. Especially after introducing the innovative idea of human preference training (RLHF: Reinforcement Learning from Human Feedback), the capabilities of large models have been significantly improved, enabling them to better understand human intentions, generate high-quality content, and meet the needs of practical applications. During the post-training process of large models, interactions among multiple models are involved, including the policy model (Actor Model), the reward model (Reward Model), the evaluation model (Critic Model), and the reference model (Reference Model). Among them, the policy model is used to generate a series of responses for a given input, the reward model scores and evaluates these responses, the evaluation model is used to predict the future cumulative rewards of the responses, and the generalized advantage estimation algorithm is used to evaluate the quality of each generation of the policy model. At the same time, the policy model and the evaluation model are trained. The reference model uses the KL divergence to prevent the policy model from being trained incorrectly and maintains the stability of the model. Large models such as openAI o1 and deepseek r1 have comprehensively outperformed human doctor experts on the benchmark GPQA-diamond for testing chemical, physical, and biological expertise, demonstrating super reasoning capabilities that previous models did not possess. The post-training stage involves multi-model interactions, and each model has a certain amount of parameter scale. It is difficult for a single computing acceleration device (such as a GPU) or even a single cluster to bear the training of multiple models. How to use distributed training technology to slice multiple models in the post-training stage of large models to multiple devices for training to improve the overall training performance has become a hot topic in the current field of deep learning.
[0003] The current distributed training of large models mainly relies on expert experience, which requires developers to deeply understand the hardware characteristics of neural network models and training devices and design the distributed training strategy of the model based on these characteristics. However, with the continuous increase in the model scale and the device cluster scale, the search space of the strategy grows exponentially, and it is difficult for the expert experience-based method to design a parallel strategy with better performance in a short time.
[0004] To improve the performance of model parallel strategies and the efficiency of strategy search, the academic and industrial communities have begun to study automated analysis methods for models. The model automatic analysis method extracts model structure and device topology information, and automatically analyzes the neural network model with the help of search algorithms, so as to simplify the design process of model parallel strategies and improve the performance of model distributed training. At present, the mainstream model automatic parallel method searches for parallel strategies for a single model, without considering the differences in the multi-model interaction scenario in the post-training scenario, including differences in model structure, differences in training and inference tasks, etc. Summary of the Invention
[0005] During the post-training process of large models, the training and inference processes of multiple models are involved, including the policy model, reward model, evaluation model, and reference model. In view of the problems in the post-training process of large models, such as a large number of models, complex neural network structures, and complex model interactions, which make it difficult to design and implement the computing power information and distributed parallel strategies allocated to each model, this invention designs a parallel strategy search method with multi-model structure awareness for the post-training of large models, which can efficiently and accurately evaluate the structures of each model, as well as their computational and communication volumes, realize load balancing between models and within models, and support the design and implementation of computing power allocation and distributed parallel strategies for multiple models. First, this invention analyzes each neural network model in the post-training of large models, and evaluates the model computational graph according to its computing and storage characteristics. Secondly, profile performance data is collected for different models (training models, inference models), including hardware information such as the single-card video memory size, the number of cluster cards, the intra-node communication bandwidth, and the inter-node communication bandwidth, as well as information such as parameter synchronization and communication between training models and inference models, and model performance data such as the activation value video memory of each layer of the model, forward and backward times. According to the computational graph and profile data, analyze the operator data flow characteristics in the computational graph, and construct the computational and video memory overheads of each model in combination with operator parameters, communication modes, etc. Then, pre-allocate the number of computing power (GPUs) for each model according to the computational and video memory overheads of each model, and perform pruning according to the load balancing between and within models. According to the model video memory and computational volume information in the profile and the cluster topology, construct a four-layer tree search space. The first layer is multiple models involved in post-training, including the policy model, reward model, evaluation model, and reference model. The second to fourth layers are the parallel strategy searches for each model, which are pipeline parallelism, data parallelism, and tensor parallelism in sequence. The parallel strategy search for each model is a subtree, which is independent of each other and does not affect each other, and can be searched concurrently. And according to whether the model performs inference to generate tasks, obtain the prefill and decode ratios according to the profile, and adopt prefill&decode separate deployment to allocate different data parallel scales for prefill and decode. Finally, based on the constructed four-layer tree search space, concurrently perform the distributed parallel strategy search for multiple models, and during the search process, prune unreasonable parallel strategies (such as OOM or more time-consuming parallel strategies) to reduce the search duration.
[0006] This invention provides a parallel strategy search method with multi-model structure awareness for the post-training of large models, and the steps are as follows:
[0007] Step 1: Analyze the network structures of various deep learning models, evaluate the model computation graphs according to the network structures, and construct the computation cost, communication cost, and video memory cost of each model based on the model computation graphs and profile measured performance data. Pre-allocate computing power (GPU) for each model according to the computation cost and video memory cost of each model.
[0008] Sub-step 1.1: Perceive the computation graph of each model in the deep learning framework and define a directed acyclic graph G i (O, E), i ∈ (1, 2,..., m), where m is the total number of models in the post-training process, and G i (O, E) is the computation graph of the i-th model. The node O represents the operator, and the directed edge E represents the data dependency relationship and the execution order of the operators.
[0009] Sub-step 1.2: The computation graph of each model shows the flow direction of the data stream. Combine the computation graph of each model and its data stream flow direction, and calculate the computation cost, communication cost, and video memory cost of each model in the post-training process according to the profile measured performance data. Obtain the cluster topology information through profile. The number of GPUs is N, the communication speed within the node is V intra , and the communication speed between nodes is V inter . Suppose there are a total of m models in the post-training process, and the performance data of each model is S i (t i , c i , m i , I i ), i ∈ (1, 2,..., m), where t i , c i , m i , I i are the computation duration, communication duration, video memory size, and model interaction duration of the i-th model, and record the performance data of each model in JSON format.
[0010] Since there are differences in the inference training tasks of models in the post-training stage of large models, the video memory and computation costs of the two models are different. The video memory composition of the inference model mainly consists of model parameters, KV cache size, activation values, etc.; the video memory composition of the training model mainly consists of model parameters, optimizers, activation values, etc. Therefore, the model video memory cost is expressed by the following formula:
[0011]
[0012] Among them, n represents the number of model parameters of the model, dtype represents the precision used by the model, usually BF16, FP16, FP32, etc., sizeof(·) is a function to obtain the data size, batch is the batch size, seq is the sequence length, l is the number of Transformer layers of the model, and h is the hidden dimension. A is the size of the activation value obtained from the actual profile. k is determined by the selected optimizer. If SGD is used, 2 times the video memory occupancy is required: the parameters of the training model + the reverse gradient; if SGD+Momentum is used, 3 times the video memory occupancy is required: the parameters of the training model + the reverse gradient + momentum; if Adam is used, 4 times the video memory occupancy is required: the parameters of the training model + the gradient + momentum + second-order momentum.
[0013] The duration of the inference model consists of the prefill and decode stages, while the duration of the training model consists of the forward and reverse stages of multiple layers. The model calculation cost is expressed by the following formula:
[0014]
[0015] Among them, I i is the time-consuming of the model interaction during the post-training process, prefill and decode are the time-consuming actually obtained from the profile, n d represents the number of decode times. t emb represents the forward and reverse time-consuming of the embedding layer, t tf_layer is the forward and reverse time-consuming of a single transformer layer, t output is the forward and reverse time-consuming of the output layer, t loss is the forward and reverse time-consuming of the loss layer.
[0016] Common communication methods include AllGather, AllReduce, P2P, etc. AllGather communication will gather the respective data from each process. After the data is aggregated, each process will get the complete data. Assuming that each process has K pieces of data and there are M processes in total, then each process will get M*K pieces of data after AllGather, and the communication volume required for each process is K*(M-1). And P2P is point-to-point communication, where the process sends and receives a single piece of data with another process. Therefore, the communication cost of each model during the post-training process can be expressed by the following formula:
[0017]
[0018] Among them, k represents the number of communications, H i represents the data volume size of the i-th communication, Type(comm i) is the number of communications for different communication types. The number of p2p communications is one, and the number of communications for AllGather is the number of processes.
[0019] Sub-step 1.3: Obtain performance data according to the profile and use it to model and generate the computational costs of intra-model and inter-model load balancing. Based on the computational costs of intra-model and inter-model load balancing as constraints, pre-allocate computational power (GPU) for each model. The inter-model load balancing is expressed by the following formula:
[0020]
[0021] where α is the conversion cost rate, and Cost compi is the computational cost of the i-th model, obtained from actual analysis. Through inter-model load balancing, the time costs of multiple models are balanced. The intra-model load balancing is expressed by the following formula:
[0022]
[0023] where β is the conversion cost rate, obtained from actual analysis. k i is the number of cards allocated to each model. Through intra-model load balancing, the computational costs of each card number within the model are balanced.
[0024] Step 2: According to the model computation graph obtained in sub-step 1.1 and the computing power allocated to each model obtained according to the computational cost, adopt prefill&decode separated deployment (PD separation) for the inference model, and allocate different data parallel scales for the full-scale inference prefill and incremental inference decode stages of the inference model.
[0025] Sub-step 2.1: During the execution of the inference model, the prefill stage is computation-intensive and has greater computational pressure. After the generation of the KV cache is completed, the prefill stage itself does not need to retain these caches. The decode stage is storage-intensive. Due to the nature of generating tokens one by one, it needs to frequently access the KV cache. The PD separation architecture can enable the prefill to focus on the prefill computation and the decode to focus on the token generation in the decode stage, optimizing the overall inference time. During the post-training process, the model inference generation stage occupies most of the time, and PD separation needs to be used to accelerate the model inference generation. Based on the time-consuming data of prefill and single decode obtained in sub-step 1.2, and the number of cards ki allocated to each model obtained in sub-step 1.3, allocate reasonable data parallel scales for the prefill and decode stages to reduce the overall inference generation time.
[0026] Step 3: Calculate the computing and video memory costs in 1.2 according to the multiple model calculation graphs obtained in sub-step 1.1. Combine the data parallel scale and adopt the tree-based multi-model parallel strategy search method to search for the distributed parallel strategy of each model during the subsequent training process. Use the tree-based multi-model parallel strategy search method to search for the distributed parallel strategy of each model during the subsequent training process. Multiple models concurrently execute their respective distributed parallel strategies. Pruning operations are performed while searching, thereby reducing the search duration.
[0027] Sub-step 3.1: According to the computing cost Cost of each model comp , video memory cost Cost mem and the computing power allocated to each model obtained according to the computing cost. Select the tree-based multi-model parallel strategy distributed parallel strategy search, with the goal of reducing the overall end-to-end time consumption overlap is the overlapping duration of computing and communication. Search in the order of pipeline parallelism PP, data parallelism DP, and tensor parallelism TP. Since the parallel strategies of each model are independent of each other, the models can concurrently execute the search, greatly reducing the search time consumption.
[0028] Sub-step 3.2 If the parallel strategy of the current model is PP i , DP i , TP i in the case of no out-of-memory OOM. If the TP in the next search is greater than the TP in the current search i , then the training or inference time consumption of this model will definitely increase, and the next search will be cancelled. Otherwise, continue the search. If OOM occurs under the current parallel strategy configuration, and the PP and TP in the next search are both smaller than the PP and TP in the current parallel configuration, then there is no need to search.
[0029] The beneficial effects of the present invention are:
[0030] Automatically analyze the structures of multiple deep neural network models in the post-training stage of large models. Based on the computing, storage, and parameter characteristics of the models, combined with the inference and training task characteristics in the post-training stage and the actual running profile performance data, evaluate and model the computing cost and video memory cost of the models. And through the load balancing within and between models for constraint, through the tree-based multi-model parallel strategy search, find the parallel strategy with the shortest overall end-to-end time consumption for each model in parallel, helping AI engineers search for distributed parallel strategies for distributed training and reducing training costs. Description of the Drawings
[0031] Figure 1 is the overall architecture diagram;
[0032] Figure 2It is a schematic diagram of multi-model structure information and data flow in the post-training stage;
[0033] Figure 3 It is a schematic diagram of tree-based multi-model parallel strategy search. Detailed implementation manners
[0034] The present invention will be further described below in conjunction with the accompanying drawings and specific implementation steps:
[0035] Such as Figures 1 to 3 shown, a distributed heterogeneous parallel strategy search method based on multi-model structure perception includes the following steps:
[0036] Step 1: Analyze the neural network structures of each model in deep learning, evaluate the model computation graph according to the network structure. And according to the model computation graph and profile data, analyze the operator data flow characteristics in the computation graph, and construct the computation cost, communication cost, and video memory cost of each model in combination with operator parameters, communication modes, etc., and pre-allocate computing power for each model according to the computation cost and video memory cost of each model.
[0037] Sub-step 1.1: Perceive the computation graph of each model in the deep learning framework, such as Figure 2 shown. For the computation graph generated by each model, define a directed acyclic graph G i (O, E), i ∈ (1, 2,..., m), m is the total number of models, G i (O, E) is the computation graph of the i-th model, the node O represents an operator, and the directed edge E represents the data dependency relationship and the execution order of the operators.
[0038] Sub-step 1.2: The computation graph of each model shows the direction of the data flow. Combine the computation graph of each model and its data flow direction, and calculate the computation cost, communication cost, and video memory cost of each model according to the profile measured performance data. Obtain the cluster topology information through profile, the number of GPUs is N, the communication speed within the node is V intra , and the communication speed between nodes is V inter . Suppose there are m models in the post-training of the large model, and the performance data of each model is S i (t i , c i , m i , I i ), i ∈ (1, 2,..., m), where, t i , c i , m i , I iLet \(t_i\), \(c_i\), \(m_i\), and \(i_i\) be the computing duration, communication duration, video memory size, and model interaction duration of the \(i\)-th model. The performance data of each model is recorded in JSON format. Taking the actor model as an example, it can be converted to \([{"embedding":{"time":"1.022210044","memory":"0.01779224","max_memory":"0.035586357"},"\(0\)":{"time":"5.975484848","memory":"0.2821211814","max_memory":"0.31946325"},"output_layer":{"time":"7.842540740","memory":"1.1420896982","max_memory":"1.13377409019470215"}},"forward_step_time":"19.4187164306", "forward_step_mem":"2.5309791564", "backward_step_time":"3.48639488220214", "backward_step_mem":"2.25961351394"]\). Among them, "embedding", "\(0\)", and "output_layer" represent the performance data of the embedding layer, one transformer layer, and the output layer in the model, and time and memory represent the time consumption and video memory occupancy of this layer.
[0039] Due to the differences in model inference and training tasks during the post-training stage of large models, the video memory and computing costs of the two models are different. The video memory composition of the inference model mainly consists of model parameters, KV cache size, activation values, etc.; the video memory composition of the training model mainly consists of model parameters, optimizers, activation values, etc. Therefore, the video memory cost of the model is expressed by the following formula:
[0040]
[0041] Among them, n represents the number of model parameters of the model. The model consists of an Embedding layer, multiple Transformer layers, and an output layer. Taking the Embedding layer as an example, it designs a matrix with a hidden size of hidden_size and a vocabulary size of vocab_size, and the number of parameters is hidden_size × vocab_size. sizeof(·) is a function to obtain the data size, and dtype represents the precision used by the model, usually BF16, FP16, FP32, etc. batch is the batch size, seq is the sequence length, l is the number of Transformer layers of the model, h is the hidden dimension, and A is the size of the activation value obtained from the actual profile. k is determined by the selected optimizer. If SGD is used, 2 times the video memory occupancy is required: model parameters + reverse gradient; if SGD + Momentum is used, 3 times the video memory occupancy is required: model parameters + reverse gradient + momentum; if Adam is used, 4 times the video memory occupancy is required: model parameters + gradient + momentum + second-order momentum.
[0042]
[0043]
[0044] The duration of the inference model consists of the prefill and decode stages, while the duration of the training model consists of the forward and backward stages of multiple layers. The time cost of the model is expressed by the following formula:
[0045]
[0046] Among them, I i is the time consumed by the model interaction during the post-training process, prefill and decode are the time consumed actually obtained from the profile, and n d represents the number of decode times. t emb represents the forward and backward time consumed by the embedding layer, tt tf_layer is the forward and backward time consumed by a single transformer layer, t output is the forward and backward time consumed by the output layer, t loss is the forward and backward time consumed by the loss layer.
[0047]
[0048]
[0049] Common communication methods include AllGather, AllReduce, P2P, etc. AllGather communication aggregates the data from each process. After the data is aggregated, each process will obtain the complete data. Assuming that each process has K pieces of data and there are M processes in total, each process will obtain M*K pieces of data after AllGather, and the communication volume required for each process is K*(M - 1). P2P is point-to-point communication, where processes perform single-piece data Send and Recv with each other. Therefore, the communication cost of each model during the post-training process can be expressed by the following formula:
[0050]
[0051] Among them, k represents the number of communication times, H i represents the data volume size of the i-th communication, dtype represents the precision used, and the common ones are FP16, BF16, etc. Type(comm i ) is the number of communication times for different communications. The P2P communication is once, and the number of communication times for AllGather is the number of processes.
[0052]
[0053] Sub-step 1.3: Obtain performance data according to the profile and use it to build a model to generate the computational cost after load balancing within and between models. Taking the computational cost of load balancing within and between models as a constraint, pre-allocate computational power (GPU) for each model. The load balancing between models is expressed by the following formula:
[0054]
[0055] Among them, α is the conversion cost rate, obtained from actual analysis. Through load balancing between models, the time-consuming costs of multiple models are balanced. The load balancing within the model is expressed by the following formula:
[0056]
[0057] Among them, β is the conversion cost rate, obtained from actual analysis. k i is the number of cards allocated to each model. Through load balancing within the model, the computational cost of each card number within the model is balanced.
[0058] Step 2: According to the deep neural network computation graph obtained in sub-step 1.1 and the computing power allocated to each model obtained according to the computational cost, adopt prefill&decode separate deployment for the inference model, and allocate different data parallel scales for the prefill and decode stages of the inference model.
[0059] Sub-step 2.1: During the execution of the inference model, the prefill stage is computationally intensive with greater computational pressure. After generating the KV cache, the prefill stage itself doesn't need to retain these caches continuously. The decode stage is memory-intensive. Due to the nature of token-by-token generation, it needs to access the KV cache frequently. The prefill&decode separation architecture allows prefill to focus on prefill computations and decode to focus on token generation in the decode stage, optimizing the overall inference time.
[0060] Step 3: Based on the multiple deep neural network computation graphs obtained in sub-step 1.1, the computation, video memory, and communication costs in 1.2, and the number of cards pre-allocated for each model in 1.3, use a tree-based multi-model parallel strategy to search for a distributed parallel strategy. Multiple models concurrently execute their respective distributed parallel strategies. And pruning operations are carried out during the search to reduce the search duration.
[0061] Sub-step 3.1: According to the computational cost Cost comp of each model, the video memory cost Cost mem and the computing power allocated to each model, select a tree-based multi-model distributed parallel strategy for searching. Search in the order of PP, DP, TP. Since the parallel strategies of each model are independent of each other, the searches can be executed concurrently, greatly reducing the search duration. The goal is to reduce the overall end-to-end duration. overlap is the overlapping duration of computation and communication.
[0062] Sub-step 3.2: During the search process, prune unnecessary parallel configurations (such as those that cause OOM or more time-consuming parallel strategies) to further reduce the search duration.
[0063]
[0064]
[0065] Suppose a cluster with 16 cards of 910B1 (64G) is used, seq_length is 1024, max_prompt_length is 512, batch_size is 8, the policy model, reward model, evaluation model, and reference model all use Llama3, hidden_size is 4096, num_attention_heads is 32, num_layer is 16, vocab_size is 128256, the data type is fp16, and the optimizer is Adam.
[0066] I. Video Memory Performance
[0067] The video memory required for model training and inference consists of model parameters, gradients, optimizers, kv, etc.
[0068] (1) The video memory required for each model parameter is:
[0069] Memoryparam = vocab_size × hidden_size + num_layer × (12 × hidden_size 2 + 4 × hidden_size) × 2 / 1024 3 = 6.98GB
[0070] (2) The video memory occupied by each model is approximately Memory model ≈ 8 × Memory param = 55.84GB
[0071] II. Computational Performance
[0072] The forward computational volume of the model is mainly composed of the transformer layer and the output layer.
[0073] (1) The computational volume of each transformer layer is:
[0074] 24 × batch_size × seq_length × hidden_size 2 + 4 × batch_size ×
[0075] seq_length 2 × hidden_size = 3.4 × 10 12 [[ID=—38]]
[0076] (2) The computational volume of the output_layer is:
[0077] 2 × batch_size × seq_length × hidden_size × vocab_size = 8.6 × 10 12
[0078] Then the forward computational volume of the model is 16 × 3.4 × 10 12 + 8.6 × 10 12 = 6.3 × 10 13
[0079] Therefore, the computational volume during inference of the policy model is 6.3 × 10 13 , and the computational volumes of the reward model and the reference model are 6.3 × 10 13. When the large model performs training, the amount of backpropagation calculation is usually twice that of the forward calculation. Therefore, when evaluating the model and the policy model perform training, it is 3×6.3×10 13 = 1.9×10 14 .
[0080] During the post-training process of the large model, the end-to-end time consumption is mainly composed of the inference generation of the policy model, occupying 60% or more of the data. Therefore, according to prior knowledge and the video memory and computing costs of each model, 8 cards are allocated for the policy model to perform inference, and 2 cards are allocated for each of the reward model, evaluation model, reference model, and policy model to perform training. Each model performs parallel policy search. When the policy model performs inference, the parallel policy is DP4 PP2 TP1, and when performing training, the inference parallel policy is DP2PP1 TP1; the evaluation model, reference model, and policy model are DP2 PP1 TP1;
[0081] III. Execution Time
[0082] (1) Policy Model
[0083] Based on the above configuration, when the policy model performs inference, one prefill takes 10ms, one decode takes 0.08ms, and the number of decode times is seq_length - max_prompt_length. Then the time consumption ratio of the model to perform inference once is:
[0084]
[0085] When the policy model performs training, the time consumption is mainly composed of the forward and backward of each layer. The forward time consumption of one layer of the transformer layer is about 0.1ms, the backward is about 0.2ms, and the communication ratio is 30%. Therefore, the model training time consumption is
[0086] T train = 1.3×(0.1 + 0.2)×num_layer = 6.24ms
[0087] (2) Reward Model and Reference Model
[0088] During the post-training process, these two models each perform one forward pass. Therefore, the time consumption is
[0089] T reward ≈T reference = 1.3×0.1×num_layer = 2.08ms
[0090] (3) Evaluation Model
[0091] During the post-training process, the evaluation model needs to continuously train to adjust its parameters. Therefore, the time consumption is approximately the same as that of the policy model, which is:
[0092] T critic = 1.3×(0.3)×num_layer = 6.24 ms
[0093] During the post-training process of the large model, there are interactions among multiple models. The model interaction time is about 2 ms. Then the overall end-to-end time consumption is as follows:
[0094] Time = T inference + T train + T reward + T critic + T refenece + 4×2 = 75.6 ms.
Claims
1. A parallel policy search method for multi-model structure perception in the post-training of large models, characterized in that It includes the following steps: Step 1: Analyze the network structures of various deep learning models, evaluate the model computation graphs, construct the computational cost, communication cost, and video memory cost of each model based on the model computation graphs and profile measured performance data, and pre-allocate computing power for each model according to the computational cost and video memory cost of each model; Step 2: According to the model computation graphs and the computing power allocated to each model, adopt prefill&decode separate deployment for the inference model, and allocate different data parallel scales for the prefill and decode stages of the inference model; Step 3: According to the model computation graphs, computational cost, and video memory cost, combined with the data parallel scale, adopt a tree-based multi-model parallel strategy search method to search for the distributed parallel strategies of each model during the post-training process, and perform pruning operations while searching.
2. The parallel policy search method for multi-model structure perception in the post-training of large models according to claim 1, wherein The specific implementation process of the above Step 1 is as follows: Step 1.1: Sense the computational graph of each model in the deep learning framework and define a directed acyclic graph G i (O, E), i ∈ (1, 2,..., m), where m is the total number of models in the reinforcement learning, and G i (O, E) is the computational graph of the i-th model, the node O represents the operator, and the directed edge E represents the data dependency relationship and the execution order of the operators; Step 1.2: Combine the computation graph of each model and its data flow direction, and calculate the computational cost, communication cost, and video memory cost of each model during the post-training process according to the profile measured performance data; Step 1.3: Obtain performance data according to the profile and use it to model and generate the computational cost after load balancing within and between models, pre-allocate computing power GPU for each model, and perform pruning with the load balancing within and between models as the constraint conditions.
3. The parallel policy search method for multi-model structure perception after training of a large model according to claim 2, wherein The specific implementation process of the computational cost, communication cost, and video memory cost of each model during the post-training process is as follows: Obtain the cluster topology information through the profile, and the number of GPUs is N; assume that there are a total of m models during the post-training process, and the performance data of each model is S i (t i ,c i ,m i ,I i ), i ∈ (1, 2,..., m), where t i 、c i 、m i 、I i are the computing duration, communication duration, video memory size, and model interaction duration of the i-th model; The video memory cost of the model is expressed by the following formula: where n represents the number of model parameters of the model, dtype represents the precision used by the model, batch is the batch size, sizeof(·) is a function to obtain the data size, seq is the sequence length, l is the number of Transformer layers of the model, h is the hidden dimension, A is the size of the activation value obtained from the actual profile, and k is determined by the selected optimizer; The duration of the inference model consists of the prefill and decode stages, while the duration of the training model consists of the forward and backward stages of multiple layers; the computational cost of the model is expressed by the following formula: Among them, prefill and decode are the actual time-consuming obtained by the profile, n d represents the number of times of decode, t emb represents the forward and backward time-consuming of the embedding layer, t tf_layer is the forward and backward time-consuming of a single transformer layer, t output is the forward and backward time-consuming of the output layer, t loss is the forward and backward time-consuming of the loss layer; The communication cost of each model during the post-training process of reinforcement learning is expressed by the following formula: where k represents the number of communications, and H i represents the data volume of the i-th communication, and Type(Commi) is the number of communications of different communication types.
4. The parallel policy search method for multi-model structure awareness after training for large models according to claim 3, wherein, The inter-model load balancing is expressed by the following formula: where α is the conversion cost rate, is the computational cost of the i-th model. Through load balancing among models, the time-consuming costs of multiple models are balanced. The intra-model load balancing is expressed by the following formula: Among them, β is the conversion cost rate, and k i is the number of cards allocated to each model. Through the load balancing within the model, the computing cost of each card number within the model is balanced.
5. The parallel policy search method for multi-model structure perception after training for large models according to claim 4, characterized in that, The specific implementation of the above Step 2 is: during the execution of the inference model, after the prefill stage completes the generation of the KV cache, the prefill stage does not continue to retain these caches; The decode stage needs to continue to access the KV cache; Based on the time-consuming data of prefill and single decode, and the number of cards ki allocated to each model, allocate data parallel scales for the prefill and decode stages to reduce the overall inference generation time.
6. The parallel policy search method for multi-model structure awareness for post-training of large models according to claim 5, wherein The specific implementation process of the distributed parallel strategies of each model during the post-training search is as follows: Step 3.1: Obtain the computational cost Cost of each model comp , the video memory cost Cost mem , and the computing power allocated to each model obtained according to the computational cost. Based on these factors, select a tree-based multi-model parallel strategy for distributed parallel strategy search with the goal of reducing the overall end-to-end latency The overlap is the duration of the overlap between computation and communication. Search in the order of pipelining parallelism PP, data parallelism DP, and tensor parallelism TP, and perform concurrent execution search between models; Step 3.2 If the parallel strategy of the current model is PP i , DP i , TP i In the case of, there is no out-of-memory (OOM). If the TP in the next search is greater than the TP in the current search i , the training or inference time of the current model increases, then cancel the next search, otherwise continue the search; If OOM occurs under the current parallel strategy configuration, and both the PP and TP in the next search are less than the PP and TP in the current parallel configuration, there is no need to search.
Citation Information
Cited By
Distributed training method of multi-modal model, electronic equipment and storage medium
CN121144857A
Distributed training method of multi-modal model, electronic device and storage medium
CN121144857B
Parallel training strategy online switching method and deep learning model training system
CN122021969A