A parallel strategy search method for efficient training of large artificial intelligence models

By building a three-dimensional hardware topology field and two-layer policy network optimization operator node, the communication overhead deviation caused by the hardware topology effect in parallel training of artificial intelligence large models is solved, efficient parallel strategy search is achieved, and training efficiency and system robustness are improved.

CN120012879BActive Publication Date: 2025-08-19西安圣瞳科技有限公司
View PDF 3 Cites 0 Cited by

Patent Information

Application Number
CN202510488064.9
Authority / Receiving Office
CN · China
Patent Type
Patents(China)
Current Assignee / Owner
Filing Date
2025-04-18
Publication Date
2025-08-19
Estimated Expiration
2045-04-18

AI Technical Summary

Technical Problem

The prior art has failed to effectively solve the problem of communication overhead estimation bias and limited strategy robustness caused by hardware topology effects in parallel training of artificial intelligence large models. Especially under the interweaving of video memory and communication bottlenecks in large-scale clusters, it is difficult to achieve the improvement of multi-strategy synergy efficiency.

Method used

By generating a three-dimensional hardware topology field, analyzing the calculation graph and mapping the operator nodes, building a hybrid integer planning model, combining a two-layer policy network for parallel policy search, dynamically adjusting task allocation, optimizing communication paths using NeRF model and hybrid integer planning, and combining Transformer and cellular automata network to achieve global-local collaborative optimization.

Benefits of technology

It significantly improves the accuracy of communication time prediction, reduces the memory fragmentation rate, improves GPU utilization, shortens the failure recovery time, and achieves a multi-objective balance between training efficiency, memory utilization and system robustness.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN120012879B_ABST
    Figure CN120012879B_ABST
Patent Text Reader

Abstract

The present invention discloses a parallel strategy search method for efficient training of large artificial intelligence models, which relates to the field of artificial intelligence technology. The method comprises the following steps: generating a three-dimensional hardware topology field by training a NeRF model; parsing the computational graph of the large model to be trained, mapping operator nodes to the three-dimensional hardware topology field, and generating a labeled computational graph; constructing a mixed integer programming model based on the labeled computational graph, and outputting a set of candidate strategies; inputting the set of candidate strategies into a two-layer strategy network, and outputting a set of parallel strategies; loading the three-dimensional hardware topology field into a simulator, verifying and adjusting the set of parallel strategies, and outputting the optimal strategy; and deploying the optimal strategy to a real cluster to trigger re-screening of strategies. The present invention models the physical-level hardware communication field based on the neural radiation field, integrates mixed integer programming with deep reinforcement learning, and realizes global-local collaborative optimization of parallel strategies, achieving a multi-objective balance among training efficiency, memory utilization, and system robustness.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence technology, and in particular to a parallel strategy search method for efficient training of large artificial intelligence models. Background Art

[0002] As the parameter size of large AI models continues to expand, parallel training based on distributed clusters has become a key technical path. Existing technologies primarily optimize basic strategies such as data parallelism, model parallelism, and pipeline parallelism, and implement policy search through heuristic algorithms or reinforcement learning. For example, Google proposed GPipe to reduce video memory usage through micro-batch pipeline partitioning, while NVIDIA Megatron-LM optimizes model partitioning based on tensor parallelism. However, these methods have significant limitations: First, traditional parallel strategy combination optimization ignores the physical topology effects of heterogeneous hardware (such as cross-rack communication attenuation and NVLink bandwidth heterogeneity), resulting in large deviations in communication overhead estimates; second, existing hierarchical optimization frameworks (such as first dividing the pipeline and then optimizing data parallelism) have difficulty modeling the synergistic effects of multiple strategies and are prone to falling into local optimality. This is especially true in large-scale clusters where video memory and communication bottlenecks are intertwined, severely limiting the robustness of the strategies. Summary of the Invention

[0003] In view of the above existing problems, the present invention is proposed.

[0004] Therefore, the present invention provides a parallel strategy search method for efficient training of large artificial intelligence models to solve the problems of low efficiency of multi-dimensional parallel strategy collaboration and poor dynamic adaptability of hardware topology in large-scale distributed training.

[0005] In order to solve the above technical problems, the present invention provides the following technical solutions:

[0006] In a first aspect, the present invention provides a parallel strategy search method for efficient training of large artificial intelligence models, which includes inputting the physical topology parameters of the cluster, training the NeRF model with multi-view bandwidth and delay data, and generating a three-dimensional hardware topology field including bandwidth density and communication attenuation coefficient;

[0007] Parse the computational graph of the large model to be trained, map the operator nodes to the three-dimensional hardware topology field, and label each operator node with a dynamic communication cost label based on the bandwidth density and communication attenuation coefficient of the operator node to generate a labeled computational graph;

[0008] Construct a mixed integer programming model based on the annotated computational graph and output a set of candidate strategies that satisfy the communication field constraints;

[0009] The candidate policy set is input into the two-layer policy network, and the output is a parallel policy set containing a global scheduling framework and local adaptive rules;

[0010] Load the 3D hardware topology field into the simulator, verify and adjust the parallel strategy set, and output the optimal strategy;

[0011] Deploy the optimal strategy to the real cluster and collect training data, provide feedback to update the NeRF model, and trigger the mixed integer programming model to re-screen the strategy.

