You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Implement Batch-Level Parallelism for Multiscale Processing #161
The current implementation processes each element in a batch sequentially during the multiscale decomposition and mixing phases. This approach doesn't leverage modern GPU architecture effectively, resulting in suboptimal training performance.
Impact
Training is slower than it could be
Inefficient GPU utilization
Limited scaling to larger batch sizes
Proposed Solution
Refactor the multiscale processing operations to support true batch-level parallelism:
Rewrite the __multi_scale_process_inputs method to operate on entire batches simultaneously
Optimize the mixing operations in MultiScaleSeasonMixing and MultiScaleTrendMixing classes
Vectorize operations where possible to avoid sequential processing
Technical Details
The issue is primarily in models/TimeMixer.py where several functions process batch elements one at a time:
# Current approach (pseudocode)foriinrange(B): # Process each sample in batch sequentially# Process sample i
...
# Proposed approach (pseudocode)# Process entire batch at once with tensor operations
...
Specific attention should be given to:
The __multi_scale_process_inputs method
future_multi_mixing implementation
Batch handling in PastDecomposableMixing
Expected Outcome
Significant speedup in training time (estimated 30-50%)
Better GPU utilization
Ability to scale to larger batch sizes with near-linear performance
Implementation Complexity
Medium - requires careful refactoring of core model operations while maintaining numerical equivalence.
Activity
kevinliu2000 commentedon Dec 5, 2025
Thanks for the advice. These implementations have been supported by CUDA.