Hybrid expert parallel training optimization method for diffusion transformer model based on video generation

Abstract

This invention discloses a hybrid expert parallel training optimization method based on the video generation diffusion Transformer model, addressing the problems of low training efficiency, high cost, and poor stability in existing technologies. This method introduces capacity constraints during the routing phase and employs a weighted truncation and backoff strategy for excess tokens, mapping them to the target device. It constructs a sending bucket and a "target device-expert-token" buffer structure, performing two full interconnect communications. For high-speed interconnect hardware features such as PCIe and IB, a hierarchical all-to-all strategy is adopted. A communication and computation overlap mechanism is constructed, achieving overlapping of distribution, computation, and merging operations through double buffering and double-stream pipelined implementation. The MoE parallel optimization subgraph is embedded into spatial blocks, constructing an alternating temporal block structure to achieve optimization. This invention effectively reduces the communication and synchronization overhead caused by expert parallelism, significantly improving the stability and overall efficiency of parallel training.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence technology, and further relates to video generation and distributed training optimization methods. Specifically, it is a hybrid expert parallel training optimization method based on the video generation diffusion Transformer model, which can be used for efficient training of large parameter video generation models, high-resolution multi-frame video generation, and acceleration of AI model training on cloud platforms. Background Technology

[0002] Diffusion models, employing a probabilistic modeling paradigm of "forward noise addition – backward denoising," are widely adopted in image and video generation tasks due to their stable training and high-quality generation. Furthermore, by introducing the Transformer into the diffusion denoising network to form the Diffusion Transformer, the model can utilize a self-attention mechanism to globally model long-range dependencies, adapting to the requirements of cross-frame consistency and cross-spatial structural correlation in video generation, thus becoming one of the important backbones of current high-performance video generation. However, video generation tasks exhibit a significant "spatiotemporal token inflation" characteristic: the number of tokens increases multiplicatively with the number of frames F and the number of spatial patches N, causing a sharp increase in the computational and communication costs of the model during training. Against this backdrop, Hybrid Experts (MoE) expands the parameter capacity without significantly increasing the computational cost per iteration through "sparse activation," becoming a common path to improve the model's expressive power. However, MoE heavily relies on Expert Parallelism (EP) in distributed training, leading to the emergence of specialized "Expert Parallelism Communication Optimization" approaches. These include optimizations around high-throughput, low-latency GPU all-to-all kernels for MoE dispatch / combime, low-precision communication, and overlapping communication and computation pipelines; as well as multi-dimensional hybrid parallel organization methods around "data parallelism / tensor parallelism / expert parallelism" to reduce global synchronization and cross-node communication costs. DeepEP provides a dedicated communication library for MoE / EP scenarios, offering high-throughput, low-latency all-to-all (dispatch / combime) communication capabilities and supporting low-precision operators to further reduce bandwidth pressure.

[0003] In summary, the key challenge of introducing MoE into the DiT model for video generation has shifted from "whether parameters can be expanded" to "how to achieve efficient and stable expert parallel training in scenarios with inflated video tokens": it's necessary to control expert load and routing stability while mitigating the communication bottleneck of EP through communication, sharding, and parallel organization strategies, making large expert parameter scale models trainable, convergent, and scalable with limited GPU resources. Currently, the Diffusion Transformer (DiT), formed by combining the diffusion model with Transformer, has become the mainstream architecture for video generation. However, video data suffers from spatiotemporal token inflation, leading to a significant increase in computational and communication costs. While Hybrid Expert (MoE) can expand model capacity through sparse activation, it relies on expert parallelism (EP), generating a large amount of all-to-all communication, becoming a training bottleneck. Existing DiT-MoE architectures are mostly geared towards image generation and cannot directly adapt to high-resolution, multi-frame video scenarios; video generation models such as VideoGPT and Latte either have excessive computational requirements or only expand capacity through intensive methods, both suffering from limitations in GPU memory and training efficiency. General-purpose MoE training systems and communication optimization libraries (such as DeepEP) are primarily geared towards language models, lacking collaborative design with video DiT-MoE structures, routing loads, and parameter sharding. This makes them unable to address routing fluctuations and communication bottlenecks caused by spatiotemporal block alternation. Other MoE solutions focus on multimodal and speech tasks, failing to optimize for the high token density and large-scale expert parallel communication issues inherent in video generation, thus failing to meet the demands for efficient and stable training of large-parameter video generation models.

[0004] Existing technologies apply hybrid expert models to sequence recognition tasks such as speech emotion recognition, such as the "Multi-task Speech Emotion Recognition Method Based on Parallel Processing Hybrid Expert Networks" proposed by Central China Normal University. This scheme uses speech signals as the main input and combines feature extraction networks with hybrid expert structures to achieve joint modeling of multiple emotion categories and multiple tasks. Its technical solution introduces parallel processing to improve the inference and training efficiency of multi-expert models in speech recognition tasks. However, the data characteristics faced by this type of scheme differ fundamentally from those of video generation tasks: speech signals exhibit one-dimensional or low-dimensional sequence features in the temporal dimension, the number of tokens is relatively controlled, and the model primarily serves discriminative tasks. In this context, its hybrid expert parallel processing focuses more on task output fusion and feature discrimination capabilities, rather than the stability and communication efficiency of expert routing in high-dimensional spatiotemporal token scenarios. Therefore, this type of scheme does not address the expert parallel communication bottleneck caused by "spatial-temporal token inflation" in video generation diffusion models, nor does it provide parallel optimization methods for the MoE training phase.