[0012] As an optimal solution of the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the physical topology parameters include rack layout, connection relationship between nodes and physical distance.

[0013] As an optimal solution for the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the operator node includes the operator's computing power, graphics memory requirements and dependencies.

[0014] As a preferred solution of the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the mixed integer programming model constructed based on the labeled computational graph is to integrate the computation time, communication time and constraints through the objective function, and enumerate the pipeline stage division and data parallel sharding strategy through the quantum annealing algorithm.

[0015] As a preferred solution of the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the candidate strategy set is input into a two-layer strategy network, and the output of the parallel strategy set including a global scheduling framework and local adaptive rules is to generate a preliminary scheduling strategy by performing global attention modeling on the spatiotemporal features of the annotated computational graph through a Transformer encoder, the cellular automaton network performs local rule adjustment based on the states of adjacent nodes, and the Transformer encoder and the cellular automaton network are jointly optimized through a differentiable interface.

[0016] As a preferred solution of the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the verification and adjustment of the parallel strategy set refers to using the low-fidelity mode to screen potential high-quality strategies, and the high-fidelity mode simulates cross-node communication paths through ray tracing, and dynamically adjusts task allocation in combination with the cellular automaton rule base.

[0017] As a preferred solution of the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the training data includes communication efficiency, video memory fluctuation data and hardware abnormal events.

[0018] As an optimal solution for the parallel strategy search method for efficient training of large artificial intelligence models described in the present invention, the training data is divided into two feedback paths. The first path inputs the NeRF model to incrementally update the three-dimensional hardware topology field parameters. The second path inputs the two-layer strategy network. By comparing the deviation between the simulated prediction value and the measured value, the dependence weight of the Transformer attention mechanism on the hardware field characteristics is adjusted, and the local rule base of the cellular automaton is evolved.

[0019] In a second aspect, the present invention provides a computer device comprising a memory and a processor, wherein the memory stores a computer program, wherein: when the computer program is executed by the processor, it implements any step of the parallel strategy search method for efficient training of large artificial intelligence models as described in the first aspect of the present invention.

[0020] In a third aspect, the present invention provides a computer-readable storage medium having a computer program stored thereon, wherein: when the computer program is executed by a processor, it implements any step of the parallel strategy search method for efficient training of large artificial intelligence models as described in the first aspect of the present invention.

[0021] The present invention achieves the following beneficial effects: modeling physical-level hardware communication fields based on Neural Radiant Fields (NeRF), integrating mixed integer programming with deep reinforcement learning to achieve global-local collaborative optimization of parallel strategies. A three-dimensional bandwidth density field is constructed using the NeRF model to quantify cross-rack link attenuation effects, significantly improving the accuracy of communication time prediction. Cluster TPI (Per Iteration Time) is reduced in GPT-3-level model training. Mixed integer programming (MIQP) is used to jointly optimize pipeline stage division, data parallel sharding, and the compute-communication overlap factor, avoiding the local optimality trap caused by traditional hierarchical optimization. Field measurements show a significant reduction in memory fragmentation and improvement in GPU utilization. Based on a Transformer-cellular automaton architecture and combined with an online incremental learning mechanism, the approach reoptimizes strategies after hardware topology changes (such as node scaling), shortens fault recovery time, and reduces system throughput fluctuation compared to static strategies. This solution overcomes the limitations of traditional methods that optimize only a single performance metric, achieving a multi-objective balance among training efficiency, memory utilization, and system robustness. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0023] Figure 1This is an operation diagram of the parallel strategy search method for efficient training of large artificial intelligence models in this embodiment.

[0024] Figure 2 This is a flowchart of the parallel strategy search method for efficient training of large artificial intelligence models in this embodiment.

[0025] Figure 3 This is an operation diagram of the calculation diagram and optimization model in this embodiment.

[0026] Figure 4 This is a diagram of the double-layer strategy network structure in this embodiment. DETAILED DESCRIPTION

[0027] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the specific embodiments of the present invention are described in detail below with reference to the accompanying drawings.

[0028] In the following description, many specific details are set forth to facilitate a full understanding of the present invention. However, the present invention may also be implemented in other ways different from those described herein. Those skilled in the art may make similar generalizations without violating the connotation of the present invention. Therefore, the present invention is not limited to the specific embodiments disclosed below.

[0029] Secondly, the term "one embodiment" or "embodiment" herein refers to a specific feature, structure, or characteristic that may be included in at least one implementation of the present invention. The phrase "in one embodiment" appearing in various places throughout this specification does not necessarily refer to the same embodiment, nor does it refer to a separate or selective embodiment that is mutually exclusive of other embodiments.

[0030] In this embodiment, refer to Figures 1 to 4 This embodiment provides a parallel strategy search method for efficient training of large artificial intelligence models, including the following steps:

[0031] S1. Input the physical topology parameters of the cluster, train the NeRF (Neural Radiated Field) model using multi-view bandwidth and delay data, and generate a three-dimensional hardware topology field that includes bandwidth density and communication attenuation coefficient.

[0032] Physical topology parameters include rack layout, inter-node connection relationships, and physical distances;

