Distributed construction method of inter-atomic potential energy prediction system

Abstract

The invention provides a distributed construction method for an inter-atomic potential energy prediction system, and the method comprises the steps: S1, obtaining distributed equipment and parameters thereof, and configuring an initial inter-atomic potential energy prediction model on the distributed equipment; s2, acquiring a plurality of atomic graphs and potential energy information of each atomic graph; s3, dividing all the atomic graphs into a plurality of sample sets, wherein the total number of atoms of all the atomic graphs in each sample set is smaller than or equal to the atomic capacity of a preset sample set; s4, preprocessing each sample set, marking the sample set without the large-scale atomic graph as a centralized training sample set, and marking the sample set with the large-scale atomic graph as a distributed training sample set; and S5, scheduling all the sample sets to distributed equipment, and performing reasoning training on the initial inter-atomic potential energy prediction model by taking the atomic graph as an input sample and the potential energy information corresponding to the atomic graph as prediction output so as to obtain an inter-atomic potential energy prediction system.

Description

A Distributed Construction Method for Interatomic Potential Energy Prediction System Technical Field

[0001] This invention relates to the field of deep learning, specifically to an efficient training method for deep learning inter-atomic potential energy surface models, and more specifically, to a distributed construction method for an inter-atomic potential energy prediction system. Background Technology

[0002] Machine Learning Interatomic Potentials (MLIPs) have become an important tool for accelerating molecular dynamics and materials simulations in recent years. MLIPs typically employ a graph neural network (GNN)-based structure, using atoms as nodes and chemical bonds as edges to construct an atomic graph as input data. In this modeling approach, each molecule or crystal sample is represented as a graph, the size of which is directly related to the number of atoms in the system. MLIP models are generally composed of stacked interaction block layers, where each layer's atoms interact with their neighboring features.

[0003] In large-scale training scenarios, graph parallelism (GP) is often employed. This involves dividing the edges of the atomic graph across multiple computational units and performing forward and backward computations via all-to-all methods such as all-gathering across devices. Compared to data parallelism (DP) and pipeline parallelism (PP), graph parallelism can handle extremely large atomic graph data that cannot be contained in the memory of a single computational unit. Therefore, it is widely used in MLIP training tasks. However, MLIP training data typically exhibits a significant long-tail distribution: most samples have a small number of atoms, while a minority of samples have an extremely large number of atoms. Existing methods use a fixed number of atomic graph samples or a fixed number of atoms in the data batch, treating the data batch as the entire large atomic graph. During graph parallelism, the data is split according to the edges of the atomic graph, regardless of the topology of the atomic graph within the data batch.

[0004] In existing distributed training frameworks, parallelization of Graph Neural Networks (GNNs) mostly adopts static partitioning and a fixed number of data samples. This approach only fixes the number of graph samples during scheduling and performs static atomic graph partitioning without differentiating between large and small graphs in the long-tail distribution of real-world atomic graphs. In particular, in the training of Machine Learning Atomic Potential Models (MLIPs), the data naturally exhibits a long-tail distribution characteristic of "most small graphs and a few large graphs." This imbalance significantly weakens the effectiveness of existing methods. Furthermore, the static data batching strategy for atomic graph data leads to an imbalance in the load between computing units. During micro-batch partitioning, since the overall iteration speed in distributed training is usually determined by the slowest device, when a micro-batch happens to contain a large number of atomic graphs, its required computation and memory usage will be significantly higher than other micro-batches, resulting in a significant increase in the execution time of the corresponding computing unit. The resulting load imbalance directly leads to a decrease in the overall training throughput.

[0005] In summary, the current GNN distributed training framework uses a parallel strategy of "static partitioning + fixed number of samples," which only divides atomic graphs according to the number of samples and does not differentiate the long-tail distribution of "most small graphs and a few large graphs." In MLIPs training, this static packaging method can easily cause some micro-batches to be mixed with large graphs, and their computation and memory overhead is much higher than that of other batches, resulting in a sharp increase in the execution time of the corresponding device. The unbalanced load directly reduces the overall training throughput.

[0006] It should be noted that the background information presented here is only for illustrating relevant information about the present invention to aid in understanding the technical solutions of the present invention, and does not imply that the relevant information is necessarily prior art. In the absence of evidence indicating that the relevant information was disclosed before the filing date of this invention, the relevant information should not be considered prior art. Summary of the Invention