[0005] Furthermore, Hangzhou Dianzi University proposed a "hybrid expert model communication optimization method based on expert load prediction and sequence parallelism." This scheme uses statistical modeling and predictive analysis of expert load to dynamically adjust the parallel scheduling strategy in the sequence model, thereby alleviating some expert overload problems and reducing the probability of communication conflicts. This method mainly focuses on optimizing the expert computation order and communication scheduling for model scenarios where sequence length and token distribution are relatively stable. It should be noted that this scheme is based on the assumption that the token size of the sequence model changes relatively smoothly and the computation graph structure is stable, making it suitable for language modeling or general sequence processing tasks. However, for the video generation diffusion Transformer model, its token size increases multiplicatively with the number of frames and spatial resolution, and frequent rearrangements occur between spatial and temporal blocks, resulting in highly dynamic routing distribution. The aforementioned expert load prediction and sequence parallelism strategies are difficult to directly adapt to such complex spatiotemporal structures and do not consider the large-scale all-to-all communication and communication-computation overlap problems required by the MoE feedforward sublayer in a multi-node environment. Therefore, they are difficult to directly apply to parallel training scenarios for video generation models.

[0006] In summary, while existing hybrid expert-related patent solutions have made some progress in areas such as multimodal data processing, speech recognition, multi-task learning, and sequence model communication optimization, their technical goals and application scenarios are not designed specifically for the training characteristics of video generation diffusion Transformer models. In particular, existing technologies fail to provide systematic solutions to problems commonly found in video generation tasks, such as high-dimensional spatiotemporal token inflation, training latency dominated by expert parallel communication, and significant cross-node all-to-all overhead. These limitations make it difficult to support stable and efficient training of large-parameter video generation models under limited computing power. Summary of the Invention

[0007] The purpose of this invention is to address the shortcomings of existing technologies by proposing a hybrid expert parallel training optimization method based on the video generation diffusion Transformer model. This method solves the problems of low efficiency, high cost, and poor stability in hybrid expert parallel training of existing video generation models. By introducing a sparse activation expert feedforward computation method within the limits of the model structure, and combining it with training organization strategies tailored to the characteristics of expert parallelism, this invention manages expert parameters, gradients, and optimizer states through sharding and on-demand communication, thereby reducing the system overhead caused by expert parallelism in multi-GPU, multi-node environments. This invention enables the video generation DiT-MoE model to maintain high training throughput and stable convergence characteristics even as the parameter scale continues to expand. Through reasonable routing control, expert load constraints, and parameter sharding and communication scheduling mechanisms, this invention effectively reduces the communication and synchronization overhead caused by expert parallelism, significantly improving the stability and overall efficiency of parallel training.

[0008] To achieve the above objectives, the technical solution of the present invention includes the following:

[0009] (1) Let the input video tensor be... Where B is the batch size, F is the number of frames, Z is the number of channels, and H and W are the height and width of a single frame image, respectively; a capacity constraint is introduced during the routing phase:

[0010] ,

[0011] in, For capacity coefficient, Indicates the hardware correlation coefficient. For the calculation period, This represents the computational cost of a single data block token. This represents the rounding up operation; let the total number of tokens received by the expert be n, and when n>C, a weighted truncation and rollback strategy is adopted for excess tokens;

[0012] (2) Map the connection of each token to the device where the target expert is located, and construct a sending bucket from each source device to the target device to reduce communication startup and small packet overhead; the sending buffer is arranged continuously in the three-level order of "target device-expert-token", and the offset is determined according to the prefix sum; perform two batch optimizations of the fully interconnected communication, the first to complete token distribution and the second to complete result collection;

[0013] (3) In view of the hardware characteristics of the high-speed serial computer expansion bus standard PCIe between supercomputing cluster nodes and the high-bandwidth IB between nodes, a hierarchical all-to-all strategy is adopted: firstly, the data sent to the same remote node is aggregated within the node, then coarse-grained all-to-all communication between nodes is carried out through IB, and finally, the data is distributed according to the image processor GPU within the remote node; expert segment computing is adopted on the target device; during the recycling phase, a fused weighted aggregation is performed, which combines the gating and scatter-add operations into a single write-back, in order to reduce the kernel and intermediate buffer and reduce the pressure on the video memory bandwidth.

[0014] (4) Construct a communication and computation overlap strategy: Divide the token distribution block into m blocks, and use a double buffering mechanism and dual-stream pipelined execution of communication and computation to make the distribution of the r-th block overlap with the expert computation of the (r-1)-th block and the merging operation of the (r-2)-th block to improve parallel efficiency.

[0015] (5) The expert parallel optimization subgraph is embedded into the spatial block of the video generation diffusion Transformer to construct a temporal block alternation structure: the tokens of the spatial block and the temporal block are rearranged to change the routing distribution. Before entering the MoE sublayer each time, the capacity C and splitting rules splits are recalculated and the layout is repackaged according to the token distribution of the current stage to ensure that the load balance and stable throughput are maintained under the temporal and spatial alternation. The model structure is only used as a carrier object. The parallel training optimization of the MoE sublayer can be reused in both the spatial block and the temporal block, which ultimately generates cumulative benefits to the training efficiency of the whole network.

[0016] Compared with the prior art, the present invention has the following advantages:

[0017] First, in the video generation DiT model, this invention introduces a hybrid expert structure and designs a training strategy oriented towards expert parallelism. This enables the model to continuously expand its parameter scale while reducing the single-card memory pressure and overall training cost through parallel training methods, thereby improving the training efficiency of the model in a multi-GPU, multi-node environment from a system level.