[0033] It should be noted that the rack layout, including rack number, location coordinates, dimensions, and internal GPU node distribution, is obtained using a laser rangefinder or RFID tag positioning system (radio frequency identification tags) from the data center infrastructure management system or the physical plan drawings during cluster deployment. Inter-node connectivity is obtained using cluster management tools (such as the Kubernetes Node Topology API) or network switch configuration tables (such as those using the Link Layer Discovery Protocol). Use commands such as lshw (list hardware) and ibstat (InfiniBand status tool) to query InfiniBand / NVLink (InfiniBand / NVIDIA high-speed interconnect technology) connection status and generate a topology diagram of direct or cross-switch connections between nodes. Physical distance: Intra-rack node distances are determined based on rack design drawings (e.g., 0.1-meter spacing between adjacent GPU nodes), and inter-rack distances are measured using a laser rangefinder. Physical distance directly influences the signal attenuation model (e.g., the attenuation coefficient is proportional to the square of the distance), and the rack layout determines the electromagnetic interference zone (e.g., signal attenuation increases by 30% at the edges of metal racks). Node connectivity determines communication path selection (e.g., NVLink direct links have a bandwidth of 256 GB / s, while inter-switch bandwidth is reduced to 100 GB / s). Physical distance is combined with physical distance to generate a bandwidth density field. Based on the physical topology, high-attenuation paths (e.g., inter-rack paths with distances greater than 3 meters) are avoided, and low-latency links within the same rack are prioritized, reducing communication time by 40%. In the event of a node failure, rapid failover to a backup node within the same rack (e.g., migrating to an adjacent GPU within 0.2 meters) is performed based on connectivity, reducing recovery time to seconds. Physical distance constrains data parallel sharding strategies (e.g., inter-rack sharding requires a bandwidth density of ≥100 Gbps / m³), avoiding communication bottlenecks caused by long-distance transmission.

[0034] Based on the cluster's physical topology parameters, a distributed probe program was deployed in the cluster to perform cross-node bandwidth testing (test tool: Network Bandwidth Test Tool Version 3). This test covered all connected paths (within the same rack, across racks, and across switches), acquiring bandwidth data from multiple perspectives. Ping (a network connectivity test tool) was used to obtain one-way latency between nodes and collect latency data.

[0035] Integrate multi-view bandwidth data and delay data to form a multi-view bandwidth-delay dataset;

[0036] The cluster physical space is divided into a three-dimensional grid (resolution 0.1 meters), and each grid point is associated with bandwidth density and communication attenuation coefficient;

[0037] Input a multi-view bandwidth-delay dataset and train the NeRF model to learn the nonlinear mapping relationship from spatial coordinates to communication performance. The optimization goal is to minimize the mean square error between the predicted bandwidth density and the actual value.

[0038] It's important to note that during training, the NeRF model learns how to map spatial coordinates to expected communication performance values. Specifically, the optimization process focuses on minimizing the mean squared error between the bandwidth density predicted by the NeRF model and the observed bandwidth density. This means the NeRF model must accurately capture the variation in communication quality at every point in space, enabling it to provide precise communication performance predictions for unobserved spatial coordinates. By continuously iteratively updating the NeRF model parameters, the predicted results are kept as close to the true values as possible, thereby optimizing network planning and performance evaluation.

[0039] Compare the NeRF model predictions with the actual test data. If the NeRF predicted bandwidth density error exceeds 5%, add probe sampling points in areas with large errors (e.g., error > 5%) (e.g., rack edges) and retrain the NeRF model.

[0040] Outputs a 3D hardware topology field (HDF5 format), including the bandwidth density, communication attenuation coefficient, and coordinate mapping table for each spatial point.

[0041] It should be noted that the NeRF model is trained using multi-view bandwidth and latency data to model the impact of physical topology. NeRF converts parameters such as rack layout and physical distance into three-dimensional bandwidth density and communication attenuation coefficients, accurately quantifying the nonlinear attenuation of cross-node communication (e.g., the error in the attenuation coefficient for cross-rack paths is less than 5%). Through an error correction mechanism (e.g., adding additional sampling points when the error exceeds 5%), the trained NeRF field can dynamically adapt to hardware changes (e.g., the addition of new nodes), preventing policy failures caused by topology changes.

[0042] S2. Analyze the computational graph of the large model to be trained, map the operator nodes to the three-dimensional hardware topology field, and label each operator node with a dynamic communication cost label based on the bandwidth density and communication attenuation coefficient of the operator node to generate a labeled computational graph.

[0043] The operator node includes the operator's computational load, memory requirements, and dependencies;

[0044] Enable the context manager in the PyTorch performance analyzer, start the advanced performance analysis mode, perform a single forward and backward pass on a large model to be trained (such as the Meta large language model), generate a trace file containing detailed operator execution information, and extract metadata for each operator node, including the operator's computational load, graphics memory requirements, and dependencies.

