Model training methods, model training devices, electronic devices, and storage media

Abstract

This disclosure provides a model training method, a model training apparatus, an electronic device, and a storage medium. The method is applied to a multi-chip unit accelerator card, which includes multiple chip units integrated in a single package. The model includes a converter decoding layer. The method includes: training the converter decoding layer using a hybrid parallel strategy; during the forward and backward propagation of the converter decoding layer, distributing the activation values ​​processed by the multi-chip unit accelerator card to the multiple chip units according to a first arrangement, and distributing the weights processed by the multi-chip unit accelerator card to the multiple chip units according to a second arrangement. This method significantly improves the training performance of large models on multi-chip unit accelerator cards without increasing inter-chip unit accelerator card communication overhead or memory resource consumption.

Description

Technical Field

[0001] The embodiments disclosed herein relate to the field of artificial intelligence, specifically to a model training method, a model training apparatus, an electronic device, and a storage medium. Background Technology

[0002] In the fields of artificial intelligence and high-performance computing, accelerator cards based on multi-die (also known as multi-grid or multi-die) architectures are gradually becoming widely used. Multi-die technology integrates two or more independent dies onto the same substrate or interposer through advanced packaging (such as 2.5D / 3D packaging or silicon interposers), enabling them to work efficiently together and present a single, logically unified, high-performance chip. When the package contains two dies, this architecture can also be called a dual-die architecture.

[0003] With the continuous growth in the scale of artificial intelligence models and the increasing computational demands, how to achieve efficient model training based on Multi-Die accelerator cards has become one of the important technical issues. Summary of the Invention

[0004] This disclosure provides at least one embodiment of a model training method applied to a multi-chip unit accelerator card, the multi-chip unit accelerator card including multiple chip units integrated in a single package, the model including a converter decoding layer, and the model training method including: training the converter decoding layer using a hybrid parallel strategy; during the forward and backward propagation of the converter decoding layer, distributing the activation values ​​processed by the multi-chip unit accelerator card to the multiple chip units according to a first arrangement, and distributing the weights processed by the multi-chip unit accelerator card to the multiple chip units according to a second arrangement, wherein the first arrangement includes splitting the data to be arranged along a preset dimension to obtain multiple data fragments, and distributing each data fragment to different chip units; the second arrangement includes copying the data to be arranged to obtain multiple data copies, and distributing each data copy to different chip units.

[0005] The model training method provided in at least one embodiment of this disclosure further includes: during the backpropagation process of the converter decoding layer, distributing the gradient processed by the multi-chip unit accelerator card to the multiple chip units according to the second arrangement.

[0006] In the model training method provided in at least one embodiment of this disclosure, the converter decoding layer includes a normalization module, an attention module, and a feedforward network module. The step of training the converter decoding layer using a hybrid parallel strategy includes: training the normalization module using a sequence parallel strategy; and training the attention module and the feedforward network module using a tensor parallel strategy.

[0007] In the model training method provided in at least one embodiment of this disclosure, the step of training the converter decoding layer using a hybrid parallel strategy further includes: performing an inter-card full aggregation operation on the activation values ​​output by the normalization module; and performing an inter-card reduced scattering operation on the activation values ​​output by the attention module and the feedforward network module.

[0008] In at least one embodiment of the model training method provided in this disclosure, the model further includes an embedding layer, which is located before the converter decoding layer. The model training method further includes: distributing the input tensors of the embedding layer to multiple chip units of the multi-chip unit accelerator card according to the second arrangement before the forward propagation process of the embedding layer; and distributing the weights and gradients processed by the multi-chip unit accelerator card to the multiple chip units according to the first arrangement during the forward and backward propagation processes of the embedding layer.

[0009] In the model training method provided in at least one embodiment of this disclosure, the activation values ​​output by the embedding layer are distributed in multiple chip units of the multi-chip unit accelerator card according to a third arrangement. The model training method further includes: converting the arrangement of the activation values ​​output by the embedding layer from the third arrangement to the first arrangement, and performing inter-card reduction scattering operation. The local tensors distributed in the multiple chip units according to the third arrangement can be reconstructed into complete tensors after being processed by a preset algorithm, and the shape of the complete tensor is the same as that of each local tensor.

[0010] In at least one embodiment of the model training method provided in this disclosure, the model further includes an output layer located after the converter decoding layer. The output layer includes a normalization module and a language model head module. The model training method further includes: training the normalization module using a sequential parallel strategy; distributing the activation values ​​processed by the multi-chip unit accelerator card to the multiple chip units according to the first arrangement during the forward and backward propagation of the normalization module; performing an inter-card full aggregation operation on the activation values ​​output by the normalization module to obtain a first activation value; and converting the arrangement of the first activation value from the first arrangement to the second arrangement for input to the language model head module.

[0011] In at least one embodiment of the model training method provided in this disclosure, the model training method further includes: during the forward and backward propagation of the language model head module, distributing the weights and gradients processed by the multi-chip unit accelerator card to the multiple chip units according to the first arrangement.

[0012] In at least one embodiment of the model training method provided in this disclosure, the model training method further includes: under the first arrangement, performing parameter optimization and updating based on the optimizer state variables and the gradient obtained through the backpropagation process to obtain a first updated weight, wherein the first updated weight is distributed according to the first arrangement; converting the first updated weight into a second updated weight, wherein the second updated weight is distributed according to the second arrangement.

[0013] In the model training method provided in at least one embodiment of this disclosure, the step of performing parameter optimization and updating based on optimizer state variables and gradients obtained through backpropagation in the first arrangement to obtain the first updated weights includes: distributing the optimizer state variables processed by the multi-chip unit accelerator card to the multiple chip units according to the first arrangement; performing a splitting operation and a gradient accumulation operation on the first gradient obtained through backpropagation to obtain a second gradient, wherein the first gradient is distributed according to the second arrangement and the second gradient is distributed according to the first arrangement; and performing parameter optimization and updating based on the optimizer state variables and the second gradient to obtain the first updated weights.

[0014] In the model training method provided in at least one embodiment of this disclosure, the second gradient and the first updated weights adopt a first precision format; the first gradient and the second updated weights adopt a second precision format, wherein the first precision is higher than the second precision.

[0015] At least one embodiment of this disclosure provides a model training apparatus applied to a multi-chip unit accelerator card. The multi-chip unit accelerator card includes multiple chip units integrated in a single package. The model includes a converter decoding layer. The model training apparatus includes: an execution module configured to train the converter decoding layer using a hybrid parallel strategy; and a distribution module configured to distribute the activation values ​​processed by the multi-chip unit accelerator card to the multiple chip units according to a first arrangement and a second arrangement during the forward and backward propagation of the converter decoding layer. The first arrangement includes splitting the data to be arranged along a preset dimension to obtain multiple data fragments and distributing each data fragment to different chip units. The second arrangement includes copying the data to be arranged to obtain multiple data copies and distributing each data copy to different chip units.

[0016] At least one embodiment of this disclosure provides an electronic device, the electronic device comprising: at least one processor; at least one memory including one or more computer program modules; wherein the one or more computer program modules are stored in the at least one memory and configured to be executed by the at least one processor, the one or more computer program modules being used to implement the model training method described in at least one embodiment of this disclosure.

[0017] At least one embodiment of this disclosure provides a non-transitory computer-readable storage medium storing computer-readable instructions thereon, wherein the computer-readable instructions, when executed by at least one processor, perform the model training method described in at least one embodiment of this disclosure.

[0018] The model training method, model training apparatus, electronic device, and storage medium provided in at least one embodiment of this disclosure, in scenarios employing a hybrid parallel strategy, arrange the weights in the forward and backward propagation processes of the converter decoding layer using a Die-BB type arrangement and the activation values ​​using a Die-SB type arrangement. This satisfies the parallel partitioning constraints and avoids communication redundancy, while also eliminating the need for any arrangement conversion operations during the computation process of the converter decoding layer. This method significantly improves the large model training performance on multi-chip unit accelerator cards without increasing inter-chip unit accelerator card communication overhead or memory resource consumption. Attached Figure Description

[0019] To more clearly illustrate the technical solutions of the embodiments of this disclosure, the accompanying drawings of the embodiments will be briefly described below. Obviously, the drawings described below only relate to some embodiments of this disclosure and are not intended to limit this disclosure.

[0020] Figure 1 This is a schematic diagram illustrating the calculation and communication process of a converter decoding layer, provided for at least one embodiment of the present disclosure.

[0021] Figure 2 A flowchart illustrating a model training method provided in at least one embodiment of this disclosure.

[0022] Figure 3 This is a schematic diagram of another converter decoding layer calculation and communication process provided for at least one embodiment of the present disclosure.

[0023] Figure 4A This is a schematic flowchart illustrating an embedded layer computing and communication method provided for at least one embodiment of the present disclosure.

[0024] Figure 4B This is a schematic flowchart illustrating a converter decoding layer calculation and communication method provided for at least one embodiment of the present disclosure.

[0025] Figure 4C This is a schematic block diagram of an attention module provided for at least one embodiment of the present disclosure.

[0026] Figure 4D This is a schematic flowchart illustrating an output layer computing and communication method provided for at least one embodiment of the present disclosure.

[0027] Figure 5 This is a schematic diagram illustrating a parameter optimization update provided for at least one embodiment of the present disclosure.

[0028] Figure 6 This is a schematic block diagram of a model training apparatus provided for at least one embodiment of the present disclosure.

[0029] Figure 7 This is a schematic block diagram of an electronic device provided for at least one embodiment of the present disclosure.

[0030] Figure 8 This is a schematic block diagram of another electronic device provided for at least one embodiment of the present disclosure.

[0031] Figure 9 This is a schematic block diagram of a non-transitory computer-readable storage medium provided for at least one embodiment of the present disclosure. Detailed Implementation

[0032] To make the objectives, technical solutions, and advantages of the embodiments of this disclosure clearer, the technical solutions of the embodiments of this disclosure will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this disclosure. All other embodiments obtained by those skilled in the art based on the described embodiments of this disclosure without creative effort are within the scope of protection of this disclosure.

[0033] Unless otherwise defined, the technical or scientific terms used in this disclosure shall have the ordinary meaning understood by one of ordinary skill in the art to which this disclosure pertains. The terms “first,” “second,” and similar terms used in this disclosure do not indicate any order, quantity, or importance, but are merely used to distinguish different components. Terms such as “comprising” or “including” mean that the element or object preceding the word encompasses the elements or objects listed following the word and their equivalents, without excluding other elements or objects. Terms such as “connected” or “linked” are not limited to physical or mechanical connections, but can include electrical connections, whether direct or indirect. Terms such as “upper,” “lower,” “left,” and “right” are used only to indicate relative positional relationships, and these relative positional relationships may change accordingly when the absolute position of the described objects changes.

