Communication method based on resource awareness and asynchronous parallel

Abstract

The invention discloses a communication method based on resource awareness and asynchronous parallel. The method comprises the following steps: calculating a key resource consumption dimension in a big language model training process; capturing dependency relationships of all layers in a model operation process to form a block-level data dependency graph; loads are distributed to all the nodes in a balanced mode; and a non-interlaced 1F1B scheduling strategy is adopted to realize parallel execution of data transmission and model calculation. According to the embodiment of the invention, by monitoring the computing power and the memory capacity of the CPU nodes and the computing complexity and the memory demand of each layer of the neural network in real time, the model is reasonably allocated to different nodes according to the calculated amount and the memory consumption by adopting the adaptive segmentation algorithm. The limitation of traditional static segmentation is broken through, and the problem of load imbalance in a heterogeneous environment is effectively solved; through an asynchronous queue and non-blocking communication technology and fusion of a double-buffer and non-staggered 1F1B scheduling strategy, data transmission and model calculation can be carried out in parallel, the defect of idle resources is effectively overcome, and communication overhead is greatly reduced.

Description

Technical Field

[0001] This invention relates to the field of communication technology, and in particular to a communication method based on resource awareness and asynchronous parallelism. Background Technology

[0002] In recent years, deep learning has continued to develop, and language models have demonstrated superior performance in tasks such as semantic understanding and text generation. This technological breakthrough has enabled language models to be deeply integrated into the field of Artificial Intelligence of Things (AIoT), showing great application potential in scenarios such as smart home voice interaction and industrial equipment fault diagnosis. AIoT integrates artificial intelligence technology into IoT systems, enabling devices not only to collect and transmit data but also to perform real-time data analysis and autonomous decision-making locally or in the cloud, thereby achieving more efficient and intelligent application scenarios. This integration breaks through the functional limitations of traditional IoT devices and is a necessary technology for the rapid modernization of today's society.

[0003] The objective of this application is to construct a distributed system for efficient training of large language models, addressing the challenges of limited CPU computing power and memory resources in the AI ​​Internet. This system utilizes advanced resource-aware technologies and parallel model training techniques to efficiently leverage edge device resources within the AI ​​Internet and reduce the training overhead of large language models.

[0004] However, applying Large Language Models (LLMs) to the field of AIoT faces several challenges: First, due to the heterogeneity of computing power in embedded CPUs (such as the differences in computing capabilities between different ARM chip models) and the uneven distribution of model computing tasks in AIoT, some CPU nodes may become "bottleneck nodes" due to undertaking too many computing tasks. Second, the large number of training parameters of large language models contrasts sharply with the actual effective communication bandwidth between AIoT devices. Unlike the hundreds of GB / s transmission rate of NVLink in GPU clusters, the communication bandwidth between AIoT devices is typically only tens to hundreds of Mbps.

[0005] Therefore, it is necessary to improve one or more of the problems existing in the above-mentioned related technical solutions.

[0006] It should be noted that this section is intended to provide background or context for the technical solutions of this disclosure as set forth in the claims. The description herein does not constitute an admission that it is prior art simply because it is included in this section. Summary of the Invention

[0007] The purpose of this invention is to provide a communication method based on resource awareness and asynchronous parallelism, thereby overcoming, at least to some extent, one or more problems caused by the limitations and defects of related technologies.

[0008] This invention first provides a communication method based on resource awareness and asynchronous parallelism, comprising: S1, calculate the key resource consumption dimensions in the training process of the large language model, including: computational complexity, parameter memory and activation memory; S2 treats the attention module as a whole unit, initializes the dependency dictionary, the preceding layer dictionary and the Transformer block dictionary, captures the dependencies of all layers during the model operation, performs sub-layer classification processing, upgrades hierarchical dependencies to block-level dependencies, and forms a block-level data dependency graph. S3. Based on the data dependency graph, modules with close data dependencies are initially assigned to the same CPU node. The number of floating-point operations of all layers on each CPU node is accumulated to obtain the computational load. The memory load is obtained by accumulating the parameter memory and active memory usage. Nodes with excessively high or low loads are identified. The load is iteratively calculated and adjusted until the load of each node is balanced. S4 divides the model training process into multiple independent stages. Each stage uses micro-batch data as the processing unit and adopts a non-interleaved 1F1B scheduling strategy to ensure that each stage only processes the task corresponding to the current micro-batch. It controls the switching of the dual buffers. When the model uses the data in the current buffer for calculation, the next batch of data is asynchronously loaded into the other buffer, so as to realize the parallel execution of data transmission and model calculation.

[0009] In this invention, in S1, computational complexity is measured by calculating the number of floating-point operations (FLOPs) for each neural network layer. The formulas for calculating FLOPs for different neural network layers are as follows:

[0010]

[0011]

[0012] in, Indicates the number of floating-point operations in the linear layer. Indicates the number of floating-point operations performed in the convolutional layer. Indicates the number of floating-point operations performed in the embedding layer. input_size Indicates the shape of the input data. output_size Indicates the shape of the output data. kernel_size Indicates the shape of the convolution kernel. in_channels Indicates the input channel. out_channels Indicates the output channel. output_length Indicates the output length. num_embed Indicates vocabulary size, embed_dim This indicates the embedding vocabulary dimension.

[0013] In this invention, in S1, the parameter memory of each neural network layer... The calculation formula is as follows:

[0014] in, Pi_num This represents the number of elements in each parameter tensor. Pi_type Indicates the format of the stored parameters.

[0015] In this invention, step S1, activating the memory computation process, includes: For fully connected layers and convolutional layers, the output shape is first calculated based on the input shape and layer parameters, and then the activation memory usage data is obtained by multiplying the total number of elements in the output shape by 4. For applications that do not generate additional significant activation memory usage, set the activation memory usage data to 0.

[0016] In this invention, in S2, the PyTorch hook function register_hook(module, name) is used to insert custom logic before / after execution in the model layer, thereby achieving dynamic capture of dependencies.

[0017] In this invention, the sub-layer classification process in S2 includes: Initialize the Transformer block dictionary, where the key is the block number in the Transformer model and the value is the set of sub-layer names contained in that block; The scattered sub-layers are categorized according to their respective Transformer block numbers; Store the categorized sub-layer names into the dictionary value of the corresponding block number to complete the sub-layer categorization.

[0018] In this invention, in step S3, multiple consecutive attention modules are allocated as a whole to the same CPU node.

[0019] The present invention further provides a communication system based on resource awareness and asynchronous parallelism, comprising: The model computation module is used to calculate the key resource consumption dimensions during the training process of a large language model. The key resource consumption dimensions include: computational complexity, parameter memory, and activation memory. The model analysis module treats the attention module as a whole unit, initializes the dependency dictionary, the preceding layer dictionary, and the Transformer block dictionary, captures the dependencies of all layers during model operation, performs sub-layer classification processing, upgrades hierarchical dependencies to block-level dependencies, and forms a block-level data dependency graph. The model segmentation module is used to initially allocate modules with close data dependencies to the same CPU node based on the data dependency graph. It accumulates the number of floating-point operations of all layers on each CPU node to obtain the computational load, accumulates the parameter memory and activation memory usage to obtain the memory load, identifies nodes with excessively high or low loads, iteratively calculates and adjusts the load until the load of each node is balanced. The asynchronous pipeline parallel module is used to divide the model training process into multiple independent stages. Each stage uses micro-batch data as the processing unit and adopts a non-interleaved 1F1B scheduling strategy to ensure that each stage only processes the task corresponding to the current micro-batch. It controls the switching of the dual buffers. When the model uses the data in the current buffer for calculation, the next batch of data is asynchronously loaded into the other buffer, realizing the parallel execution of data transmission and model calculation.

[0020] The technical solution provided by this invention may include the following beneficial effects: This invention presents a resource-aware and asynchronous parallel communication method, employing a resource-aware model segmentation approach: a multi-dimensional resource matching model is constructed. By real-time monitoring of the CPU node's computing power, memory capacity, and the computational complexity and memory requirements of each layer of the neural network, an adaptive segmentation algorithm is used to rationally allocate the model to different nodes based on computational load and memory consumption. This method overcomes the limitations of traditional static segmentation and effectively solves the load imbalance problem in heterogeneous environments. An asynchronous pipelined parallel strategy is adopted: To address the issues of long training times and high communication overhead, an asynchronous non-interleaved pipeline strategy is proposed. This method utilizes asynchronous queues and non-blocking communication techniques, along with a fusion of double buffering and a non-interleaved 1F1B scheduling strategy, to enable parallel data transmission and model computation. Compared to traditional synchronous communication, this method effectively overcomes the drawback of idle computing resources during data transmission and significantly reduces communication overhead. This has significant value for improving the training performance of AIoT devices. Attached Figure Description

[0021] The accompanying drawings, which are incorporated in and form part of this specification, illustrate embodiments consistent with this disclosure and, together with the description, serve to explain the principles of this disclosure. It is obvious that the drawings described below are merely some embodiments of this disclosure, and those skilled in the art can obtain other drawings based on these drawings without any inventive effort.

[0022] Figure 1 A flowchart illustrating a resource-aware and asynchronous parallel communication method in an exemplary embodiment of this disclosure is shown. Figure 2 This diagram illustrates the segmentation flowchart of a resource-aware load balancing model in an exemplary embodiment of this disclosure. Figure 3This diagram illustrates a non-interleaved pipeline in an exemplary embodiment of this disclosure. Figure 4 A flowchart illustrating the asynchronous pipelined communication optimization scheme in an exemplary embodiment of this disclosure is shown. Figure 5 A comparative diagram of segmentation methods in exemplary embodiments of this disclosure is shown; Figure 6 A comparison diagram of different parallel methods in exemplary embodiments of this disclosure is shown; Figure 7 The diagram shows the ablation experiment results in an exemplary embodiment of this disclosure; Figure 8 A structural diagram of a resource-aware and asynchronous parallel communication system according to an exemplary embodiment of this disclosure is shown. Figure 9 This diagram illustrates the overall architecture of a resource-aware and asynchronous parallel communication system in an exemplary embodiment of this disclosure. Detailed Implementation

[0023] Exemplary embodiments will now be described more fully with reference to the accompanying drawings. However, these exemplary embodiments can be implemented in many forms and should not be construed as limited to the examples set forth herein; rather, they are provided so that this disclosure will be more comprehensive and complete, and will fully convey the concept of the exemplary embodiments to those skilled in the art. The described features, structures, or characteristics may be combined in any suitable manner in one or more embodiments.