[0045] It should be noted that the computational cost is calculated by directly reading the calculated value. If the computation graph is not provided, it is dynamically calculated by the operator type (such as two-dimensional convolution layer, matrix multiplication) and input dimension (such as the computational cost of the MatMul operator = 2×Batch×M×N×K, where M is the row of the left matrix, N is the column of the right matrix, K is the column of the left matrix, and Batch is the number of data samples), and normalized. The memory requirement is calculated by recording the input and output tensor dimensions of the operator, calculating the number of bytes according to the data type, and using the sum of the bytes of all input tensors to calculate the input memory requirement. The sum of the bytes of all output tensors is used to calculate the output memory requirement. The peak memory requirement is calculated using the maximum value of the memory temporarily allocated during the execution of the operator. The creation and release timestamps of each tensor are recorded, and the memory lifecycle is obtained by subtracting the earliest creation time from the last release time (unit: microseconds); the dependency relationship includes the following steps: loading the trace file generated by the profiler, extracting all operator nodes and thread / stream IDs in the order of execution, including predecessor operators and successor operators. The predecessor operator is the operator that has not been completed before the current operator in the same thread, and the successor operator is the operator that starts executing earliest after the current operator in the same thread. For multi-GPU / multi-thread scenarios, different execution flows are distinguished. If the output tensor of operator A is used by operator B (located on a different GPU), it is marked as a cross-device dependency. Circular dependencies in the computational graph are detected through depth-first search. If they exist, an exception is thrown (large models are usually directed acyclic graphs). The longest path from input to output is found, all operators on the path are marked, and a dependency table is output, which contains the predecessor list, successor list, and whether each operator is on the critical path.

[0046] Bind model layers to physical GPU nodes using a loop-based allocation strategy in layer order (e.g., 64 layers of LLaMA are allocated to 8 GPUs, with 8 layers per GPU).

[0047] Query the bandwidth density and communication attenuation coefficient of each operator node in the 3D hardware topology field, and add a communication cost label to each operator node. The communication cost calculation formula is: Communication cost = Transmitted data volume / (Bandwidth density × Attenuation compensation factor), where the attenuation compensation factor = 1 / (1 + Accumulated value of path attenuation coefficient). The path attenuation coefficient comes from the accumulated attenuation value of the communication path in the 3D hardware field.

[0048] Construct a 32-dimensional joint feature vector using computational load, memory requirements, and communication cost;

[0049] It should be noted that the computational load (8 dimensions) includes operator type encoding (4 dimensions), logarithmically normalized computational effort (1 dimension), degree of parallelism (1 dimension), and dependency depth (2 dimensions); the video memory requirement (8 dimensions) includes input / output tensor size (2 dimensions), video memory peak value (1 dimension), lifecycle (2 dimensions), and reuse count (1 dimension); the communication cost (16 dimensions) includes cross-device count (1 dimension), maximum single communication cost (1 dimension), and critical path marker (1 dimension), etc.

[0050] Outputs an annotated computation graph with hardware features (JSON file containing 32-dimensional joint feature vectors and physical binding information).

[0051] It should be noted that computation graphs are parsed and hardware features are integrated to annotate operator communication costs. A 32-dimensional joint feature vector (computational load + memory requirement + communication cost) is used to establish a mapping between operators and hardware resources. This improves the scheduling matching of operators on critical paths by 45% in GPT-3-level models. Dynamically adjusting communication priorities based on a decay compensation factor (1 / (1 + path decay)) reduces data transmission on high-decay paths, resulting in a 38% reduction in communication conflict rates.

[0052] S3. Construct a mixed integer programming model based on the labeled computational graph and output a set of candidate strategies that satisfy the communication field constraints.

[0053] The mixed integer programming model based on the annotated computational graph is constructed by integrating computation time, communication time and constraints through the objective function, and enumerating pipeline stage division and data parallel sharding strategies through the quantum annealing algorithm;

[0054] Define discrete and continuous variables. Discrete variables include the number of pipeline stages and the number of data parallelism. The number of pipeline stages is an integer variable with a value range of [1,16], indicating that the mixed integer programming model is divided into 1 to 16 pipeline stages. The number of data parallelism is an enumeration value {1,2,4,8}, indicating the number of data shards (the enumeration value must satisfy ≤ the total number of GPUs). The continuous variable [0,1] indicates the time overlap ratio between computation and communication through the computation-communication overlap factor (0 means no overlap, 1 means complete overlap).

[0055] The objective function is to minimize the single iteration time, which is expressed as:

[0056] ;

[0057] in, represents a single iteration time, Indicates the The total computational amount of the stage (extracted from the annotated computational graph), Indicates the The amount of data communicated across devices during the phase (calculated based on the operator binding location), Represents the first query from the 3D hardware topology field Current path bandwidth density at the current stage (unit: Gbps / m³), The index variable representing the stage division in the pipeline parallelism is used to traverse the time consumption of all stages. Indicates the total number of stages of parallel partitioning of the pipeline, The adjustment factor representing the overlap of computation and communication time directly affects the actual proportion of communication overhead. When it is 1, it means that the communication time is completely covered by the computing time (ideally, there is zero communication overhead). When it is 0, it means that the calculation and communication are completely serial (no overlap, and the communication time is all extra consumption). represents the attenuation compensation factor, =1 / (1+accumulated value of path attenuation coefficient), the path attenuation coefficient comes from the accumulated attenuation value of the communication path in the three-dimensional hardware field, The GPU computing power is the measured floating-point computing power (FLOP / s) based on the MLPerf benchmark. The real-time utilization is obtained by running standard operators such as GEMM (general matrix multiplication).

[0058] Constraints include memory constraints and communication path constraints. The memory constraint is that the memory usage of a single graphics card is ≤ 40GB (based on the cumulative memory requirements of the annotated computation graph). The communication path constraint is that the bandwidth density of the cross-shard communication path is ≥ 100Gbps / m³ (verified from the 3D hardware topology field).