[0018] Secondly, this invention addresses the spatiotemporal modeling characteristics of video generation tasks by designing a parallel training approach for the combination of DiT and MoE. This enables the hybrid expert structure to operate stably in video generation scenarios through reasonable parallel organization and scheduling, thereby improving the controllability of expert parallel training while ensuring temporal consistency and generation quality.

[0019] Third, this invention alleviates the problems of expert load imbalance and routing collapse by introducing routing control, expert load constraints, and parameter sharding and communication scheduling strategies for the parallel characteristics of experts during the parallel training process, reduces the overhead of expert parallel communication, and improves the training throughput and stability in multi-GPU and multi-node environments, enabling the hybrid expert structure to complete training efficiently and scalably in the video generation DiT model. Attached Figure Description

[0020] Figure 1 This is a flowchart illustrating the overall implementation of the method of the present invention;

[0021] Figure 2 This is a schematic diagram of the overall architecture of the MoE-based video generation diffusion Transformer model;

[0022] Figure 3 for Figure 2 A schematic diagram of the hybrid expert MoE layer structure in the model;

[0023] Figure 4 This is a schematic diagram of the video result generated using the method of the present invention in an embodiment of the present invention. Detailed Implementation

[0024] The present invention will now be further described with reference to the accompanying drawings.

[0025] Example 1: Refer to Appendix Figure 1 This invention proposes a hybrid expert parallel training optimization method based on a video generation diffusion Transformer model. The specific implementation steps include the following:

[0026] Step 1) Let the input video tensor be... Where B is the batch size, F is the number of frames, Z is the number of channels, and H and W are the height and width of a single frame image, respectively; a capacity constraint is introduced during the routing phase:

[0027] ,

[0028] in, For capacity coefficient, Indicates the hardware correlation coefficient. For the calculation period, This represents the computational cost of a single data block token. This represents the rounding up operation; let the total number of tokens received by the expert be n, and when n>C, a weighted truncation and rollback strategy is adopted for excess tokens;

[0029] The capacity coefficient mentioned above in this embodiment The frame complexity is dynamically adjusted based on the current video tensor: when the edge density of a single frame exceeds a preset threshold, Automatically boost to increase expert token processing capacity; when frame complexity falls below a threshold, To reduce computational resource waste, the hardware correlation coefficient k is positively correlated with the memory bandwidth and CPU core count of the computing device, and the value of k is determined through performance testing during the pre-training phase. The above-mentioned weighted truncation and rollback strategy for excess tokens specifically involves retaining the top C tokens within expert e in order of sorting, and prioritizing rollback of the remaining tokens to the next-best expert of the same token. If overload still occurs, rollback to the shared backup branch to ensure gradient path continuity, suppress route collapse, and achieve stable convergence. When excess tokens roll back to the next-best expert or the shared backup branch, gradient information of the original expert is preserved through gradient weighted compensation. That is, the gradient of the rolled-back token is superimposed on the gradient update of the original expert according to a preset weight ratio to avoid gradient breakage caused by route adjustment.

[0030] Step 2) Map the connection of each token to the device where the target expert is located, and construct a sending bucket from each source device to the target device to reduce communication startup and small packet overhead; the sending buffer is arranged continuously in the three-level order of "target device-expert-token", and the offset is determined according to the prefix sum; perform two batch optimizations of the fully interconnected communication, the first to complete token distribution and the second to complete result collection.

[0031] In this embodiment, the aforementioned sending buffer is arranged continuously in a three-level order of "target device-expert-token". The offset is determined based on the prefix sum. Specifically, it is arranged continuously using a three-level index structure of "target device-expert-token", and the storage offset of each token is determined by dual calculation using a prefix hash table and a suffix offset. The prefix hash table is used to quickly locate the device segment where the target expert is located, and the suffix offset is used to accurately locate the physical address of the token within that device segment. After the sending bucket is filled with data, communication is triggered in the form of batch atomic operations. The amount of data in a single batch operation is not less than the minimum transmission unit and does not exceed the hardware cache limit.

[0032] Step 3) Targeting the hardware characteristics of the high-speed serial computer expansion bus standard PCIe between supercomputing cluster nodes and the high-bandwidth IB between nodes, a hierarchical all-to-all strategy is adopted: First, data destined for the same remote node is aggregated within the node, then coarse-grained all-to-all communication is performed between nodes via IB, and finally, data is distributed within the remote node according to the image processor (GPU); expert segment-based computation is used on the target device; during the recycling phase, a fused weighted aggregation is performed, which combines gating and scatter-add operations into a single write-back to reduce the kernel and intermediate buffers and reduce the pressure on video memory bandwidth.

[0033] In this embodiment, the hierarchical all-to-all strategy described above specifically includes the following three stages:

[0034] (3.1) Intra-node aggregation stage: Using the high-speed serial computer expansion bus standard PCIe array or NV-Link high-speed interconnect within the node, the token data sent to the same remote node is reduced and aggregated between the local GPU / graphics processor to form data blocks;

[0035] (3.2) Inter-node transmission phase: Coarse-grained all-to-all communication is performed through the high-bandwidth IB network between nodes, and only the aggregated data blocks are transmitted to reduce the number of transmissions across physical links between nodes;

[0036] (3.3) Remote distribution stage: Inside the remote node, the data blocks are again distributed to the GPU / graphics processor where the corresponding target expert is located using a PCIe array or NV-Link; wherein the aggregation granularity of the data blocks is matched with the maximum transmission unit (MTU) of the IB network to achieve high bandwidth and low latency optimization for cross-node communication.

