Accelerating process of training hybrid expert model

Abstract

This disclosure describes techniques for accelerating a process of training a hybrid expert (MOE) model. A sequence in training data is divided into a plurality of segments. The plurality of segments are input in parallel into a plurality of devices. Attention calculations of the layers are performed in parallel by the plurality of devices. The lems from the attention calculations of the layer are assigned to different devices among the plurality of devices, and expert calculations of the layer are implemented by these different devices. The traffic is reduced by reserving at least a portion of the lexical elements from each of these different devices on the same device for performing attention calculations for subsequent layers after completing expert calculations for the layers.

Description

Background Technology

[0001] Machine learning models are increasingly being used across various industries to perform a wide range of tasks. Improved technologies are needed to efficiently scale these models. Attached Figure Description

[0002] The following detailed embodiments can be better understood when read in conjunction with the accompanying drawings. For illustrative purposes, exemplary embodiments of various aspects of this disclosure are shown in the drawings; however, the invention is not limited to the specific methods and tools disclosed.

[0003] Figure 1 An example system for accelerating the training process of a hybrid expert (MoE) model according to this disclosure is shown.

[0004] Figure 2 An example system for sequence parallel attention according to this disclosure is shown.

[0005] Figure 3 An example attention block according to this disclosure is shown.

[0006] Figure 4 An example of an overlapping communication technique according to this disclosure is shown.

[0007] Figure 5 An example system for out-of-order parallelism according to this disclosure is shown.

[0008] Figure 6A and Figure 6B An example load balancing method utilizing out-of-order parallelism is shown according to this disclosure.

[0009] Figure 7 An example process for accelerating the training process of a MoE model according to this disclosure is shown.

[0010] Figure 8 An example process for accelerating the training process of a MoE model according to this disclosure is shown.

[0011] Figure 9 An example process for accelerating the training process of a MoE model according to this disclosure is shown.

[0012] Figure 10 An example process for accelerating the training process of a MoE model according to this disclosure is shown.

[0013] Figure 11 An example process for accelerating the training process of a MoE model according to this disclosure is shown.

[0014] Figure 12 An example process for accelerating the training process of a MoE model according to this disclosure is shown.

[0015] Figure 13 An example table showing the evaluation results according to this disclosure is provided.

[0016] Figure 14 An example table showing the evaluation results according to this disclosure is provided.

[0017] Figure 15A An example table showing the evaluation results according to this disclosure is shown.

[0018] Figure 15B An example table showing the evaluation results according to this disclosure is shown.

[0019] Figure 16 Example graphs are shown to illustrate the evaluation results according to this disclosure.

[0020] Figure 17A Example graphs are shown to illustrate the evaluation results according to this disclosure.

[0021] Figure 17B Example graphs are shown to illustrate the evaluation results according to this disclosure.

[0022] Figure 18 Example charts are shown illustrating the evaluation results according to this disclosure.

[0023] Figure 19 An example computing device is shown that can be used to perform any of the techniques disclosed herein. Detailed Implementation

[0024] This paper describes an improved technique for accelerating the training process of Hybrid Expert (MoE) models. In recent years, Large Language Models (LLMs) have become a cornerstone of modern artificial intelligence research, demonstrating unparalleled capabilities in generating human-like text, understanding complex queries, and driving breakthroughs across numerous fields. The increasing importance of LLMs is highlighted by their growing role in a wide range of applications, from augmenting natural language processing tasks to driving innovation in generative AI technologies. As expectations for these models continue to grow, the scale of their training systems is also expanding. This increased training scale makes efficiency improvements not only desirable but also crucial; even small improvements in training efficiency can lead to significant reductions in computational resource consumption and time, profoundly impacting the feasibility and sustainability of developing state-of-the-art LLMs.

[0025] The MoE mechanism is a sophisticated approach designed to improve the performance and efficiency of transformer models, which are becoming increasingly crucial in the LLM field. At its core, the MoE mechanism diversifies the transformer architecture by incorporating multiple specialized network components (e.g., experts) into a feedforward network (FFN) component. Unlike traditional transformer models that uniformly process all data through the same layers, MoE models dynamically route input terms to the most relevant experts based on the nature of the input. This routing is typically managed by a trainable gating mechanism that determines which experts are best suited for each piece of data. This architectural innovation allows for a significant expansion of the capacity of MoE models without proportionally increasing the computational cost for inference, as only a subset of experts is activated for each input. The MoE mechanism provides a more flexible and efficient way to improve model performance, rather than simply increasing network size.

[0026] Within the LLM architecture, the MoE model stands out due to its sparse activation architecture, which dynamically routes input terms to selected expert groups rather than all components. This design allows the required computational budget (e.g., FLOPs) to scale sublinearly with model size, significantly reducing computational costs. Despite MoE's inherently lower training cost, two significant challenges remain during its training. These challenges differ from those encountered when training dense models.

[0027] The first challenge encountered during MoE training stems from the significant differences between the characteristics of the attention mechanism and the characteristics of the FFN components. For example, to facilitate sparse computation, FFN components require two additional full-to-full communications during both forward and backward computations, which typically hinders ongoing computations and consumes a large portion of the training step time. This difference highlights the necessity of customized optimization strategies that can effectively address the unique needs of each component, thereby improving overall efficiency and further reducing the training cost of the MoE model.

[0028] The second challenge encountered during MoE training is the shifting balance between computation and communication. While model architectures have advanced, hardware capabilities have also rapidly evolved, with graphics processing units (GPUs) becoming increasingly faster. Simultaneously, training accuracy is decreasing in pursuit of more efficient and cost-effective training. These trends have led to a situation where the raw processing time available for computation is decreasing, making the relative communication overhead between computational units a more significant bottleneck. For example, it has been observed that in some cases, simply scaling existing intra-layer parallelism methods to multiple nodes can result in communication overhead exceeding 50%. Therefore, optimizing communication is crucial for maintaining and enhancing the scalability of large-scale MoE models, especially in distributed training environments where data must be synchronized across multiple devices.

[0029] This paper describes a system optimized for efficient large-scale MoE training on high-performance GPU clusters. The system described is a dedicated LLM training system tailored for MoE models, enabling full utilization of the computational power of high-performance GPUs. The key system principle described is the co-design of a dedicated parallel strategy and communication-computation overlap, which addresses the unique challenges posed by attention and FFN components in MoE layers.