[0007] Therefore, the purpose of this invention is to overcome the shortcomings of the prior art and provide a distributed construction method for an interatomic potential energy prediction system.

[0008] The objective of this invention is achieved through the following technical solution:

[0009] According to a first aspect of the present invention, a distributed construction method for an inter-atomic potential energy prediction system is proposed. The inter-atomic potential energy prediction system is used to predict the potential energy information of an atomic graph based on an atomic graph, where the atomic graph is a graph with chemical atoms as nodes and chemical bonds as edges. The method includes: step S1, acquiring a distributed device and its parameters, and configuring an initial inter-atomic potential energy prediction model on the distributed device, wherein the parameters are the atomic capacity of each computing unit in the distributed device; step S2, acquiring multiple atomic graphs and the potential energy information of each atomic graph, wherein the potential energy information of the atomic graph includes the energy of the atomic graph and the force information of each atom; step S3, dividing all atomic graphs into multiple sample sets according to a preset partitioning method, wherein the total number of atoms in all atomic graphs in each sample set is less than or equal to the preset atomic capacity of the sample set; step S4... 4. Preprocess each sample set, marking the sample set without large-scale atomic graphs as the centralized training sample set, and the sample set with large-scale atomic graphs as the distributed training sample set. The large-scale atomic graph is defined as an atomic graph whose number of atoms exceeds a preset threshold based on the atomic capacity configuration of the computing unit. Step S5. Schedule all preprocessed sample sets to a distributed device and train the initial inter-atomic potential energy prediction model using atomic graphs as input samples and the potential energy information corresponding to the atomic graphs as the prediction output to obtain the inter-atomic potential energy prediction system. Specifically, when processing the centralized training sample set, the distributed device trains the inter-atomic potential energy prediction model using centralized training; when processing the distributed training sample set, the distributed device trains the inter-atomic potential energy prediction model using distributed training.

[0010] Optionally, the preset partitioning method is as follows: initialize a preset number of empty sample sets, wherein the preset number is the integer value obtained by dividing the total number of atoms of all acquired atomic graphs by the preset sample set atomic capacity and rounding up; arrange all atomic graphs in descending order of the number of atoms; and sequentially assign all the arranged atomic graphs to the sample set with the smallest current total number of atoms that has not reached the preset sample set atomic capacity.

[0011] Optionally, in step S4, the preset threshold is: ,in, This represents the atomic capacity of the computational unit, where T represents a preset parameter, and the value of T ranges from [value missing]. .

[0012] Optionally, between step S4 and step S5, the following step is also included: randomly arranging the order of the atomic graphs in each preprocessed sample set.

[0013] Optionally, in step S5, the training of the initial interatomic potential energy prediction model includes multiple training batches, wherein each training batch uses a portion of the preprocessed sample set to perform multiple rounds of iterative training on the interatomic potential energy prediction model.

[0014] Optionally, each training batch includes the following steps: Step S51, dividing the computing units in the distributed device into multiple computing groups on an average basis; Step S52, extracting multiple sample sets without replacement from the preprocessed full sample set, wherein the number of sample sets extracted is a positive integer multiple of the number of computing groups; Step S53, distributing all sample sets extracted in Step S52 equally to each computing group, wherein the computing units in each computing group select a centralized training method or a distributed training method to iteratively train the interatomic potential energy prediction model based on the label of the allocated sample set.

[0015] Optionally, in step S53: with the constraint of the balanced total number of atoms of all atomic graphs allocated to each computing unit, all atomic graphs allocated to each computing group are evenly distributed to each computing unit in the computing group. For large-scale atomic graphs, they are evenly divided into multiple sub-atomic graphs and distributed to computing units in units of sub-atomic graphs. For non-large-scale atomic graphs, they are distributed to computing units in units of atomic graphs.

[0016] Optionally, each computing group can choose between centralized training or distributed training to iteratively train the inter-atomic potential energy prediction model in the following ways: When the sample set allocated to the computing group is the centralized training sample set, the centralized training method is used to iteratively train the inter-atomic potential energy prediction model. The centralized training method is as follows: each atomic graph is processed independently by one computing unit, and no feature exchange is performed between computing units. When the sample set allocated to the computing group includes the distributed training sample set, the distributed training method is used to iteratively train the inter-atomic potential energy prediction model. The distributed training method is as follows: for sub-atomic graphs, computing units exchange features according to a preset communication method to collaboratively complete the calculation of the large-scale atomic graph to which the sub-atomic graph belongs; for non-sub-atomic graphs, each non-sub-atomic graph is processed independently by one computing unit.