[0037] The aforementioned expert-segmented calculation specifically involves: the receive buffer being segmented by expert, and each expert only performs local feedforward on its own segment's input.

[0038] ,

[0039] in, For the first The output tensor of each expert It is a feedforward neural network. For the first Each expert has a segmented input tensor; and the feedforward computation for each expert is completed only within the corresponding segment's memory area, without accessing the buffers of other experts across segments.

[0040] Step 4) Construct a communication and computation overlap strategy: Divide the token distribution block into m blocks, and use a double-buffering mechanism to perform dual-stream pipelined execution of communication and computation. This ensures that the dispatch of the r-th block overlaps with the expert computation of the (r-1)-th block and the merge operation of the (r-2)-th block, thereby improving parallel efficiency. In this embodiment, the double-buffering mechanism in the communication and computation overlap strategy described in this step is specifically as follows: the first buffer is used for the dispatch of the r-th block of tokens and the merge of the results of the (r-2)-th block; the second buffer is used for the expert computation of the (r-1)-th block of tokens; and when r≤2, an idle filling mechanism is used to fill in the pipeline stages, ensuring that at least two parallel tasks are executing at any given time.

[0041] Step 5) Embed the expert parallel optimization subgraph into the spatial block of the video generation diffusion Transformer to construct a temporal block alternation structure: the tokens of the spatial block and temporal block are rearranged to change the routing distribution. Before entering the MoE sublayer each time, the capacity C and splitting rules splits are recalculated, and the tokens are repackaged and laid out according to the current stage to ensure that the load balance and stable throughput are maintained under the temporal and spatial alternation. The model structure is only used as a carrier object. The parallel training optimization of the MoE sublayer can be reused in both spatial and temporal blocks, which ultimately generates cumulative benefits to the training efficiency of the entire network.

[0042] In this embodiment, the recalculation of the capacity C and splitting rules before each entry into the MoE sublayer in the above-mentioned time block alternation structure specifically includes: firstly, statistically analyzing the spatial-temporal distribution of tokens in the current stage, calculating the expected reception volume of each expert, and dynamically updating the capacity coefficient accordingly. Then, based on the current GPU memory usage and total token amount, the number of sub-blocks and the size of each block are determined, and new splitting rules (splits) are generated. After each recalculation, the token sequence is repackaged and laid out according to the new splits, and then routing and distribution are performed. Before the spatial block enters the MoE sub-layer, the intra-frame spatial distribution is used as the statistical benchmark, and before the temporal block enters the MoE sub-layer, the cross-frame temporal distribution is used as the statistical benchmark. The recalculation process structure of the two is the same and fully reused. Only the statistical dimension of the input token distribution is different, so as to maintain the load balance and stable throughput of each expert during the spatiotemporal alternation execution.

[0043] Example 2: The hybrid expert parallel training optimization method proposed in this example abstracts the MoE feedforward sub-layer into a schedulable "communication-computation composite subgraph," and constructs expert parallel topology, communication timing, and load control strategies around it. The training is organized using a "routing-distribution-expert segmented computation-collection aggregation" framework as the basic execution structure. Unlike existing schemes that only describe the MoE structure or rely solely on default all-to-all communication, this invention explicitly incorporates the communication topology, data layout, and scheduling timing in the expert parallel process into the training algorithm, ensuring stable throughput and scalability even in video generation scenarios (where tokens grow rapidly with frame rate and resolution). This architecture covers the organization of MoE sub-layer distribution, computation, and collection during forward and backward propagation in the training phase. Now, in conjunction with the appendix... Figure 1-3 The implementation process of the present invention will be described in further detail below.

[0044] Reference Figure 1After the input video is embedded and tokenized, the model is trained synchronously on GPU groups corresponding to multiple data parallel replicas. In each data parallel replica, routing and gating computation are first performed on the tokens, and the Top-k experts are selected and conditional information is broadcast. Then, expert capacity and backoff control are introduced to limit the maximum token load of a single expert and avoid load imbalance. On this basis, the tokens that need to be sent across GPUs are packed, bucketed, and offset, and expert parallel scheduling and mapping information is generated according to the distribution relationship of experts on GPUs. Then, tokens are distributed (Dispatch) through hierarchical All-to-All communication, and the tokens are sent to the GPUs where the corresponding experts are located. Each GPU performs feedforward computation on the tokens corresponding to the local experts in segments, and hides the communication latency through a dual-stream overlap mechanism of communication and computation. The computation results are then collected (Combine) through hierarchical All-to-All communication and weighted aggregation is completed. Finally, the output tokens enter the subsequent alternating spatial block and temporal block structure, and the above parallel process is repeated after token rearrangement.

[0045] Reference Figure 2 The model takes a multi-frame sequence of video as input. After processing by the embedding and positional encoding modules, diffusion time step information and positional information are injected into the feature representation. Subsequently, the features are input to a denoising backbone network composed of multiple DiT-MoE modules. The backbone network uses an alternating stacking of spatial modeling and temporal modeling modules to model the intra-frame spatial structure and inter-frame temporal dependencies of the video, respectively. At the end of the network, through linear mapping and reconstruction operations, the latent variables are restored to video pixel or latent space representations, thereby completing the generation of the video sequence. This structure can simultaneously take into account spatial consistency and temporal continuity, and is the core network architecture of this invention.