[0030] For attention mechanisms in large-scale MoE training, tensor parallelism is typically applied to self-attention, and sequence parallelism is commonly applied to the LayerNorm and Dropout operators. This deployment aims to alleviate computationally intensive tasks and minimize activation memory footprint, respectively. However, this deployment introduces necessary full-collection and reduce-scatter communication along the critical path. With increasing GPU computing power and decreasing training accuracy, this relative communication overhead becomes unsustainable. The system described in this paper utilizes Sequence Parallel Attention (SPA), which partitions the entire attention computation along the sequence dimension, effectively eliminating collection and scatter operations from the critical path. The system also employs an overlap strategy that decomposes the projections of queries, keys, and values ​​to hide the key-value collection overhead in forward propagation, and a hierarchical parameter synchronization method that considers both intra-node and inter-node bandwidth in backpropagation.

[0031] For the FFN component, when the model parameters exceed the memory limit of a single GPU, the system described in this paper employs expert parallelism. In this scheme, the FFN component is distributed across multiple GPUs as a separate expert. Due to the inherent sparsity in the MoE model, full-to-full communication operations are required before and after expert computation, which has become a known major bottleneck. To overcome this bottleneck, the system described in this paper utilizes out-of-order expert parallelism, which retains a portion of the terms after computation on the expert side, thereby reducing the full-to-full communication volume by 1 / (2×top-k) while maintaining computational consistency. In addition, the system described in this paper utilizes an in-layer pipeline method to maximize the overlap between computation and full-to-full communication.

[0032] Figure 1 An example system 100 for accelerating the training process of MoE models is shown. System 100 includes multiple devices 102a to 102d. Each of the multiple devices 102a to 102d may include a GPU and / or a network interface controller (NIC). Although Figure 1 Only four devices are shown, but it should be understood that multiple devices may alternatively include any other number of devices.

[0033] System 100 employs SPA to address the challenges posed by attention blocks. SPA can partition the entire attention computation along the sequence dimension. This approach can significantly reduce communication overhead by leveraging a grouped query attention architecture. To adopt SPA, each sequence in the training data (e.g., sequence 103) can be divided into multiple segments. Figure 1 In the example, the first segment includes the lexical units "my" and "cat". The second segment includes the lexical units "slept" and "on". The third segment includes the lexical units "the" and "cozy". The fourth segment includes the lexical units "sofa" and ".". These multiple segments can be input in parallel to multiple devices 102a to 102d. For example, the first segment can be input to device 102a, the second segment can be input to device 102b, the third segment can be input to device 102c, and the fourth segment can be input to device 102d. Attention computation for a layer (e.g., the i-th layer) can be performed in parallel by multiple devices (e.g., 102a to 102d).

[0034] Device 102a can perform attention calculations associated with the first segment, device 102b can perform attention calculations associated with the second segment, device 102c can perform attention calculations associated with the third segment, and device 102d can perform attention calculations associated with the fourth segment. For example, device 102a can perform attention calculations associated with the first segment to generate the word unit "my". a (my a)” and “cat” a (Cat a)". Device 102b can perform attention calculations associated with the second segment to generate the lexical "slept". a (sleep a) and "on" a (above a)". Device 102c can perform attention calculations associated with the third segment to generate the lexical "the". a (the a) and "cozy" a (comfortable a)". Device 102d can perform attention calculations associated with the fourth segment to generate the lexical "sofa". a (Sofa a) and ". a (。a)”.

[0035] System 100 employs out-of-order expert parallelism to address the challenges posed by the FFN component. In an FFN block, expert parallelism eliminates the need to send all tokens back to their originating device. Out-of-order expert parallelism avoids sending all tokens back to their originating device, thus reducing communication costs. Compared to traditional expert parallelism, out-of-order expert parallelism also introduces a more balanced load across multiple experts.

[0036] To employ out-of-order expert parallelism, system 100 can assign the lexical units from the attention computation of a layer (e.g., layer i) to different devices among multiple devices (e.g., 102a to 102d). All-to-all (A2A) communication can be used to assign the lexical units from the attention computation of a layer to different devices based on the selected expert. For example, the lexical unit "my" a (my a)” and “cat” a "(cat a)" can be assigned to devices 102d and 102c respectively. The lexical "slept" a (sleep a) and "on" a "(on top of)" can be assigned to devices 102a and 102c respectively. The lexical "the" a (the a) and "cozy" a (comfortable a) can be assigned to device 102b and device 102a respectively. a (Sofa a) and ". a(。a)” can be assigned to device 102d and device 102b respectively.

[0037] Expert computation at each layer (e.g., the i-th layer) can be performed by different devices. For example, device 102a can perform computation with the lexical "cozy". a (comfortable a) and "slept" a Expert calculations associated with "(sleep a)" were used to generate the word morpheme "cozy". f (comfortable f) and "slept" f (Sleeping f)". Device 102b can implement with the word "". a (。a)” and “the” a (The a) is associated with expert calculations to generate lexical units. f (。f)” and “the f (the f)". Device 102c can implement the word "on" a (above a) and "cat" a Expert calculations associated with "(cat a)" were used to generate the word unit "on". f (above f) and "cat" f (cat f). Device 102d can implement with the word "my" a (my a) and "sofa" a Expert calculations associated with "(sofa a)" were used to generate the word unit "my" f (my f) and "sofa" f (Sofa f)

[0038] After expert computation is completed for a layer (e.g., the i-th layer), at least a subset of lexical units from each of these different devices can be retained on the same device for attention computation in subsequent layers (e.g., the i+1-th layer). For example, the lexical unit "cozy" can be used... f (comfortable f) and "slept" f (sleeping f) is retained on device 102a for attention calculation in subsequent layers (e.g., the i+1th layer), instead of storing the lexical "cozy". f (comfortable f) and "slept" f (sleeping f)” is sent back to its original device (e.g., device 102c and device 102b respectively). Similarly, the word character “” can be sent back. f (。f)” and “the f (The f) is retained on device 102b for attention calculations in subsequent layers, rather than for word units. f (。f)” and “the f(The f) is sent back to its original device (e.g., device 102d and device 102c respectively). The word "on" can be used. f (above f) and "cat" f (cat f) is retained on device 102c for attention calculations in subsequent layers, instead of storing the lexical "on". f (above f) and "cat" f (cat f)” is sent back to its original device (e.g., device 102b and device 102a respectively). Finally, the word character “” can be sent back. f (。f)” and “the f (The f) is retained on device 102d for attention calculations in subsequent layers, instead of storing the lexical "my". f (my f) and "sofa" f (Sofa f)” is sent back to its original device (e.g., device 102a and device 102d respectively). The total amount of communication can be reduced by keeping at least a portion of the lexical units from each of these different devices on the same device for implementing attention computation in subsequent layers.

