Read MiniMax-01 paper.

Summary

Introducing MiniMax-01:

  • A new series of foundational models, including MiniMax-Text-01 (language model) and MiniMax-VL-01 (vision-language model)
  • Designed to achieve top-tier performance on standard benchmarks while excelling in long-context processing, handling context windows of up to 4 million tokens
  • Challenges the prevailing assumption that state-of-the-art language models must rely on traditional, computationally expensive attention mechanisms

Key Innovations of MiniMax-Text-01:

  • Hybrid Architecture: Combines the efficiency of lightning attention, a novel linear attention variant, with the strengths of softmax attention.
  • Lightning Attention: Achieves linear complexity through a tiling technique that avoids costly cumsum operations, making it ideal for long sequences.
  • Mixture of Experts (MoE): Enhances scalability and efficiency by activating only a subset of parameters for each token.
  • Global Routing Strategy: Ensures balanced token allocation across experts in MoE layers, preventing routing collapse and improving training stability.
  • Focus on Data Quality and Diversity: Utilizes a meticulously curated training corpus with diverse sources, including academic literature, code, web content, and books, enhanced through a reward-based quality assessment.
  • Hierarchical Reward System and Multi-Stage Training: Refines model performance, long-context capabilities, and real-world applicability through supervised fine-tuning, offline and online reinforcement learning.

Highlights of MiniMax-VL-01:

  • Impressive Performance in Vision-Language Tasks: Excels in visual question answering and long-context comprehension.
  • Large-Scale Image-Description Dataset: Trained on 100 million images paired with fine-grained descriptions, boosting its ability to align and understand visual and textual information.

Optimization Strategies:

  • Training Optimizations: Data-packing, Varlen Ring Attention, Improved Linear Attention Sequence Parallelism (LASP+), Three-Stage Training for Long-Context Extension, Multi-Stage Post-Training.
  • Inference Optimizations: Batched Kernel Fusion, Separated Prefill and Decoding Execution, Multi-level Padding, StridedBatchedMatmul Extension.

Impact and Future Directions:

  • Challenges the Limitations of Traditional Attention Mechanisms: Demonstrates the potential of linear attention for building highly efficient and scalable long-context models.
  • Pushes the Boundaries of Context Window Size: Enables the processing of significantly longer sequences, opening up new possibilities for applications like document analysis, code generation, and conversational AI.
  • Highlights the Need for More Advanced Evaluation Benchmarks: Existing datasets primarily focus on simplified long-context scenarios, calling for the development of more complex and realistic evaluation methods.

MoE Architecture

The model’s hidden size is configured to 6144, and each layer incorporates 32 experts with a top-2 (2 active experts for each token) routing strategy. The feed-forward network (FFN) within each expert has a hidden dimension of 9216. Given that the total model size is 456 billion parameters and 45.9 billion are activated for each token, the size of each expert can be estimated to be approximately 14 billion parameters (456 billion / 32 experts). The MoE architecture in MiniMax-01 employs several unique strategies to enhance its effectiveness:

  • Hybrid Attention Mechanism: This model uniquely combines both linear attention (specifically lightning attention) and softmax attention in a structured pattern. It uses one transformer block with softmax attention after every 7 transnormer blocks with lightning attention. This hybrid approach allows the model to leverage the efficiency of linear attention for long sequences while maintaining the strengths of softmax attention for retrieval and in-context learning.
  • Global routing strategy: To address the challenge of routing collapse, where token allocation becomes concentrated and can hinder training stability, MiniMax-01 introduces a global routing strategy. This strategy builds upon the auxiliary loss mechanism from GShard, a technique for distributing model parameters across multiple devices. This global routing ensures better load balancing across the experts within the MoE layers.
  • Token-drop strategy: The model utilizes a token-drop approach during training for improved efficiency. This strategy sets a capacity limit for each expert, discarding any tokens exceeding that limit. This allows for more streamlined training without processing every token by every expert.
  • Optimization for reduced communication overhead: The authors have implemented a token-grouping-based overlap scheme to minimize communication overhead during MoE training. This involves overlapping the all-to-all communication within expert parallel groups with the processing of tokens from different groups, significantly improving training efficiency.
  • Decoupling of parallel strategies: The MoE component’s parallel strategies are decoupled from those of non-MoE components through expert parallel (EP) and expert tensor parallel (ETP). This configuration grants flexibility in expert distribution, weight partitioning, and the application of the Zero Redundancy Optimizer (ZeRO) algorithm.
  • EP-ETP overlap strategy: Further minimizing communication overhead, an overlap strategy between EP and ETP is implemented to maximize the utilization of network and computational resources.
  • Integration of DeepNorm: To enhance overall performance, the architecture incorporates DeepNorm, a normalization technique designed to improve training stability and model convergence.

These optimizations and integration strategies enable the MoE architecture in MiniMax-01 to achieve a balance between performance, memory usage, and computational efficiency, allowing for the effective handling of long sequences with minimal overhead.