[0059] The combinatorial optimization problem of discrete and continuous variables is converted into a quadratic unconstrained binary optimization model. This model is then embedded in a D-Wave quantum computer (quantum annealing hardware) and the quantum annealing algorithm is run to screen candidate solutions that meet both memory and bandwidth constraints. The top 50 quantum annealing solutions are then extracted as the initial solutions for the classical optimizer.

[0060] For each quantum annealing solution, the discrete variables and continuous variables are fixed, and the continuous variables are optimized to minimize the single iteration time. The algorithm is terminated when the change in the continuous variable is less than 0.01 or when 100 iterations are reached.

[0061] It's important to note that when performing quantum annealing calculations to solve an optimization problem, a series of possible solutions is first obtained. For each solution obtained from the quantum annealing process, the values of the discrete and continuous variables are first identified and fixed. This fixing means temporarily treating these variables as constants during further optimization. Next, the focus is on optimizing the continuous variables marked as adjustable, aiming to reduce the time required for a single iteration and thus speed up the overall solution process. Careful analysis is required to identify which continuous variables in the current solution significantly affect iteration time. Appropriate methods (such as local search and gradient descent) are then used to find the optimal values for these variables, minimizing iteration time while maintaining solution quality. The ultimate goal is to improve the overall efficiency and practicality of the quantum annealing algorithm. It is important to note that in practice, the optimization process is repeated to gradually approach the optimal solution while continuously optimizing iteration time.

[0062] Eliminate memory overruns and communication timeouts. Specifically, remove policies with a single GPU memory usage > 40GB and single iteration times > 500ms (based on high-fidelity simulation predictions).

[0063] Sort in ascending order by single iteration time, and select the strategy with the shortest single iteration time. If there are duplicate discrete or continuous variables in the top 200 strategies, fill in the suboptimal solution based on the single iteration time.

[0064] Output a set of candidate strategies that satisfy the communication field constraints (CSV file, each line contains the following fields: number of pipeline stages, number of data parallelism, overlap factor, single iteration time (ms), video memory usage (GB), shard binding list).

[0065] It should be noted that a mixed integer programming model is constructed, and quantum annealing is used to solve Nash optimal strategies. The quantum annealing algorithm supports traversing the combinatorial space with a 16-stage pipeline and 8-shard data parallelism (traditional brute force search only supports 4 stages x 4 shards), expanding the range of candidate strategies by 16 times. Constraints of video memory usage ≤ 40GB and bandwidth density ≥ 100Gbps / m³ are used to eliminate invalid strategies (such as those caused by memory overflow or communication bottlenecks).

[0066] S4. Input the candidate strategy set into the two-layer strategy network and output a parallel strategy set containing the global scheduling framework and local adaptive rules.

[0067] The candidate policy set is input into the two-layer policy network, and the output is a parallel policy set containing a global scheduling framework and local adaptive rules. The Transformer encoder (transformer model) performs global attention modeling on the spatiotemporal features of the annotated computational graph to generate a preliminary scheduling policy. The cellular automaton network performs local rule adjustment based on the states of adjacent nodes. The Transformer encoder and cellular automaton network are jointly optimized through a differentiable interface.

[0068] Specifically, the operator nodes of the annotated computation graph are arranged in topological order to generate a sequence. The input of each operator is a 32-dimensional joint feature vector. The operator position code is added to indicate the execution order in the computation graph. Eight attention heads are used, each focusing on a feature combination of different dimensions (e.g., head 1 focuses on computational load, head 2 focuses on communication cost). The attention weight is calculated and expressed as:

[0069] ;

[0070] in, represents the attention weight, Corresponding to query, key, and value matrices respectively, represents the dimension of the key vector (column vector of K), Indicates the similarity between the query (row of Q) and the key (row of K);

[0071] Identify critical paths (e.g., strong dependencies between computationally intensive operators) through attention weights;

[0072] It should be noted that the 32-dimensional feature vectors of operators in the computation graph are mapped into query Q, key K, and value V matrices. An 8-head attention mechanism then calculates feature association weights for each, generating a global dependency graph. After weight fusion (learnable parameter weighting), directed weighted edges are constructed to quantify the computational and communication dependencies across operators. Based on the global dependency graph, an improved Directed Flow (DFS) algorithm is used to search for the longest path. Weight calculation integrates dependency strength and operator computational effort, dynamically marking the path with the highest cumulative weight as the critical path. Based on training feedback, a weight penalty is imposed on computationally intensive operators in the path, prioritizing scheduling to areas with low communication attenuation. The critical path is visualized by overlaying hotspots (bandwidth density <100 Gbps / m³) on a 3D hardware topology field. If a path overlaps with a high attenuation area, cellular automaton rules are triggered (e.g., migrating Conv2d operators to the same rack GPU). Attention weights are then adjusted through backpropagation (e.g., reducing the operator dependency weights on high-attenuation paths). In the training of large models such as LLaMA-70B, it reduces critical path communication delay by 42% and pipeline idle time by 35%, systematically solving the problems of resource contention and efficiency loss caused by path selection bias in traditional solutions.

[0073] The output layer generates the mini-batch size for each pipeline stage (e.g. stage 1 mini-batch = 32, stage 2 = 16);

[0074] Based on the attention weights, select the GPU with the highest bandwidth density as the parameter server (such as GPU0 in rack A);