[0017] According to a second aspect of the present invention, a material structure screening method is proposed for screening material structures that meet target performance indicators from a variety of material structures. The method includes: obtaining atomic diagrams of a variety of material structures and target performance indicators; constructing an interatomic potential energy prediction system using the method described in any one of claims 1-8 to predict the potential energy information of the atomic diagram of each material structure, and evaluating the performance indicators of the material structure based on the potential energy information of the atomic diagram of each material structure; and screening material structures that meet the target performance indicators based on the evaluated performance indicators of each material.

[0018] Compared with the prior art, the advantages of the present invention are as follows:

[0019] This invention optimizes graph-parallel training in long-tailed distribution scenarios. At the data parallelism and pipeline parallelism levels, sorting and packaging ensure load balancing; at the graph parallelism level, random shuffling avoids long-term imbalance; and at the communication level, the atomic graph topology-aware mechanism reduces redundant operations. These three aspects complement each other, significantly improving end-to-end training throughput, fully utilizing GPUs, and effectively alleviating cross-device communication pressure. Furthermore, this invention is not only applicable to MLIPs model training but can also be extended to other scientific computing and artificial intelligence tasks employing graph neural networks, providing an efficient and reliable load balancing and communication optimization solution for large-scale distributed training. Attached Figure Description

[0020] The embodiments of the present invention will be further described below with reference to the accompanying drawings, wherein:

[0021] Figure 1 is a schematic diagram of a distributed construction method for an interatomic potential energy prediction system according to an embodiment of the present invention;

[0022] Figure 2 is a schematic diagram of the process of dividing the sample set and configuring sample set labels according to an embodiment of the present invention. Detailed Implementation

[0023] To make the objectives, technical solutions, and advantages of this invention clearer, the invention is further described in detail below through specific embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention.

[0024] As mentioned in the background section, the current GNN distributed training framework uses a parallel strategy of "static partitioning + fixed number of samples", which only divides atomic graphs according to the number of samples and does not differentiate the long-tail distribution of "most small graphs and a few large graphs". In MLIPs training, this static packaging method is prone to mixing some micro-batch into large graphs. Their computation and memory overhead is much higher than that of other batches, resulting in a sharp increase in the execution time of the corresponding device. The unbalanced load directly reduces the overall training throughput.

[0025] To address the aforementioned issues, this invention proposes a scheme for constructing an inter-atomic potential energy prediction system. This scheme divides the atomic graphs in the training samples into multiple sample sets, constrained by atomic number balance and atomic graph size balance. This ensures that the total number of atoms in the atomic graphs contained in each sample set is consistent, and atomic graph samples of different sizes are evenly distributed among the various sample sets. Furthermore, this scheme combines the sample sets with the atomic capacity of the computing unit to design a topology-aware execution mode for atomic graphs. Based on the number of atoms in the atomic graphs in the sample sets, a centralized or distributed training method is selected to train the inter-atomic potential energy prediction model.

[0026] According to an embodiment of the present invention, a distributed construction method for an inter-atomic potential energy prediction system is proposed. Referring to Figure 1, the figure illustrates the steps of the method, which can be summarized as steps S1, S2, S3, S4, and S5. Step S1 involves acquiring a distributed device and its parameters, and configuring an initial inter-atomic potential energy prediction model on the distributed device, wherein the parameters are the atomic capacity of each computing unit in the distributed device; Step S2 involves acquiring multiple atomic graphs and the potential energy information of each atomic graph, the potential energy information of the atomic graph including the energy of the atomic graph and the force information of each atom; Step S3 involves dividing all atomic graphs into multiple sample sets according to a preset partitioning method, wherein the total number of atoms in all atomic graphs in each sample set is less than or equal to the preset atomic capacity of the sample set; Step S4... 4. Preprocess each sample set, marking the sample set without large-scale atomic graphs as the centralized training sample set, and the sample set with large-scale atomic graphs as the distributed training sample set. The large-scale atomic graph is defined as an atomic graph whose number of atoms exceeds a preset threshold based on the atomic capacity configuration of the computing unit. Step S5. Schedule all preprocessed sample sets to a distributed device and train the initial inter-atomic potential energy prediction model using atomic graphs as input samples and the potential energy information corresponding to the atomic graphs as the prediction output to obtain the inter-atomic potential energy prediction system. Specifically, when processing the centralized training sample set, the distributed device trains the inter-atomic potential energy prediction model using centralized training; when processing the distributed training sample set, the distributed device trains the inter-atomic potential energy prediction model using distributed training.