[0046] Reference Figure 3 The MoE layer starts with the input latent variables and sequentially includes a first normalization layer (LayerNorm1), a self-attention module (Attention), a second normalization layer (LayerNorm2), and a hybrid expert feedforward sublayer (MoE). The hybrid expert feedforward sublayer consists of multiple parallel expert subnetworks, each an independent feedforward neural network structure. Internally, a routing unit calculates expert selection weights based on the feature information of the input token and selects a subset of experts to participate in the computation according to a preset strategy. The outputs of the selected experts are weighted and fused, then superimposed with the residual branch to obtain the output latent variables of the module. This structural design expands the model parameter capacity within a single Transformer Block while avoiding the drastic increase in computational complexity caused by all experts participating in the computation simultaneously.

[0047] This invention addresses the issues of expert parallel communication dominance, load imbalance, and decreased throughput that arise after introducing Hybrid Experts (MoE) in the training of Video Generation Diffusion Transformer. It proposes an optimized expert parallel training method based on "routing-distribution-expert computation-collection aggregation," adapted to the spatiotemporal token organization of Video Generation Diffusion Transformer. Let the input video tensor be... Where B is the batch size, F is the number of frames, C is the number of channels, and H and W are the spatial resolutions; after Patch Embedding, the resulting token is represented as:

[0048] ,

[0049] Where N is the number of patches per frame, and D is the hidden dimension. This invention treats the MoE feedforward sub-layer as a communication computation subgraph that needs to be optimized in parallel, i.e., an expert parallel optimization subgraph; and constructs expert parallel topology, communication timing, and load control strategies around it. In each MoE feedforward sub-layer, the gated routing calculation is as follows:

[0050] ,

[0051] For each token, select a Top-k set of experts. .

[0052] The optimizations made for hybrid expert parallel training are described in detail step by step below:

[0053] Step 1. To address the issues of expert overload and amplified communication hotspots and tail latency caused by a large number of tokens and easily biased routing distribution during video generation, this invention introduces an expert capacity constraint and backoff redistribution mechanism in the gated routing stage to ensure stable training even when the number of video tokens surges. First, the upper limit of the token capacity that an expert can receive is determined based on the current batch token size and the number of experts. Then, tokens exceeding the capacity limit are truncated according to the gated weights, and backoff is prioritized to the next-best expert for the same token. If necessary, further backoff is performed to a shared backup branch to ensure gradient path continuity. The constraint effect of capacity control and backoff redistribution on parallel training efficiency in this method is reflected in: eliminating overloaded experts before communication, avoiding extreme imbalance in buckets in all-to-all communication, thereby reducing the probability of communication time being dragged down by the slowest expert and improving overall throughput.

[0054] Let the desired load balancing be Define the expert capacity limit as:

[0055] ,

[0056] in, For capacity coefficient, Indicates the hardware correlation coefficient. For the calculation period, This represents the computational cost of a single data block token. This represents the rounding up operation; let the total number of tokens received by the experts be n. When n>C, a weighted truncation and backoff strategy is adopted for excess tokens: the top C tokens within expert e are retained in order, and the remaining tokens are backoffed to the next-choice expert of the same token. If it is still overloaded, it is backoffed to the shared backup branch to ensure the continuity of the gradient path, thereby suppressing route collapse and achieving stable convergence.

[0057] Step 2. After completing routing and capacity constraints, this invention addresses cross-GPU data exchange in parallel MoE expert mode. Instead of sending token-by-token, it proposes an all-to-all packet layout method: tokens to be sent are bucketed and arranged continuously according to the order "target device → expert → token," and a sending buffer is constructed using prefixes and offsets, generating matching segment information (splits). Based on this, a two-stage batch full interconnect communication is employed: the first stage, dispatch, distributes tokens to the GPUs holding the corresponding experts; the second stage, combine, reclaims the expert outputs to the source GPUs and restores the token order. The core of this protection point is that the packet layout and the "two-stage all-to-all + unpacking" training implementation belong to parallel training optimization methods, not network structures; it can cover equivalent implementations under different communication libraries / implementations (NCCL, etc.), as long as the organizational relationship of "packing—distribution—reclaiming—restoration" is satisfied.

[0058] In this embodiment, the specific step involves mapping the connection of each token to the device where the target expert resides, and constructing a sending bucket from each source device to the target device. To reduce communication startup and small packet overhead, the sending buffer is arranged continuously in a three-layer order of target device-expert-token, with the offset determined by a prefix sum. Subsequently, two batch full interconnect communication operations are performed: the first is token distribution (dispatch), and the second is result retrieval (combine), represented as follows:

[0059]

[0060] The number of splits / splits2 is determined by the size of each bucket. This "packaging-two-stage all-to-all-unpacking" structure makes the MoE computation path controllable and schedulable at the parallel level, avoiding the passive dominance of training latency by the all-to-all approach.

[0061] Step 3. Addressing the cluster characteristics of high-speed interconnects within nodes and inter-node IB, a hierarchical all-to-all approach (intra-node aggregation → inter-node exchange → intra-node distribution) is adopted to reduce cross-node small packet and tail latency. The target GPU performs local feedforward computation according to expert segmentation, and at the recycling end, gated weighting and aggregated write-back are fused to reduce additional kernel and intermediate buffers and lower memory bandwidth pressure. This combination of "hierarchical communication + segmented computation + fused write-back" is a key point in parallel training optimization. It includes: first, aggregating data destined for the same remote node within the node; then performing coarse-grained all-to-all inter-node communication via IB; and finally, distributing data within the remote node according to GPU. This strategy reduces the number of cross-node small packets and lowers tail latency, making cross-node bandwidth more efficient for large-block transmission; it also allows for secondary distribution within the node using higher bandwidth links. On the target device, expert segmented computation is used to reduce invalid memory access: the receive buffer is segmented by expert, and expert e only performs local feedforward on its own segment input.

[0062]