[0039] In an embodiment, system 100 can utilize overlap techniques to minimize communication overhead in the SPA. To employ overlap, system 100 can overlap computation and communication. The projection of the query, key, and value can be decomposed into three separate matrix multiplication operations, unlike conventional methods that typically employ a single matrix multiplication operation for this purpose. This strategic decomposition facilitates the concurrent execution of query projection computation with the full collection communication process of the key and value components. By facilitating this overlap, communication overhead on the critical path can be significantly reduced, practically approaching zero.

[0040] Figure 2An example system 200 for sequence-parallel attention is shown. During the training of MoE models, tensor parallelism is typically employed to efficiently parallelize computationally intensive attention operations, while operations such as LayerNorm and DropOut are parallelized along the sequence dimension to save GPU memory. However, tensor-parallel attention introduces unavoidable communication for collecting and distributing activations along the critical path. While increasing the number of GPUs and utilizing lower-precision computations significantly reduces the computational burden of attention mechanisms, the relative increase in communication overhead becomes a more pressing issue. Techniques such as Multi-Query Attention (MQA) and Grouped Query Attention (GQA), where multiple queries share the same keys and values, can exacerbate this problem, resulting in suboptimal performance. Primarily, the increased communication rate can offset the benefits of parallelizing attention mechanisms across GPUs. Furthermore, communication time exceeding computation time means that communication overhead cannot be effectively overlapped and hidden.

[0041] The Sequence Parallel Attention (SPA) disclosed herein overcomes these limitations associated with tensor parallelism. SPA can be based on a GQA architecture. As shown in system 200, SPA efficiently partitions all computations of the attention mechanism across the sequence dimension. Self-attention is not readily parallelizable along the sequence dimension due to the necessary interactions between queries, keys, and values ​​of terms. SPA partitions queries across devices and performs a full collection operation on keys and values ​​before self-attention, thereby maintaining computational consistency. Compared to tensor parallelism, utilizing a GQA architecture allows for a significant reduction in communication while reducing computation at the same rate.

[0042] Figure 3 An example attention block 300 is shown. Attention block 300 may include, for example... Figure 1Any attention block shown (e.g., Attn SP0, Attn SP1, Attn SP2, Attn SP3, etc.). Each attention block may include five operations: key projection operation 302, value projection operation 304, query projection operation 306, attention operation 308, and output projection 310. These five operations can be performed by a GPU (e.g., one of multiple devices 102a to 102d). After performing key projection operation 302, a full collection operation 312 for keys can be performed concurrently with (e.g., by a NIC) performing value projection operation 304. After performing value projection operation 304, a full collection operation 314 for values ​​can be performed concurrently with (e.g., by a NIC) performing query projection operation 306. By performing full collection operations 312 and 314 concurrently with value projection operation 304 and query projection operation 306, the training process of the MOE model can be accelerated.

[0043] Figure 4 An example of an overlapped communication technique 400 is illustrated. As described above, system 100 can utilize overlapped techniques to minimize communication overhead in a Special Application (SPA). To employ overlapped techniques, system 100 can overlap computation (e.g., computation performed by a GPU) and communication (e.g., performed by a NIC). Each attention operation can be divided into two chunks. Figure 4 In the example, an attention block (such as attention block 300) can be divided into two chunks: attention chunk 402 and attention chunk 404. Similarly, each FFN component can be divided into two chunks: FFN chunk 406 and FFN chunk 408.

[0044] After performing the operation associated with attention block 402, the operation associated with attention block 404 can be performed simultaneously with A2A communication between all attention blocks. After performing the operation associated with attention block 404, the operation associated with FFN block 406 can be performed simultaneously with A2A communication between all attention blocks. After performing the operation associated with FFN block 406, the operation associated with FFN block 408 can be performed.

[0045] Figure 5 An example system 500 for out-of-order parallelism according to this disclosure is shown. Expert parallelism is a common parallelization strategy used in MoE models. This strategy involves distributing experts across different devices for parallel processing. Using all-to-all communication, tokens are assigned to different devices based on the selected (multiple) experts(s). After expert computation is completed, another all-to-all communication is used to send the tokens back to the original device for further processing.

[0046] However, after expert computation (i.e., post-computation), it is not necessary to restore all lexical units to their original positions. Given the case where top-k equals 1, lexical units can be retained on their assigned devices for subsequent attention layer computations. While this may seem challenging due to the necessary lexical interactions in attention computation, the application of SPAs (where lexical computation is performed independently) mitigates this problem, allowing for uninterrupted progress.

[0047] For cases where the top-k value exceeds 1, the traditional collection operation at the end of expert computation introduces complexity due to the need for weighted aggregation of lexical components across devices. However, by retaining a subset of lexical components on the current device and aggregating the rest onto that device, we can effectively reduce the total amount of all-to-all communication by 1 / (2×top-k). Figure 5 An example of out-of-order expert parallelism 500 is shown, where top-k equals 2. The communication from expert computation to the next layer of self-attention is reduced by half. Given that the top-k value in the MoE model mainly ranges between 1 and 2, this reduction in communication brought about by out-of-order expert parallelism has a significant impact.

[0048] Figure 6A and Figure 6B An example load balancing method utilizing out-of-order parallelism according to this disclosure is shown. Besides reducing communication overhead, using out-of-order expert parallelism also helps with load balancing in SPAs. When using SPAs, computation is partitioned based on the query dimension. Causal masks can lead to uneven distribution of computation across multiple devices, such as... Figure 6A Example distribution 600 is shown. Subsequent ranks typically handle a higher computational load because these queries are computed against most of the preceding keys and values. This computational imbalance between ranks can lead to a straggler effect, where slower devices delay synchronization points during training.