[0027] To better understand the present invention, each step will be described in detail below with reference to specific embodiments.

[0028] According to one embodiment of the present invention, in step S1, the distributed device is a device comprising multiple computing units used to execute pipelined training tasks. The atomic capacity of each computing unit is used to determine whether the atomic graph is a large-scale atomic graph that needs to be partitioned. In scenarios where atomic graphs of different sizes coexist, traditional fixed graph sample allocation methods easily lead to uneven load distribution among different data batches, making individual computing units performance bottlenecks. According to one embodiment of the present invention, all atomic graphs are first sorted in descending order according to the number of nodes or the number of atoms. After sorting, a greedy allocation strategy is used to partition the samples. Specifically, each time, the graph with the largest current size is allocated to the sample set with the smallest cumulative number of atoms. As the traversal proceeds, the large-scale graph is evenly distributed in different micro-batches (or sample sets), while the smaller graphs are used to fill the remaining capacity, making the number of atoms in each sample set as close as possible. Through this method, the size difference between sample sets is significantly reduced, the workload among computing units is more balanced, and the performance bottleneck caused by large-scale graph sets is avoided. For example, in a training task with eight sample sets, if random allocation is used directly, two sample sets may contain multiple large-scale atomic graphs, resulting in a computation time far exceeding the average level. However, by using the sorting and greedy packing strategy of this invention, large-scale atomic graphs are distributed among various sample sets, thereby effectively improving the real-time balance of the training process and laying the foundation for subsequent optimization steps.

[0029] According to an embodiment of the present invention, in step S3, the preset partitioning method is as follows: Initialize a preset number of empty sample sets, wherein the preset number is the integer value obtained by rounding up the quotient of the total number of atoms in all acquired atomic graphs divided by the preset atomic capacity of the sample set; sequentially traverse and acquire all atomic graphs in descending order of the number of atoms, wherein when the total number of atoms in a sample set is greater than the preset atomic capacity of the sample set, no more atomic graphs are added to that sample set, and each traversed atomic graph is added to the sample set with the smallest total number of atoms, until all atomic graphs have been traversed. The above partitioning method is a partitioning strategy based on a combination of atomic graph size sorting and a greedy algorithm. This strategy can distribute larger atomic graphs across different sample sets, while smaller atomic graphs are evenly distributed throughout the sample sets, thereby making the total number of atoms in each sample set close to the target average (i.e., the preset atomic capacity of the sample set), achieving an approximate balance of computational load. This can significantly reduce the load skew between computational units caused by the long-tail distribution of samples, providing a more balanced input data foundation for subsequent model training. Although dividing the atomic graph into multiple sample sets with a balanced number of atoms in the foregoing embodiments of the present invention can achieve good load balancing, the large-scale atomic graphs in the sample sets are all arranged at the beginning of the sample sets. However, during long-term training, this can lead to certain computing units bearing the computation of large-scale atomic graphs for an extended period. To avoid this "long-term imbalance," according to one embodiment of the present invention, a random shuffling mechanism is introduced to randomly shuffle the order of atomic graphs in each sample set. This prevents certain devices from continuously bearing more atomic graph samples during training, thereby maintaining load balancing for graph parallelism during long-term training.

[0030] According to one embodiment of the present invention, in order to reduce the communication volume between computing units, the present invention labels each data set according to the size of the atomic graph in each data set. The labeling is divided into two types: one is comm-free labeling, where the sample set with comm-free labeling adopts a centralized training method, that is, each atomic graph is calculated within a computing unit without communicating with other computing units; the other is dist labeling, where the sample set with dist labeling adopts a distributed training method, that is, the large-scale atomic graph is divided and distributed to multiple computing units, and the multiple computing units jointly complete the calculation. This method can effectively reduce redundant communication while ensuring the correctness of training, and the load of different devices is more balanced.

[0031] To facilitate understanding of the process of dividing the sample set and configuring labels for the sample set in this invention, the solution of this invention will be further explained below in conjunction with pseudocode and accompanying drawings.

[0032] According to one embodiment of the present invention, all acquired atomic graphs are divided into multiple sample sets, and the pseudocode for configuring labels for each sample set is shown below:

[0033] 1.Sort B in desending order by size(g)

[0034] 2.Initialize MB_list = {0, 0, …, 0} of length