[0063] During the recycling phase, a fused weighted aggregation is performed at the receiving end, merging gated addition and scatter-add into a single write-back, reducing additional kernel and intermediate buffers and lowering memory bandwidth pressure.

[0064] Step 4. To reduce communication latency, an all-to-all communication timing and computation overlap strategy is designed: the token distribution is divided into m blocks, and double buffering / dual streams are used to achieve pipeline overlap of dispatch, expert computation, and combination, thereby hiding some of the all-to-all latency in the computation to hide some communication latency; at the same time, when the token is rearranged due to the alternation of spatial blocks and temporal blocks, the routing statistics, capacity, and splits are dynamically updated and the packing layout is rebuilt to ensure that load balancing and stable throughput are maintained under the alternation of spatial and temporal blocks.

[0065] Step 5. Embed the expert parallel optimization subgraph into the spatial-temporal block alternation structure of the video generation diffusion Transformer: The token rearrangement of spatial and temporal blocks will change the routing distribution. Therefore, before entering the MoE sublayer each time, the capacity C and splits are recalculated, and the layout is repackaged according to the token distribution of the current stage to ensure that load balancing and stable throughput can still be maintained under the alternation of space and time. The model structure only serves as a carrier. The parallel training optimization of the MoE sublayer can be reused in both spatial and temporal blocks, thereby generating cumulative benefits to the training efficiency of the entire network.

[0066] The effects of the present invention will be further explained below with reference to simulation experiments.

[0067] 1. Experiment Content:

[0068] The experimental training and testing of this invention were completed in a distributed environment at the High-Performance Computing Center of Xi'an University of Electronic Science and Technology. Four computing nodes were used for collaborative training, each with eight NVIDIA RTX 6000 Ada (48GB) GPUs, totaling 32 GPUs. The GPUs within each node were interconnected via PCIe 5.0, and cross-node communication was achieved through InfiniBand (IB). This hardware topology matches the hierarchical all-to-all parallel strategy of this invention, which involves "intra-node aggregation—inter-node IB full interconnection—intra-node distribution," and is used to objectively reflect the applicability and benefits of expert parallel training optimization in a real cluster environment.

[0069] 2. Analysis of Experimental Results:

[0070] like Figure 4 As shown, representative video frame sequences generated under different datasets and video categories are presented, including various content types such as faces (FFS dataset), natural scenes (SkyTimelapse dataset), and motion behavior (UCF101 dataset). It can be seen that the generated videos exhibit good temporal continuity in frame transitions, maintain high consistency and clarity in content structure, and do not show obvious temporal jitter or structural collapse. This result intuitively demonstrates the improved ability of this invention to model complex spatiotemporal distributions in video generation tasks, verifying the effectiveness of the proposed technical solution.

[0071] While ensuring efficient and stable training of MoE in the video diffusion Transformer, this invention outperforms the comparison methods in terms of generation quality on multiple datasets. FVD results are shown in Table 1:

[0072] Table 1: Comparison of generation quality of the present invention and existing methods on multiple datasets

[0073]

[0074] As shown in the table, this invention achieves lower FVD on datasets such as FaceForensics, UCF101, and Taichi-HD, and has a significant advantage over VideoGPT, indicating that this invention performs better in terms of temporal consistency and overall distribution approximation. It is important to emphasize that this performance improvement does not solely rely on changes in model structure, but is built upon the expert parallel training optimization strategy proposed in this invention. This strategy enables the training process of MoE in the video diffusion Transformer to be stably executed and converged in a 32-card multi-node environment, thereby fully releasing the modeling capabilities of the hybrid expert structure.

[0075] In existing technologies, after introducing MoE into the diffusion Transformer, the training bottleneck is often dominated by expert parallel communication: unbalanced routing leads to overload of some experts, differences in bucket size in all-to-all cause tail latency, and training throughput decreases with scale, especially in video tasks where tokens grow linearly with the number of frames, further amplifying communication and scheduling overhead. This invention optimizes the system by using an expert parallel training framework of "routing-distribution-expert computation-collection aggregation" and treating the MoE feedforward sub-layer as a communication-computation composite subgraph: in the routing stage, capacity constraints and backoff redistribution are introduced to reduce expert overload and hotspot probability; in the communication stage, an all-to-all-oriented packing layout and two-stage dispatch / combine batch communication are adopted to reduce small packet and startup overhead; in multi-node scenarios, a layered all-to-all (intra-node aggregation / distribution + inter-node IB exchange) is used to reduce the number of cross-node small packets and tail latency; in the execution stage, dual-stream and dual-buffered pipeline overlap is used to hide some communication waiting in expert computation. Compared to existing implementations that rely solely on the default all-to-all approach or simply perform conventional DP / FSDP, this invention is better suited for stably scaling expert parallel training in multi-node, multi-GPU environments, significantly improving training throughput and resource utilization, thereby reducing the engineering threshold and time cost of training large-scale video generation models.

[0076] The above simulation analysis proves the correctness and effectiveness of the method proposed in this invention.