[0024] Furthermore, the accompanying drawings are merely illustrative diagrams of embodiments of this disclosure and are not necessarily drawn to scale. The same reference numerals in the drawings denote the same or similar parts, and therefore repeated descriptions of them will be omitted. Some block diagrams shown in the drawings are functional entities and do not necessarily correspond to physically or logically independent entities.

[0025] This invention aims to construct an optimized solution for efficient AI model operation in CPU cluster environments by combining innovative resource scheduling and parallel computing technologies with the resource requirements of AI model execution. This provides high resource utilization and short communication latency support for AI model deployment in clusters. The solution design considers resource adaptability, computational parallelism, and communication efficiency to meet the high-efficiency operation requirements of AI models in large-scale data processing scenarios. In designing the optimized AI model operation solution for CPU clusters, this invention needs to consider the uneven distribution of resources among nodes in the CPU cluster and the collaborative relationship between model computation and data transmission. The core objective of this solution is to improve the resource utilization efficiency of the CPU cluster through optimized model partitioning and parallel scheduling strategies, while reducing communication latency during model execution, thereby improving the overall performance of AI models in cluster environments.

[0026] The method of this invention mainly includes two aspects: First, it creates a load balancing model segmentation method specifically designed for resource awareness. This method can intelligently divide the AI ​​model into multiple sub-modules based on the computing, storage, and network resources of each node in the CPU cluster. Each sub-module is assigned to the node most suitable for its processing, thus fully utilizing the resources of each node and avoiding excessive concentration and idleness of resources, laying a resource foundation for efficient model operation. Second, it proposes an asynchronous pipelined parallel strategy method. This method achieves parallel execution of data transmission and model computation through the integration of non-blocking communication technology and a dual-buffer and non-interleaved 1F1B scheduling strategy. Data flows sequentially through each stage like a pipeline, and each stage uses asynchronous queues and non-blocking communication to allow computation and communication operations to be performed simultaneously, maximizing the overlap of computation and communication time and effectively reducing the impact of communication waiting on model running efficiency. The comprehensive application of these methods will provide key technical support for the efficient operation of AI models in CPU clusters. This example implementation first provides a communication method based on resource awareness and asynchronous parallelism. Please refer to [link / reference]. Figure 1 This method may include: S1-S4, as follows: S1, calculate the key resource consumption dimensions in the training process of the large language model, including: computational complexity, parameter memory and activation memory; S2 treats the attention module as a whole unit, initializes the dependency dictionary, the preceding layer dictionary and the Transformer block dictionary, captures the dependencies of all layers during the model operation, performs sub-layer classification processing, upgrades hierarchical dependencies to block-level dependencies, and forms a block-level data dependency graph. S3. Based on the data dependency graph, modules with close data dependencies are initially assigned to the same CPU node. The number of floating-point operations of all layers on each CPU node is accumulated to obtain the computational load. The memory load is obtained by accumulating the parameter memory and active memory usage. Nodes with excessively high or low loads are identified. The load is iteratively calculated and adjusted until the load of each node is balanced. S4 divides the model training process into multiple independent stages. Each stage uses micro-batch data as the processing unit and adopts a non-interleaved 1F1B scheduling strategy to ensure that each stage only processes the task corresponding to the current micro-batch. It controls the switching of the dual buffers. When the model uses the data in the current buffer for calculation, the next batch of data is asynchronously loaded into the other buffer, so as to realize the parallel execution of data transmission and model calculation.

[0027] In this embodiment, a resource-aware model segmentation method is adopted: a multi-dimensional resource matching model is constructed. By real-time monitoring of the computing power and memory capacity of CPU nodes, as well as the computational complexity and memory requirements of each layer of the neural network, an adaptive segmentation algorithm is used to rationally allocate the model to different nodes according to computational load and memory consumption. This method breaks through the limitations of traditional static segmentation and effectively solves the problem of load imbalance in heterogeneous environments. An asynchronous pipelined parallel strategy is adopted: To address the issues of long training times and high communication overhead, an asynchronous non-interleaved pipeline strategy is proposed. This method utilizes asynchronous queues and non-blocking communication techniques, along with a fusion of double buffering and a non-interleaved 1F1B scheduling strategy, to enable parallel data transmission and model computation. Compared to traditional synchronous communication, this method effectively overcomes the drawback of idle computing resources during data transmission and significantly reduces communication overhead. This has significant value for improving the training performance of AIoT devices.

[0028] The specific process of each step in the above embodiments will be described below.

[0029] S1, to achieve efficient adaptation and operational optimization of AI models on CPU clusters, innovatively proposes a model analysis scheme of "resource awareness + data dependency graph based on attention modules". This scheme provides data support for subsequent load balancing model segmentation and asynchronous pipeline parallel strategies by accurately calculating key resource consumption dimensions during model training. The core focuses on three dimensions: computational complexity, parameter memory, and activation memory. The specific process is as follows: (1) Computational complexity is the core indicator for evaluating the computational pressure of each layer of the model, and it directly affects the rationality of sub-module allocation. This invention quantifies computational complexity by monitoring the FLOPs (floating-point operations) of each neural network layer during model training in real time. Furthermore, it designs differentiated FLOPs calculation formulas for different types of neural network layers (such as fully connected layers, convolutional layers, pooling layers, etc.). For example, convolutional layers need to consider the impact of parameters such as the number of input channels, kernel size, and stride on the computational load, while fully connected layers need to combine the number of input features and the number of output features to calculate the computational load. This ensures the accuracy of computational complexity evaluation for different layers and provides a basis for matching "computational load - node resources" when dividing subsequent sub-modules. The formula is as follows:

[0030]

[0031]

[0032] in, Indicates the number of floating-point operations in the linear layer. Indicates the number of floating-point operations performed in the convolutional layer. Indicates the number of floating-point operations performed in the embedding layer. input_size Indicates the shape of the input data. output_size Indicates the shape of the output data. kernel_size Indicates the shape of the convolution kernel. in_channels Indicates the input channel. out_channels Indicates the output channel. output_length Indicates the output length. num_embed Indicates vocabulary size, embed_dim This indicates the embedding vocabulary dimension.

[0033] (2) Parameter memory usage is a key consideration for CPU cluster node storage resource adaptation. This invention clarifies the calculation logic of parameter memory usage for each layer of the neural network through a formula, as follows:

[0034] in, Pi_num This represents the number of elements in each parameter tensor. Pi_type Indicates the format of the stored parameters.

[0035] Core Computational Dimension Definition: First, clarify the core elements of parameter calculation, including... Pi (The learnable parameters corresponding to the neural network layers, such as convolution kernel parameters, weight matrices of fully connected layers, etc.) Pi_num (each) Pi (Total number of elements in the corresponding parameter tensor) Pi_type (The storage format of the parameters determines the number of bytes of memory occupied by a single element.) Storage format and byte count mapping: Taking the mainstream PyTorch framework as an example, parameters are stored as 32-bit floating-point numbers (float32) by default, in which case Pi_type is 4 (i.e., a single parameter element occupies 4 bytes of memory); if other formats are used (such as 16-bit floating-point numbers (float16), then... Pi_type The corresponding value is adjusted to 2 to ensure a precise correspondence between the storage format and the number of bytes. Single-layer parameter memory computation: The memory usage is calculated using the formula "total number of parameter tensor elements × number of bytes occupied by a single element", i.e., parameter memory usage = Pi_num × Pi_type The memory consumption of each layer of the neural network is calculated, thereby understanding the storage resource requirements of each layer for nodes.

[0036] (3) Activation memory is the memory temporarily occupied by data during model operation, and its size is directly related to the layer type and the shape of the input and output. This invention designs differentiated activation memory calculation rules for different types of neural network layers. The core rule of activation memory calculation generally follows the principle of "differentiating layer type calculations, with non-critical layers occupying zero memory by default": For fully connected layers (nn.Linear) and one-dimensional convolutional layers (nn.Conv1d): the output shape needs to be calculated based on the input shape and layer parameters, and then the activation memory usage is obtained by multiplying the total number of elements in the output shape by 4 (the number of bytes occupied by a single element of type float32). For other layers (such as pooling layer nn.AvgPool1d, activation function layer nn.ReLU, etc.): Since they do not generate additional significant activation memory consumption, their activation memory usage is set to 0.

[0037] S2, Attention Module-Based Data Dependency Graph: To simplify the construction process of AI model data dependency graphs and reduce the operational complexity of subsequent model segmentation, this invention innovatively treats the attention module (common in Transformer models) as a whole unit from the perspective of computational and memory co-optimization, rather than breaking it down into its internal sub-layers. The core advantage of this design is that when constructing the data dependency graph, only the attention module needs to be regarded as a single node, focusing on the interaction relationship of "module input (original data or previous module output) - module output (processed feature vector)," ignoring the detailed dependencies of each sub-layer within the module (such as multi-head attention sub-layers, layer normalization sub-layers, etc.), thereby significantly reducing the number of nodes and edge complexity of the data dependency graph, providing a clear and efficient structural foundation for subsequent model segmentation based on the dependency graph.

[0038] First, by initializing three types of core dictionaries, a "container" for storing dependency relationships is built to ensure the orderly capture and organization of subsequent dependency data, as detailed below: (1) Dependency dictionary depend: Uses a default dictionary (values ​​are set types), where the key is the name of the current neural network layer (e.g., "transformer.layer.0.attention"), and the value is a set of names of the preceding layers that this layer directly depends on. The choice of set type can automatically avoid storing duplicate dependencies and ensure the uniqueness of the dependency relationship.

[0039] (2) Prev_layers dictionary: Initialized as an empty dictionary, the key is the output tensor of a certain layer (torch.Tensor type), and the value is the set of names of the current layer that produced the output tensor (usually a single layer, so it is represented by the set {name}). Its core function is to "reverse associate the output tensor with the layer that produced the tensor", providing a "tensor-layer" mapping bridge for subsequent capture of dependencies. (3) Transformer block dictionary transformer_blocks: This uses a default dictionary (values ​​are set types). The keys are the block numbers in the Transformer model (e.g., 0, 1, 2, corresponding to the 1st, 2nd, and 3rd encoder / decoder blocks of the Transformer), and the values ​​are the set of names of all sub-layers contained in that block (e.g., "transformer.layer.0.attention", "transformer.layer.0.linear", etc.). This dictionary is used to categorize scattered sub-layers by "blocks", providing a foundation for subsequent integration of "block-level dependencies".

[0040] To automatically and in real-time capture the input-output relationships (i.e. dependencies) of each layer during model execution, this application defines the register_hook(module, name) function. This function uses PyTorch's hook mechanism to insert custom logic before / after the execution of the model layer, thereby achieving dynamic capture of dependencies.

[0041] After completing the dependency capture of all layers (the depend dictionary stores hierarchical dependencies) and sub-layer classification (the transformer_blocks dictionary stores block-sub-layer mappings), the "hierarchical dependency" is further upgraded to "block dependency", which is in line with the design idea of ​​"holistic attention module".

[0042] The sub-layer classification process includes: (1) Initialize the Transformer block dictionary, with the key being the block number in the Transformer model and the value being the set of sub-layer names contained in the block.

[0043] (2) Classify the scattered sub-layers according to their respective Transformer block numbers.

[0044] (3) Store the classified sub-layer names into the dictionary value of the corresponding block number to complete the sub-layer classification.

[0045] S3, please refer to Figure 2 Model partitioning based on load balancing: This step follows the logic of "initial partitioning foundation - load calculation and evaluation - dynamic transfer optimization - iterative balancing and finalization". It uses the data dependency graph as the structural guide and CPU cluster load balancing as the goal. Through multiple rounds of dynamic adjustment, it achieves efficient deployment of the model on the cluster nodes.

[0046] First, the constructed data dependency graph clearly shows the flow of data in different parts of the model and the dependencies in the computation process. When determining the initial model partitioning scheme, guided by the data dependency graph, modules with close data dependencies are assigned to the same CPU node. This ensures efficient data processing within the node and reduces data transfer overhead across nodes. For example, for the attention module, due to its complex global and cyclic dependencies, allocating multiple consecutive attention modules as a whole to the same CPU avoids the impact of frequent cross-node communication on training efficiency. This initial partitioning based on the data dependency graph lays the foundation for subsequent load balancing optimization, ensuring the consistency and rationality of the training process after model partitioning.

[0047] Initial segmentation is the fundamental step in model segmentation. Its core objective is to "divide modules based on data dependencies to reduce cross-node communication costs." The specific process is as follows: Based on the data dependency graph constructed above, which uses Transformer blocks (including attention modules) as core nodes, we can clearly understand the flow of data between different modules of the model (e.g., "block 0 output → block 1 input") and computational dependencies (e.g., "block 2 depends on the outputs of blocks 0 and 1"), providing a basis for "dependency constraints" for module partitioning. The core principle of module partitioning is to follow the rule of "closely dependent modules on the same node," allocating modules with close data dependencies (e.g., consecutive Transformer blocks, layer combinations with direct data flow) to the same CPU node. This design ensures that data flows directly within the node, avoiding latency and bandwidth consumption caused by cross-node data transmission. Special handling for attention modules: For global dependencies (e.g., key-value pair interactions in multi-head attention) and cyclic dependencies (e.g., inter-layer residual connections) within attention modules, splitting them to different nodes would lead to frequent cross-node communication. Therefore, multiple consecutive attention modules (usually contained in the same or adjacent Transformer blocks) are allocated as a whole to the same CPU node, reducing the impact of communication on training efficiency from the source. The value of initial segmentation: Initial segmentation constrained by data dependency graphs not only ensures the continuity of data flow and computational dependencies during model training, but also provides a basic solution with "reasonable structure and low communication overhead" for subsequent load balancing optimization, avoiding excessive adjustment costs caused by chaotic initial structure in subsequent optimization.

[0048] Secondly, based on the calculated FLOPs, parameter memory usage, and activation memory usage for each layer, and in conjunction with the current model segmentation scheme, the computational load and memory load of each CPU are calculated. The computational load is obtained by summing the FLOPS of each neural network layer allocated to that CPU, as shown in the following formula, reflecting the CPU's pressure on the computational task. The memory load is determined by the sum of the parameter memory usage and activation memory usage of each layer, as shown in the following formula, reflecting the CPU's memory resource usage. By calculating the load of each CPU, the resource usage status of each node can be clearly understood, identifying CPU nodes with excessively high or low loads, providing a basis for subsequent model re-segmentation.

[0049] After the initial partitioning is completed, the load of each CPU node needs to be quantitatively calculated to identify nodes that are "overloaded and need to be offloaded" and "underloaded and can be taken over," providing optimization targets for dynamic migration. The specific calculation logic and dimensions are as follows: (1) Core calculation dimensions and formulas Load calculation revolves around two core dimensions: "computational load" and "memory load," both of which are obtained by accumulating the corresponding consumption of all neural network layers on the node.

[0050] Computational load calculation: Based on the FLOPs (floating-point operations) of each layer, the FLOPs of all layers allocated to the CPU node are accumulated. The formula can be expressed as: Computational load = Σ (FLOPs of all layers on a certain CPU node). This indicator directly reflects the computational task pressure of the CPU node in model training. The higher the total FLOPs, the greater the computational load.

[0051] Memory load calculation: Taking into account both parameter memory usage and activation memory usage of each layer, the two types of memory consumption of all layers on a node are summed. The formula can be expressed as: Memory load = Σ (Parameter memory usage of all layers on a CPU node + Activation memory usage of that layer). Among them, parameter memory usage is calculated by "total number of parameter tensor elements × number of bytes of storage format" (as detailed above), and activation memory usage is calculated by layer type and input / output shape. The sum of the two reflects the intensity of memory resource usage of the node.

[0052]

[0053] in, part_flots This indicates the number of floating-point operations for a specific part of the model, and `part_mem` indicates the memory usage for that part of the model. Flotsi Indicates the first i The number of floating-point operations per layer flots_all This represents the sum of floating-point operations for all levels. parami This represents the memory of the parameters in the i-th layer. activateiIndicates the first i Activation memory of the layer, param_all The sum of the memory usage for all layers. activate_all This represents the sum of active memory across all layers.

[0054] (2) Core objective of load computing Through the above calculations, the "computational load - memory load" data for each CPU node can be obtained, thereby: identifying nodes with excessively high loads: nodes whose computational or memory loads far exceed the cluster average level. These nodes are prone to computational blocking or memory overflow, and some layers need to be transferred to reduce the load; identifying nodes with excessively low loads: nodes whose computational or memory loads are far below the average level. These nodes have idle resources and can take over layers transferred from other nodes to improve resource utilization; providing clear "source node (excessively high load)" and "target node (excessively low load)" pointers for subsequent "dynamic layer transfer", ensuring accurate optimization direction.

[0055] Then, after obtaining the load status of each CPU, the neural network layers are dynamically transferred according to the load balancing goal. The specific strategy is as follows: for CPU nodes with excessive load, some of their computationally complex or memory-intensive neural network layers are transferred to nodes with lower loads; simultaneously, it is ensured that the transfer operation does not disrupt the established dependencies in the data dependency graph. For example, if a CPU node has an excessively high computational load, some of its convolutional layers can be transferred to other CPUs with lower computational loads. During the transfer process, the data dependency graph is used to verify whether the transferred solution conforms to the data flow and computational dependency requirements.

[0056] During the dynamic transfer process, an iterative optimization approach is adopted. Each iteration includes four main steps: "load calculation sorting - balance check - candidate block selection - block transfer execution". The load of each CPU is continuously recalculated and the partitioning scheme is adjusted until the computational and memory loads of each CPU reach a balanced state. Through this dynamic re-partitioning, the resources of the CPU cluster can be fully utilized, resource waste can be avoided, and the overall efficiency of model training can be improved.

[0057] S4, please refer to Figure 3 and Figure 4 Asynchronous pipelined parallel processing: This step uses a 1F1B non-interleaved pipelined parallelism. By dividing the model training process into multiple stages and processing data in micro-batch units, it has a unique advantage in ensuring training stability. Simultaneously, a double-buffered data loading strategy is incorporated, allowing the model to asynchronously process the next batch of data while processing the current batch, achieving deep parallelism between communication and computation, significantly improving model training performance.

[0058] 1F1B Non-Interleaved Pipeline Parallelism: The model training process is logically divided into multiple independent stages according to the sequence "data input → forward computation → backpropagation → gradient propagation" (e.g., the initial stage handles data preprocessing and distribution, the intermediate stage handles core layer computation, and the final stage handles final output and gradient summarization). Each stage uses "micro-batch data" as the processing unit (a complete batch is split into multiple micro-batches, e.g., a batch size of 32 is split into 4 micro-batches, each with a size of 8). The non-interleaved design ensures that each stage only processes the corresponding task of the current micro-batch, avoiding cross-execution of tasks from different micro-batches within a stage, thus guaranteeing the stability of the training process and data consistency.

[0059] Double-buffered data loading strategy: Two independent buffers are set up (controlled by buffer_idx, buffer_idx=0 or 1 corresponds to two buffers respectively). When the model uses data from the currently active buffer (e.g., buffer_idx=0) for computation, the next batch of data is asynchronously loaded into the other buffer (e.g., buffer_idx=1). This "parallel computation and data loading" design completely eliminates the idle time of traditional "computation waiting for data loading," while providing a data foundation for "parallel communication and computation."

[0060] To verify the effectiveness of the proposed method, this application validates the method on the Distil-Bert distillation model. In Distil-Bert, each CPU has its own CPU model, which is trained locally in a pipelined parallel mode. Training data of different sizes were randomly generated and compared with other methods to test the effectiveness of the proposed method.

[0061] All experiments in this application were implemented using Python 3.11.2 and run on a Raspberry Pi 4B with 8GB of RAM. The BERT distillation version of the model was used for training. The hyperparameters in the experiments were set as follows: 3 epochs per round, Adamw as the optimizer, and an initial learning rate (lr) of 0.00005. This application sets up an edge environment and simulates four CPU-centric edge devices using the four CPUs of the Raspberry Pi 4B. This application was validated using data of different scales and different micro-batch sizes. The evaluation metric for this application is: load variance. and training time The experiment verifies the method using the following two formulas:

[0062]

[0063] in, For maximum load, For minimum load, This is the end time of training. This is the start time of training.

[0064] The experimental results are as follows: Figure 5 The image shows the results of comparing the model segmentation method of this application with traditional model segmentation algorithms (data parallelism, micro-batch pipeline, naive pipeline). Figure 5 It can be seen that this application focuses more on CPU load optimization. Compared with the traditional model segmentation load difference of 75%, the load difference after the model re-segmentation in this application is reduced to 23%, which is an optimization of nearly 69%, which proves the effectiveness of model re-segmentation.

[0065] Figure 6 This compares data parallelism, micro-batch pipeline parallelism, naive pipeline parallelism, and asynchronous non-interleaved pipeline (the method of this application). The horizontal axis represents data size, the vertical axis represents training time, and the batch size is set to 4. The micro_batch size for micro-batch pipeline parallelism and asynchronous non-interleaved pipeline parallelism is set to 2. Table 1 compares the training time of the three methods (data parallelism, micro-batch pipeline, and naive pipeline) with the method of this application. Positive values ​​represent the percentage improvement in training time compared to the comparison methods, and negative values ​​represent the percentage increase in training time compared to the comparison methods. From... Figure 6 As can be seen from Table 1, in a CPU environment, pipelined parallelism has more advantages than data parallelism. At the same time, compared with micro-batch pipelined parallelism and naive pipelined parallelism, the method of this application has an improvement of about 50% and 3%, respectively.

[0066] Table 1 Comparison results of the four algorithms

[0067] Figure 7 This is an ablation experiment comparing the method in this application with asynchronous non-interleaved pipelined parallelism (without double buffers), while progressively increasing their micro-batch size, i.e., progressively increasing their communication time by increasing the amount of data transmitted. Figure 7 It can be seen that the training time decreases with the increase of the micro-batch size (micro_batchsize). Table 2 shows the training time optimization results of this application compared with asynchronous non-interleaved pipelines and naive pipelines. Similarly, positive values ​​represent the percentage improvement in training time compared to the comparison methods, and negative values ​​represent the percentage increase in training time compared to the comparison methods. From... Figure 7 As shown in Table 2, the training time of this application decreases more rapidly, which proves the effectiveness of the double buffer.

[0068] Table 2 Comparison results of the three algorithms

[0069] Please refer to Figure 8 and Figure 9 This example implementation further provides a communication system based on resource awareness and asynchronous parallelism, including: The model computation module is used to calculate the key resource consumption dimensions during the training process of a large language model. The key resource consumption dimensions include: computational complexity, parameter memory, and activation memory. The model analysis module treats the attention module as a whole unit, initializes the dependency dictionary, the preceding layer dictionary, and the Transformer block dictionary, captures the dependencies of all layers during model operation, performs sub-layer classification processing, upgrades hierarchical dependencies to block-level dependencies, and forms a block-level data dependency graph. The model segmentation module is used to initially allocate modules with close data dependencies to the same CPU node based on the data dependency graph. It accumulates the number of floating-point operations of all layers on each CPU node to obtain the computational load, accumulates the parameter memory and activation memory usage to obtain the memory load, identifies nodes with excessively high or low loads, iteratively calculates and adjusts the load until the load of each node is balanced. The asynchronous pipeline parallel module is used to divide the model training process into multiple independent stages. Each stage uses micro-batch data as the processing unit and adopts a non-interleaved 1F1B scheduling strategy to ensure that each stage only processes the task corresponding to the current micro-batch. It controls the switching of the dual buffers. When the model uses the data in the current buffer for calculation, the next batch of data is asynchronously loaded into the other buffer, realizing the parallel execution of data transmission and model calculation.

[0070] In this embodiment, the model computation module and model analysis module are the foundation for deploying large language model training, providing key basis for model segmentation and performance optimization, and are responsible for analyzing the computational resource requirements of each layer of the model and sensing all available resources in the system. First, the model's neural network layers are analyzed by dynamically quantifying the number of floating-point operations (FLOPs), parameter memory usage, and activation memory usage of each layer. Second, a data dependency graph is constructed to describe the flow of data in each layer of the model, clarifying the dependencies in the computation process and providing a basis for model segmentation strategies. The model segmentation module dynamically deploys the model to different computational resources using the data dependency graph. The asynchronous pipelined parallel module orderly divides model training into multiple stages, processing data in micro-batch units.

[0071] Regarding the system in the above embodiments, the specific ways in which each module performs operations have been described in detail in the embodiments related to the method, and will not be elaborated here.

[0072] In summary, this invention creates a resource-aware load balancing model segmentation method. This method intelligently divides the AI ​​model into multiple sub-modules based on the computing, storage, and network resources of each node in the CPU cluster. Each sub-module is assigned to the most suitable node for processing, thus fully utilizing node resources and avoiding excessive concentration and idleness. This invention proposes an asynchronous pipelined parallel strategy method. This method achieves parallel execution of data transmission and model computation through the fusion of non-blocking communication technology and a dual-buffer and non-interleaved 1F1B scheduling strategy, thereby reducing communication waiting time. Data flows sequentially through each stage like a pipeline, with each stage using asynchronous queues and non-blocking communication, allowing computation and communication operations to occur simultaneously, maximizing the overlap of computation and communication time.

[0073] This invention aims to construct a distributed system for efficient training of large language models. It deeply analyzes the problems of insufficient computing power, limited memory, and low communication efficiency in training large language models in the context of artificial intelligence and the Internet, and proposes an optimization method combining resource-aware model segmentation and asynchronous parallel communication. This invention creates a resource-aware model segmentation method and innovatively integrates a double-buffer and 1F1B non-interleaved scheduling strategy, thereby improving the efficiency of large language model training.

[0074] It should be noted that although several modules of the system for executing actions are mentioned in the detailed description above, this division is not mandatory. In fact, according to embodiments of the present invention, the features and functions of two or more modules described above can be embodied in one module. Conversely, the features and functions of one module described above can be further divided into multiple modules for embodiment. Components shown as modules may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of the present invention according to actual needs. Those skilled in the art can understand and implement this without any inventive effort.

[0075] The embodiments of the present invention have been described above with reference to the accompanying drawings. However, the present invention is not limited to the specific embodiments described above. The specific embodiments described above are merely illustrative and not restrictive. Those skilled in the art can make many other forms under the guidance of the present invention without departing from the spirit and scope of the claims. All of these forms are within the protection scope of the present invention.

Claims

1. A communication method based on resource awareness and asynchronous parallelism, characterized in that, include: S1, calculate the key resource consumption dimensions in the training process of the large language model, including: computational complexity, parameter memory and activation memory; S2 treats the attention module as a whole unit, initializes the dependency dictionary, the preceding layer dictionary and the Transformer block dictionary, captures the dependencies of all layers during the model operation, performs sub-layer classification processing, upgrades hierarchical dependencies to block-level dependencies, and forms a block-level data dependency graph. S3. Based on the data dependency graph, modules with close data dependencies are initially assigned to the same CPU node. The number of floating-point operations of all layers on each CPU node is accumulated to obtain the computational load. The memory load is obtained by accumulating the parameter memory and active memory usage. Nodes with excessively high or low loads are identified. The load is iteratively calculated and adjusted until the load of each node is balanced. S4 divides the model training process into multiple independent stages. Each stage uses micro-batch data as the processing unit and adopts a non-interleaved 1F1B scheduling strategy to ensure that each stage only processes the task corresponding to the current micro-batch. It controls the switching of the dual buffers. When the model uses the data in the current buffer for calculation, the next batch of data is asynchronously loaded into the other buffer, so as to realize the parallel execution of data transmission and model calculation.

2. The communication method based on resource awareness and asynchronous parallelism according to claim 1, characterized in that, In S1, computational complexity is measured by the number of floating-point operations (FLOPs) for each neural network layer. The formulas for calculating FLOPs for different neural network layers are as follows: in, Indicates the number of floating-point operations in the linear layer. Indicates the number of floating-point operations performed in the convolutional layer. Indicates the number of floating-point operations performed in the embedding layer. input_size Indicates the shape of the input data. output_size Indicates the shape of the output data. kernel_size Indicates the shape of the convolution kernel. in_channels Indicates the input channel. out_channels Indicates the output channel. output_length Indicates the output length. num_embed Indicates vocabulary size, embed_dim This indicates the embedding vocabulary dimension.

3. The communication method based on resource awareness and asynchronous parallelism according to claim 1, characterized in that, In S1, the parameter memory of each neural network layer The calculation formula is as follows: in, Pi_num This represents the number of elements in each parameter tensor. Pi_type Indicates the format of the stored parameters.

4. The communication method based on resource awareness and asynchronous parallelism according to claim 1, characterized in that, In S1, the activation of the memory computation process includes: For fully connected layers and convolutional layers, the output shape is first calculated based on the input shape and layer parameters, and then the activation memory usage data is obtained by multiplying the total number of elements in the output shape by 4. For applications that do not generate additional significant activation memory usage, set the activation memory usage data to 0.

5. The communication method based on resource awareness and asynchronous parallelism according to claim 1, characterized in that, In S2, PyTorch hooks are used with the register_hook(module, name) function to insert custom logic before / after execution in the model layer, thereby achieving dynamic capture of dependencies.

6. The communication method based on resource awareness and asynchronous parallelism according to claim 1, characterized in that, In S2, the sub-layer classification process includes: Initialize the Transformer block dictionary, where the key is the block number in the Transformer model and the value is the set of sub-layer names contained in that block; The scattered sub-layers are categorized according to their respective Transformer block numbers; Store the categorized sub-layer names into the dictionary value of the corresponding block number to complete the sub-layer categorization.

7. The communication method based on resource awareness and asynchronous parallelism according to claim 1, characterized in that, In S3, multiple consecutive attention modules are allocated as a whole to the same CPU node.

8. A communication system based on resource awareness and asynchronous parallelism, characterized in that, include: The model computation module is used to calculate the key resource consumption dimensions during the training process of a large language model. The key resource consumption dimensions include: computational complexity, parameter memory, and activation memory. The model analysis module treats the attention module as a whole unit, initializes the dependency dictionary, the predecessor layer dictionary, and the Transformer block dictionary, captures the dependencies of all layers during the model's operation, performs sub-layer classification, upgrades hierarchical dependencies to block-level dependencies, and forms a block-level data dependency graph. The model segmentation module is used to initially allocate modules with close data dependencies to the same CPU node based on the data dependency graph. It accumulates the number of floating-point operations of all layers on each CPU node to obtain the computational load, accumulates the parameter memory and activation memory usage to obtain the memory load, identifies nodes with excessively high or low loads, iteratively calculates and adjusts the load until the load of each node is balanced. The asynchronous pipeline parallel module is used to divide the model training process into multiple independent stages. Each stage uses micro-batch data as the processing unit and adopts a non-interleaved 1F1B scheduling strategy to ensure that each stage only processes the task corresponding to the current micro-batch. It controls the switching of the dual buffers. When the model uses the data in the current buffer for calculation, the next batch of data is asynchronously loaded into the other buffer, realizing the parallel execution of data transmission and model calculation.

Images