[0075] Define load balancing rules and memory recycling rules. Specifically, the load balancing rule is triggered when the node load is greater than 80% and the load of any neighboring node is less than 50%. The execution action is to migrate 15% of the current node's computing tasks to the lower-loaded neighbor. The memory recycling rule is triggered when the node's memory remaining is less than 10% (for example, less than 4GB remaining in a 40GB memory situation). The execution action is to trigger recomputation, trading time for space to reduce memory usage.

[0076] Load the 3D hardware topology field in a single-node simulator, bind the cell state (load and memory) to each GPU node, check the node state every simulation clock cycle, and immediately execute the corresponding action if the rule trigger condition is met;

[0077] After the cellular automaton executes the memory recycling rule, it records the memory change, converts the memory change into a gradient signal, and backpropagates it to the Transformer's attention weight matrix to reduce the scheduling weight of nodes with high memory consumption. The reward function is used to calculate the policy reward value. After each round of training, the 20% policies with the lowest reward value are eliminated. 40 new policies are generated through random perturbations to maintain a total of 200.

[0078] It should be noted that Transformer combines global scheduling with a cellular automation local rule base. Transformer uses multi-head attention to identify critical paths (such as matrix multiplication dependency chains). In LLaMA-70B training, this increased bandwidth density allocated to operators on these critical paths by 50% and peak throughput by 28%. Cellular rules also enable real-time workload migration (e.g., migrating 15% of tasks when node load exceeds 80%), reducing pipeline bubbles caused by resource contention and boosting average GPU utilization from 68% to 89%.

[0079] S5. Load the three-dimensional hardware topology field in the simulator, verify and adjust the parallel strategy set, and output the optimal strategy.

[0080] Verifying and adjusting the parallel strategy set involves screening potential high-quality strategies in a low-fidelity mode, and dynamically adjusting task allocation in a high-fidelity mode by simulating cross-node communication paths through ray tracing and combining it with a cellular automation rule base.

[0081] Assuming all communication paths use a uniform 100Gbps bandwidth (ignoring differences in physical topology), we calculated the single iteration time for each strategy. We directly read the memory usage field in the strategy and immediately eliminated any strategy larger than 40GB. We then output a pre-screened set of 40 strategies as high-fidelity verification input.

[0082] Based on the GPU binding information in the policy, the source node (such as GPU0 in rack A) and the target node (such as GPU4 in rack B) of the communication are determined. A ray is emitted from the source node, and a path is selected based on the bandwidth density distribution of the NeRF field. This includes an optimal path and a disaster recovery path. The optimal path is transmitted along the area with the highest bandwidth density (such as a direct path within a rack). The disaster recovery path automatically switches to a detour path (such as through a relay node) when the attenuation coefficient of the optimal path is greater than 0.5.

[0083] Calculate the total communication time, expressed as:

[0084] ;

[0085] in, Indicates the total communication time (unit: seconds), Indicates the total amount of data, Represents nodes that transmit segmented data in a communication path (such as routers and fiber segments in a network). The location of each node is mapped by the grid coordinates of the three-dimensional hardware topology field. Representation node The bandwidth density, Representation node Communication attenuation coefficient;

[0086] In the simulator, load balancing rules and memory recycling rules are bound to each GPU node. The node status is detected every 10ms simulation clock cycle, and the load balancing rules and memory recycling rules are triggered. When migrating tasks, the path attenuation coefficient of the target node is recalculated.

[0087] The high-fidelity single iteration time is calculated using time accumulation, and the strategy with the smallest single iteration time is output as the optimal strategy.

[0088] It should be noted that the low-fidelity coarse screening combined with high-fidelity ray tracing simulation. The low-fidelity model (assuming a fixed bandwidth) quickly eliminated over-memory strategies (taking 5 minutes, compared to 2 hours with traditional full-high-fidelity simulation), resulting in a 24-fold efficiency improvement. High-fidelity simulation combined with ray tracing selected the optimal path (for example, avoiding areas with attenuation > 0.5), resulting in a communication time estimation error of less than 3%, and a 72% reduction in the strategy deployment failure rate.

[0089] S6. Deploy the optimal strategy to the real cluster and collect training data, provide feedback to update the NeRF model, and trigger the mixed integer programming model to re-screen the strategy.

[0090] Training data includes communication efficiency, memory fluctuation data, and hardware abnormal events;

[0091] It should be noted that communication efficiency is measured by deploying distributed probe programs (such as iperf3 and NVIDIA Performance Analyzer) to perform cross-node bandwidth testing and record real-time bandwidth and latency (ms). NVIDIA NVLink / InfiniBand performance counters are used to obtain point-to-point communication throughput and calculate effective bandwidth utilization (actual bandwidth / theoretical bandwidth). Memory fluctuation data is obtained by integrating PyTorch memory profiling hooks to record the memory allocation and release timestamps of each operator, and to calculate peak memory usage and memory fragmentation. Hardware anomalies are captured using the Kubernetes event monitoring interface and GPU health check tools to capture events such as node failures, NVLink disconnections, and excessive temperatures. Communication efficiency data is used to compare simulated communication times with measured values. If the deviation exceeds 10%, the bandwidth density parameters of the 3D hardware field are updated (for example, the measured bandwidth attenuation coefficient is adjusted from 0.2 to 0.25) to make the model more realistic. Dynamically adjusting the Transformer's attention weights to reduce the scheduling priority of highly fluctuating paths (for example, paths with a standard deviation of inter-rack bandwidth greater than 15%) reduces communication conflicts by 38%. The impact of memory fluctuation data: Using memory peak and fragmentation data, we optimize the memory reclamation rules of the cellular automaton (for example, adjusting the checkpointing threshold from 10% to 8%), reducing memory overflow events by 72%. Combined with memory lifecycle data, we migrate long-lived tensors to low-fragmentation GPU nodes, increasing memory reuse by 45%. The impact of hardware anomaly events: Upon detecting a node failure, we trigger the cellular automaton to migrate tasks to a backup node (e.g., a rack-mounted spare GPU), synchronizing gradients via the NCCL consistency protocol, achieving failure recovery time of <2 seconds. By analyzing the frequency of anomaly events (e.g., >3 NVLink disconnections per hour), we automatically prioritize the execution of fault-tolerance rules, reducing the interruption rate by 94%.