Inference Optimizations

Key optimization strategies for lightning attention inference:

  • Batched Kernel Fusion: This technique reduces intermediate result storage and memory access by fusing multiple memory-bound kernels and extending support for all batch inputs. In the “prefill” phase, a kernel fusion handles the processing of Q, K, and V tensors, which involves tasks like padding, partitioning, layout adjustment, and decay value calculation. Similarly, another kernel fusion is applied during the “decoding” phase to compute K and update the prefix K cache. This optimization significantly enhances memory access efficiency and leads to a 10% reduction in end-to-end latency, especially for decoding and short-text input scenarios.
  • Separated Prefill and Decoding Execution: Recognizing that the decoding phase often deals with single tokens (length = 1), the authors propose using separate kernels and CUDA streams to handle tokens with lengths of 1 and greater than 1 in parallel. This separation improves efficiency, especially when dealing with batches containing a mix of short and long inputs. For instance, in a batch with a few long sequences and many single-token sequences, this approach can reduce latency to a level comparable to processing only the long sequences.
  • Multi-level Padding: This strategy involves padding the Q, K, and V tensors along the sequence dimension to enable decomposition into identical matrix multiplications. This decomposition allows for efficient use of the StrideBatchedMatmul interface, maximizing parallel processing. The authors initially used a fixed block size of 256 for padding but found it inefficient when dealing with shorter sequences common in prefix caching scenarios. To address this, they introduced additional segmentation options (32, 64, and 128) to allow dynamic selection of the optimal computational scale based on input length, minimizing padding overhead and maximizing resource utilization.
  • StridedBatchedMatmul Extension: Leveraging the optimized cublasGemmStridedBatchedEx function from the NVIDIA cuBLAS Library for StridedBatchedMatmul operations, this approach ensures high performance and architectural versatility. The implementation incorporates techniques like warpgroup-wide WGMMA instructions for efficient handling of 256x256 GEMM operations and the asynchronous capabilities of the Tensor Memory Accelerator (TMA) to further enhance memory access and processing efficiency. The goal is to achieve dynamic regulation of pipeline stages for optimal performance across different GPU architectures like H20 and H800.

These optimizations collectively contribute to achieving a Model Flops Utilization (MFU) exceeding 75% on the H20 GPU for end-to-end inference. They significantly reduce the latency contribution of lightning attention compared to softmax attention in long-sequence scenarios. Furthermore, the optimizations enable the model to efficiently handle heterogeneous batch inputs with varying sequence lengths, crucial for real-world applications.

Training Optimizations

Optimizing both the training process and the underlying architecture of MiniMax-01 to handle long sequences effectively. Here are the key training optimizations employed:

  • Data-packing: To avoid computational waste from padding in sequences of varying lengths, the training process utilizes “data packing.” This technique concatenates different samples end-to-end along the sequence dimension, minimizing padding and maximizing computational efficiency, especially important for training at the 1M sequence length scale.
  • Varlen Ring Attention: Building on data-packing, Varlen Ring Attention eliminates the need for excessive padding by directly applying the ring attention algorithm to the entire concatenated sequence. It distinguishes the offset of each sequence’s attention mask within the computation, enabling efficient processing of variable-length inputs without wasted computation.
  • Improved Linear Attention Sequence Parallelism (LASP+): Addressing the sequential dependency bottleneck in the original LASP algorithm for lightning attention, the paper introduces an improved version called LASP+. This enhancement replaces the send-recv operations with an AllGather operation across all parallel computing ranks, allowing for parallel computation of key-value (KV) blocks and significantly speeding up the training process. While LASP+ increases communication volume and temporary memory usage, the performance benefits outweigh the added overhead.
  • Three-Stage Training for Long-Context Extension: The model’s context window is gradually expanded to 1M tokens using a three-stage training procedure. This approach progressively upsamples long-context data while maintaining the distribution of shorter contexts to ensure stable performance across different sequence lengths.
  • Multi-Stage Post-Training: The model undergoes a multi-stage post-training process (detailed in Table 7) designed to improve long-context handling while maintaining performance on shorter sequences. This involves stages of Supervised Fine-Tuning (SFT), Offline and Online Reinforcement Learning (RL) with varying sequence lengths and batch sizes.
  • Curriculum Learning for Reward Model Training: The reward model training utilizes curriculum learning, starting with easier tasks and gradually increasing the difficulty. This approach improves the stability and effectiveness of reward model training, crucial for aligning the model with human preferences.
  • Adaptive Batch Size Scaling: The batch size is dynamically adjusted throughout training based on a power-law relationship between training loss and the critical batch size. The batch size increases when the loss reaches a certain threshold, ensuring an optimal balance between training time and data efficiency.
  • Source-Specific Weight Interpolation for Stability: During transitions between different data distributions in the long-context extension process, linear interpolation of source-specific weights is employed. This technique mitigates potential instabilities from distributional shifts by facilitating a gradual and controlled evolution of the data distribution, ensuring stable and convergent training.