[0035] 3. For i from 1 to M do:

[0036] 4. = argminLoad(MB_list)

[0037] 5. =

[0038] 6. for j from 1 to N_mb do:

[0039] 7. shuffle ( )

[0040] 8. =

[0041] 9. If

[0042] 10. Type(MB_j) = comm_free

[0043] 11. else

[0044] 12. Type(MB_j) = dist

[0045] 13. Return MB_list

[0046] The input to this pseudocode is a global data batch. Quantity of micro-batch Dimensionality of graph parallelism ,parameter The output is There are several sample sets, each labeled with either "comm-free" or "dist" tags. "Comm-free" tags represent centralized training labels, while "dist" tags represent distributed training labels. The pseudocode above will be explained line by line below. Lines 1-5 of the algorithm represent dividing all atomic graphs into... Each micro-batch, specifically, the first line sorts all samples in the input batch in descending order by the number of atoms; the second line initializes... Given a sample set, the third line iterates through the sorted atomic graphs; the fourth line calculates the cumulative atomic count of all current micro-batches and selects the micro-batch with the smallest cumulative atomic count. Line 5, the atomic diagram Add to This method ensures that large images are preferentially and evenly distributed across different micro-batches, while smaller images are used to fill empty spaces, thus maintaining a relatively balanced total size across all micro-batches. Lines 6-13 of the algorithm represent randomly shuffling the order of atomic images in each micro-batch and marking each micro-batch. Specifically, line 6 iterates through all micro-batches; line 7 shuffles the order of atomic images in each micro-batch; and line 8 determines a preset threshold parameter. It is configured according to user needs. For any number within the range, determine the atomic capacity of the computational unit. Preset threshold Lines 9-12: If the number of atoms in all atomic maps of a micro-batch is less than or equal to a preset threshold... If the value of a micro-batch exceeds the preset threshold, it is determined to be a "small graph batch" and does not require cross-device splitting. This micro-batch is marked as comm-free, i.e., a centralized training sample set. Otherwise, if any atomic graph exceeds the preset threshold, it is marked as distributed execution (dist), i.e., a distributed training sample set.

[0047] According to an example of the present invention, referring to Figure 2, it is assumed that the obtained atomic graph consists of 7 graphs, of which 4 are small graphs with 3 atoms, 2 are medium graphs with 4 atoms, and 1 is a large graph (large-scale atomic graph) with 7 atoms. It is necessary to divide the above atomic graphs into two sample sets (data batches). It should be understood that "small graph", "medium graph" and "large graph" are only descriptive labels set for the convenience of distinguishing relative scale. In real training samples, the number of atoms in an atomic graph can range from a few to thousands or even more, and is continuously distributed, exhibiting a typical long-tail pattern (i.e., most atomic graphs have fewer atoms, and a few atomic graphs have a very large number of atoms). Therefore, the actual partitioning and scheduling strategy should be based on dynamic load balancing based on the continuous number of atoms, rather than relying on manual threshold classification. Figure 2 illustrates the process of partitioning and labeling the sample sets, which includes three steps: First, sorting the atomic graphs from largest to smallest according to their size (number of atoms); second, allocating atomic graphs according to the minimum load principle, i.e., traversing the sorted atomic graphs and assigning each atomic graph to the sample set with the smallest total number of atoms, resulting in two initial sample sets, sample set 0 and sample set 1. Sample set 0 includes two atomic graphs with 3 atoms and two atomic graphs with 4 atoms, while sample set 1 includes two atomic graphs with 3 atoms and one atomic graph with 7 atoms; third, shuffling the data within each batch (sample set). (i.e., random rearrangement) and assign a label to each data batch. Specifically, assuming a preset threshold of 5, the order of atomic graphs in each sample set is randomly shuffled, and the sample set with atomic graphs having 5 or more atoms is configured as a distributed training sample set. This set needs to average the large-scale atomic graphs with more than the preset threshold and perform iterative training using a distributed training method. Conversely, the sample set with atomic graphs having less than 5 atoms is configured as a centralized training sample set. This sample set assigns each atomic graph to a computing unit, and each atomic graph is computed by a computing unit, thereby avoiding redundant communication between computing units.