[0049] However, by implementing out-of-order expert parallelism (where lexical units are distributed across devices based on selected experts and processed directly in subsequent attention layers), causal masks are effectively shuffled along the query dimension. Causal masks make the computation more evenly distributed across devices, such as... Figure 6B Example distribution 601 is shown. This brings a certain degree of load balancing, thereby mitigating the imbalance that may be caused by SPA.

[0050] Figure 7 An example procedure 700 is shown to accelerate the training process of the MoE model. Although in Figure 7 The operations are described as a series of operations, but those skilled in the art will understand that various embodiments may add, remove, reorder or modify the operations described.

[0051] According to the Sequence Parallel Attention (SPA) disclosed herein, the entire attention computation can be partitioned along the sequence dimension. This approach can significantly reduce communication overhead by utilizing a grouped query attention architecture. At 702, the sequence in the training data (e.g., sequence 103) can be divided into multiple segments. For example, the sequence can be divided into a first segment, a second segment, a third segment, and a fourth segment. At 704, these multiple segments can be input in parallel into multiple devices (e.g., multiple devices 102a to 102d). For example, the first segment can be input into the first device among the multiple devices, the second segment into the second device among the multiple devices, the third segment into the third device among the multiple devices, and the fourth segment into the fourth device among the multiple devices.

[0052] At point 706, the attention computation of a layer (e.g., the i-th layer) can be performed in parallel by multiple devices. For example, a first device can perform the attention computation associated with a first segment, a second device can perform the attention computation associated with a second segment, a third device can perform the attention computation associated with a third segment, and a fourth device can perform the attention computation associated with a fourth segment. For example, the first device can perform the attention computation associated with the first segment to generate attention lexicons A and B; the second device can perform the attention computation associated with the second segment to generate attention lexicons C and D; the third device can perform the attention computation associated with the third segment to generate attention lexicons E and F; and the fourth device can perform the attention computation associated with the fourth segment to generate attention lexicons G and H.

[0053] At point 708, attention terms from the layer can be assigned to different devices across multiple devices. A2A communication can be used to assign these terms to different devices based on a selected expert. For example, attention terms A and B can be assigned to the second and third devices, respectively. Attention terms C and D can be assigned to the first and third devices, respectively. Attention terms E and F can be assigned to the second and first devices, respectively. Attention terms G and H can be assigned to the fourth and second devices, respectively.

[0054] The expert computation for each layer can be performed by different devices. For example, a first device can perform expert computation associated with attention lexicons F and C to generate expert lexicons F and C, respectively. A second device can perform expert computation associated with attention lexicons H and E to generate expert lexicons H and E, respectively. A fourth device can perform expert computation associated with attention lexicons D and B to generate expert lexicons D and B, respectively. A fifth device can perform expert computation associated with attention lexicons A and G to generate expert lexicons A and G, respectively.

[0055] At point 710, at least a portion of the lexical units from each of these different devices can be retained on the same device for attention computation in subsequent layers (e.g., the (i+1)th layer). For example, expert lexical units F and C can be retained on the first device for attention computation in subsequent layers, instead of sending them back to their original devices (e.g., the third and second devices). Similarly, expert lexical units H and E can be retained on the second device for attention computation in subsequent layers, instead of sending them back to their original devices (e.g., the fourth and third devices, respectively). Expert lexical units D and B can be retained on the third device for attention computation in subsequent layers, instead of sending them back to their original devices (e.g., the second and first devices, respectively). Finally, expert lexical units A and G can be retained on the fourth device for attention computation in subsequent layers, instead of sending them back to their original devices (e.g., the first and fourth devices, respectively). By keeping at least a portion of the lexical units from each of these different devices on the same device for use in attention computation in subsequent layers, the total amount of communication can be reduced.

[0056] Figure 8 An example procedure 800 is shown to accelerate the training process of the MoE model. Although in Figure 8 The operations are described as a series of operations, but those skilled in the art will understand that various embodiments may add, remove, reorder or modify the operations described.

[0057] SPAs can efficiently partition all computations of the attention mechanism across the sequence dimension. Since there are necessary interactions between queries, keys, and values ​​of terms, self-attention is not easily parallelized along the sequence dimension. At 802, queries can be partitioned across multiple devices. At 804, full collection operations for keys and values ​​can be performed before attention. Each full collection operation can include communication operations for gathering information from multiple devices. Performing full collection operations for keys and values ​​before self-attention helps maintain computational consistency. At 806, attention computation can be implemented. Attention computation can be implemented in parallel based on the query dimension.

[0058] Figure 9 An example procedure 900 is shown to accelerate the training process of the MoE model. Although in Figure 9 The operations are described as a series of operations, but those skilled in the art will understand that various embodiments may add, remove, reorder or modify the operations described.

[0059] The system (e.g., system 100) can utilize overlap techniques to minimize communication overhead in the SPA. To employ overlap, the system can overlap computation and communication. At 902, the projections for queries, keys, and values ​​can be decomposed into separate matrix multiplication operations. This differs from conventional methods, which typically employ a single matrix multiplication operation for this purpose. This strategic decomposition facilitates the concurrent execution of query projection computation with the full collection communication process for key and value components. At 904, query projection computation can be performed concurrently with the full collection operation performed for keys and values ​​to accelerate the training process of the MOE model. By facilitating this overlap, communication overhead on the critical path can be significantly reduced, practically approaching zero.

[0060] Figure 10 Example procedure 1000 is shown to accelerate the training process of the MoE model. Although in Figure 10 The operations are described as a series of operations, but those skilled in the art will understand that various embodiments may add, remove, reorder or modify the operations described.

[0061] At point 1002, the tokens computed from the attention calculation of a layer (e.g., layer i) can be assigned to different devices among multiple devices (e.g., multiple devices 102a to 102d). The tokens computed from the attention calculation of a layer can be assigned to different devices based on a selected expert. Full-to-full communication can be used to assign the tokens computed from the attention calculation of a layer. At point 1004, full-to-full communication can be hidden. Full-to-full communication can be hidden through overlapping computation and communication to accelerate the training process of the MOE model. Hiding full-to-full communication can effectively hide the communication overhead in SPA.

[0062] Figure 11 An example procedure 1100 is shown to accelerate the training process of the MoE model. Although in Figure 11 The operations are described as a series of operations, but those skilled in the art will understand that various embodiments may add, remove, reorder or modify the operations described.