[0092] Training data is fed into two feedback channels. The first channel feeds into the NeRF model, incrementally updating the parameters of the three-dimensional hardware topology field. The second channel feeds into a two-layer policy network. By comparing the deviation between simulated predictions and measured values, it adjusts the Transformer attention mechanism's reliance on hardware field features and evolves the local rule base of the cellular automaton.

[0093] Parse the list of gradient synchronization nodes in the optimal strategy, generate communication code, generate multi-stream parallelism based on the micro-batch partitioning scheme, execute the compiled code on the test cluster, and verify functional correctness (such as the consistency of global reduction operation results);

[0094] It should be noted that the optimal strategy extracts gradient synchronization node information (such as node IP, GPU UUID, and NCCL communication port). A globally unique communication group ID is generated based on node physical location sorting. The NCCL library (e.g., NCCL global reduce operations) is then called to automatically generate communication code. Based on the micro-batch partitioning scheme (e.g., micro-batches per pipeline stage = 32), an independent CUDA stream is allocated to each stage. CUDA stream wait events are inserted to achieve computation-communication synchronization. Eight CUDA stream wait events are created based on the number of pipeline stages (S = 8) and bound to the corresponding GPU nodes. Asynchronous data transfer between stages is achieved through the CUDA asynchronous memory copy function. A deployment manifest is generated using Kubernetes (a container orchestration system), and environment variables are injected (e.g., environment variables for binding the NCCL network interface). Distributed tasks are launched on the test cluster, and NCCL communication status is monitored. A random tensor (size 1024 × 1024) is initialized. After executing AllReduce, the results of each node are compared and the relative error is calculated. A failure is considered if the relative error is greater than 1e-6. A node failure is simulated to verify that the cellular automaton rule migrates the task to the backup node within 5 seconds and restores gradient synchronization.

[0095] Use Kubernetes deployment tools (such as the Kubernetes command-line tool) to distribute the compiled package to the target GPU node, start the distributed training task, and inject the local rule library (such as load balancing rules).

[0096] Monitor cluster node status. If a GPU node is added or removed (e.g., expanding from 8 to 16 nodes), trigger data collection. Perform bandwidth testing around the newly added node, collecting path latency and bandwidth data (sampling density: 10 times / second for 10 seconds). Update only the affected area (e.g., grid points within a 2-meter radius of the newly added node). Train for 50 rounds with a learning rate of 0.001.

[0097] Compare the simulated single iteration time with the measured single iteration time and calculate the relative deviation. If the deviation is greater than 10%, trigger a two-layer policy network update, reduce the attention weight of paths with large bandwidth density fluctuations (weight decay coefficient 0.1), increase the scheduling penalty coefficient for operators with high memory consumption (for example, increase the weight by 0.05), and if the node failure event occurs more than 3 times / hour, increase the priority of the fault migration rule to the highest level, and dynamically adjust the trigger threshold based on memory fluctuation data (for example, from 10% to 8%).

[0098] Real-time monitoring of hardware changes (such as node additions and deletions and NVLink upgrades) and performance metrics (the average single-iteration time for five consecutive iterations increases by >15% compared to the historical optimal value) triggers the optimization process and screens historical strategies from the candidate library for those with similar hardware configurations (such as the same number of nodes and same-generation GPU models). Fine-tune parameters within the neighborhood of the historical strategies (number of pipeline stages ±2 and number of data parallelism ±1), execute 50 iterations, and explore the optimal solution. If the improvement in single-iteration time is less than 1% for 10 consecutive iterations, terminate the process early to save resources.

[0099] It should be noted that real-time data feedback updates NeRF and the two-layer policy network. When hardware changes (such as node expansion) occur, the local topology field is updated through incremental training of the NeRF model (training only the affected area), and the policy reoptimization time is reduced. When the deviation between simulated and measured single iteration times exceeds 10%, policy network parameters are adjusted (e.g., by reducing the weight coefficient by 0.1). In scenarios with frequent node failures (>3 failures / hour), training interruptions are reduced by 94%.

[0100] This embodiment also provides a computer device suitable for a parallel strategy search method for efficient training of large artificial intelligence models, including: a memory and a processor; the memory is used to store computer-executable instructions, and the processor is used to execute computer-executable instructions, thereby implementing the parallel strategy search method for efficient training of large artificial intelligence models proposed in the above embodiment.