[0048] According to an embodiment of the present invention, in step S5, the inference training includes multiple training batches, wherein each training batch uses a portion of the preprocessed sample set from all sample sets to perform multiple rounds of iterative training on the inter-atomic potential energy prediction model on a distributed device. Specifically, in each training batch, multiple rounds of iterative training are performed as follows: Step S51: Divide the computing units in the distributed device into multiple computing groups on an equal basis. It should be noted that the operation performed in step S51 can be pre-configured on the distributed device or configured during the model training phase; Step S52: Select a preset number of sample sets from all processed sample sets as training samples, wherein the preset number is a positive integer multiple of the number of computing groups; Step S53: Distribute all sample sets selected in step S52 equally to each computing group, and each computing group performs iterative training on the inter-atomic potential energy prediction model based on the assigned sample set.

[0049] It should be noted that communication overhead is particularly critical when training atomic potential energy models on distributed devices, especially when large-scale graphs are split across multiple devices. Boundary nodes and their gradients need to be frequently exchanged across devices via operations such as All-Gather or Reduce-Scatter, and the communication cost often exceeds the computational cost, becoming the main performance bottleneck for model training. To reduce unnecessary communication, this invention proposes a communication-aware execution mode partitioning mechanism. Different training modes are adopted for sample sets with different labels. Specifically, when the sample sets allocated to the computing group are all centralized training sample sets, the atomic graphs in the allocated sample sets are evenly distributed to each computing unit in a preset manner. Each computing unit then iteratively trains the inter-atomic potential energy prediction model based on the allocated atomic graphs. Since there are no cross-device connection edges between the atomic graphs held by each computing unit, cross-device feature communication is not required at the interaction layer, and the complete computation can be completed within a single device. When the sample set allocated to the computation group includes a distributed training sample set, each large-scale atomic graph is divided into multiple sub-atomic graphs according to the average edge. Using a preset method, all sub-atomic graphs and other atomic graphs (excluding the large-scale atomic graphs) in the sample set allocated to the computation group are evenly distributed to each computation unit. Each computation unit then iteratively trains the inter-atomic potential energy prediction model based on its allocated atomic graphs and sub-atomic graphs. When a computation unit trains using sub-atomic graphs as samples, the large-scale atomic graphs are distributed across different computation units. The computation unit needs to perform feature exchange with other computation units through an All-Gather operation to complete the iterative training of the large-scale atomic graph samples. Through this communication-aware execution mode partitioning mechanism, cross-device communication overhead can be significantly reduced while ensuring computational correctness, thereby improving the overall efficiency of model training.

[0050] According to one embodiment of the present invention, a large-scale atomic graph is divided into multiple sub-atomic graphs according to the average number of edges as follows: a hypergraph segmentation technique is used to segment the large-scale atomic graph into multiple sub-atomic graphs with the same number of edges. For example, it is assumed that the number of edges in the atomic graph is... Divide the edges in the atomic graph into d sub-atomic graphs, that is, each sub-atomic graph includes Each edge and the atoms connected to it can be used to divide a large-scale atomic graph into sub-atomic graphs, which can be distributed among multiple computational units to perform computations together. It should be noted that the number of sub-atomic graphs after the division of each large-scale atomic graph is consistent with the number of computational units in its computational group (i.e., the dimension of graph parallelism). Although dividing a large-scale atomic graph introduces feature communication, it can significantly balance the load of each computational unit, thereby achieving efficient and stable model training.

[0051] According to one embodiment of the present invention, in order to achieve load balancing among each computing unit, the present invention aims to balance the total number of atoms in all atomic graphs allocated to each computing unit, and distributes all atomic graphs allocated to each computing group equally to each computing unit in the computing group, wherein large-scale atomic graphs are allocated in units of sub-atomic graphs after being divided.

[0052] According to an embodiment of the present invention, the present invention proposes a material structure screening method based on the above-mentioned interatomic potential energy prediction system. The method includes: obtaining atomic maps of multiple material structures to be screened; using the interatomic potential energy prediction system to predict the interatomic potential energy of the atomic map of each material structure to be screened, and evaluating the performance of the material structure to be screened based on the interatomic potential energy of the atomic map of each material structure to be screened, so as to screen out material structures that meet the target performance.