[0063] In addition to reducing all-to-all communication, each micro-batch can be split into two, and the computation of a micro-batch can begin immediately after the communication phase of the previous micro-batch starts. At 1102, each micro-batch can be split into two sub-micro-batches. For example, each attention operation can be divided into a first attention block and a second attention block (e.g., attention block 402 and attention block 404). Similarly, each FFN component can be divided into a first FFN block and a second FFN block (e.g., FFN block 406 and FFN block 408).

[0064] At point 1104, when the current sub-micro-batch begins its communication phase, the computation of a new sub-micro-batch can be initiated. For example, after performing operations associated with the first attention block, operations associated with the second attention block can be performed simultaneously with A2A communication between all attention blocks. After performing operations associated with the second attention block, operations associated with the first FFN block can be performed simultaneously with A2A communication between all attention blocks. After performing operations associated with the first FFN block, operations associated with the second FFN block can be performed.

[0065] Figure 12 An example procedure 1200 is shown to accelerate the training process of the MoE model. Although in Figure 12 The operations are described as a series of operations, but those skilled in the art will understand that various embodiments may add, remove, reorder or modify the operations described.

[0066] At point 1202, the tokens from the attention computation of the layer (e.g., the i-th layer) can be assigned to different devices among multiple devices (e.g., multiple devices 102a to 102d). All-to-all communication can be used to assign the tokens from the attention computation of the layer to different devices based on the selected expert. The expert computation of the layer can be performed by different devices.

[0067] At point 1204, at least a portion of the lexical units from each of these different devices can be retained on the same device for attention computation in subsequent layers (e.g., layer i+1). For example, lexical units generated by the first device during expert computation can be retained on the first device for attention computation in subsequent layers, instead of sending them back to their original devices(s). Similarly, lexical units generated by the second device during expert computation can be retained on the second device for attention computation in subsequent layers, instead of sending them back to their original devices(s). Lexical units generated by the third device during expert computation can be retained on the third device for attention computation in subsequent layers, instead of sending them back to their original devices(s). Lexical units generated by the fourth device during expert computation can be retained on the fourth device for attention computation in subsequent layers, instead of sending them back to their original devices(s). At point 1206, the computational load can be balanced. The computational load can be balanced by performing attention computation in subsequent layers on different devices.

[0068] To demonstrate the effectiveness of sequence parallel attention, a detailed theoretical analysis was conducted, where b represents the micro-batch size, P represents the parameter size of the attention block, s represents the sequence length, h represents the hidden dimension size, d represents the data parallel size, e represents the expert parallel size, n represents the model parallel size (i.e., tensor or sequence parallel size) in the attention block, and m represents the ratio between the number of query headers and the number of key-value headers.

[0069] The attention mechanism mainly involves QKV (query, key, value) projection, self-attention, and output projection. The total FLOPs required by the GQA mechanism consist of the following four parts: 1) QKV projection: 2bsh²(1+2 / m) / n FLOPs; 2) QK matrix multiplication: 2bs²h / n FLOPs; 3) Attention over values: 2bs²h / n FLOPs; 4) Output projection: 2bsh² / n FLOPs. Summing these components, the total computational cost of the attention block is 4bsh(h+s+h / m) / n FLOPs.

[0070] For communication, when using tensor parallelism, the communication amount for each full collection or reduce-distribute operation is bsh(n-1) / n elements. In the case of sequence parallel attention, the communication amount for each full collection is reduced to bsh(n-1) / n / m elements.

[0071] We can assume the model is trained with a peak performance of 1979 TFLOPS and a bandwidth of 450 GB / s, and all computations are performed with FP8 precision, where full collection is FP8 communication and reduce-divide is BF16 communication, as reduce-divide requires higher precision. Furthermore, we can assume computational and communication utilization are both 60% in the settings (b = 1, s = 2, h = 12288), and the self-attention model parallel size is n = 8 with GQA coefficients m = 12 (i.e., 12 query headers share 1 key / value header). Performance is shown in... Figure 13 Table 1300 provides two observations. First, on high-performance hardware such as the H100 GPU, the communication time for tensor parallel attention can significantly exceed the computation time. Second, by employing sequence parallel attention, the communication volume can be drastically reduced to 1 / m of its original size while maintaining the same computational cost.

[0072] The performance of the techniques described in this paper was evaluated. Ablation analysis was performed to assess the effectiveness of various model parallelism strategies and overlap methods. In the experiments, a setup consisting of 32 GPUs was used, with each GPU managing one of the 32 experts. The focus was primarily on analyzing the communication exposure ratio and model FLOP utilization (MFU) during single-layer forward propagation of the training procedure. Similar conclusions were reached because the reverse procedure is the inverse of the forward procedure, and its communication time is consistent, but its computation time is approximately twice that of the forward procedure. Figure 14 As shown in Table 1400, the initial attempt was to simply extend model parallelism to a multi-node setup. However, as Experiment 1 (e.g., Experiment Index 1) demonstrates, this approach is ineffective due to the large amount of tensor parallel communication overhead (which constitutes a significant portion of the overall training time). The parallelism strategy for the attention mechanism is fixed at intra-node tensor parallelism and inter-node data parallelism, and various parallelism strategies, such as EP32, EP4TP8, and out-of-order EP32, are experimented for the MLP portion (Experiment Indexes 2, 3, and 4). The results confirm that the out-of-order execution strategy minimizes communication overhead, consistent with theoretical predictions.

[0073] Further modifications were made by changing the parallel strategy of the attention component from intra-node tensor parallelism to intra-node sequence parallelism attention, as shown in Experiment Indices 5 and 6. This adjustment significantly reduced the communication overhead in the attention mechanism and markedly improved MFU. Finally, by applying the designed overlap method, the exposed communication overhead was reduced to zero, resulting in MFU scores of 0.65 and 0.9 under bfloat16 and float8 conditions, respectively, as shown in Experiment Indice 7. This experiment also demonstrates the importance of communication optimization. When communication time is dominant, switching computation from bf16 to fp8 does not significantly improve performance. However, once the communication overhead is fully optimized, the benefits of using fp8 become very apparent.