[0077] This invention addresses the engineering bottleneck of rapidly increasing token size in video generation diffusion Transformers under high resolution and multi-frame input, providing a training paradigm that can be directly deployed on multi-GPU / multi-node clusters. By introducing load control (capacity constraints and backoff), all-to-all-oriented packet layout, hierarchical communication topology (intra-node aggregation / diffusion + inter-node IB exchange), and communication-computation overlap scheduling into the MoE sublayer, it significantly alleviates the communication-dominated and tail-latency problems in expert parallel training, improves throughput and resource utilization, and makes the video generation model more convergent, scalable, and reproducible under limited cluster resources. This paradigm can be integrated into existing distributed training frameworks as an "optimization component for video generation MoE training systems," possessing clear engineering application value and potential for widespread adoption. This invention belongs to the training system / parallel implementation layer improvement and does not protect the diffusion modeling formula or the generation algorithm itself. Furthermore, this invention can directly serve the training deployment needs of various types of video generation and content production, such as film and television special effects generation, virtual character / digital human animation synthesis, motion demonstration and instructional video generation, advertising creativity, and automatic short video generation. The aforementioned applications typically require training iterations at high resolutions and long frame sequences. While the MoE structure possesses stronger expressive power, its training efficiency heavily relies on expert parallel communication and scheduling. The hierarchical all-to-all and pipelined overlap strategy provided by this invention better adapts to the communication characteristics of practical supercomputing platforms (high-speed interconnects within nodes, inter-node IB), reducing cross-node communication overhead and improving end-to-end training efficiency. This shortens the model training cycle, reduces computing costs, and increases deployment iteration speed. Furthermore, this invention is also applicable to the development process of video generation models that require continuous expansion of data and training scale, providing a technical path for universities and enterprises to build scalable and deployable video generation systems in multi-node GPU environments.

[0078] The parts of this invention not described in detail are common knowledge to those skilled in the art. The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Obviously, those skilled in the art, after understanding the content and principles of the present invention, may make various modifications and changes in form and details without departing from the principles and structure of the present invention. For example: 1) Equivalent alternative to the expert parallel communication topology: The dispatch / combine two-stage fully interconnected communication of the present invention can be replaced by other equivalent cross-GPU exchange methods, including multi-level communication with "first local aggregation and then cross-node exchange", establishing communication groups according to spatial blocks and time blocks respectively, and adopting different communication granularities and grouping strategies for different token sources (spatial tokens and time tokens). In summary, after the alternation of spatial blocks and time blocks causes token rearrangement, the distribution of "token → expert GPU" and the recycling of "expert output → source GPU" are still completed, and the number of cross-node small packets and tail latency are reduced. These can all be used as alternative solutions to achieve throughput improvement and stability in video generation training. 2) Equivalent alternatives to routing load balancing strategies: The capacity cap and fallback redistribution strategies of this invention can be replaced by other load balancing methods for video token distribution, including setting capacity caps by frame grouping, and separately calculating and limiting load by spatial and temporal blocks. In summary, any method that can suppress the phenomenon of all-to-all trailing caused by overload of a few experts in video scenarios and avoid training instability caused by routing collapse can be considered an equivalent alternative. 3) Equivalent alternatives to packing layout and bucketing methods: The packing layout of this invention can be replaced by an organization method oriented towards the spatiotemporal locality of video, including layered packing for target GPUs by spatial location blocks and target experts, and establishing continuous buffers by spatial and temporal blocks to reduce non-contiguous memory access and small packet communication after token rearrangement. In summary, methods that batch and contiguously distribute and recycle video tokens during the distribution / recycling phase and reduce communication startup overhead and memory bandwidth pressure can be considered an equivalent alternative to this invention. 4) Equivalent alternative to communication-computation overlap scheduling: The dual-stream, dual-buffered block pipeline of the present invention can be replaced by other overlap mechanisms that can hide communication waiting in video training, including setting block size and pipeline depth according to spatial block stage / temporal block stage, prioritizing cross-node communication overlap in temporal blocks, prioritizing intra-node aggregation and expert computation in spatial blocks, etc. In the alternating training of spatial blocks and temporal blocks, the time interleaving of communication and expert computation is continuously realized, thereby improving end-to-end throughput and maintaining convergence stability, which can be regarded as an equivalent alternative to the present invention.The aforementioned alternatives all use "spatiotemporal token organization and alternating modeling of video generation diffusion Transformer" as the application constraint. Without changing the core idea of ​​expert parallel training optimization in this invention, they can still achieve the inventive objectives of reducing communication dominance, alleviating load imbalance, improving training efficiency, and achieving stable convergence. However, these modifications and changes based on the ideas of this invention are still within the scope of protection of the claims of this invention.

Claims

1. A hybrid expert parallel training optimization method for generating a diffusion Transformer model based on a video, characterized in that, The implementation steps are as follows: (1) Let the input video tensor be... Where B is the batch size, F is the number of frames, Z is the number of channels, and H and W are the height and width of a single frame image, respectively; a capacity constraint is introduced during the routing phase: ， in, For capacity coefficient, Indicates the hardware correlation coefficient. For the calculation period, This represents the computational cost of a single data block token. This represents the rounding up operation; let the total number of tokens received by the expert be n, and when n>C, a weighted truncation and rollback strategy is adopted for excess tokens; (2) Map the connection of each token to the device where the target expert is located, and construct a sending bucket from each source device to the target device to reduce communication startup and small packet overhead; the sending buffer is arranged continuously in the three-level order of "target device-expert-token", and the offset is determined according to the prefix sum; perform two batch optimizations of the fully interconnected communication, the first to complete token distribution and the second to complete result collection; (3) In view of the hardware characteristics of the high-speed serial computer expansion bus standard PCIe between supercomputing cluster nodes and the high-bandwidth IB between nodes, a hierarchical all-to-all strategy is adopted: firstly, the data sent to the same remote node is aggregated within the node, then coarse-grained all-to-all communication between nodes is carried out through IB, and finally, the data is distributed according to the image processor GPU within the remote node; expert segment computing is adopted on the target device; during the recycling phase, a fused weighted aggregation is performed, which combines the gating and scatter-add operations into a single write-back, in order to reduce the kernel and intermediate buffer and reduce the pressure on the video memory bandwidth. (4) Construct a communication and computation overlap strategy: Divide the token distribution block into m blocks, and use a double buffering mechanism and dual-stream pipelined execution of communication and computation to make the distribution of the r-th block overlap with the expert computation of the (r-1)-th block and the merging operation of the (r-2)-th block to improve parallel efficiency. (5) The expert parallel optimization subgraph is embedded into the spatial block of the video generation diffusion Transformer to construct a temporal block alternation structure: the tokens of the spatial block and the temporal block are rearranged to change the routing distribution. Before entering the MoE sublayer each time, the capacity C and splitting rules splits are recalculated and the layout is repackaged according to the token distribution of the current stage to ensure that the load balance and stable throughput are maintained under the temporal and spatial alternation. The model structure is only used as a carrier object. The parallel training optimization of the MoE sublayer can be reused in both the spatial block and the temporal block, which ultimately generates cumulative benefits to the training efficiency of the whole network.

