AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
Yucheng Guo * 1 Yongjian Guo * 2 1 Zhong Guan 3 1 Haoran Sun 4 1 Wen Huang 2 1 Wanting Xu 1 Jing Long 4 1 Shuai Di 1 Junwu Xiong 1
arXiv:2605.17923v1 [cs.DC] 18 May 2026
Abstract
et al., 2023) must support highly heterogeneous data shapes, ranging from single-frame static images to long-duration, high-frame-rate complex video sequences. This diversity in data dimensionality introduces severe Bucket Management challenges. In standard distributed training frameworks, the synchronization of model parameters relies on the AllReduce (Patarasuk & Yuan, 2009) collective communication operation at the end of each step. Consequently, the global computational efficiency is strictly constrained by the “Long-tail Effect” (Constantinides et al., 2021; Guan et al., 2026), where the global step latency is stochastically determined by the slowest GPU in the cluster:
In video generation models, particularly world models, training large-scale video diffusion Transformers (such as DiT and MMDiT) poses significant computational challenges due to the extreme variance in sequence lengths within mixed-mode datasets. Existing bucket-based data loading strategies typically rely on ”equal token length” constraints. This approach fails to account for the quadratic complexity of self-attention mechanisms, leading to severe load imbalance and underutilization of GPU resources. This paper proposes AdaptiveLoad, an integrated optimization framework consisting of two core components: (1) A dual-constraint adaptive load balancing system, which eliminates long-sequence bottlenecks by simultaneously limiting memory consumption and computational load (B × S p ≤ Mcomp ); (2) A fused LayerNorm-Modulate CUDA kernel, which utilizes a D-tile coalesced reduction strategy to increase throughput and alleviate memory pressure. Experimental results on the Wan 2.1 world model demonstrate that our method reduces the computational imbalance rate from 39% to 18.9%, improves peak VRAM utilization efficiency by 22.7%, and achieves an overall training throughput increase of 27.2%.
Tsync =
max
i∈{1,...,N }
{Ti }
(1)
where Ti denotes the execution time of the i-th GPU. However, the “Equal-token” bucket strategy (Diao et al., 2025) commonly adopted in current industrial pipelines severely overlooks the non-linear computational characteristics of the core self-attention mechanism. While maintaining a constant token count per batch provides a firstorder approximation of memory usage, it fails to account for the quadratic complexity scaling of Transformers. Experimental observations in our study indicate that the correlation between training step latency and the total number of tokens is merely R ≈ 0.35, whereas it exhibits an extremely strong correlation (R ≈ 0.92) with the square of the sequence length. This systematic bias in load estimation causes long-sequence buckets to become significant “computational stragglers” (Yan et al., 2026). This imbalance forces GPUs processing long-sequence data into a state of chronic computational overload, while other nodes face idle wait times (synchronization bubbles) of several seconds, resulting in a substantial waste of high-value computing power in large-scale GPU clusters. This gap between linear token counting and quadratic computational reality motivates our first core contribution: a dual-constraint scheduling logic that reconciles memory safety with computational parity.
1. Introduction The paradigm of generative artificial intelligence is undergoing a profound transformation from static image synthesis to high-fidelity dynamic video generation, an evolution primarily driven by the exceptional scalability of Diffusion Transformer (DiTs) architectures (Peebles & Xie, 2023). Unlike image models that process fixed resolutions, modern video generators (Skorokhodov et al., 2022; Khachatryan * Equal contribution 1 JDT AI Infra, Beijing, China 2 Tsinghua University, Beijing, China 3 Tianjin University, Tianjin, China 4 Peking University, Beijing, China. Correspondence to: Junwu Xiong <[email protected]>.
Parallel to the scheduling bottleneck, the architectural evolution towards Multimodal DiT (MMDiT) modules (Reuss et al., 2024)—exemplified by SD3 (Esser et al., 2024) and Wan 2.1 (Wan et al., 2025)—has introduced new lowlevel performance hurdles. In these architectures, Adaptive
Preprint. May 19, 2026.
1
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
Layer Normalization (AdaLN) serves as the critical bridge connecting timestep embeddings with multimodal feature flows (Zhang et al., 2024), being invoked hundreds of times per training iteration (Team et al., 2025). Although memoryefficient attention kernels like FlashAttention (Dao et al., 2022; Dao, 2023) have significantly mitigated the core attention bottleneck, peripheral auxiliary operators have emerged as the new primary performance ceiling due to the “Memory Wall” (Wulf & McKee, 1995) problem. Standard AdaLN implementations typically rely on a sequence of discrete CUDA kernel calls. In the forward pass, this leads to frequent, redundant round-trips of intermediate tensors between the GPU’s Streaming Multiprocessors (SMs) and the High Bandwidth Memory (HBM). This memory-bound behavior is further exacerbated during backpropagation. The gradient reduction operations for modulation parameters (scale and shift) are hindered by inefficient strided memory access patterns. Because the hidden dimension D is the contiguous dimension in memory, accumulating gradients across the sequence dimension N fails to trigger hardware-level coalesced access (Harris et al., 2007). This leads to effective memory bandwidth utilization far below the theoretical peak of modern HBM3/HBM3e systems. The fundamental inefficiency of these “auxiliary” operators in the MMDiT backbone necessitates a hardwareaware redesign, shifting the paradigm from discrete operator execution to fused, bandwidth-optimized kernels.
Figure 1. Comparison of Naive Access and D-tile Coalesced Access Patterns
• We provide a rigorous quantification of the load mismatch phenomenon in video diffusion training and propose a dual-constraint adaptive bucket scheduling algorithm. This approach effectively eliminates the synchronization stragglers caused by long-sequence quadratic complexity, reducing the load imbalance rate by up to 50%. • We design and implement an end-to-end training bottleneck quantization and feedback closed-loop system. This shifts distributed training configurations from being heuristic-driven to data-driven, utilizing automated cost-model fitting via our Shape Benchmark to optimize cluster-wide resource utilization.
To address these multi-level challenges, this paper proposes a full-stack optimization framework named AdaptiveLoad. Our approach synergizes data-driven load balancing with hardware-aware operator fusion to reclaim lost training efficiency. First, we transcend the traditional linear load assumption by introducing a dual-constraint adaptive load balancing strategy. This mechanism dynamically optimizes the batch size of each bucket by simultaneously enforcing a linear memory capacity boundary and a polynomial computational complexity boundary. Second, we design an automated shape benchmarking and cost-fitting system. By capturing execution traces in a live distributed environment and utilizing non-linear regression, the system achieves precise modeling of computational costs, effectively eliminating the sub-optimality inherent in manual empirical tuning. Finally, addressing the AdaLN efficiency gap, we developed a Fused AdaLN CUDA kernel and an innovative D-tile coalesced reduction strategy, as show in Figure 1. This innovation fundamentally restructures memory access patterns—transforming strided access into coalesced HBM transactions—thereby significantly enhancing the arithmetic intensity and maximizing effective bandwidth utilization during gradient reduction.
• We develop a high-performance Fused AdaLN CUDA kernel specifically optimized for MMDiT backbones. By introducing D-tile coalesced reduction, we resolve the long-standing strided access bottleneck in backpropagation, increasing training throughput by over 25% while simultaneously reducing the activation memory footprint, thereby providing a robust foundation for scaling ultra-long sequence video generation models.
2. Related Work 2.1. Memory Optimization Techniques In large-scale model training, memory capacity remains the core bottleneck restricting the scaling of sequence lengths. Existing optimization research primarily focuses on three dimensions: first, memory-efficient attention operators (such as the FlashAttention series (Dao et al., 2022; Dao, 2023; Shah et al., 2024)) effectively reduce the spatial complexity of self-attention from O(S 2 ) to O(S) through tiling and operator restructuring; second, gradient and parameter sharding techniques (such as the ZeRO series (Rajbhandari et al., 2020)) eliminate redundant model state occu-
The main contributions of this paper are summarized as follows:
2
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
pancy by leveraging distributed storage; and finally, activation checkpointing and offloading strategies (such as DeepSpeed-Ulysses (Jacobs et al., 2023)) alleviate peak memory pressure by trading computation for space or utilizing CPU memory as a secondary cache. However, most of these general-purpose techniques are designed for fixeddimension sequences and fail to fully account for the dynamic shape characteristics of multimodal data in video generation models. Particularly in complex scenarios involving mixed training of video and images, fixed-stride optimization strategies struggle to cope with drastic fluctuations in memory demand, leading to constrained computational efficiency for long-sequence samples (Han et al., 2026; Liu et al., 2025).
such as LayerNorm—is relatively increasing as attention efficiency improves, gradually becoming a new performance bottleneck. Existing fusion schemes typically overlook the specific strided memory access patterns inherent in AdaLN during backpropagation. This non-contiguous access behavior severely hinders the GPU’s coalescing properties; especially when processing ultra-long sequences, low bandwidth utilization in reduction operations has become a critical obstacle preventing further increases in training throughput.
2.2. Load Balancing in Distributed Training
The training framework proposed in this paper addresses the dual challenges of computational heterogeneity and memory constraints inherent in large-scale video diffusion Transformers (e.g., DiT, MMDiT), as shown in Figure 2. We have constructed a full-stack optimization system spanning from high-level distributed scheduling to low-level kernel execution.
3. AdaptiveLoad: Dual-Constraint Load Balancing and Fused Modulated Kernels 3.1. Overview
Load balancing in distributed deep learning has traditionally focused on uniform batch processing in data parallelism. For Transformer architectures, the industry commonly adopts a ”constant token budget” heuristic strategy, which predicts computational load by constraining B × S to a constant. However, as sequence lengths S in video generation tasks scale to the order of 105 , the quadratic complexity O(S 2 ) of the self-attention mechanism renders this linear-assumptionbased heuristic completely ineffective. While existing research such as Megatron-LM (Shoeybi et al., 2019) proposes sophisticated 3D parallelism partitioning, or addresses data heterogeneity through dynamic task scheduling (e.g., ByteScale (Ge et al., 2025)), these solutions mostly focus on a single linguistic modality. When dealing with world models like Wan 2.1 (Wan et al., 2025), the significant differences in computational patterns between multi-modal encoding and language modeling tasks prevent existing schedulers from effectively resolving load mismatches between modalities, resulting in severe GPU idle waste at global synchronization points (Barriers). To mitigate this, ByteScale (Ge et al., 2025) introduces a flexible framework to handle the mismatch between data heterogeneity (Huang et al., 2024; Zhang et al., 2025) and static mesh. Similarly, MegaScaleData (Zhao et al., 2026) addresses the workload imbalance specifically in multisource foundation model training by optimizing the distributed dataloader architecture.
The core of this approach is the Dual-Constraint Adaptive Load Balancing strategy, which moves away from traditional fixed-token allocation in favor of an empirical complexity power-law distribution. This mitigates the ”longtail” latency issues caused by O(S 2 ) complexity in longsequence video data buckets. To ensure this strategy aligns with actual hardware performance, we employ a data-driven Shape Benchmarking workflow that fits a parameterized overhead model with real-world synchronization telemetry, enabling precise calibration of computational boundaries. Building upon this macro-system equilibrium, we further implement Fused AdaLN CUDA kernels to optimize the micro-architectural execution logic of the MMDiT backbone. By fusing normalization and modulation operations and introducing a D-tile Coalesced Reduction strategy for gradient computation, the system significantly alleviates HBM bottlenecks and eliminates redundant memory access. Together, these components form a closed-loop optimization system that dynamically adapts to diverse input morphologies while maintaining near-peak hardware utilization across distributed GPU clusters.
2.3. Hardware-Aware Operator Fusion 3.2. Dual-Constraint Adaptive Load Balancing & Cost Fitting
As a key technology to break the ”Memory Wall” (Gholami et al., 2024) bottleneck, operator fusion significantly enhances arithmetic intensity by reducing redundant reads and writes of intermediate results in the GPU’s HBM. Representative works include FlashAttention-2 (Dao, 2023) for attention mechanisms and FusedSoftmaxCrossEntropy (Ge et al., 2025) for loss functions. Although these solutions have achieved success in core operators, the overhead of ”auxiliary” operators within actual Transformer Blocks—
To address the computational synchronization bottlenecks faced by video diffusion Transformers when processing heterogeneous data, we implement a dual-boundary strategy to determine the batch size B for any given bucket shape defined by sequence length S. Traditional ”equal-token” allocation schemes (which constrain B × S = Constant) over-
3
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
Figure 2. The data pipeline for joint image-video training with adaptive bucketing.
look the O(S 2 ) computational complexity of the core Transformer attention operators, causing long-sequence buckets to become significant ”long-tail” bottlenecks in distributed environments. Experimental observations show that the correlation between training step latency and total token count is only 0.35, whereas the correlation with B × S p reaches as high as 0.92.
Figure 3. Schematic diagram of forward fusion kernel data flow
To eliminate this imbalance, the system first calculates the logical sequence length S = Stext + Svisual for each data shape (nframe , H, W ) after VAE encoding, where Svisual is compressed according to temporal and spatial downsampling factors (8 and 16, respectively). Subsequently, we simultaneously apply a linear memory constraint and a polynomial computational constraint to each bucket, taking their intersection as the final batch size Bshape : Bshape = max(1, min(⌊Mmem /S⌋, ⌊Mcomp /S p ⌋))
racy and cost, the system employs a Throughput Sweep mode, prioritizing multi-level batch size tests for longsequence buckets where S ≥ 20, 000 to capture performance characteristics in the compute-intensive regime. Based on the collected benchmark data, we construct a parameterized cost model: step time sync ≈ a + b × B × S p . By performing a grid search for optimization within the interval p ∈ [1.6, 2.4], we select the p̂ that maximizes the coefficient of determination R2 as the model parameter. Subsequently, based on a preset training target step latency target sync, we back-derive the computational load upper bound Mcomp = (target sync − a)/b. This process establishes a closed-loop optimization framework: it monitors the waiting time wait sync of each GPU in real-time, identifies the primary bottleneck using bottleneck analysis tools, and dynamically recalibrates bucket configurations.
(2)
where Mmem represents the upper memory bound determined by GPU capacity and static model overhead, Mcomp is the upper bound for computational load, and p denotes the empirical exponent of attention complexity. Under this strategy, short-sequence buckets are governed by memory limits to maintain high throughput, while long-sequence buckets trigger the computational constraint. By actively reducing the batch size, we ensure that GPUs holding longsequence data do not slow down the global synchronization time step time sync due to O(S 2 ) load surges.
3.3. Fused LayerNorm-Modulate CUDA Kernels
To precisely quantify the computational constraint parameter Mcomp and the exponent p, the system incorporates a data-driven shape benchmarking system. This system operates within a real distributed environment (e.g., FSDP communication paths) and measures the mapping matrix (B, S) → step time sync via synthetic pixel scans that exclude data-loading I/O jitter. To balance measurement accu-
In the MMDiT architecture, AdaLN is invoked with high frequency. Traditional decoupled operator implementations introduce substantial bandwidth waste and scheduling overhead due to frequent HBM reads/writes of intermediate tensors (such as the normalized xnorm ). This paper designs and implements a fully fused CUDA kernel aimed at completing the entire normalization and modulation process in 4
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
a single kernel launch. In the forward pass, the kernel uses the Token as the unit of parallelism, utilizing registers and Shared Memory to temporarily store intermediate statistics, as show in Figure 3. Mean µ and variance σ 2 are calculated via two-stage reduction (Warp-level Shuffle and Block-level Reduction), and linear transformations are written out directly in combination with modulation parameters during a third traversal of dimension D. This eliminates approximately 131 GB/step of redundant HBM access in a 40-layer MMDiT model, significantly alleviating memory pressure.
Figure 4. Schematic diagram of grad shift/grad scale reduction process.
In the backpropagation stage, specifically for the computation of modulation parameter gradients ∇shift and ∇scale , we propose an innovative D-tile Coalesced Reduction strategy.Traditional gradient reduction occurs along the sequence dimension N . Since the feature dimension D is the contiguous storage dimension, this results in strided, noncontiguous memory access, which severely limits memory bandwidth. The D-tile strategy swaps the loop hierarchy, partitioning the Grid into a (dtile , ntile ) layout. Each thread is fixed to a feature index d and performs vertical accumulation along the sequence tiles, as show in Figure 4. In this mode, the physical addresses accessed by the 32 threads within the same Warp are perfectly contiguous, thereby boosting memory bandwidth utilization to near-theoretical peaks. Furthermore, the kernel caches computed statistics in global memory for subsequent reuse, avoiding redundant reduction calculations and ensuring natural compatibility and numerical stability in distributed sequence-parallel scenarios.
the framework’s global memory management system. This redundancy elimination at the computational graph level provides a theoretical ”negative entropy” for video generation models, serving as critical infrastructure for supporting ultra-long sequence lengths (Seqlen).
4. Experimental Evaluation 4.1. Experimental Setup We adopt Wan 2.1 (Wan et al., 2025) as our primary architecture, implemented within the DiffSynth-Studio (ModelScope, 2024) ecosystem to ensure compatibility with largescale generative benchmarks (Fu et al., 2025; Qiu et al., 2026; Yang et al., 2025; He et al., 2025). The system is stress-tested using a mixed corpus of 10 million samples from WebDataset (Face, 2025) and Koala-36m (Wang et al., 2025), creating extreme sequence length variance. To evaluate efficiency, we employ three metrics: ThroughB×[ F −1 +1]× W × H
3.4. Activation Computational Graph Simplification and Memory Recovery
λ γ η put Efficiency (Θ = ), reflecting latent T units processed per unit time; Load Balancing Efficiency −lenmin (CVstep = lenmax ), measuring synchronization overlenmax head; and Physical Load Pressure (O = B × S 2 ), which models the quadratic complexity of self-attention.
From the perspective of the computational graph, operator fusion not only reduces hardware-level memory access overhead but also achieves significant simplification of the graph within the deep learning framework’s automatic differentiation (Autograd) engine. In traditional discrete operator implementations, each constituent operation (e.g., Mean, Var, Standardize, Mul, Add) generates an independent node in the computational graph. This forces the system to retain the forward outputs of each node as ”activations” to facilitate gradient computation during the backward pass. Consequently, the activation footprint scales linearly with the length of the operator chain.
4.2. Overall System Throughput and Scalability The experimental results provide empirical evidence for the significant superiority of our integrated optimization strategy in enhancing system throughput. By mitigating computational imbalances and kernel-level redundancies, we achieve substantial gains across different cluster scales. Throughput Gains in 8-GPU Cluster: As illustrated in Figure 5a, the initial introduction of operator fusion techniques (e.g., Fused AdaLN) effectively reduced redundant HBM accesses. However, the most significant leap occurred after deploying the dual-constraint adaptive load balancing strategy. The mean training throughput surged from 14,383 tokens/sec to 18,069 tokens/sec, representing a substantial efficiency gain of 25.6%. Notably, the AdaptiveLoad curve (dark green) consistently maintains a higher floor compared to the Baseline (grey), indicating that even in worst-case
By integrating the entire path of AdaLN into a single fused CUDA kernel, we collapse the computational subgraph—originally consisting of 5–8 discrete nodes—into a single atomic node. Under this paradigm, only the final output of the kernel needs to be stored in the activation checkpoint. All intermediate statistics, including means, variances, and intermediate scaling factors, complete their entire lifecycle within registers instantaneously, bypassing 5
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
(a) 8-GPU Configuration
(b) 16-GPU Configuration
Figure 5. Throughput (tokens/sec) comparison between Baseline and AdaptiveLoad. The shaded areas represent the raw fluctuations per step, while the solid lines indicate the moving average.
sequence distributions, our system preserves high hardware utilization. Scaling Advantage in 16-GPU Cluster: This improvement is further amplified in the 16-GPU scaling experiments (Figure 5b). The average throughput rose from 30,170 tokens/sec to 38,372 tokens/sec, achieving an optimization margin of 27.18%. The widening gap between the two configurations as the cluster scale doubles highlights a critical insight: as parallel scale expands, the synchronization overhead caused by the “long-tail effect” of heterogeneous sequences becomes the primary bottleneck. Analysis of the “Long-Tail” Suppression: The Baseline throughput exhibits frequent downward spikes (notably around Step 200 and Step 450 in Figure 5b), where massive computational imbalances force high-performance nodes into prolonged idle states. AdaptiveLoad successfully flattens these throughput dips. By utilizing precise cost modeling and dynamic bucketing, our method liberates the latent computing power previously restricted by synchronization barriers. This robust scalability demonstrates that AdaptiveLoad is particularly well-suited for large-scale training on heterogeneous video datasets, where sequence length variance is inherently high.
Figure 6. Comparison of CVstep distribution between Baseline and AdaptiveLoad across 8-GPU and 16-GPU configurations.
of iterations operate under highly synchronized conditions. The optimization effect becomes more pronounced as the cluster scales. In the 16-GPU configuration, while the Baseline mean CVstep rises to 18.7% due to increased fragmentation, AdaptiveLoad effectively suppresses this trend, maintaining a robust mean of 10.4%. This improvement is attributed to the dynamic bucketing mechanism, which intelligently re-aligns input dimensions in real-time to mitigate the “straggler effect,” ensuring near-uniform peak memory usage and maximizing effective hardware throughput.
4.3. Multidimensional Load Balancing Evaluation To analyze load balancing in distributed training, we performed deep quantification across two dimensions: load balancing efficiency and computational complexity.
4.3.2. C OMPUTATIONAL C OMPLEXITY AND L OAD A LIGNMENT
4.3.1. LOAD BALANCING EFFICIENCY
The empirical results in Figure 7 reveal a substantial optimization in hardware utilization. In a high-pressure 16-GPU environment, the variation coefficient of computational complexity plummeted from a baseline mean of 39.0% to 18.9% with AdaptiveLoad. This reduction is critical for several reasons:
As illustrated in Figure 6, AdaptiveLoad demonstrates a superior capability in harmonizing cross-GPU workloads. In the 8-GPU configuration, the mean CVstep drops sharply from 15.9% to 8.9%, with the distribution exhibiting a significantly narrower body. This shift indicates a transition from a high-variance, unpredictable scheduling state to a tightly regulated, load-balanced regime where the majority
• Mitigating Quadratic Imbalance: As shown by the 6
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
+ Modulate), Query/Key Normalization (Q-Norm + K-Norm), and Gated Normalization (Gate + Norm). By consolidating these discrete logics into unified CUDA kernels, we aim to fundamentally eliminate the high-frequency round-trips of intermediate activations to High Bandwidth Memory (HBM), thereby increasing computational density while reclaiming precious memory resources for longsequence training. Table 1. System-level Comparison of Performance and Memory Metrics Before and After Operator Fusion
Figure 7. Comparison of Compute CV (%) between Baseline and AdaptiveLoad. The metric reflects the variance in quadratic computational pressure (B × S 2 ) across 16 GPUs.
grey curve in Figure 7, the Baseline frequently exhibits extreme spikes in Compute CV (exceeding 55%), where the quadratic nature of attention S 2 exacerbates even minor differences in sequence distribution. AdaptiveLoad (blue curve) successfully flattens these peaks by actively perceiving the S p weight of each bucket and executing a refined batch size reduction strategy for long-sequence clusters.
Metric Metric
Baseline (Max)
Fused (same)
Fused (Max)
Allocated Mem (GB) Reserved Mem (GB) Peak Memory (GB) Frames Step Time (s) Seq. Length Batch Size Throughput
88 128 139 233 62 48k 3 2,322
87.9 134 136 (↓3) 233 56 (↓10.7%) 48k 3 2,571 (↑10.7%)
93.81 136.78 139 257 68 52.8k (↑10%) 3 2,328
The quantitative results, summarized in Table 1, clearly illustrate significant gains in both throughput and load capacity. In our benchmark, the maximum sequence length supported by the GPU without operator fusion was capped at 48,000 tokens. Upon introducing operator fusion, the single-step execution time for the same load was significantly reduced from 62 seconds to 56 seconds, representing a 10.7% throughput increase. A critical breakthrough is observed in the expansion of load boundaries: the optimized system can stably support sequence lengths up to 52,800 tokens—a 10% improvement. Physically, this optimization saves approximately 3 GB of peak memory under identical loads.
• Eliminating Synchronization Bubbles: The drastic narrowing of the CV gap directly translates to minimized load deviation at synchronization points (Barriers). By ensuring that each GPU processes a similar total O per step, AdaptiveLoad effectively fills the idle bubbles that traditionally plague distributed training. This transformation converts previously wasted “sync-waiting periods” into effective “floating-point operation periods.” • Deterministic Throughput: Unlike the Baseline, which fluctuates wildly across Sample Steps, AdaptiveLoad maintains a consistently low Compute CV. This stability ensures predictable iteration times and prevents thermal throttling or power surges often caused by highly imbalanced, bursty computational workloads.
To further dissect these system-level gains, we conducted a micro-level kernel benchmark as shown in Table 2. By analyzing the execution of a single Fused AdaLN operator across varying sequence lengths (N ), we observe a consistent Pareto Improvement in both computational and storage dimensions.
Through this active weight-aware scheduling, AdaptiveLoad moves beyond simple padding-minimization to achieve true hardware-level load symmetry.
Table 2. Performance and Memory Comparisonxx N
4.4. Evaluation of Operator Fusion and Memory Efficiency To further validate the superior performance of operator fusion strategies in mitigating “Memory Wall” bottlenecks, we implemented deep kernel fusion optimizations for critical operator chains within the MMDiT architecture. This experimental suite focuses on core forward and backward propagation operator combinations, including Adaptive Layer Normalization and Modulation (Norm 7
Forward (s) Nat
Spd
Backward (s) Fsd
Nat
Spd
Mem (MB)
(Seq)
Fsd
Fsd
Nat
8k 16k 24k 32k 40k 48k 56k 64k
0.227 0.709 3.12× 1.051 0.774 0.74× 376 0.409 1.360 3.33× 1.373 1.487 1.08× 752 0.594 2.000 3.37× 1.706 2.169 1.27× 1125 0.783 2.645 3.38× 2.067 2.878 1.39× 1500 0.973 3.291 3.38× 2.381 3.592 1.51× 1877 1.161 3.933 3.39× 3.320 4.256 1.28× 2253 1.356 4.583 3.38× 3.647 4.973 1.36× 2626 1.547 5.240 3.39× 3.986 5.668 1.42× 3001
1072 2012 2956 3942 4930 5912 6892 7878
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
egy, which utilizes float32 accumulation for critical gradient paths, preserves sufficient numerical precision to avoid divergence or accuracy degradation. • Distribution Consistency: The adaptive bucketing and dynamic batch size adjustment strategy does not disrupt the statistical distribution of the training data. Despite the intra-step re-alignment of sequences, the effective gradients across the cluster remain unbiased compared to the Baseline’s stochastic sampling. • Enhanced Robustness: Interestingly, during the late training stage (approx. 500–900 steps), the AdaptiveLoad curve (dark green) exhibits slightly smoother convergence with fewer abrupt spikes compared to the Baseline. This suggests that by mitigating extreme sequence length imbalances, AdaptiveLoad may reduce the occurrence of ”gradient noise” typically caused by highly skewed batch compositions.
Figure 8. Training loss curves of Wan 2.1 model for Baseline and AdaptiveLoad. The shaded curves represent raw loss values, while the solid lines show the smoothed trend.
Micro-architectural Analysis In the forward pass, the fusion strategy achieves a stable speedup of 3.21× to 3.39×, effectively eliminating ∼70% of computation time by transforming memory-bound operations into compute-bound ones. More importantly, in the backward pass, our proposed D-tile coalesced reduction strategy demonstrates its scalability; as N increases to 64,000, the backward speedup steadily grows to 1.42×, resolving the strided memory access bottleneck inherent in discrete gradient accumulation.
The AdaptiveLoad framework seamlessly accelerates the training of the Wan 2.1 model while perfectly preserving its convergence robustness and final model quality.
5. Conclusion and Future Work This paper presents AdaptiveLoad, a full-stack optimization framework that addresses computational imbalance and memory access inefficiencies in large-scale video diffusion Transformer training. By implementing a dual-constraint adaptive load balancing mechanism driven by O(S 2 ) complexity modeling, we reduce the load Coefficient of Variation (CV) by nearly 50%, effectively eliminating synchronization bottlenecks. Complementing this, our fused LayerNorm-Modulate CUDA kernel—leveraging D-tile coalesced reduction—boosts training throughput by 28.63% on a 16-GPU cluster while perfectly maintaining the numerical stability of the Wan 2.1 model.
From a storage perspective, the fusion kernel reduces activation memory by approximately 61.9% across all tested sequence lengths (e.g., at N = 64, 000, saving 4.8 GB). By “collapsing” the computational graph, intermediate tensors are intercepted within registers rather than written to HBM, providing a theoretical “negative entropy” for memory management. This micro-level efficiency directly translates to the 10% expansion of global load limits observed in Table 1, allowing the model to bypass physical hardware constraints and process longer temporal sequences without relying on additional parallelization strategies.
Future research will focus on scaling these optimizations to thousand-GPU clusters by integrating with hybrid parallelism strategies and generalizing cost-fitting models for emerging architectures like State Space Models (SSMs). Furthermore, we aim to utilize compiler technologies such as OpenAI Triton for automated kernel generation and incorporate real-time power monitoring to pursue a more powerefficient and sustainable approach to large-scale model training.
4.5. Training Stability and Convergence Verification While achieving substantial throughput gains, maintaining numerical stability and convergence quality is a fundamental prerequisite for any system-level optimization. We conducted rigorous monitoring of the training loss trajectories to ensure that our optimizations did not compromise the model’s learning behavior. As illustrated in Figure 8, the loss curve of the AdaptiveLoad-optimized training remains highly congruent with the Baseline throughout the first 1,000 steps. This observation yields several critical conclusions:
References Constantinides, G., Dahlqvist, F., Rakamarić, Z., and Salvia, R. Rigorous roundoff error analysis of probabilistic floating-point computations. In International Conference on Computer Aided Verification, pp. 626–650. Springer, 2021.
• Numerical Fidelity: The highly overlapping trajectories validate that our D-tile coalesced reduction strat8
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
Huang, J., Zhang, Z., Zheng, S., Qin, F., and Wang, Y. {DISTMM}: Accelerating distributed multimodal model training. In 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24), pp. 1157– 1171, 2024.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022.
Jacobs, S. A., Tanaka, M., Zhang, C., Zhang, M., Song, S. L., Rajbhandari, S., and He, Y. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023.
Diao, S., Yang, Y., Fu, Y., Dong, X., Su, D., Kliegl, M., Chen, Z., Belcak, P., Suhara, Y., Yin, H., et al. Nemotronclimb: Clustering-based iterative data mixture bootstrapping for language model pre-training. arXiv preprint arXiv:2504.13161, 2025.
Khachatryan, L., Movsisyan, A., Tadevosyan, V., Henschel, R., Wang, Z., Navasardyan, S., and Shi, H. Text2videozero: Text-to-image diffusion models are zero-shot video generators. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 15954–15964, 2023.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.
Liu, Z., Cheng, S., Tan, G., You, Y., and Tao, D. Elasticmm: Efficient multimodal llms serving with elastic multimodal parallelism. arXiv preprint arXiv:2507.10069, 2025.
Face, H. Webdataset. Hugging Face Hub Documentation, 2025. URL https://huggingface.co/docs/ hub/datasets-webdataset.
ModelScope. Diffsynth-studio. GitHub Repository, 2024. URL https://github.com/modelscope/ diffsynth-studio.
Fu, X., Huang, N., Li, J., Guo, J., Li, X., and Lee, T.-Y. Polyart: Customizable multilingual movie poster generation via diffusion transformer. In Proceedings of the SIGGRAPH Asia 2025 Posters, pp. 1–3. 2025.
Patarasuk, P. and Yuan, X. Bandwidth optimal all-reduce algorithms for clusters of workstations. Journal of Parallel and Distributed Computing, 69(2):117–124, 2009.
Ge, H., Feng, J., Huang, Q., Fu, F., Nie, X., Zuo, L., Lin, H., Cui, B., and Liu, X. Bytescale: Communication-efficient scaling of llm training with a 2048k context length on 16384 gpus. In Proceedings of the ACM SIGCOMM 2025 Conference, pp. 963–978, 2025.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195–4205, 2023.
Gholami, A., Yao, Z., Kim, S., Hooper, C., Mahoney, M. W., and Keutzer, K. Ai and memory wall. IEEE Micro, 44 (3):33–39, 2024.
Qiu, Z., Wang, B., Chen, X., He, Y., and Wang, Z. Emovid: A multimodal emotion video dataset for emotion-centric video understanding and generation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pp. 8612–8620, 2026.
Guan, Z., Sun, H., Guo, Y., Di, S., Bai, X., Long, J., Zhao, T., Luo, M., Zhou, C., Guo, Y., et al. Rl-vla 3: Reinforcement learning vla accelerating via full asynchronism. arXiv preprint arXiv:2602.05765, 2026.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC20: international conference for high performance computing, networking, storage and analysis, pp. 1–16. IEEE, 2020.
Han, Y., Guo, M.-H., Liu, Z., Chen, W., and Hu, S.-M. Making llms optimize multi-scenario cuda kernels like experts. arXiv preprint arXiv:2603.07169, 2026. Harris, M. et al. Optimizing parallel reduction in cuda. Nvidia developer technology, 2(4):70, 2007.
Reuss, M., Yağmurlu, Ö. E., Wenzel, F., and Lioutikov, R. Multimodal diffusion transformer: Learning versatile behavior from multimodal goals. arXiv preprint arXiv:2407.05996, 2024.
He, Z., Qu, X., Li, Y., Zhu, T., Huang, S., and Cheng, Y. Diffthinker: Towards generative multimodal reasoning with diffusion models. arXiv preprint arXiv:2512.24165, 2025.
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658–68685, 2024. 9
AdaptiveLoad: Towards Efficient Video Diffusion Transformer Training
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-lm: Training multibillion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
multimodal large language models. In Proceedings of the ACM SIGCOMM 2025 Conference, pp. 24–38, 2025. Zhao, J., Lu, Q., Jia, W., Wan, B., Zuo, L., Feng, J., Jiang, J., Chen, Y., Cao, S., He, J., Jiang, K., Hu, Y., Nong, S., Peng, Y., Lin, H., and Wu, C. Megascale-data: Scaling dataloader for multisource large foundation model training, 2026. URL https://arxiv.org/abs/2504. 09844.
Skorokhodov, I., Tulyakov, S., and Elhoseiny, M. Styleganv: A continuous video generator with the price, image quality and perks of stylegan2. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 3626–3636, 2022. Team, K., Chen, J., Ci, Y., Du, X., Feng, Z., Gai, K., Guo, S., Han, F., He, J., He, K., et al. Kling-omni technical report. arXiv preprint arXiv:2512.16776, 2025. Wan, T., Wang, A., Ai, B., Wen, B., Mao, C., Xie, C.-W., Chen, D., Yu, F., Zhao, H., Yang, J., et al. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv:2503.20314, 2025. Wang, Q., Shi, Y., Ou, J., Chen, R., Lin, K., Wang, J., Jiang, B., Yang, H., Zheng, M., Tao, X., et al. Koala-36m: A large-scale video dataset improving consistency between fine-grained conditions and video content. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 8428–8437, 2025. Wulf, W. A. and McKee, S. A. Hitting the memory wall: Implications of the obvious. ACM SIGARCH computer architecture news, 23(1):20–24, 1995. Yan, Z., Bai, H., Yao, X., Liu, D., Liu, T., Liu, H., Li, P., Wu, E., Fan, S., Tao, L., Zhang, R., Wang, Y., Xu, S., Chang, J., Chen, X., Li, K., Bai, Y., Deng, G., Zheng, N., Korthikanti, V. A., Khattar, A., He, E., Govande, S., Lym, S., Zhu, Z., Zhang, Q., Yuan, H., Ren, X., Fu, D., Ma, T., Zhang, S., Shao, J., Wang, R., Rengasamy, V., Garg, R., Bhavani, S., Li, X., Zhou, C., Wu, D., Wei, Y., Aithal, A., Andersch, M., Shoeybi, M., Yao, J., and Yang, J. Scalable training of mixture-of-experts models with megatron core, 2026. URL https://arxiv.org/ abs/2603.07685. Yang, S., Wang, Y., Wang, Y., Zhu, L., and Zheng, Z. Towards scalable video anomaly retrieval: A synthetic video-text benchmark. arXiv preprint arXiv:2506.01466, 2025. Zhang, F., Ji, N., Gao, F., Zhao, B., Wu, J., Jiang, Y., Du, H., Ye, Z., Zhu, J., Zhong, W., et al. Dim-gesture: Co-speech gesture generation with adaptive layer normalization mamba-2 framework. arXiv preprint arXiv:2408.00370, 2024. Zhang, Z., Zhong, Y., Jiang, Y., Hu, H., Sun, J., Ge, Z., Zhu, Y., Jiang, D., and Jin, X. Disttrain: Addressing model and data heterogeneity with disaggregated training for 10