[0074] In a weakly scalable setup, the scalability of the single-layer converter across multiple devices with 1 to 64 GPUs was evaluated. This setup means that the workload of each worker node remains constant, while the total system workload grows linearly. To address this, the micro-batch size within each node and the number of DP units across nodes were increased. Simultaneously, the number of experts and expert parallelism in the MLP portion were scaled proportionally to the number of GPUs. MFUs were reported separately for both the forward and backward processes. Figure 15A Table 1500 (which shows the weak extension performance of BF16 accuracy) and Figure 15B The results in Table 1501 (which shows the weak scaling performance at FP8 precision) indicate that MFU remains consistently high, with near-linear scaling observed. Even at a scale of 64 GPUs, the proportion of communication exposed during the forward and backward phases is zero, suggesting that communication overhead is minimal and likely cannot be further optimized. The total runtime consists primarily of GEMM operations and other miscellaneous operations, with the slight decrease in MFU attributed to the non-linear scaling of these miscellaneous operations.

[0075] Subsequently, an evaluation was conducted under a strongly scalable configuration. Strong scaling presents a greater challenge because the total system workload remains constant, while the workload assigned to each worker node continuously decreases. Two configurations were employed to address this. Within a single node, the micro-batch size remained constant while the number of sequence-parallel units was increased. Across multiple machines, the number of data-parallel units was increased, and the micro-batch size was decreased. In the strongly scalable scenario, the primary focus was on whether adding more worker nodes could reduce task execution time. Figure 16 As shown in graph 1600, as the number of working nodes increases, the latency, including both forward and reverse operations, continues to decrease across different settings.

[0076] As mentioned above, the difference between SP attention and TP attention lies in the altered parameter synchronization mode. TP attention requires synchronizing parameters of size P / n across d DP numbers. In contrast, SP attention requires synchronizing the full-size P parameters across n×d numbers. Theoretically, by utilizing a hierarchical architecture of both intra-node and inter-node networks, the time costs associated with these synchronization processes can be roughly equal. Experiments were conducted to verify this theory. In the experiments, the parameter synchronization communication latency between TP8 and SP8 was evaluated across 32 and 64 GPU settings. The data size increased from 384MB to 1536MB. Figures 17A to 17B The experimental results shown in graphs 1700 and 1701 demonstrate that the delays of TP8 and SP8 remain comparable, with no significant difference observed. This observation confirms the hypothesis that TP8 and SP8 exhibit similar performance characteristics in terms of data parallel communication delay.

[0077] Under a weakly scaled setting, the MoE training performance of System 100 and mainstream frameworks (e.g., Megatron) across configurations of 1 to 64 GPUs was compared using intra-layer model parallelism. Mainstream frameworks employ TP and EP to partition each transformer layer. For example... Figure 18 As shown in Figure 1800, the latency of System 100 (e.g., AdvMoE) remains relatively stable, while the latency of mainstream frameworks gradually increases. With 64 GPUs, System 100 can achieve up to 2.5 times the execution speed of mainstream frameworks. Several issues were found when implementing mainstream frameworks: (1) In the case of intra-layer TP, communication involves full collection and reduce-distribute operations, which are time-consuming on Hopper GPUs; (2) The lack of overlap in MoE training affects performance; (3) In the FP8 configuration, only QKVO GEMM uses the FP8 data type, while FFN still uses BF16, thus limiting the speedup benefits of FP8.

[0078] Figure 19 The exhibition showcased computing devices that can be used in various fields, such as Figures 1 to 5 The (multiple) models, components, and / or devices depicted herein. Regarding... Figures 1 to 5 Any or all components can be freely separated. Figure 19 To implement this using one or more instances of the computing device 1900. Figure 19 The computer architectures shown illustrate conventional server computers, workstations, desktop computers, laptop computers, tablet computers, network facilities, PDAs, e-readers, digital cellular phones, or other computing nodes, and can be used to perform any aspect of the computers described herein, such as implementing the methods described herein.

[0079] The computing device 1900 may include a substrate or “motherboard”, which is a printed circuit board to which multiple components or devices may be connected via a system bus or other electrical communication paths. One or more central processing units (CPUs) 1904 may operate in conjunction with a chipset 1906. The CPUs(multiple) 1904 may be standard programmable processors that perform the arithmetic and logic operations required to perform the operation of the computing device 1900.

[0080] Multiple CPU 1904 processors can perform necessary operations by manipulating switching elements to transition from one discrete physical state to the next. These switching elements can distinguish and change these states. Switching elements typically include electronic circuitry (such as flip-flops) that holds one of two binary states, and electronic circuitry (such as logic gates) that provides an output state based on a logical combination of the states of one or more other switching elements. These basic switching elements can be combined to create more complex logic circuits, including registers, adder-subtractor units, arithmetic logic units, floating-point units, etc.

[0081] The (multiple) CPUs 1904 can be enhanced or replaced by other processing units (such as (multiple) GPUs 1905). The (multiple) GPUs 1905 may include processing units specifically designed for, but not necessarily limited to, highly parallel computing (such as graphics and other visualization-related processing).

[0082] Chipset 1906 provides an interface between CPU(s) 1904 and the remaining components and devices on the substrate. Chipset 1906 provides an interface to random access memory (RAM) 1908, which is used as the main memory in computing device 1900. Chipset 1906 may also provide an interface to computer-readable storage media, such as read-only memory (ROM) 1920 or non-volatile RAM (NVRAM) (not shown), for storing basic routines that help boot computing device 1900 and transfer information between various components and devices. ROM 1920 or NVRAM may also store other software components required for the operation of computing device 1900 according to the various aspects described herein.

[0083] Computing device 1900 can operate in a networked environment using a logical connection to remote computing nodes and computer systems via a local area network (LAN). Chipset 1906 may include functionality for providing network connectivity via a network interface controller (NIC) 1922 (such as a Gigabit Ethernet adapter). NIC 1922 can connect computing device 1900 to other computing nodes via network 1918. It should be understood that multiple NICs 1922 can be present in computing device 1900, thereby connecting the computing device to other types of networks and remote computer systems.

[0084] Computing device 1900 can be connected to mass storage device 1928, which provides non-volatile storage for the computer. Mass storage device 1928 can store system programs, application programs, other program modules, and data, which are described in more detail herein. Mass storage device 1928 can be connected to computing device 1900 via storage controller 1924, which is connected to chipset 1906. Mass storage device 1928 can consist of one or more physical storage units. Mass storage device 1928 may include management unit 1910. Storage controller 1924 can interface with physical storage units via a Serial Amount SCSI (SAS) interface, a Serial Advanced Technology Attachment (SATA) interface, a Fibre Channel (FC) interface, or other types of interfaces used for physical connections and data transfer between the computer and physical storage units.