2. The method according to claim 1, characterized in that: The capacity coefficient mentioned in step (1) The frame complexity is dynamically adjusted based on the current video tensor: when the edge density of a single frame exceeds a preset threshold, Automatic promotion to increase the token processing capacity of experts; When the frame complexity is below the threshold To reduce the waste of computing resources; the hardware correlation coefficient k is positively correlated with the video memory bandwidth and the number of CPU cores of the computing device, and the value of k is determined by performance testing during the pre-training phase.

3. The method according to claim 1, characterized in that: The weighted truncation and backoff strategy for excess tokens described in step (1) involves retaining the top C tokens within expert e in order of sorting, and backoffing the remaining tokens to the next best expert of the same token. If the load is still too high, the tokens are backoffed to the shared backup branch to ensure the continuity of the gradient path, suppress route collapse and achieve stable convergence.

4. The method according to claim 3, characterized in that: When the excess token is rolled back to the secondary expert or the shared backup branch, the gradient information of the original expert is preserved through gradient weighted compensation. That is, the gradient of the rolled-back token is superimposed on the gradient update of the original expert according to a preset weight ratio to avoid gradient breakage caused by routing adjustment.

5. The method according to claim 1, characterized in that: In step (2), the sending buffer is arranged continuously in a three-level order of "target device-expert-token". The offset is determined according to the prefix sum. Specifically, it is arranged continuously in a three-level index structure of "target device-expert-token". The storage offset of each token is determined by double calculation using the prefix hash table and the suffix offset. The prefix hash table is used to quickly locate the device segment where the target expert is located, and the suffix offset is used to accurately locate the physical address of the token in the device segment. After the data is filled, the sending bucket triggers communication in the form of batch atomic operations. The amount of data in a single batch operation is not less than the minimum transmission unit and does not exceed the upper limit of the hardware cache.

6. The method according to claim 1, characterized in that: The hierarchical all-to-all strategy described in step (3) specifically includes the following three stages: (3.1) Intra-node aggregation stage: Using the high-speed serial computer expansion bus standard PCIe array or NV-Link high-speed interconnect within the node, the token data sent to the same remote node is reduced and aggregated between the local GPU / graphics processor to form data blocks; (3.2) Inter-node transmission phase: Coarse-grained all-to-all communication is performed through the high-bandwidth IB network between nodes, and only the aggregated data blocks are transmitted to reduce the number of transmissions across physical links between nodes; (3.3) Remote distribution stage: Inside the remote node, the data blocks are again distributed to the GPU / graphics processor where the corresponding target expert is located using a PCIe array or NV-Link; wherein the aggregation granularity of the data blocks is matched with the maximum transmission unit (MTU) of the IB network to achieve high bandwidth and low latency optimization for cross-node communication.

7. The method according to claim 1, characterized in that: The expert segmented calculation described in step (3) specifically involves: the receiving buffer being segmented by expert, and each expert performing local feedforward only on the input of their current segment. ， in, For the first The output tensor of each expert It is a feedforward neural network. For the first Each expert has a segmented input tensor; and the feedforward computation for each expert is completed only within the corresponding segment's memory area, without accessing the buffers of other experts across segments.

8. The method according to claim 1, characterized in that: In the communication and computation overlap strategy described in step (4), the double buffer mechanism is as follows: the first buffer is used for the dispatch of the r-th token and the combination of the results of the (r-2)-th block; the second buffer is used for expert computation of the (r-1)-th token; and when r≤2, the idle filling mechanism is used to fill the pipeline stage, so that at any time there are at least two parallel tasks being executed.

9. The method according to claim 1, characterized in that: In step (5), the recalculation of the capacity C and splitting rules before each entry into the MoE sublayer in the time block alternation structure specifically includes: First, we statistically analyze the spatial-temporal distribution of tokens at the current stage, calculate the expected reception volume for each expert, and then dynamically update the capacity coefficient accordingly. Then, based on the current GPU memory usage and total token amount, the number of sub-blocks and the size of each block are determined, and new splitting rules (splits) are generated. After each recalculation, the token sequence is repackaged and laid out according to the new splits, and then routing and distribution are performed. Before the spatial block enters the MoE sub-layer, the intra-frame spatial distribution is used as the statistical benchmark, and before the temporal block enters the MoE sub-layer, the cross-frame temporal distribution is used as the statistical benchmark. The recalculation process structure of the two is the same and fully reused. Only the statistical dimension of the input token distribution is different, so as to maintain the load balance and stable throughput of each expert during the spatiotemporal alternation execution.

Images