[0034] The present disclosure will now be described through several specific embodiments. To keep the following description of the embodiments of the present disclosure clear and concise, detailed descriptions of known functions and known components may be omitted. When any component of an embodiment of the present disclosure appears in more than one drawing, the component is represented by the same or similar reference numerals in each drawing.

[0035] The accelerator card described in at least one embodiment of this disclosure may include a graphics processing unit (GPU), a general-purpose graphics processing unit (GPGPU), a tensor processing unit (TPU), a deep learning processing unit (DPU), an accelerated processing unit (APU), a neural network processing unit (NPU), an application-specific integrated circuit (ASIC), or a field-programmable gate array (FPGA), etc.

[0036] In multi-die accelerator cards, data can be arranged across multiple dies in various ways, including Die-SB (Split-Broadcast), Die-BB (Broadcast-Broadcast), and Die-PP (Partial-Partial). This classification originates from the Split-Broadcast-Partial (SBP) mechanism proposed by the OneFlow deep learning framework. The SBP mechanism can automatically select the optimal data distribution strategy by abstracting logical and physical views, defining data mapping rules (i.e., Split / Broadcast / Partial), and predefined operator behavior rules.

[0037] For example, Die-SB means splitting the complete tensor along a preset dimension to obtain multiple data fragments, and distributing each data fragment into different dies. In other words, in a Multi-Die accelerator card, the complete tensor is split and stored in multiple dies.

[0038] For example, Die-BB means copying the complete tensor to obtain multiple copies of the data and distributing these copies across different dies. For instance, data can be synchronized from one die to another via a Die-to-Die broadcast operation, ensuring that each die holds a complete copy of the data.

[0039] For example, Die-PP means that each die stores a local tensor with the same shape as the complete tensor, but its value is only a part of the complete tensor. A preset algorithm (such as performing partial sum, maximum, minimum, etc. on the local tensor) is needed to reconstruct the logically complete tensor.

[0040] When using multi-die accelerator cards to perform model training and inference tasks, the arrangement of the model's activation values, weights, gradients, and optimizer states among the dies directly affects computational performance and memory utilization. However, the data arrangement problem of multi-die accelerator cards spans multiple interdisciplinary fields such as chip architecture design, parallel computing strategies, and memory management, and currently, a mature and practical system-level solution is still lacking.

[0041] During model training, various parallel strategies can be employed to conserve GPU memory and improve computational efficiency, such as data parallelism (DP), pipeline parallelism (PP), tensor parallelism (TP), and sequence parallelism (SP). These strategies partition computation and storage from different dimensions to adapt to hardware resource constraints. While each of these parallel strategies has its advantages, when dealing with large language models (LLMs) with a large number of parameters, a single strategy often struggles to meet the comprehensive optimization requirements of GPU memory, computation, and communication. To further improve training efficiency and accelerator card hardware utilization, hybrid parallel strategies, that is, a combination of multiple parallel strategies, can also be used.

[0042] For example, a data parallel strategy divides the input batch data into multiple sub-batches, each processed by a different accelerator card. Each accelerator card holds a complete copy of the model, and after completing forward and backward propagation computations, gradients are synchronized through an all-reduce operation between the accelerator cards. For multi-Die accelerator cards, since each accelerator card executes the complete computation process, the data parallel strategy does not affect the data arrangement within the multi-Die accelerator card, i.e., it does not involve tensor rearrangement within the card.

[0043] For example, the pipelined parallel strategy divides the model vertically into several stages according to network layers, with different accelerator cards responsible for the computation of each stage. Data of the same batch is passed sequentially between accelerator cards in a pipelined manner in the form of micro-batches, reducing the memory pressure on a single card by trading time for space. For multi-Die accelerator cards, since each stage usually contains several consecutive, isomorphic transformer decoder layers, each stage constitutes a complete computational structure on its respective accelerator card. The access patterns of its weight matrix and activation value matrix behave consistently in different transformer decoder layers. Therefore, adopting the pipelined parallel strategy does not affect the data arrangement within the multi-Die accelerator card.

[0044] For example, tensor parallelism partitions the weight matrix within a single network layer (e.g., by row or column), with multiple accelerator cards collaboratively performing a matrix multiplication operation. The computation results are aggregated via communication, thus distributing the memory and computational load of a single layer across multiple accelerator cards. For multi-die accelerator cards, tensor parallelism causes the computational tensors of the model (e.g., weights) to be partitioned, thereby affecting how the data is arranged within the card (across multiple dies).

[0045] For example, sequence parallelism, designed for long sequence scenarios, splits the input data along the sequence length dimension, with each accelerator card processing only the activation values ​​corresponding to a subset of tokens. Before operations requiring global context (such as attention modules), a full-gather operation between accelerator cards is used to aggregate the complete sequence information, which is suitable for alleviating memory bottlenecks in long text training. For multi-Die accelerator cards, similar to tensor parallelism, sequence parallelism also leads to the splitting of the model's computational tensors (such as activation values), thus affecting the arrangement of data within the card (across multiple dies).

[0046] For tensor parallel and sequence parallel modes, rearranging tensor data typically requires two methods: inter-accelerator card communication and intra-accelerator card communication (i.e., inter-die communication). For example, in a multi-die accelerator card, the inter-die interconnect bandwidth can reach the order of 10TB / s, while the inter-accelerator card interconnect bandwidth is usually only 1.8TB / s. Since the bandwidth of inter-die communication is significantly higher than that of inter-card communication, when there is redundancy in the data transmitted across cards on two dies within a single accelerator card, redundant data should be eliminated first within the accelerator card using high-bandwidth inter-die interconnects. This reduces the amount of data that needs to be transmitted across cards and improves overall communication efficiency.

[0047] Therefore, in order to achieve the optimization goals of reducing memory usage, balancing computational load, and reducing communication overhead, it is necessary to rationally determine the arrangement of the model's activation values, weights, and gradients on the Multi-Die accelerator card, while satisfying parallel partitioning constraints and avoiding communication redundancy. How to specifically set this arrangement becomes a problem that needs to be solved.

[0048] The following section uses the training of the Transformer Decoder Layer as an example to propose a method for model training in a distributed system composed of multiple Multi-Die accelerator cards using a hybrid parallel strategy. This disclosure does not limit the structure of the Transformer Decoder Layer; an example can be found in the decoding layer design of the Llama2-70B model.

[0049] Figure 1 This is a schematic diagram illustrating the calculation and communication process of a converter decoding layer, provided for at least one embodiment of the present disclosure.

[0050] like Figure 1 As shown, the converter decoding layer includes a normalization module, an attention module, and a feed-forward network (FFN) module. The normalization module can be implemented using layer normalization (LayerNorm), root mean square layer normalization (RMSNorm), or similar methods. The attention module can employ self-attention or cross-attention mechanisms, typically including operations such as linear projection, multi-attention head splitting, attention score calculation, and weighted aggregation. The feed-forward network usually consists of one or more fully connected layers, with non-linear activation functions (e.g., GELU, SwiGLU) introduced in between. This disclosure does not limit the specific implementation of the above network modules.

[0051] Figure 1 Two normalization modules are shown, referred to as the first normalization module (Norm1) and the second normalization module (Norm2) for easy distinction. Their network structures can be identical. That is, Figure 1 The converter decoding layer shown includes a first normalization module, an attention module, a second normalization module, and a feedforward network module connected in sequence. Figure 1 In the example, the normalization module is trained using a sequence parallel strategy, while the attention module and the feedforward network module are trained using a tensor parallel strategy.

[0052] Specifically, during the training of the normalization module, since the normalization module performs position-wise operations, it only relies on the activation statistics of a single token and does not require cross-token information. Therefore, the input sequence can be split along the sequence length dimension and assigned to different accelerator cards for computation. The input sequence is usually represented as a three-dimensional tensor with the shape [batch size, sequence length, hidden layer size]. The batch size represents the number of samples processed in one forward propagation, the sequence length represents the number of tokens in each sample, and the hidden layer size represents the feature vector dimension of each token. Assuming the input sequence length is 2048 and the sequence parallelism (i.e., the number of accelerator cards participating in sequence parallelism) is 4, then accelerator card 0 is responsible for processing tokens 0-511, accelerator card 1 is responsible for processing tokens 512-1023, accelerator card 2 is responsible for processing tokens 1024-1535, and accelerator card 3 is responsible for processing tokens 1536-2047.

[0053] During the training of the attention module and the feedforward network module, since global interactions across words or high-dimensional matrix operations are involved, the model weights can be divided according to a preset dimension, and each weight slice can be assigned to different accelerator cards for collaborative computation to improve computational efficiency.

[0054] Because there is a switching of parallel strategies between modules, communication operations need to be inserted at the boundaries to align the data arrangement. Figure 1 In the forward propagation, g represents the All-Gather operation and the Reduce-Scatter operation; g' represents the Reduce-Scatter operation during forward propagation and the All-Gather operation during backward propagation.

[0055] For a single Multi-Die accelerator card, the output of the first normalization module is the local activation value corresponding to the local sequence; that is, each Multi-Die accelerator card participating in sequence parallelism only holds a fragment of the complete activation value. However, the attention mechanism needs to model the global dependency between any two lexical units, and its computation requires each accelerator card to access the complete activation value. Therefore, before feeding the activation value output by the first normalization module into the attention module, each accelerator card participating in sequence parallelism needs to perform a full aggregation operation between the cards, pooling the activation value fragments they hold and reconstructing the complete activation value locally. Since the attention module adopts a tensor parallelism strategy, its weights are split and distributed across the accelerator cards. Therefore, each accelerator card participating in tensor parallelism only performs partial computation based on its local weight fragment and the complete activation value; the activation value output by the attention module is a local result based on the complete input but limited by the local weights. The second normalization module runs in sequence parallelism mode and only needs to process its assigned local sequence fragment. Therefore, it is necessary to perform inter-card reduction and scattering operations on the local results output by each accelerator card participating in tensor parallelism. It should be understood that the reduction and scattering operation includes reduction and scattering operations: first, the local results of each accelerator card participating in tensor parallelism are aggregated through reduction operations to restore the global activation value, and then the results are split according to the sequence dimension and redistributed to each accelerator card participating in sequence parallelism (that is, the scattering operation), thereby obtaining the shape of the local activation value required for sequence parallelism.

[0056] Therefore, the activation values ​​output by the first normalization module need to undergo a full aggregation operation before being input into the attention module, while the activation values ​​output by the attention module need to undergo a reduction scattering operation before being input into the second normalization module. Similarly, the activation values ​​output by the second normalization module need to undergo a full aggregation operation before being input into the feedforward network module, while the activation values ​​output by the feedforward network module need to undergo a reduction scattering operation before being input into the subsequent module (still a normalization module trained using a sequence parallel strategy, not shown in the figure).