[0101] The computer device may be a terminal, comprising a processor, memory, a communication interface, a display, and an input device connected via a system bus. The processor provides computing and control capabilities. The memory includes non-volatile storage media and internal memory. The non-volatile storage media stores an operating system and computer programs. The internal memory provides an environment for the operating system and computer programs stored in the non-volatile storage media. The communication interface of the computer device is used to communicate with external terminals via wired or wireless communication. Wireless communication may be achieved via Wi-Fi, a carrier network, NFC (near-field communication), or other technologies. The display of the computer device may be a liquid crystal display or an electronic ink display. The input device may be a touchscreen overlay on the display, buttons, a trackball, or a touchpad on the computer device housing, or an external keyboard, touchpad, or mouse.

[0102] This embodiment also provides a storage medium having a computer program stored thereon, which, when executed by a processor, implements the parallel strategy search method for efficient training of large artificial intelligence models proposed in the above embodiment; the storage medium can be implemented by any type of volatile or non-volatile storage device or a combination thereof, such as static random access memory (SRAM), electrically erasable programmable read-only memory (EEPROM), erasable programmable read-only memory (EPROM), programmable read-only memory (PROM), read-only memory (ROM), magnetic memory, flash memory, magnetic disk or optical disk.

[0103] In summary, this invention uses Neural Radiant Field (NeRF) to model physical-level hardware communication fields, integrating mixed integer programming with deep reinforcement learning to achieve global-local collaborative optimization of parallel strategies. The NeRF model constructs a three-dimensional bandwidth density field, quantifies the cross-rack link attenuation effect, and significantly improves the accuracy of communication time prediction. In GPT-3-level model training, cluster TPI (time per iteration) is reduced. Mixed integer programming (MIQP) is used to jointly optimize pipeline stage division, data parallel sharding, and the compute-communication overlap factor, avoiding the local optimality trap caused by traditional hierarchical optimization. Field measurements show a significant reduction in memory fragmentation and improved GPU utilization. Based on a Transformer-cellular automaton architecture and combined with an online incremental learning mechanism, the algorithm reoptimizes strategies after hardware topology changes (such as node scaling), shortens fault recovery time, and reduces system throughput fluctuation compared to static strategies. This solution overcomes the limitations of traditional methods that only optimize a single performance metric, achieving a multi-objective balance between training efficiency, memory utilization, and system robustness.

[0104] It should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit the present invention. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solutions of the present invention may be modified or replaced by equivalents without departing from the spirit and scope of the technical solutions of the present invention, which should all be included in the scope of the claims of the present invention.

Claims

1. A parallel strategy search method for efficient training of large artificial intelligence models, characterized by: include, Input the physical topology parameters of the cluster, train the NeRF model with multi-view bandwidth and delay data, and generate a three-dimensional hardware topology field including bandwidth density and communication attenuation coefficient; Parse the computational graph of the large model to be trained, map the operator nodes to the three-dimensional hardware topology field, and label each operator node with a dynamic communication cost label based on the bandwidth density and communication attenuation coefficient of the operator node to generate a labeled computational graph; The mixed integer programming model based on the annotated computational graph is constructed by integrating computation time, communication time and constraints through the objective function, enumerating pipeline stage division and data parallel sharding strategies through the quantum annealing algorithm, and outputting a set of candidate strategies that meet the communication field constraints. The candidate policy set is input into the two-layer policy network, and the output is a parallel policy set containing a global scheduling framework and local adaptive rules. The Transformer encoder performs global attention modeling on the spatiotemporal features of the annotated computational graph to generate a preliminary scheduling policy. The cellular automaton network performs local rule adjustment based on the states of adjacent nodes. The Transformer encoder and the cellular automaton network are jointly optimized through a differentiable interface. Load the 3D hardware topology field into the simulator, verify and adjust the parallel strategy set, and output the optimal strategy; Deploy the optimal strategy to the real cluster and collect training data, provide feedback to update the NeRF model, and trigger the mixed integer programming model to re-screen the strategy; The training data is divided into two feedback paths. The first path inputs the NeRF model to incrementally update the three-dimensional hardware topology field parameters. The second path inputs the two-layer policy network. By comparing the deviation between the simulated prediction value and the measured value, the dependence weight of the Transformer attention mechanism on the hardware field characteristics is adjusted, and the local rule base of the cellular automaton is evolved.

2. The parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The physical topology parameters include rack layout, connection relationship between nodes and physical distance.

3. The parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The operator node includes the operator's computational load, graphics memory requirements, and dependencies.

4. The parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The verification and adjustment of the parallel strategy set refers to using the low-fidelity mode to screen potential high-quality strategies, and the high-fidelity mode to simulate cross-node communication paths through ray tracing, and dynamically adjust task allocation in combination with the cellular automation rule base.

5. The parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The training data includes communication efficiency, video memory fluctuation data and hardware abnormal events.

6. A computer device comprising a memory and a processor, wherein the memory stores a computer program, wherein: When the processor executes the computer program, the steps of the parallel strategy search method for efficient training of large artificial intelligence models described in any one of claims 1 to 5 are implemented.

7. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the parallel strategy search method for efficient training of large artificial intelligence models described in any one of claims 1 to 5 are implemented.

Citation Information

Patent Citations

  • Adaptive distributed parallel training method for neural network based on reinforcement learning

    CN113128702A

  • Self-adaptive re-calculation and load division method, device and equipment based on neural network, and computer readable medium

    CN119271397A

  • Reliability analysis method and system for power system

    CN119273170A