[0053] To demonstrate the beneficial effects of this invention, the inventors conducted experiments on a multi-node GPU cluster. The experimental platform consisted of 32 NVIDIA A100 40GB GPUs, interconnected via InfiniBand at high speed. The software environment included CUDA 12.4, PyTorch 2.6, and the NCCL communication library. Regarding models, representative machine learning inter-atomic potential energy models (MLIPs) were selected, including UMA and eSEN. Both UMA and eSEN, based on the basic model structure, were expanded to 1.2B and 2.3B parameter scale versions by increasing the number of layers and adjusting the hidden width; eSEN was further developed into 100M and 220M parameter versions. In terms of experiments, a greedy packing strategy with a fixed number of atoms was employed, i.e., traversing all data and sequentially combining it into batches with a fixed number of atoms, along with a default graph parallelism strategy, i.e., a static graph partitioning method. Simultaneously, this invention also integrates data parallelism and pipeline parallelism mechanisms into the system. Compared with existing technologies, the proposed inter-atom potential energy prediction system construction method exhibits significant advantages in actual training. Firstly, when processing inter-atom potential energy model graph data with long-tailed distributions, this invention effectively alleviates the load differences between different sample sets, making the computation time of each stage of the pipeline more balanced, thereby significantly reducing pipeline cavitation caused by graph size differences. Experimental results show that when training the UMA-1.2B model on 32 GPUs, this invention improves the overall training throughput by an average of 1.11 times (approximately 10.8%). Secondly, by randomly shuffling graph data within micro-batches, this invention reduces the communication imbalance problem caused by ultra-large-scale graph partitioning. For example, on the UMA-1.2B model, when the pipeline parallelism dimension is 8 (corresponding to 8 GPUs) and the graph parallelism and data parallelism dimensions are 2, this invention reduces communication time by approximately 38%, significantly improving the stability and efficiency of cross-device communication. Furthermore, under other parallel configurations (pipeline parallelism dimension of 4, graph parallelism and data parallelism dimension of 1), this invention demonstrates improved load balancing of pipeline parallelism, reducing training time by approximately 11%. The method proposed in this invention reduces communication time and achieves greater load balancing among different devices while ensuring computational correctness, thus realizing efficient and stable training processes.

[0054] In summary, this invention optimizes graph-parallel training in long-tailed distribution scenarios. At the data parallelism and pipeline parallelism levels, sorting and packaging ensure load balancing; at the graph parallelism level, random shuffling avoids long-term imbalance; and at the communication level, the atomic graph topology-aware mechanism reduces redundant operations. These three aspects complement each other, significantly improving end-to-end training throughput, fully utilizing GPUs, and effectively alleviating cross-device communication pressure. Furthermore, this invention is not only applicable to MLIPs model training but can also be extended to other scientific computing and artificial intelligence tasks employing graph neural networks, providing an efficient and reliable load balancing and communication optimization solution for large-scale distributed training.

[0055] It should be noted that although the steps are described in a specific order above, it does not mean that the steps must be executed in the above specific order. In fact, some of these steps can be executed concurrently, or even in a different order, as long as the required function can be achieved.

[0056] This invention can be a system, method, and / or computer program product. A computer program product may include a computer-readable storage medium having computer-readable program instructions loaded thereon for causing a processor to implement various aspects of the invention.

[0057] Computer-readable storage media can be tangible devices that hold and store instructions for use by an instruction execution device. Computer-readable storage media can include, for example, but not limited to, electrical storage devices, magnetic storage devices, optical storage devices, electromagnetic storage devices, semiconductor storage devices, or any suitable combination thereof. More specific examples (a non-exhaustive list) of computer-readable storage media include: portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), static random access memory (SRAM), portable compact disc read-only memory (CD-ROM), digital multifunction disc (DVD), memory sticks, floppy disks, mechanical encoding devices, such as punch cards or recessed protrusions storing instructions thereon, and any suitable combination thereof.

[0058] The various embodiments of the present invention have been described above. These descriptions are exemplary and not exhaustive, nor are they limited to the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is chosen to best explain the principles, practical application, or technical improvements to the embodiments in the market, or to enable others skilled in the art to understand the embodiments disclosed herein.

Claims