[0057] like Figure 1 As shown, in a single Multi-Die accelerator card, the weights and gradients of each module are of Die-SB type ( Figure 1 The array is represented as sb, and the activation values ​​of the attention module and the feedforward network module are both of Die-BB type ( Figure 1The normalization module's activation values ​​and the inputs and outputs of the communication components (including full aggregation and reduction scattering operations) are arranged in Die-SB format. This arrangement can also be called the Multi-Die TP (Multi-Die Tensor Parallel) arrangement strategy. To meet the requirements of parallel partitioning constraints and avoid communication redundancy, a Die-BB to Die-SB reordering operation is performed on the activation values ​​before each reduction scattering operation; correspondingly, a Die-SB to Die-BB reordering operation is also performed on the activation values ​​after each full aggregation operation. Therefore, a total of 8 reordering operations are required during a complete computation (including forward and backward propagation) of each converter decoding layer. It should be noted that the total number of reordering operations will increase linearly with the number of model layers (the model may include multiple stacked converter decoding layers) and the number of mini-batches to be processed in each training iteration.

[0058] Taking the Llama2-70B model trained using a hybrid parallel strategy (e.g., sequence parallelism + tensor parallelism + pipeline parallelism) as an example, assuming the following configuration: the actual micro-batch size processed in each forward / backward propagation is 8, the global batch size processed in one iteration of the entire model is 1024, and a single pipeline stage includes 20 transformer decoding layers. Under this configuration, the number of micro-batch sizes processed in each training iteration is Globalbatch size / Micro-batch size = 1024 / 8 = 128.

[0059] In the Multi-Die TP deployment strategy described above, the number of deployment mode conversion operations performed by a single accelerator card in each iteration is num_reorder_per_layer × num_layers × num_micro_batches = 8 × 20 × 128 = 20,480. Here, num_reorder_per_layer represents the number of deployment mode conversion operations during the computation of each converter's decoding layer (8 times in the example above), num_layers represents the number of converter decoding layers in a single pipeline stage, and num_micro_batches represents the number of mini-batches processed in each training iteration. As the number of layers or the number of mini-batches increases, the deployment mode conversion overhead will increase significantly, becoming a key factor affecting training efficiency. Each deployment mode conversion operation involves data movement across dies. Such high-frequency data movement not only introduces significant communication overhead and reduces training efficiency but also increases the adaptation complexity of the model in distributed systems.

[0060] It should be noted that if the data parallelism strategy is further enabled, the 128 mini-batches calculated above will be evenly distributed among the various DP parallel groups. In this case, the number of mini-batches that a single accelerator card needs to process is 128 / DP_Size, where DP_Size represents the DP parallelism (i.e., the number of accelerator cards participating in data parallelism). At this time, a single accelerator card will perform the following data reordering operations in each iteration: num_reorder_per_layer × num_layers × num_micro_batches = 8×20×128 / DP_Size = 20,480 / DP_Size.

[0061] Regarding memory usage, when using the Multi-Die TP layout strategy described above, model weights, gradients, and optimizer states are stored as follows: weights in forward and backward propagation are stored with BF16 (16-bit floating-point) precision and arranged in Die-SB format; the optimizer's FP32 (32-bit single-precision floating-point) precision gradients, FP32 precision weight copies, and optimizer state variables (including first-order momentum and second-order momentum) are also arranged in Die-SB format. Gradients are accumulated using a gradient accumulation strategy, directly accumulating at FP32 precision, eliminating the need for additional storage of BF16 precision gradients. Therefore, the memory overhead per parameter includes 2 bytes of BF16 precision weights, 4 bytes of FP32 precision weight copies (for optimizer updates), 4 bytes of FP32 precision gradients, and 8 bytes of optimizer state variables, meaning the overall memory usage is approximately 18 times the number of parameters (in bytes). For example, for a model with 70B parameters, the weights and optimizer alone would require 70B × 18 = 1.26 TB of video memory. Assuming tensor parallelism is 8 and pipeline parallelism is 4, then a single accelerator card would require 70B × 18 / (8 × 4) = 78.75 GB of video memory.

[0062] While the aforementioned Multi-Die TP deployment strategy supports model splitting and computational parallelism, its frequent data movement introduces significant communication overhead. Therefore, it is necessary to further reduce the number of data movement operations without increasing GPU memory usage, thereby improving training efficiency.

[0063] At least one embodiment of this disclosure provides a model training method applied to a multi-chip unit accelerator card. The multi-chip unit accelerator card includes multiple chip units integrated in a single package, and the model includes a converter decoding layer. The model training method includes: training the converter decoding layer using a hybrid parallel strategy; during the forward and backward propagation of the converter decoding layer, distributing the activation values ​​processed by the multi-chip unit accelerator card to the multiple chip units according to a first arrangement, and distributing the weights processed by the multi-chip unit accelerator card to the multiple chip units according to a second arrangement. The first arrangement includes splitting the data to be arranged along a preset dimension to obtain multiple data fragments, and distributing each data fragment to different chip units. The second arrangement includes copying the data to be arranged to obtain multiple data copies, and distributing each data copy to different chip units.

[0064] The model training method provided in at least one embodiment of this disclosure, in a scenario employing a hybrid parallel strategy, arranges the weights in the forward and backward propagation processes of the converter's decoding layer using a Die-BB type arrangement and the activation values ​​using a Die-SB type arrangement. This satisfies the parallel partitioning constraints and avoids communication redundancy, while also eliminating the need for any arrangement conversion operations during the computation process of the converter's decoding layer. This method significantly improves the training performance of large models on multi-chip unit accelerator cards without increasing inter-chip unit accelerator card communication overhead or memory resource consumption.

[0065] Figure 2 This is a flowchart illustrating a model training method provided in at least one embodiment of the present disclosure. This model training method can be applied to any of multiple multi-chip unit accelerator cards (i.e., the Multi-Die accelerator card described above, hereinafter also referred to as an accelerator card for simplified description), where the multi-chip unit accelerator card includes multiple chip units (dies) integrated in a single package. The model training method can be applied to a model to be trained, including a converter decoding layer, such as a model composed of multiple stacked converter decoding layers. The model to be trained can be, for example, a large language model, i.e., a deep learning model trained on a large amount of text data, capable of understanding and generating natural language text. Large language models are typically based on the Transformer architecture and learn the statistical regularities and semantic information of language from massive amounts of text data through self-supervised learning, and are commonly used for tasks such as text generation, translation, question answering, and summarization. The embodiments of this disclosure do not limit the specific implementation of the model to be trained.

[0066] For example, such as Figure 2 As shown, the model training method provided in at least one embodiment of this disclosure may include steps S101 to S102.

[0067] Step S101: Train the converter decoding layer using a hybrid parallel strategy.

[0068] Step S102: During the forward and backward propagation of the converter decoding layer, the activation values ​​processed by the multi-chip unit accelerator card are distributed to multiple chip units according to the first arrangement, and the weights processed by the multi-chip unit accelerator card are distributed to multiple chip units according to the second arrangement.

[0069] In step S101, the converter decoding layer is the computational module in the Transformer architecture used for the decoding task. It achieves autoregressive or conditional generation of the input sequence by stacking multiple layers with the same structure. The converter decoding layer typically includes a normalization module, an attention module, a feedforward network module, etc., and may also include mechanisms for training optimization such as residual connections and dropout layers. It should be understood that the specific structure of the converter decoding layer in this disclosure is not limited; it is merely an example and can be referred to... Figure 1 The converter decoding layer structure is shown. A detailed description of the hybrid parallel strategy can be found above and will not be repeated here. The parallel strategy employed can be flexibly determined according to actual needs, such as sequence parallelism + tensor parallelism, sequence parallelism + tensor parallelism + pipelined parallelism, sequence parallelism + tensor parallelism + pipelined parallelism + data parallelism, etc. This disclosure does not limit the specific combination of parallel strategies in its embodiments.

[0070] In the model training method provided in at least one embodiment of this disclosure, the converter decoding layer includes a normalization module, an attention module, and a feedforward network module. An example of step S101 may include the following steps S111 to S112.

[0071] Step S111: Train the normalization module using a sequence parallel strategy.

[0072] Step S112: Train the attention module and the feedforward network module using a tensor parallel strategy.

[0073] An example of step S101 may also include steps S113 to S114.

[0074] Step S113: Perform inter-card full aggregation operation on the activation values ​​output by the normalization module.

[0075] Step S114: Perform inter-card reduction scattering operation on the activation values ​​output by the attention module and the feedforward network module.

[0076] For steps S111 to S114 above, please refer to... Figure 1 The specific details are not elaborated here.

[0077] In step S102, the first arrangement method includes dividing the data to be arranged along a preset dimension to obtain multiple data fragments, and distributing each data fragment to different chip units; the second arrangement method includes copying the data to be arranged to obtain multiple data copies, and distributing each data copy to different chip units. The first arrangement method can be the Die-SB arrangement method described above, and the second arrangement method can be the Die-BB arrangement method described above. In other words, for a single multi-chip unit accelerator card, the weights in the forward and backward propagation processes of the converter decoding layer are arranged in Die-BB type among multiple chip units, and the activation values ​​are arranged in Die-SB type among multiple chip units.

[0078] "Activation values ​​processed by the multi-chip unit accelerator card" can be intermediate activation tensors that the accelerator card locally computes or holds based on its allocated input data and model parameters during forward or backward propagation; "Weights processed by the multi-chip unit accelerator card" can be a subset of model parameters that the accelerator card locally stores and uses for computation under the current parallel configuration. For example, under a tensor parallel strategy, the complete weights are partitioned along a specific dimension, and each accelerator card holds only a portion of them; under other parallel strategies, each accelerator card may hold the complete weights.

[0079] For example, the preset dimensions for data segmentation between chip units can be set according to actual needs, typically depending on the data dimensions of the tensor and the type of operator associated with it. The Die-SB arrangement can be further defined as Die-S0B, Die-S1B, Die-S2B, etc., representing the distributed arrangement after segmentation along the 0th, 1st, or 2nd dimension of the tensor, respectively. For example, the above segmentation method can employ uniform segmentation.

[0080] For example, assuming a single multi-chip unit accelerator card integrates M chip units, for a tensor of shape [b, s, h] (where b represents the batch size, s represents the sequence length, and h represents the hidden layer size), if a Die-S0B arrangement is used, the tensor shape held by each chip unit can be [b / M, s, h]; if a Die-S1B arrangement is used, the tensor shape held by each chip unit can be [b, s / M, h]; and if a Die-S2B arrangement is used, the tensor shape held by each chip unit can be [b, s, h / M]. In sequence parallelization schemes, the first two partitioning methods are usually used, that is, partitioning is performed in the batch size dimension or the sequence length dimension. For example, if a Die-BB arrangement is used, the tensor shape held by each chip unit is [b, s, h].

[0081] The model training method provided in at least one embodiment of this disclosure may further include step S103.

[0082] Step S103: During the reverse propagation process of the converter decoding layer, the gradient processed by the multi-chip unit accelerator card is distributed to multiple chip units according to the second arrangement.

[0083] For example, in step S103, the "gradient processed by the multi-chip unit accelerator card" can be the gradient tensor calculated by the accelerator card based on its locally held weights and activation values ​​during the model backpropagation process. In other words, for a single multi-chip unit accelerator card, the gradients during the backpropagation process of the converter decoding layer are arranged in a Die-BB type among multiple chip units.

[0084] Figure 3 This is a schematic diagram of another converter decoding layer calculation and communication process provided for at least one embodiment of the present disclosure.

[0085] Figure 3 This is a specific example of steps S101 to S102 above. Figure 3 The converter decoding layer shown includes a first normalization module (Norm1), an attention module, a second normalization module (Norm2), and a feedforward network module connected in sequence.

[0086] For example, such as Figure 3 As shown, in a single Multi-Die accelerator card, the weights and gradients of each module in the converter decoding layer are arranged in Die-BB type, while the activation values ​​of the attention module and the feedforward network module are arranged in Die-SB type. Figure 3 The normalization module's activation values ​​and the inputs and outputs of the communication section (including full aggregation and reduction scattering operations) are arranged in Die-SB type. This arrangement can also be called a multi-die data parallel (Multi-Die DP) arrangement strategy. Figure 1 Compared to the Multi-Die TP layout strategy shown, the Multi-Die DP layout strategy eliminates the need for any layout conversion during the forward and backward propagation of the converter's decoding layer, while still satisfying the parallel partitioning constraints and avoiding communication redundancy. Therefore, the Multi-Die DP layout strategy significantly improves the training performance of large models on Multi-Die accelerator cards without increasing inter-card communication overhead or memory resource consumption.

[0087] In the model training method provided in at least one embodiment of this disclosure, the model to be trained may further include an embedding layer. The embedding layer uses a learnable embedding matrix to convert discrete symbols (e.g., words in the input sequence) into continuous, dense vector representations so that subsequent attention modules and feedforward network modules can process them effectively.

[0088] The embedding layer precedes the converter decoding layer. When the model includes multiple stacked converter decoding layers, the embedding layer is placed before the first converter decoding layer. Depending on the parallel strategy employed, communication operations (such as inter-card reduction scattering operations) may need to be inserted between the execution of the embedding layer computation and the execution of subsequent converter decoding layer computations to split the output of the embedding layer along the sequence dimension and distribute it to multiple multi-chip unit accelerator cards, thereby adapting to the sequence parallel layout of the normalization module in the converter decoding layer.

[0089] The model training method provided in at least one embodiment of this disclosure may further include steps S201 to S202.

[0090] Step S201: Before the forward propagation process of the embedding layer, the input tensor of the embedding layer is distributed to multiple chip cells of the multi-chip unit accelerator card according to the second arrangement.

[0091] Step S202: During the forward and backward propagation of the embedded layer, the weights and gradients processed by the multi-chip unit accelerator card are distributed to multiple chip units according to the first arrangement.

[0092] For example, in step S201, since the embedding layer is typically the first step in processing model input, the input tensor of the embedding layer can be understood as the input sequence of the entire model. This input sequence is, for example, a tensor of shape [b, s] (where b represents the batch size and s represents the sequence length), where each element is an index of a corresponding term in the vocabulary, forming a representation of the text to be processed in the discrete symbol space. The vocabulary is a set of all possible terms, where each term is assigned a unique index. In this example, for a single multi-chip accelerator card, the input tensor is arranged in a Die-BB type across multiple chip units.

[0093] For example, in step S202, for a single multi-chip unit accelerator card, the weights and gradients of the embedding layer are arranged in a Die-SB type across multiple chip units. Since the total size of the weight matrix of the embedding layer is v×h (where v represents the vocabulary size and h represents the hidden layer size), the weight size of the embedding layer increases with the vocabulary size. Storing the weights of the embedding layer using a Die-BB type would result in significant waste of GPU memory resources. Therefore, in this embodiment, the weights of the embedding layer are arranged in a Die-SB type, meaning each chip unit only holds a portion of the weights processed by the accelerator card, thereby reducing additional GPU memory usage and improving training efficiency. The gradient arrangement is similar. In this example, a tensor parallel strategy can also be used to train the embedding layer; for example, the weight matrix can be divided along the vocabulary size dimension, with each accelerator card holding a portion of the complete weights.

[0094] After calculation by the embedding layer, the activation values ​​output by the embedding layer are distributed in multiple chip cells of the multi-chip unit accelerator card according to the third arrangement. Among them, the local tensors distributed in multiple chip cells according to the third arrangement can be reconstructed into complete tensors after being processed by a preset algorithm (such as partial summation, maximum value, minimum value, etc.), and the shape of the complete tensor is the same as that of each local tensor.

[0095] For example, the third arrangement is the Die-PP arrangement described above. For a single multi-chip unit accelerator card, since the input tensor of the embedding layer uses a Die-BB arrangement while its weights use a Die-SB arrangement, the activation value calculated by the embedding layer in each chip unit only corresponds to a portion of the complete activation result. As a natural result of this calculation mode, the activation values ​​output by the embedding layer are arranged in a Die-PP manner among multiple chip units, thereby ensuring the integrity of the embedded layer output data. To obtain the complete activation tensor, a preset algorithm (e.g., an accumulation operation in this case) can be performed on the local activation values ​​from each chip unit to reconstruct the complete activation values.

[0096] After step S202, the model training method provided in at least one embodiment of this disclosure may further include step S203.

[0097] Step S203: Change the arrangement of the activation values ​​output by the embedding layer from the third arrangement to the first arrangement, and perform inter-card reduction scattering operation.

[0098] Because the local tensors arranged in a Die-PP configuration have the same shape as the complete tensor, only some elements are valid activation values, while the remaining positions are usually filled with invalid data. Therefore, a large amount of redundant information is transmitted during subsequent inter-card communication, resulting in wasted communication resources. To avoid this problem, in step S203, an arrangement conversion operation can be performed to convert the Die-PP arrangement to a Die-SB arrangement. This ensures that each chip unit holds valid data fragments, reducing the amount of data on each chip unit and significantly improving the efficiency of subsequent inter-card communication. This arrangement conversion can be achieved through tensor reconstruction and dimension repartitioning, that is, redistributing the local tensors under the Die-PP arrangement according to the Die-SB partitioning method.

[0099] Since the output of the embedding layer will serve as the input to the normalization module in the subsequent converter's decoding layer, and the training of the normalization module employs a sequence parallel strategy, an inter-card reduction scattering operation is required. This operation splits the rearranged activation values ​​and distributes them to the various multi-chip unit accelerator cards participating in the sequence parallelism. Assuming that before performing the inter-card reduction scattering operation, the activation values ​​held on a single accelerator card have a shape of [b, s, h], after performing the inter-card reduction scattering operation along the sequence length dimension, each accelerator card will retain only one slice of the activation value tensor, with a shape of [b, s / t', h], where t' represents the sequence parallelism, i.e., the number of accelerator cards participating in tensor parallel computation.

[0100] The arrangement conversion operation involved in step S203 is outside the converter decoding layer, so it is only executed once at the beginning of each batch. It does not increase linearly with the number of small batches (num_micro_batch) that a single accelerator card needs to process, which greatly reduces data movement across chip units, reduces communication overhead, and thus improves training efficiency.

[0101] Figure 4A This is a schematic flowchart illustrating an embedded layer computing and communication method provided in at least one embodiment of the present disclosure. Figure 4A This is an example of steps S201 to S203 described above. Figure 4A In the diagram, dashed boxes represent the shape and arrangement of activation values, while rounded rectangles represent the shape and arrangement of weights.

[0102] For example, such as Figure 4A As shown, the input tensor has a shape of [b, s] (where b represents the batch size and s represents the sequence length) and uses the Die-BB type. Figure 4A The embedded layer (represented as bb) is arranged among multiple chip units. Training of the embedded layer employs a tensor parallel strategy with a tensor parallelism degree of t (i.e., t accelerator cards participate in tensor parallel computation). The shape of the embedded layer weights in each accelerator card participating in tensor parallelism is [h, v / t]. The weights of the embedded layer adopt the Die-SB type (…). Figure 4A The values ​​(represented as sb) are arranged among multiple chip cells. After the embedding layer calculation is completed, the activation values ​​output by the embedding layer are expressed as Die-PP (…). Figure 4AThe activation values ​​are arranged in a pp (or pp) manner among multiple chip cells. After the arrangement conversion operation, the activation values ​​output by the embedding layer are arranged in a Die-SB manner among multiple chip cells. Next, an inter-card reduction scattering operation is performed: first, the activation values ​​(of shape [b, s, h]) on each accelerator card participating in tensor parallelism are aggregated through a reduction operation to recover the global activation values. Then, the values ​​are split according to the sequence dimension and redistributed to each accelerator card participating in sequence parallelism, thereby obtaining the local activation value shape [b, s / t', h] required for sequence parallelism, which serves as the input to the subsequent converter decoding layer. Here, t' represents the sequence parallelism, that is, the number of accelerator cards participating in sequence parallel computation. The tensor parallelism t and the sequence parallelism t' can be the same or different, depending on the actual needs.

[0103] Figure 4B This is a schematic flowchart illustrating a converter decoding layer calculation and communication method provided for at least one embodiment of the present disclosure.

[0104] Figure 4B It can be Figure 4A The example shown is a concrete implementation of "converter decoding layer + communication". When the model includes multiple stacked converter decoding layers, Figure 4B This could be an example of any of the converter's decoding layer computation and communication methods. In other words, Figure 4B The input can be connected to Figure 4A The output of the "reduction scattering" operation can also be connected to the output of another "converter decoding layer + communication" structure. Correspondingly, Figure 4B The output can also be connected to the input of another "converter decoding layer + communication" structure. Figure 4B In the diagram, dashed boxes represent the shape and arrangement of activation values, while rounded rectangles represent the shape and arrangement of weights.

[0105] For example, such as Figure 4B As shown, the training of the first normalization module adopts a sequence parallel strategy with a sequence parallelism degree of t', meaning that t' accelerator cards participate in the sequence parallel computation. For each accelerator card participating in the sequence parallel computation, the shape of the input tensor of the first normalization module (i.e., the activation value received by the first normalization module) is [b, s / t', h], and it adopts the Die-SB type ( Figure 4B The weights of the first normalization module (represented as sb) are arranged among multiple chip units. To meet computational requirements, the weights of the first normalization module are arranged in a Die-BB type with a shape of [h,]. The activation values ​​output by the first normalization module are also arranged in a Die-SB type with a shape of [b, s / t', h].

[0106] For a single accelerator card, the output of the first normalization module is the local activation value corresponding to the local sequence. That is, each accelerator card participating in sequence parallelism holds only a fragment of the complete activation value, with the shape [b, s / t', h] and arranged in a Die-SB format. However, the attention mechanism needs to model the global dependency between any two tokens, and its computation requires each accelerator card to access the complete activation value. Before feeding the activation value output by the first normalization module into the attention module, each accelerator card participating in sequence parallelism needs to perform a full aggregation operation between the cards, pooling the activation value fragments they hold and reconstructing the complete activation value locally, with the shape [b, s, h] and arranged in a Die-SB format. Since the attention module adopts a tensor parallelism strategy, its weights are split and distributed across each accelerator card. Therefore, each accelerator card participating in tensor parallelism only performs partial computation based on the local weight fragment and the complete activation value. The activation value output by the attention module is a local result based on the complete input but limited by the local weights, with the shape [b, s, h] and arranged in a Die-SB format. The second normalization module runs in sequence parallel mode and only needs to process the local sequence slices it is assigned to. Therefore, it is necessary to perform inter-card reduction scattering operation on the local results output by each accelerator card: first, the local results of each accelerator card participating in tensor parallelism are aggregated through reduction operation to recover the global activation value, and then the results are split according to the sequence dimension and redistributed to each accelerator card participating in sequence parallelism, thereby obtaining the local activation value shape [b,s / t',h] required for sequence parallelism, which still adopts the Die-SB type arrangement.

[0107] For a single accelerator card, the output of the second normalization module is the local activation value corresponding to the local sequence. That is, each accelerator card participating in sequence parallelism holds only one fragment of the complete activation value, with a shape of [b, s / t', h] and arranged in Die-SB type. Before feeding the activation value output by the second normalization module into the feedforward network module, each accelerator card participating in sequence parallelism needs to perform an inter-card full aggregation operation to gather the activation value fragments held by each card and reconstruct the complete activation value locally, with a shape of [b, s, h] and arranged in Die-SB type. Since the feedforward network module adopts a tensor parallelism strategy, its weights are split and distributed among the accelerator cards. Therefore, each accelerator card participating in tensor parallelism only performs partial computation based on the local weight fragment and the complete activation value. The activation value output by the feedforward network module is a local result based on the complete input but limited by the local weights, with a shape of [b, s, h] and arranged in Die-SB type. Next, the activation values ​​output by the feedforward network module need to undergo a reduction scattering operation before they can be input into the subsequent module (which is still a normalization module trained using a sequence parallel strategy).

[0108] like Figure 4BAs shown, for a single multi-chip unit accelerator card, the weights in the forward and backward propagation processes of the converter decoding layer are arranged in Die-BB type across multiple chip units, and the activation values ​​are arranged in Die-SB type across multiple chip units. For example, the weights of gate_proj, up_proj, and down_proj in the feedforward network module are all arranged in Die-BB type. This fully utilizes the parallel acceleration capability of the multi-chip unit accelerator card for high-dimensional matrix operations while maintaining the unchanged arrangement type of input and output activation values ​​during the calculation process. For example, the intermediate activation values ​​ctx_mul_fwd_out, ctx_silu_fwd_out, ctx_gate_fwd_out, and ctx_up_fwd_out in the feedforward network module are all arranged in Die-SB type, with a shape of [b, s, f / t], where f is the dimension of the intermediate hidden layer of the feedforward network. Under this arrangement strategy, both card-to-card converged communication and card-to-card protocol scattering communication can directly use Die-SB type as input and output without additional arrangement mode conversion operations, thereby achieving zero redundancy in card-to-card communication.

[0109] Figure 4C This is a schematic block diagram of an attention module provided in at least one embodiment of the present disclosure. Figure 4C For example, Figure 4B An example of the attention module shown.

[0110] exist Figure 4C In the diagram, dashed boxes represent the shape and arrangement of activation values, while rounded rectangles represent the shape and arrangement of weights. For example... Figure 4C As shown, the attention module includes a linear projection unit, a fusion splitting unit, a query rotational position embedding (Rope) unit, a key rotational position encoding unit, a scaled dot-product attention unit, a shape transformation unit, and an output linear projection unit.

[0111] The linear projection unit transforms the input hidden state using a linear transformation to simultaneously generate three tensors: query (Q), key (K), and value (V). Its weight shape is [(head_num+2×kv_g)×head_dim / t, h], and its output activation value shape is [b, s, (head_num+2×kv_g)×head_dim / t]. The fusion and segmentation unit segments the output of the linear projection unit into three parts: Q, K, and V. Its output activation value shape is [b×kv_g / t, s,head_dim / head_num]. The query rotation position encoding unit and the key rotation position encoding unit apply rotation position encoding to Q and K respectively to introduce relative position information. Their output activation value shapes are [b×head_num / t, s, head_dim / head_num] and [b×kv_g / t, s, h], respectively. The scaling dot product attention unit performs attention computation, and its output activation value has the shape [b×head_num / t, s, head_dim / head_num]. The shape transformation unit rearranges the dimensions of the attention output to fit the multi-head structure, and its output activation value has the shape [b, s, head_dim / t]. The output linear projection unit maps the multi-head attention result back to the original hidden dimensions, and its weights have the shape [h, head_dim / t], and its output activation value has the shape [b, s, h]. Here, b represents the batch size, s represents the sequence length, h represents the hidden layer size, head_num represents the number of query heads, kv_g represents the number of key / value heads, head_dim represents the dimension of each attention head, and t represents the tensor parallelism.

[0112] like Figure 4C As shown, for a single multi-chip unit accelerator card, the weights in the forward and backward propagation processes of the attention module both adopt the Die-BB type. Figure 4C The values ​​(represented as bb) are arranged among multiple chip units, and the activation values ​​all adopt the Die-SB type ( Figure 4C The values ​​are represented as sb) and arranged among multiple chip units. For example, the weights of the linear projection unit and the output linear projection unit are arranged in Die-BB type; the activation values ​​output by the linear projection unit, the fusion segmentation unit, the query rotation position encoding unit, the key rotation position encoding unit, the scaling dot product attention unit, the shape transformation unit, and the output linear projection unit are all arranged in Die-SB type.

[0113] Thanks to the aforementioned arrangement, the inter-card convergence communication preceding the attention module can directly use the Die-SB arrangement as input and output, avoiding data redundancy and saving on additional arrangement conversion operations. Similarly, the inter-card reduction scattering communication following the attention module can also use the Die-SB arrangement as input and output, thereby further improving communication efficiency.

[0114] In the model training method provided in at least one embodiment of this disclosure, the model to be trained may further include an output layer. The output layer is used to convert the hidden states output by the last decoding layer of the model into a probability distribution on the vocabulary, thereby supporting the execution of downstream tasks (such as text generation).

[0115] The output layer follows the converter decoding layer. When the model includes multiple stacked converter decoding layers, the output layer is placed after the last converter decoding layer. The output layer may include a normalization module and a Language Model Head (LM Head) module.

[0116] The normalization module included in the output layer is used to normalize the hidden states output by the last converter decoding layer to stabilize the training process, alleviate gradient problems, and improve the model's generalization ability. Its structure can be the same as the normalization module included in the converter decoding layer, or a different structure can be used; this disclosure does not limit this. The language model head is used to map the normalized hidden vectors to the vocabulary dimension, generating the raw prediction score (logits) for each word.

[0117] The output layer may also include a loss calculation module, which is used to calculate the loss value (e.g., cross-entropy loss, which is not limited in this embodiment) based on the original predicted score and the ground truth, in order to guide the backpropagation and parameter optimization of the model.

[0118] The model training method provided in at least one embodiment of this disclosure may further include steps S301 to S304. Steps S301 to S304 are applicable to the training of the normalization module included in the output layer.

[0119] Step S301: Train the normalization module using a sequence parallel strategy.

[0120] Step S302: During the forward and backward propagation of the normalization module, the activation values ​​processed by the multi-chip unit accelerator card are distributed to multiple chip units according to the first arrangement.

[0121] Step S303: Perform inter-card full aggregation operation on the activation value output by the normalization module to obtain the first activation value.

[0122] Step S304: Change the arrangement of the first activation values ​​from the first arrangement to the second arrangement, and input them into the language model head module.

[0123] Regarding step S301, please refer to the description of step S111 above; this embodiment will not repeat the details.

[0124] In step S302, for a single multi-chip unit accelerator card, the activation values ​​of the normalization module of the output layer are arranged in Die-SB type among multiple chip units during the forward and backward propagation processes.

[0125] In step S303, before sending the activation value output by the normalization module of the output layer into the language model head module, each accelerator card participating in the sequence parallelism performs a full aggregation operation between cards to gather the activation value fragments held by each card and reconstruct the complete activation value (i.e., the first activation value) locally.

[0126] Similar to the weight arrangement in the embedding layer, to reduce weight storage overhead, the weights of the language model head module can be arranged in a Die-SB manner. Therefore, in step S304, an arrangement conversion operation needs to be performed to change the arrangement of the first activation value from Die-SB to Die-BB to match subsequent computational requirements. Under the Die-SB arrangement, there is no need to redundantly store the complete language model head weights on each chip unit, thus significantly reducing GPU memory usage.

[0127] Correspondingly, the model training method provided in at least one embodiment of this disclosure may further include step S305.

[0128] Step S305: During the forward and backward propagation of the language model head module, the weights and gradients processed by the multi-chip unit accelerator card are distributed to multiple chip units according to the first arrangement.

[0129] For example, in step S305, for a single multi-chip unit accelerator card, the weights and gradients of the language model head module during forward and backward propagation are arranged in a Die-SB type among multiple chip units.

[0130] Figure 4D This is a schematic flowchart illustrating an output layer computation and communication method provided in at least one embodiment of the present disclosure. Figure 4D This is an example of steps S301 to S305 above.

[0131] exist Figure 4DIn the diagram, dashed boxes represent the shape and arrangement of activation values, while rounded rectangles represent the shape and arrangement of weights. To distinguish it from the normalization module in the converter's decoding layer, the normalization module in the output layer is referred to here as the third normalization module. Assume... Figure 4B This refers to the last converter decoding layer. Figure 4D The input of the "Third Normalization Module" can be connected to Figure 4B The output of the "converter decoding layer + communication" structure shown.

[0132] For example, such as Figure 4D As shown, the training of the third normalization module adopts a sequence parallel strategy with a sequence parallelism degree of t', meaning that t' accelerator cards participate in the sequence parallel computation. For each accelerator card participating in the sequence parallel computation, the shape of the input tensor of the third normalization module (i.e., the activation value received by the third normalization module) is [b, s / t', h], and it adopts the Die-SB type ( Figure 4D The weights of the third normalization module (represented as sb) are arranged among multiple chip units. To meet computational requirements, the weights of the third normalization module are arranged in a Die-BB type with a shape of [h,]. The activation values ​​output by the third normalization module are also arranged in a Die-SB type with a shape of [b, s / t', h].

[0133] For a single accelerator card, the output of the third normalization module is the local activation value corresponding to the local sequence. That is, each accelerator card participating in the sequence parallelism holds only one fragment of the complete activation value, with the shape [b, s / t', h]. Before sending the activation value output by the third normalization module into the language model head module, each accelerator card participating in the sequence parallelism performs an inter-card full aggregation operation, merging the activation value fragments it holds and reconstructing the complete activation value (i.e., the first activation value) locally, with the shape [b, s, h].

[0134] To reduce weight storage overhead, the weights in the language model head module are stored using a Die-SB layout and trained using a tensor parallel strategy. Specifically, the weights are divided into t weight slices along the vocabulary dimension and distributed to the various accelerator cards participating in tensor parallelism. Therefore, the weight shape held by each accelerator card participating in tensor parallelism is [v / t, h]. To meet computational requirements, a layout conversion operation is performed to change the layout of the first activation value from Die-SB to Die-BB, while the shape of the first activation value on each accelerator card remains [b, s, h]. Subsequently, each accelerator card participating in tensor parallelism performs language model head computation based on its local weight slices and the first activation value, obtaining local logits with a shape of [b, s, v / t], arranged in a Die-SB type. These local logits serve as input to the loss calculation module, used to calculate the loss value between the model prediction and the true label, and update the model parameters accordingly.

[0135] The model training method provided in at least one embodiment of this disclosure may further include steps S401 to S402. Steps S401 to S402 involve parameter optimization and updating during the training process, and are typically performed after forward propagation and back propagation are completed. In the forward propagation phase, the model processes the input data based on the current parameters, generates output, and calculates the loss; in the back propagation phase, the gradient of the loss with respect to the parameters of each layer is calculated using the chain rule. Based on the calculated gradients, steps S401 to S402 further optimize and update the parameters to gradually reduce the loss function value and improve model performance.

[0136] Step S401: Under the first arrangement, the optimization parameters are updated based on the optimizer state variables and the gradients obtained through the backpropagation process to obtain the first updated weights, wherein the first updated weights are distributed according to the first arrangement.

[0137] Step S402: Convert the first updated weights into the second updated weights, wherein the second updated weights are distributed according to the second arrangement.

[0138] For example, in step S401, parameter optimization updates are performed based on the optimizer state variables and gradients distributed in a Die-SB manner, resulting in updated weights distributed in a Die-SB manner. This is because optimizer calculations are all pointwise operations and do not change the data arrangement. For example, optimizer state variables may include first-order momentum and second-order variance.

[0139] For example, in step S402, for the converter decoding layer, for the updated weights (first updated weights) distributed in a Die-SB arrangement, an arrangement conversion operation can be performed to convert them from a Die-SB arrangement to a Die-BB arrangement, so that the updated weights adapt to the weight arrangement requirements of the calculation process in the next iteration (that is, the weights in the converter decoding layer calculation process need to use a Die-BB arrangement). The updated weights distributed in a Die-SB arrangement (i.e., the first updated weights) use a first precision format, and the updated weights distributed in a Die-BB arrangement (i.e., the second updated weights) use a second precision format, with the first precision being higher than the second precision. This is because the updated weights distributed in a Die-SB arrangement (as weight copies described later) directly participate in the optimizer update and are more sensitive to numerical precision, while the updated weights distributed in a Die-BB arrangement are used in the forward and backward propagation calculation processes, which can improve bandwidth efficiency while tolerating a certain loss of precision. After each iteration, the second updated weights will be automatically passed to the next iteration.

[0140] An example of step S401 may include steps S411 to S413.

[0141] Step S411: Distribute the optimizer state variables processed by the multi-chip unit accelerator card to multiple chip units according to the first arrangement.

[0142] Step S412: Perform a splitting operation and a gradient accumulation operation on the first gradient obtained through the backpropagation process to obtain the second gradient, wherein the first gradient is distributed according to the second arrangement and the second gradient is distributed according to the first arrangement.

[0143] Step S413: Perform parameter optimization and update based on the optimizer state variables and the second gradient to obtain the first updated weights.

[0144] In step S411, for a single multi-chip unit accelerator card, the optimizer state variables are arranged in a Die-SB type across multiple chip units. For example, a splitting operation can be performed on the optimizer state variables on a single multi-chip unit accelerator card to arrange them in a Die-SB type. A description of the splitting operation can be found in the ZeRO Splitting section below.

[0145] In step S412, the partitioning operation is, for example, ZeRO partitioning. ZeRO partitioning refers to using the Zero Redundancy Optimizer (ZeRO) strategy to partition the data, that is, the data that would normally need to be fully redundantly stored on each chip unit is divided into N parts according to the number of chip units in the accelerator card, and each chip unit retains only one part. For example, if the first gradient is obtained through backpropagation in a Die-BB arrangement, ZeRO partitioning can be performed on the first gradient, and the partitioning result can be directly accumulated into the pre-allocated gradient buffer (i.e., gradient accumulation operation) to obtain the second gradient in a Die-SB arrangement. Therefore, there is no need to store the first gradient in a Die-BB arrangement, effectively reducing data redundancy and memory usage. It should be noted that the above operation does not involve additional data movement and is not considered as the arrangement conversion operation described above. Here, the first gradient uses the second precision format, while the second gradient and gradient buffer both use the first precision format. The first precision is higher than the second precision. For example, the first precision is FP32, and the second precision is BF16. This disclosure does not limit this.

[0146] In step S413, based on the optimizer state variables arranged in the Die-SB type and the gradient execution parameter updates arranged in the Die-SB type, the updated weights, which also adopt the Die-SB type arrangement, can be obtained.

[0147] In the model training method provided in at least one embodiment of this disclosure, the high-precision weights, gradients, and optimizer state variables are all arranged in Die-SB type, which effectively reduces memory usage while ensuring computational correctness. Furthermore, through gradient accumulation, low-precision gradients are directly added to high-precision gradients, eliminating the need for additional storage of low-precision gradients and further saving memory.

[0148] Figure 5 This is a schematic diagram illustrating a parameter optimization update provided for at least one embodiment of the present disclosure.

[0149] For example, such as Figure 5As shown, in the Multi-Die DP layout strategy provided in at least one embodiment of this disclosure, the weights, gradients, and optimizer states of the converter decoding layer are stored as follows: the weights in forward and backward propagation are stored with BF16 precision and arranged in Die-BB type; the FP32 precision gradients, FP32 precision weight copies, and FP32 precision optimizer state variables (including first-order momentum and second-order momentum) in the optimizer part are arranged in Die-SB type. For the BF16 precision gradients obtained directly through the backward propagation process, a gradient accumulation strategy is used to directly accumulate them onto the FP32 precision gradients, eliminating the need for additional storage of the BF16 precision gradients. After obtaining the updated weights with FP32 precision through the parameter optimization update process, a layout conversion operation is performed to convert them from Die-SB layout to Die-BB layout to meet the requirements of forward and backward propagation calculations for weight layout, and they are automatically passed to the next iteration.

[0150] The following examples compare the communication overhead and memory usage of the Multi-Die DP and Multi-Die TP layout strategies provided in at least one embodiment of this disclosure.

[0151] exist Figure 4B The converter decoding layer shown, and Figure 4C The attention module shown involves seven weights arranged in a Die-BB configuration, including the weights of the first normalization module, the weights of the second normalization module, the weights of gate_proj, up_proj, and down_proj in the feedforward network module, and the weights of the linear projection unit and the output linear projection unit in the attention module. Based on the above analysis, in the parameter optimization update process provided in at least one embodiment of this disclosure, the first-precision weight copies of the above weights all require an arrangement conversion operation, changing the Die-SB arrangement to a Die-BB arrangement. However, no arrangement conversion operation is involved in the forward and backward propagation processes of the converter's decoding layer. Therefore, a total of seven arrangement conversion operations are required during the training (including forward propagation, backward propagation, and parameter optimization update) of each converter's decoding layer.

[0152] Taking the Llama2-70B model trained using a hybrid parallel strategy (e.g., sequence parallelism + tensor parallelism + pipeline parallelism) as an example, assuming the following configuration: the actual micro-batch size processed in each forward / backward propagation is 8, the global batch size processed in one iteration of the entire model is 1024, and a single pipeline stage includes 20 transformer decoding layers. Under this configuration, the number of micro-batch sizes processed in each training iteration is Globalbatch size / Micro-batch size = 1024 / 8 = 128.

[0153] In the Multi-Die DP layout strategy provided in at least one embodiment of this disclosure, for the converter decoding layer, the number of layout mode conversion operations performed by a single accelerator card in each iteration is num_reorder_per_layer × num_layers = 7 × 20 = 140. Here, num_reorder_per_layer represents the number of layout mode conversion operations in the calculation process of each converter decoding layer (7 times in the example above), and num_layers represents the number of converter decoding layers in a single pipeline stage. Outside of the converter decoding layer, the remaining part includes two layout mode conversion operations, located in... Figure 4A After the embedding layer is calculated and Figure 4D After the full aggregation operation, each operation is executed once in both forward and backward propagation, for a total of four layout transformation operations. In summary, for the entire model structure, a single accelerator card will perform 140 + 4 = 144 layout transformation operations in each iteration. In this example, the total number of layout transformation operations is independent of the batch size, significantly reducing the number of data moves. Compared to the 20480 layout transformation operations of the Multi-Die TP layout strategy, the Multi-Die DP layout strategy saves 20336 layout transformation operations, significantly reducing the performance loss caused by data moves during training for multi-chip accelerator cards.

[0154] Regarding memory usage, the Multi-Die DP layout strategy provided in at least one embodiment of this disclosure effectively avoids data redundancy caused by using the Die-BB type by splitting the FP32 precision data of the optimizer and storing it in Die-SB type. In this example, model weights, gradients, and optimizer states are stored as follows: weights in forward and backward propagation are stored in BF16 precision and arranged in Die-BB type; the FP32 precision gradients, FP32 precision weight copies, and optimizer state variables (including first-order momentum and second-order momentum) in the optimizer part are arranged in Die-SB type. For gradients, a gradient accumulation strategy is adopted, which directly accumulates at FP32 precision, eliminating the need for additional storage of BF16 precision gradients. Therefore, the memory overhead per parameter includes 4 bytes of BF16 precision weights, 4 bytes of FP32 precision weight copies (for optimizer updates), 4 bytes of FP32 precision gradients, and 8 bytes of optimizer state variables, meaning the overall memory usage is approximately 20 times the number of parameters (in bytes). For example, for a model with 70 parameters, assuming tensor parallelism of 8 and pipeline parallelism of 4, the weight and optimizer portion requires 70B × 20 / (8 × 4) = 87.5 GB of video memory on a single accelerator card. Although the Multi-Die DP layout strategy uses 87.5 – 78.75 = 8.75 GB more video memory for the weight and optimizer portion compared to the 78.75 GB used by the Multi-Die TP layout strategy, the memory savings in terms of activation value memory usage are as follows:

[0155] Memory saving = micro_batch × seq_length × hidden_size × datatype × layer_nums × Die_BB_nums = 8 × 4096 × 8192 × 2 × 20 × 2 = 20GB per stage, which means 20GB of memory is saved per stage.

[0156] Here, `micro_batch` represents the actual batch size processed in each forward / backward propagation, the same as the configuration of the Multi-Die TP layout strategy, with a value of 8; `seq_length` represents the sequence length, assumed to be 4096; `hidden_size` represents the hidden layer size, assumed to be 8192; `datatype` represents the size of the data type in bytes, here the data type is BF16, so it is assumed to be 2; `num_layers` represents the number of converter decoding layers in a single pipeline stage, the same as the configuration of the Multi-Die TP layout strategy, with a value of 20; and `Die-BB_nums` represents the number of chip units in a single multi-chip unit accelerator card, assumed to be 2. Therefore, under the same hybrid training scheme, using the Multi-Die DP layout strategy can save approximately 11.25GB of GPU memory on a single accelerator card compared to the Multi-Die TP layout strategy.

[0157] The above comparison clearly demonstrates that the model training method provided in at least one embodiment of this disclosure satisfies both the parallel partitioning constraints and the avoidance of communication redundancy, while also eliminating the need for any arrangement conversion operations in the converter's decoding layer computation process. This method significantly improves the training performance of large models on Multi-Die accelerator cards without increasing inter-card communication overhead or memory resource consumption.

[0158] It should also be noted that the execution order of the various steps of the model training method in the various embodiments of this disclosure is not limited. Although the execution process of each step has been described in a specific order above, this does not constitute a limitation on the embodiments of this disclosure. The various steps in the model training method can be executed sequentially or in parallel, which can be determined according to actual needs.

[0159] For example, compared to the above description, the model training method provided in at least one embodiment of this disclosure may include more or fewer steps, and the embodiments of this disclosure do not limit this.

[0160] Figure 6 This is a schematic block diagram of a model training apparatus provided in at least one embodiment of the present disclosure. The model training apparatus is applied to a multi-chip unit accelerator card, which includes multiple chip units integrated in a single package, and the model includes a converter decoding layer.

[0161] For example, such as Figure 6As shown, the model training apparatus 600 provided in at least one embodiment of this disclosure includes an execution module 601 and a distribution module 602. The execution module 601 is configured to train the converter decoding layer using a hybrid parallel strategy; the distribution module 602 is configured to distribute the activation values ​​processed by the multi-chip unit accelerator card to multiple chip units according to a first arrangement method and the weights processed by the multi-chip unit accelerator card to multiple chip units according to a second arrangement method during the forward and backward propagation of the converter decoding layer. The first arrangement method includes splitting the data to be arranged along a preset dimension to obtain multiple data fragments and distributing each data fragment to different chip units. The second arrangement method includes copying the data to be arranged to obtain multiple data copies and distributing each data copy to different chip units.

[0162] In some examples, the distribution module 602 is further configured to distribute the gradients processed by the multi-chip unit accelerator card to multiple chip units in a second arrangement during the back propagation of the converter decoding layer.

[0163] In some examples, the converter decoding layer includes a normalization module, an attention module, and a feedforward network module. The execution module 601 includes an execution unit and a communication unit. The execution unit is configured to train the normalization module using a sequence parallel strategy and to train the attention module and the feedforward network module using a tensor parallel strategy. The communication unit is configured to perform inter-card full aggregation operations on the activation values ​​output by the normalization module and to perform inter-card reduction scattering operations on the activation values ​​output by the attention module and the feedforward network module.

[0164] In some examples, the model also includes an embedding layer located before the converter decoding layer. The distribution module 602 is further configured to distribute the input tensors of the embedding layer to multiple chip cells of the multi-chip unit accelerator card according to a second arrangement before the forward propagation process of the embedding layer; and to distribute the weights and gradients processed by the multi-chip unit accelerator card to multiple chip cells according to a first arrangement during the forward and backward propagation processes of the embedding layer.

[0165] In some examples, the activation values ​​output by the embedding layer are distributed across multiple chip cells of the multi-chip cell accelerator card in a third arrangement. The execution module 601 also includes a conversion unit configured to convert the arrangement of the activation values ​​output by the embedding layer from the third arrangement to the first arrangement. The communication unit is configured to perform inter-card reduction scattering operations, wherein the local tensors distributed across multiple chip cells in the third arrangement can be reconstructed into complete tensors after processing by a preset algorithm, and the shape of the complete tensor is the same as that of each local tensor.

[0166] In some examples, the model also includes an output layer located after the converter decoding layer. This output layer includes a normalization module and a language model head module. The execution unit is further configured to train the normalization module using a sequential parallel strategy. The distribution module is further configured to distribute the activation values ​​processed by the multi-chip unit accelerator card to multiple chip units according to a first arrangement during the forward and backward propagation of the normalization module. The communication unit is further configured to perform an inter-card full aggregation operation on the activation values ​​output by the normalization module to obtain a first activation value. The conversion unit is further configured to convert the arrangement of the first activation value from the first arrangement to a second arrangement for input to the language model head module.

[0167] In some examples, the distribution module is further configured to distribute the weights and gradients processed by the multi-chip unit accelerator card to multiple chip units in a first arrangement during the forward and backward propagation of the language model head module.

[0168] In some examples, the execution unit is further configured to optimize and update the execution parameters based on the optimizer state variables and the gradient obtained through the backpropagation process in the first arrangement, to obtain the first updated weights, wherein the first updated weights are distributed according to the first arrangement; the transformation unit is further configured to convert the first updated weights into second updated weights, wherein the second updated weights are distributed according to the second arrangement.

[0169] In some examples, the distribution module is further configured to distribute the optimizer state variables processed by the multi-chip unit accelerator card to multiple chip units according to a first arrangement; the execution unit is further configured to: perform a splitting operation and a gradient accumulation operation on the first gradient obtained through the backpropagation process to obtain a second gradient, wherein the first gradient is distributed according to a second arrangement and the second gradient is distributed according to a first arrangement; perform parameter optimization and update based on the optimizer state variables and the second gradient to obtain the first updated weights.

[0170] In some examples, the second gradient and the first updated weights are in a first-precision format; the first gradient and the second updated weights are in a second-precision format, where the first precision is higher than the second precision.

[0171] It should be noted that the above-mentioned modules can be implemented by software, hardware, firmware or any combination thereof. For example, the execution module can be implemented as an execution circuit, and the distribution module can be implemented as a distribution circuit. The embodiments of this disclosure do not limit the specific implementation methods.

[0172] It should be understood that the model training apparatus provided in at least one embodiment of this disclosure can be used to implement the aforementioned model training method and can also achieve similar technical effects as the aforementioned model training method, which will not be elaborated here.

[0173] It should be noted that, in the embodiments of this disclosure, the model training device may include more or fewer modules or units, and the connection relationship between the various modules or units is not limited and can be determined according to actual needs. The specific configuration of each module or unit is not limited; it can be constructed from analog devices according to circuit principles, or from digital chips, or in other suitable ways.

[0174] Figure 7 This is a schematic block diagram of an electronic device provided for at least one embodiment of the present disclosure.

[0175] For example, such as Figure 7 As shown, the electronic device 700 includes at least one processor 701 and at least one memory 702. The at least one memory 702 includes one or more computer program modules. These computer program modules are stored in the memory 702 and configured to be executed by the at least one processor 701. The one or more computer program modules include instructions for performing the model training method described above. When executed by the at least one processor 701, they can perform one or more steps of the model training method provided in at least one embodiment of this disclosure. The memory 702 and the processor 701 can be interconnected via a bus system and / or other forms of connection mechanisms (not shown).

[0176] For example, processor 701 can be a central processing unit (CPU), digital signal processor (DSP), graphics processing unit (GPU), general-purpose graphics processing unit (GPGPU), artificial intelligence (AI) accelerator, or other form of processing unit with data processing and / or program execution capabilities, such as a field-programmable gate array (FPGA); for example, the central processing unit (CPU) can be an x86, ARM, or RISC-V architecture. Processor 701 can be a general-purpose processor or a special-purpose processor, capable of controlling other components in electronic device 700 to perform desired functions.

[0177] For example, memory 702 may include any combination of one or more computer program products, which may include various forms of computer-readable storage media, such as volatile memory and / or non-volatile memory. Volatile memory may include, for example, random access memory (RAM) and / or cache memory. Non-volatile memory may include, for example, read-only memory (ROM), hard disk, erasable programmable read-only memory (EPROM), portable compact disc read-only memory (CD-ROM), USB memory, flash memory, etc.

[0178] Figure 8This is a schematic block diagram of another electronic device provided for at least one embodiment of the present disclosure.

[0179] The electronic devices in at least one embodiment of this disclosure may include, but are not limited to, mobile terminals such as mobile phones, laptops, digital broadcast receivers, personal digital assistants (PDAs), tablet computers (PADs), portable multimedia players (PMPs), in-vehicle terminals (e.g., in-vehicle navigation terminals), wearable electronic devices, and fixed terminals such as digital TVs and desktop computers. Figure 8 The electronic device shown is merely an example and should not be construed as limiting the functionality and scope of the embodiments disclosed herein.

[0180] The electronic device includes at least one processor and a memory. The processor may be referred to as processing device 801 as described below, and the memory may include at least one of ROM 802, RAM 803, and storage device 808 as described below. The memory is used to store programs for performing the methods described in the various method embodiments above; the processor is configured to execute the programs stored in the memory. The processor may include a central processing unit (CPU) or other forms of processing unit having data processing capabilities and / or instruction execution capabilities, and may control other components in the electronic device to perform desired functions.

[0181] like Figure 8 As shown, the electronic device 800 may include a processing unit 801 (e.g., a central processing unit, a graphics processor, etc.), which can perform various appropriate actions and processes according to a program stored in ROM 802 or a program loaded from storage device 808 into RAM 803. RAM 803 also stores various programs and data required for the operation of the electronic device 800. The processing unit 801, ROM 802, and RAM 803 are interconnected via bus 804. Input / output (I / O) interfaces are also connected to bus 804.

[0182] Typically, the following devices can be connected to I / O interface 805: input devices 806 including, for example, touchscreens, touchpads, keyboards, mice, cameras, microphones, accelerometers, gyroscopes, etc.; output devices 807 including, for example, displays, speakers, vibrators, etc.; storage devices 808 including, for example, magnetic tapes, hard disks, etc.; and communication devices 809. Communication device 809 allows electronic device 800 to communicate wirelessly or wiredly with other devices to exchange data. Although Figure 8 An electronic device 800 with various devices is shown; however, it should be understood that it is not required to implement or possess all of the devices shown. More or fewer devices may be implemented or possessed alternatively.

[0183] In particular, according to at least one embodiment of this disclosure, the processes described above with reference to the flowcharts can be implemented as computer software programs. For example, at least one embodiment of this disclosure includes a computer program product comprising a computer program carried on a non-transitory computer-readable medium, the computer program containing program code for performing the methods shown in the flowcharts. In such an embodiment, the computer program can be downloaded and installed from a network via a communication device 809, or installed from a storage device 808, or installed from a ROM 802. When the computer program is executed by a processing device 801, it performs the functions defined in the methods of at least one embodiment of this disclosure.

[0184] It should be noted that the computer-readable medium described above in this disclosure can be a computer-readable signal medium or a computer-readable storage medium, or any combination of the two. A computer-readable storage medium can be, for example, but not limited to, an electrical, magnetic, optical, electromagnetic, infrared, or semiconductor system, apparatus, or device, or any combination thereof. More specific examples of a computer-readable storage medium may include, but are not limited to: an electrical connection having one or more wires, a portable computer disk, a hard disk, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), optical fiber, portable compact disk read-only memory (CD-ROM), optical storage device, magnetic storage device, or any suitable combination thereof. In at least one embodiment of this disclosure, a computer-readable storage medium can be any tangible medium containing or storing a program that can be used by or in conjunction with an instruction execution system, apparatus, or device. In at least one embodiment of this disclosure, a computer-readable signal medium can include a data signal propagated in baseband or as part of a carrier wave, carrying computer-readable program code. Such propagated data signals can take various forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination thereof. A computer-readable signal medium may be any computer-readable medium other than a computer-readable storage medium, which can send, propagate, or transmit a program for use by or in connection with an instruction execution system, apparatus, or device. The program code contained on the computer-readable medium can be transmitted using any suitable medium, including but not limited to: wires, optical fibers, radio frequency (RF), etc., or any suitable combination thereof.

[0185] The aforementioned computer-readable medium may be included in the aforementioned electronic device 800; or it may exist independently and not assembled into the electronic device 800.

[0186] Figure 9 This is a schematic block diagram of a non-transitory computer-readable storage medium provided for at least one embodiment of the present disclosure.

[0187] For example, such as Figure 9 As shown, a non-transitory computer-readable storage medium 900 stores computer-readable instructions 901, which, when executed by at least one processor, perform one or more steps of the model training method described above.

[0188] For example, the storage medium may include a memory card for a smartphone, a storage component for a tablet computer, a hard drive for a personal computer, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM), portable compact disc read-only memory (CD-ROM), flash memory, or any combination of the above storage media, or other suitable storage media. For example, the readable storage medium may also be... Figure 7 The memory 702 in the memory is described in the foregoing content and will not be repeated here.

[0189] Although the present disclosure has been described in detail above with general descriptions and specific embodiments, modifications or improvements can be made to the embodiments of the present disclosure, which will be obvious to those skilled in the art. Therefore, all such modifications or improvements made without departing from the spirit of the present disclosure are within the scope of protection claimed by the present disclosure.

[0190] The following points should be noted regarding this disclosure:

[0191] (1) The accompanying drawings of the embodiments of this disclosure only involve the structures involved in the embodiments of this disclosure. Other structures can be referred to the general design.

[0192] (2) For clarity, the thickness of layers or regions in the drawings used to describe embodiments of the present disclosure is enlarged or reduced, i.e., these drawings are not drawn to actual scale.

[0193] (3) Where there is no conflict, the embodiments of this disclosure and the features in the embodiments can be combined with each other to obtain new embodiments.

[0194] The above description is merely a specific embodiment of this disclosure, but the scope of protection of this disclosure is not limited thereto. The scope of protection of this disclosure should be determined by the scope of protection of the claims.

Claims

1. A model training method, characterized in that, The model training method is applied to a multi-chip unit acceleration card including a plurality of chip units integrated in a single package, the model includes a transducer decoder layer, and the model training method includes: training the transducer decoder layer using a hybrid parallel strategy; in the forward propagation and backward propagation processes of the transducer decoder layer, distributing the activation values processed by the multi-chip unit acceleration card to the plurality of chip units according to a first arrangement mode and distributing the weights processed by the multi-chip unit acceleration card to the plurality of chip units according to a second arrangement mode, wherein the first arrangement mode includes splitting the data to be arranged along a preset dimension to obtain a plurality of data shards, and distributing each data shard to a different chip unit, the second arrangement mode includes copying the data to be arranged to obtain a plurality of data copies, and distributing each data copy to a different chip unit.

2. The model training method of claim 1, wherein, The model training method further includes: in the backward propagation process of the transducer decoder layer, distributing the gradients processed by the multi-chip unit acceleration card to the plurality of chip units according to the second arrangement mode.

3. The model training method of claim 1, wherein, The transducer decoder layer includes a normalization module, an attention module, and a feedforward network module, and the training of the transducer decoder layer using a hybrid parallel strategy includes: training the normalization module using a sequence parallel strategy; training the attention module and the feedforward network module using a tensor parallel strategy.

4. The model training method of claim 3, wherein, The training of the transducer decoder layer using a hybrid parallel strategy further includes: performing an inter-card all-reduce operation on the activation values output by the normalization module; performing an inter-card all-gather operation on the activation values output by the attention module and the feedforward network module.

5. The model training method of claim 1, wherein, The model further includes an embedding layer located before the transducer decoder layer, and the model training method further includes: before the forward propagation process of the embedding layer, distributing the input tensor of the embedding layer to the plurality of chip units of the multi-chip unit acceleration card according to the second arrangement mode; in the forward propagation and backward propagation processes of the embedding layer, distributing the weights and gradients processed by the multi-chip unit acceleration card to the plurality of chip units according to the first arrangement mode.

6. The model training method of claim 5, wherein, The activation values output by the embedding layer are distributed in the plurality of chip units of the multi-chip unit acceleration card according to a third arrangement mode, and the model training method further includes: converting the arrangement mode of the activation values output by the embedding layer from the third arrangement mode to the first arrangement mode, and performing an inter-card all-reduce operation, wherein a local tensor distributed in the plurality of chip units according to the third arrangement mode can be reconstructed into a complete tensor after being processed by a preset algorithm, and the shape of the complete tensor is the same as that of each local tensor.

7. The model training method of claim 1, wherein, The model further includes an output layer located after the transducer decoder layer, the output layer includes a normalization module and a language model head module, and the model training method further includes: training the normalization module using a sequence parallel strategy; In the forward propagation and the backward propagation of the normalization module, the activation values processed by the multi-chip unit acceleration card are distributed to the plurality of chip units according to the first arrangement mode; An inter-card full-aggregation operation is performed on the activation values output by the normalization module to obtain first activation values; The arrangement mode of the first activation values is converted from the first arrangement mode to the second arrangement mode to be input to the language model head module.

8. The model training method according to claim 7, characterized in that, The model training method further includes: In the forward propagation and the backward propagation of the language model head module, the weights and the gradients processed by the multi-chip unit acceleration card are distributed to the plurality of chip units according to the first arrangement mode.

9. The model training method of claim 1, wherein, The model training method further includes: In the first arrangement mode, parameter optimization updates are performed based on the optimizer state variables and the gradients obtained through the backward propagation process to obtain first updated weights, wherein the first updated weights are distributed according to the first arrangement mode; The first updated weights are converted into second updated weights, wherein the second updated weights are distributed according to the second arrangement mode.

10. The model training method of claim 9, wherein, The model training method further includes: In the first arrangement mode, parameter optimization updates are performed based on the optimizer state variables and the gradients obtained through the backward propagation process to obtain first updated weights, wherein the first updated weights are distributed according to the first arrangement mode; The optimizer state variables processed by the multi-chip unit acceleration card are distributed to the plurality of chip units according to the first arrangement mode; A splitting operation and a gradient accumulation operation are performed on the first gradients obtained through the backward propagation process to obtain second gradients, wherein the first gradients are distributed according to the second arrangement mode, and the second gradients are distributed according to the first arrangement mode; 11. The model training method of claim 10, wherein, Parameter optimization updates are performed based on the optimizer state variables and the second gradients to obtain the first updated weights.

12. A model training apparatus, comprising: The second gradients and the first updated weights adopt a first precision format; the first gradients and the second updated weights adopt a second precision format, and the first precision is higher than the second precision. The model training device is applied to a multi-chip unit acceleration card including a plurality of chip units integrated in a single package, and the model includes a transducer decoder layer, and the model training device includes: An execution module configured to train the transducer decoder layer using a hybrid parallel strategy; A distribution module configured to distribute, in the forward propagation and the backward propagation of the transducer decoder layer, activation values processed by the multi-chip unit acceleration card to the plurality of chip units according to a first arrangement mode and weights processed by the multi-chip unit acceleration card to the plurality of chip units according to a second arrangement mode, wherein the first arrangement mode includes splitting to-be-arranged data along a preset dimension to obtain a plurality of data shards and distributing each data shard to a different chip unit, 13. An electronic device, comprising: and the second arrangement mode includes copying to-be-arranged data to obtain a plurality of data copies and distributing each data copy to a different chip unit. The electronic device includes: at least one processor; at least one memory including one or more computer program modules; The one or more computer program modules are stored in the at least one memory and configured to be executed by the at least one processor, and the one or more computer program modules are used to implement the model training method in any one of claims 1-11.

14. A non-transitory computer readable storage medium, characterized in that, The non-transitory computer readable storage medium has stored computer readable instructions, wherein the computer readable instructions are executed by at least one processor to perform the model training method in any one of claims 1-11.

Images