[0085] The computing device 1900 can store data on the mass storage device 1928 by changing the physical state of the physical storage units to reflect the stored information. The specific changes in the physical state can depend on various factors and different implementations described herein. Examples of such factors may include, but are not limited to, the technology used to implement the physical storage units and whether the mass storage device 1928 is characterized as a primary or secondary storage device.

[0086] For example, computing device 1900 can issue instructions via storage controller 1924 to change the magnetic properties of a specific location within a disk drive unit, the reflection or refraction properties of a specific location within an optical storage unit, or the electrical properties of a specific capacitor, transistor, or other discrete component in a solid-state storage unit, thereby storing information in mass storage device 1928. Other transformations of the physical medium are also possible without departing from the scope and spirit of this specification, wherein the foregoing examples are provided merely to assist this specification. Computing device 1900 can also read information from mass storage device 1928 by detecting the physical state or characteristics of one or more specific locations within a physical storage unit.

[0087] In addition to the mass storage device 1928 described above, the computing device 1900 may also access other computer-readable storage media to store and retrieve information, such as program modules, data structures, or other data. Those skilled in the art will understand that a computer-readable storage medium can be any available medium that provides non-transitory data storage and can be accessed by the computing device 1900.

[0088] By way of example and not limitation, computer-readable storage media may include volatile and non-volatile computer-readable storage media, transient computer-readable storage media and non-transitory computer-readable storage media, and removable and non-removable media implemented in any method or technology. Computer-readable storage media include, but are not limited to, RAM, ROM, erasable programmable ROM (“EPROM”), electrically erasable programmable ROM (“EEPROM”), flash memory or other solid-state memory technologies, compact disc ROM (“CD-ROM”), digital versatile optical disc (“DVD”), high-definition DVD (“HD-DVD”), BLU-RAY, or other optical storage devices, magnetic tape cassettes, magnetic tape, disk storage devices, other magnetic storage devices, or any other medium that can be used to store desired information in a non-transitory manner.

[0089] High-capacity storage devices (such as) Figure 19 The mass storage device 1928 depicted herein can store an operating system used to control the operation of the computing device 1900. The operating system may include a version of the LINUX operating system. The operating system may include a version of the WINDOWS SERVER operating system from Microsoft Corporation. Depending on other aspects, the operating system may include a version of the UNIX operating system. Various mobile phone operating systems, such as iOS and Android, may also be utilized. It should be understood that other operating systems may also be utilized. The mass storage device 1928 can store other systems, applications, and data used by the computing device 1900.

[0090] Mass storage device 1928 or other computer-readable storage media may also be encoded with computer-executable instructions that, when loaded into computing device 1900, transform the computing device from a general-purpose computing system into a special-purpose computer capable of implementing the various aspects described herein. These computer-executable instructions transform computing device 1900 by specifying the state transitions of CPU(s) 1904, as described above. Computing device 1900 can access computer-readable storage media storing computer-executable instructions that, when executed by computing device 1900, can perform the methods described herein.

[0091] Computing devices (e.g.) Figure 19The computing device 1900 depicted may also include an input / output controller 1932 for receiving and processing input from multiple input devices (such as a keyboard, mouse, touchpad, touchscreen, electronic stylus) or other types of input devices. Similarly, the input / output controller 1932 may provide output to a display (such as a computer monitor, flat panel display, digital projector, printer, plotter) or other types of output devices. It should be understood that the computing device 1900 may not include... Figure 19 All components shown may include Figure 19 Other components not explicitly shown in the document, or those that can be utilized with Figure 19 The architecture shown is completely different.

[0092] As described in this article, a computing device can be a physical computing device, such as... Figure 19 The computing device 1900. A computing node may also include virtual machine host processes and one or more virtual machine instances. Computer-executable instructions can be indirectly executed by the physical hardware of the computing device by interpreting and / or executing instructions stored and executed in the context of a virtual machine.

[0093] It should be understood that these methods and systems are not limited to any particular method, component, or implementation. It should also be understood that the terminology used herein is for the purpose of describing particular embodiments only and is not intended to be limiting.

[0094] Unless the context clearly specifies otherwise, the singular forms “a,” “an,” and “the” as used in this specification and the appended claims include a plural of indicators. A range herein may be expressed as from “about” a specific value, and / or to “about” other specific values. When such a range is expressed, another embodiment includes from one specific value and / or to another specific value. Similarly, when a value is expressed as an approximation using the antecedent “about,” it will be understood that the specific value forms another embodiment. It will also be understood that the endpoints of each range are significant relative to and independent of the other endpoints.

[0095] "Optional" or "optionally" means that the event or situation described below may or may not occur, and the description includes instances where the event or situation occurs or does not occur.

[0096] Throughout the description and claims of this specification, the word "comprise" and variations thereof (such as "comprising" and "comprises") mean "including but not limited to" and are not intended to exclude, for example, other components, integrals, or steps. "Exemplary" means "an example of..." and is not intended to convey indications of preferred or ideal embodiments. "For example" is not used in a limiting sense but for illustrative purposes.

[0097] Components that can be used to perform the described methods and systems are described. When describing combinations, subsets, interactions, groups, etc., of these components, it should be understood that while specific references to each of the various individual and collective combinations and arrangements of these components may not be explicitly described, each component is specifically conceived and described herein for all methods and systems. This applies to all aspects of this application, including but not limited to operations in the described methods. Therefore, if various additional operations exist that can be performed, it should be understood that each of these additional operations can be performed using any specific embodiment or combination of embodiments of the described methods.

[0098] The method and system can be more readily understood by referring to the following detailed description of preferred embodiments and examples included therein, as well as the accompanying drawings and their descriptions.

[0099] As will be understood by those skilled in the art, the methods and systems may take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, the methods and systems may take the form of a computer program product on a computer-readable storage medium having computer-readable program instructions (e.g., computer software) embodied within the storage medium. More specifically, the methods and systems may take the form of computer software implemented on the web. Any suitable computer-readable storage medium may be utilized, including hard disks, CD-ROMs, optical storage devices, or magnetic storage devices.