1. A distributed construction method for an interatomic potential energy prediction system, wherein the interatomic potential energy prediction system is used to predict the potential energy information of an atomic graph based on an atomic graph, wherein the atomic graph is a graph with chemical atoms as nodes and chemical bonds as edges, characterized in that, The method includes: Step S1, acquiring a distributed device and its parameters, and configuring an initial inter-atomic potential energy prediction model on the distributed device, wherein the parameters are the atomic capacity of each computing unit in the distributed device; Step S2, acquiring multiple atomic graphs and the potential energy information of each atomic graph, wherein the potential energy information of the atomic graph includes the energy of the atomic graph and the force information of each atom; Step S3, dividing all atomic graphs into multiple sample sets according to a preset partitioning method, wherein the total number of atoms in all atomic graphs in each sample set is less than or equal to the preset atomic capacity of the sample set; Step S4, preprocessing each sample set, and marking the sample set without large-scale atomic graphs as a centralized training sample set. Step S5: Mark the sample set containing large-scale atomic graphs as the distributed training sample set, wherein the large-scale atomic graph is an atomic graph whose number of atoms exceeds a preset threshold based on the atomic capacity configuration of the computing unit; Step S6: Schedule all preprocessed sample sets to a distributed device and train the initial inter-atomic potential energy prediction model with atomic graphs as input samples and the potential energy information corresponding to the atomic graphs as prediction output to obtain the inter-atomic potential energy prediction system. When processing the centralized training sample set, the distributed device trains the inter-atomic potential energy prediction model in a centralized training manner, and when processing the distributed training sample set, the distributed device trains the inter-atomic potential energy prediction model in a distributed training manner.

2. The method according to claim 1, characterized in that, The preset partitioning method is as follows: initialize a preset number of empty sample sets, wherein the preset number is the integer value obtained by dividing the total number of atoms of all acquired atomic graphs by the preset sample set atomic capacity and rounding up; arrange all atomic graphs in descending order of the number of atoms; and sequentially assign all the arranged atomic graphs to the sample set with the smallest current total number of atoms that has not reached the preset sample set atomic capacity.

3. The method according to claim 1, characterized in that, In step S4, the preset threshold is , in, This represents the atomic capacity of the computational unit, where T represents a preset parameter, and the value of T ranges from [value missing]. 。 4. The method according to claim 3, characterized in that, Between step S4 and step S5, the following step is also included: in each preprocessed sample set, the order of the atomic graphs is randomly arranged.

5. The method according to claim 4, characterized in that, In step S5, the training of the initial interatomic potential energy prediction model includes multiple training batches, wherein each training batch uses a portion of the preprocessed sample set to perform multiple rounds of iterative training on the interatomic potential energy prediction model.

6. The method according to claim 5, characterized in that, Each training batch includes the following steps: Step S51, dividing the computing units in the distributed device into multiple computing groups on an average basis; Step S52, extracting multiple sample sets without replacement from the preprocessed full sample set, wherein the number of sample sets extracted is a positive integer multiple of the number of computing groups; Step S53, distributing all the sample sets extracted in Step S52 equally to each computing group, and the computing units in each computing group select either centralized training or distributed training based on the labels of the allocated sample sets to iteratively train the interatomic potential energy prediction model.

7. The method according to claim 6, characterized in that, In step S53: With the constraint of the balanced total number of atoms of all atomic graphs allocated to each computing unit, all atomic graphs allocated to each computing group are evenly distributed to each computing unit in the computing group. For large-scale atomic graphs, they are evenly divided into multiple sub-atomic graphs and distributed to computing units in units of sub-atomic graphs. For non-large-scale atomic graphs, they are distributed to computing units in units of atomic graphs.

8. The method according to claim 7, characterized in that, Each computing group selects either centralized training or distributed training to iteratively train the inter-atomic potential energy prediction model in the following manner: when the sample set allocated to the computing group is the centralized training sample set, the centralized training method is used to iteratively train the inter-atomic potential energy prediction model. The centralized training method is as follows: each atomic graph is processed independently by a computing unit, and no feature exchange is performed between computing units. When the sample set allocated to the computing group includes a distributed training sample set, the inter-atomic potential energy prediction model is iteratively trained using a distributed training method. The distributed training method is as follows: for sub-atomic graphs, computing units exchange features according to a preset communication method to collaboratively complete the calculation of the large-scale atomic graph to which the sub-atomic graph belongs; for non-sub-atomic graphs, each non-sub-atomic graph is processed independently by a computing unit.

9. A material structure screening method for selecting material structures that meet target performance indicators from a variety of material structures, characterized in that, The method includes: acquiring atomic maps of various material structures and target performance indicators; constructing an interatomic potential energy prediction system using the method described in any one of claims 1-8 to predict the potential energy information of the atomic map of each material structure, and evaluating the performance indicators of the material structure based on the potential energy information of the atomic map of each material structure; and selecting material structures that meet the target performance indicators based on the evaluated performance indicators of each material.

10. An electronic device, characterized in that, include: One or more processors; And a memory, wherein the memory is used to store executable instructions; the one or more processors are configured to implement the steps of the method of any one of claims 1 to 9 by executing the executable instructions.

Images