[0100] Embodiments of the methods and systems are described below with reference to block diagrams and flowcharts illustrating methods, systems, apparatuses, and computer program products. It should be understood that each block in the block diagrams and flowcharts, as well as combinations of blocks in the block diagrams and flowcharts, can be implemented by computer program instructions. These computer program instructions can be loaded onto a general-purpose computer, a special-purpose computer, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute on the computer or other programmable data processing apparatus, create means for implementing the functions specified in one or more blocks of the flowchart.

[0101] These computer program instructions may also be stored in a computer-readable storage medium, which may instruct a computer or other programmable data processing apparatus to operate in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of writing comprising computer-readable instructions for implementing the functions specified in one or more boxes of a flowchart. The computer program instructions may also be loaded onto a computer or other programmable data processing apparatus to cause a series of operational steps to be performed on the computer or other programmable apparatus to produce a computer-implemented process, such that the instructions, which execute on the computer or other programmable apparatus, provide steps for implementing the functions specified in one or more boxes of a flowchart.

[0102] The various features and processes described above can be used independently of each other or combined in various ways. All possible combinations and sub-combinations are intended to fall within the scope of this disclosure. Additionally, in some embodiments, certain method or process blocks may be omitted. The methods and processes described herein are not limited to any particular order, and blocks or states associated with these methods and processes may be performed in other suitable orders. For example, described blocks or states may be performed in an order different from that specifically described, or multiple blocks or states may be combined in a single block or state. Example blocks or states may be performed serially, in parallel, or in some other manner. Blocks or states may be added to or removed from the described example embodiments. The example systems and components described herein may be configured differently than described. For example, elements may be added, removed, or rearranged compared to the described example embodiments.

[0103] It should also be understood that the items are shown to be stored in memory or storage devices when in use, and these items, or portions thereof, may be transferred between memory and other storage devices for memory management and data integrity purposes. Alternatively, in other embodiments, some or all of the software modules and / or systems may execute in memory on another device and communicate with the illustrated computing system via inter-computer communication. Furthermore, in some embodiments, some or all of the systems and / or modules may be implemented or provided in other ways, such as at least in part as firmware and / or hardware, including but not limited to one or more application-specific integrated circuits (“ASICs”), standard integrated circuits, controllers (e.g., by executing appropriate instructions, and including microcontrollers and / or embedded controllers), field-programmable gate arrays (“FPGAs”), complex programmable logic devices (“CPLDs”), etc. Some or all of the modules, systems, and data structures may also be stored (e.g., as software instructions or structured data) on computer-readable media, such as hard disks, memory, networks, or portable media articles to be read by appropriate devices or via appropriate connections. The system, modules, and data structures can also be transmitted as generated data signals (e.g., as part of a carrier or other analog or digital propagation signal) over various computer-readable transmission media, including wireless and wired / cable-based media, and can take various forms (e.g., as part of a single or multiplexed analog signal, or as multiple discrete digital packets or frames). In other embodiments, such a computer program product can also take other forms. Therefore, the invention can be practiced with other computer system configurations.

[0104] Although the methods and systems have been described in conjunction with preferred embodiments and specific examples, this is not intended to limit the scope to the specific embodiments illustrated, as the embodiments herein are intended to be illustrative rather than restrictive in all respects.

[0105] Unless otherwise expressly stated, it should not be construed that any method described herein requires its operations to be performed in a particular order. Therefore, no order is intended to be inferred in any way where a method claim does not actually describe the order in which its operations are followed, or where the claims or specification do not otherwise specifically state that the operations are limited to a particular order. This applies to any possible non-express basis for interpretation, including: logical questions concerning the arrangement of steps or the flow of operations; direct meaning derived from grammatical organization or punctuation; and the number or type of embodiments described in the specification.

[0106] It will be apparent to those skilled in the art that various modifications and variations can be made without departing from the scope or spirit of this disclosure. Other embodiments will be apparent to those skilled in the art in light of the description and practice described herein. This specification and example figures are to be considered exemplary only, and the true scope and spirit are indicated by the appended claims.

Claims

1. A method for accelerating the training process of a hybrid expert MOE model, the method comprising: Divide the sequences in the training data into multiple segments; The multiple segments are input into multiple devices in parallel; The attention calculation of the layer is performed in parallel by the multiple devices; The attention-based lexical units from the layer are assigned to different devices among the multiple devices, and the layer is computed by experts on the different devices. as well as To reduce communication volume, after the expert calculations of the layer are completed, at least a portion of the lexical units from each of the different devices are retained on the same device for use in the attention calculations of subsequent layers.

2. The method of claim 1, wherein the attention calculation of the layer performed in parallel by the plurality of devices comprises: The query is divided across the multiple devices; as well as The attention computation is performed in parallel based on the query dimension.

3. The method according to claim 2, further comprising: Before self-attention, a full collection operation for keys and values ​​is performed, wherein each of the full collection operations includes a communication operation for collecting information from the plurality of devices.

4. The method of claim 1, wherein the attention calculation of the layer performed in parallel by the plurality of devices comprises: The projections for queries, keys, and values ​​are decomposed into separate matrix multiplication operations.

5. The method according to claim 4, further comprising: Query projection computation and full collection operations for keys and values ​​are performed concurrently to accelerate the process of training the MOE model.

6. The method according to claim 1, further comprising: Using full-to-full communication, the lexical units computed from the attention of the layer are assigned to the different devices based on the selected experts; as well as The whole-to-whole communication is hidden by overlapping computation and communication to accelerate the process of training the MOE model.

7. The method according to claim 6, further comprising: Each micro-batch is split into two sub-micro-batches; as well as The computation of the new sub-micro-batch begins when the previous sub-micro-batch starts its communication phase.

8. The method according to claim 1, further comprising: The computational load is balanced by distributing the lexical units across the different devices for expert computation, and then performing the attention computation directly in the subsequent layers.

9. A system for accelerating the training process of a hybrid expert MOE model, the system comprising: At least one processor; as well as At least one memory, communicatively coupled to the at least one processor and including computer-readable instructions that, when executed by the at least one processor, cause the at least one processor to perform the method according to claims 1-8.

10. A non-transitory computer-readable storage medium storing computer-readable instructions that, when executed by a processor, cause the processor to perform the method according to claims 1-8.

Images