LLM Learning Roadmap
✨ Awesome-Anything
🖼️ Digital Image Processing
🍃 LLM Components
🌱 LLM Pre-training
☘️ LLM Post-Training
🍀 LLM Popular Models
🪴 LLM Applications
🌿 LLM Optimization
🌾 LLM Compression
🌵 LLM Hands-on Practice
🌴 LLM Must-read Papers
🌳 LLM Q&A
🐝 VLM Image Encoders
📝 MISC.
🔗9.1.3 MiniMax-01
Features
- Ultra-long context window: 1 million tokens during training and 4 million tokens during inference by extrapolation.
- Using a hybrid architecture, mainly including lightning attention and an I/O-aware implementation of a linear attention variant.
- Uses a MoE (Mixture of Experts) architecture to maximize parameter and computation capacity
- Implement the all-to-all communication in MoE using expert parallel (EP) and expert tensor parallel (ETP), to minimize the overhead associated with inter-GPU communication.
- Design varlen ring attention to reduce the redundancy in computation, and the improved version of Linear Attention Sequence Parallelism (LASP) to utilize the device’s parallel capabilities.

MoE Architecture
This study adopts the MoE structure and linear attention over the traditional softmax attention used in standard transformers, aiming to optimize performance under limited resources and improve the handling of longer sequences.

We also discussed the MoE architecture in Section 4.7.4 Deepseek-V3 and 4.8 MOE Series.
MoE Computation
For a given input token , the corresponding output hidden state is computed as:
Here:
- is the total number of experts,
- is the weight of the gate,
- denotes the feed-forward network of the expert, and
- selects the top 𝑘 scores across all experts, setting the rest to , thereby zeroing them out after the softmax.
Problem Encountered: When scaling up to larger models, this study faces the challenge of routing collapse, which happens because the tokens are mostly sent to only a few experts instead of being spread out.
Global Router: Based on GShard auxiliary Loss, a global token dispatching strategy across EP groups was implemented to solve the expert loading imbalance problem. Specifically, an all-gather step was added to synchronize the number of tokens each expert needs to process before sending tokens across different EP groups. With the same capacity limits, this global routing method helps lower the overall token drop rate and keeps training stable.
MoE Token Dropping & Load Balancing
- Token Dropping in MoE Models: During training, each input token is routed to TopK experts based on gating scores. If too many tokens are routed to the same expert (i.e., it becomes overloaded), some tokens are dropped due to fixed capacity. Dropped tokens don’t get processed, but models avoid training crashes using fail-safe mechanisms like zero outputs, and it might skip loss/backprop for those tokens.
- Why Overload Happens: In multi-GPU training, gating decisions are made independently. Popular experts may receive tokens from many batches on different GPUs. Without mechanisms like all-gather, this causes imbalance and token drops.
- Solutions in DeepSeek-V3: Auxiliary load-balancing losses are added to the training objective to encourage more even expert usage.
Linear Attention
Linear Attention Can’t Forget
“Forgetting” in self-attention comes from the softmax normalization: for each query we compute a fresh probability distribution over all keys, and we can push irrelevant keys’ weights arbitrarily close to zero. In contrast, linear attention builds up two accumulators once and for all, so there’s no way to “unaccumulate” or “forget” a past key/value pair.
- Standard Self-Attention
- Fresh weights at each step.
- Zeroing out irrelevant tokens.
At time , we have query and all past keys . The output for this query is:
If a past token is irrelevant at , the model can learn , so . We re-compute a brand-new distribution that “forgets” the token at .
- Linear Attention: Memory Without Forgetting
- Running accumulators. Define feature map . At each step, we update two states
- One‐shot readout. Then
- No forgetting. Once or is added, it never decays, there’s no gate or subtraction.
Low-Rank Bottlenecks in Linear Attention
The expressiveness of attention mechanisms is crucial for their performance. Standard self-attention (using softmax) can adapt its focus based on the query, while linear attention can suffer from a "low-rank bottleneck,". Let’s use a simple example to illustrate this.
Suppose we have a hidden dimension and are considering a sequence up to its current length, which is 3. We are now generating the 4th output token (so the effective sequence length for this step). The query, key, and value matrices are . With and , these become . We can denote them as:
Since , each is a scalar. We are interested in the attention computation for the 4th query.
Standard Self-Attention: Dynamic and Expressive
- Score Computation for the 4th Query: The raw scores are computed between the 4th query and each key :
The vector of scores for is . This score vector is a scaled version of the key vector .
- Softmax Weights: These scores are then passed through a softmax function to obtain attention weights :
- Query-Dependent: Crucially, the query does not cancel out from the expression for . This means the attention weights change dynamically based on the query . As varies, the attention distribution shifts, allowing the model to focus on different parts of the sequence.
- Expressivity:
- For a fixed and varying scalar (since ), the vector traces a 1-dimensional curve within the 3-simplex (the space of 4-element probability distributions. It is a 3-simplex geometry because and ). While this single query with doesn't allow to be any arbitrary distribution in the simplex, it can still produce a rich variety of attention patterns that are important to .
- More broadly, the non-linear softmax function "inflates" the rank of the overall attention mechanism. If we consider the full score matrix (which is rank 1 in this example), the resulting full attention matrix (after exp and softmax) can be of much higher rank (up to n). This higher effective rank allows the model to capture more complex relationships.
The set of weights forms a probability distribution over the keys (i.e., and ).
Properties of Standard Attention Weights:
- Output: The output for is a weighted sum of the value vectors:
Because depends on , the output is also dynamically shaped by the query.
Linear Attention: A Potential Bottleneck
Linear attention methods approximate the softmax. A common approach involves using feature maps . In our simple scalar case (), let's use a scalar feature map .
- Approximated Unnormalized Weights: The core idea is to approximate with a factorized form:
So, the unnormalized weights for the query are:
- Normalized Weights: Normalizing these weights gives :
- Query-Independent: Notice that cancels out completely. The resulting attention weights depend only on the keys (via ) and not on the query .
- Fixed Attention Pattern: For a given set of keys , the attention distribution is fixed, regardless of what is. The model always attends to past tokens with the same pattern of importance. This is the "low-rank bottleneck."
- Rank Collapse: If we were to form an attention matrix B where (the weight query gives to key ), every row of this matrix would be identical: . Such a matrix has a rank of 1. This severely restricts the model's ability to learn diverse attention patterns.
Suppose is non-zero, it can be factored out of the sum in the denominator.
This cancellation happens only because the feature map is 1-dimensional. When we move to , the query no longer cancels, but the mechanism is still much less expressive than soft-max, because its rank can never exceed the feature-vector length.
Properties of These Linear Attention Weights:
- Output: The output for is:
Since is independent of , the way values are combined is fixed and does not adapt to the query.
Linear Attention can be treated as a type of RNN-like solution
Recently, methods resembling RNNs have been emerging frequently. Researchers prefer to refer to them as State Space Models (SSMs), such as Mamba, rwkv, RetNet, etc. In the inference stage, RNN-like methods maintain constant complexity per token, regardless of sequence length . In the training stage, they incorporate various designs to address challenges like parallelism. From this perspective, Linear Attention or Lightning Attention belongs to this category of methods, though Linear Attention is derived from the formula of Softmax Attention.
Given all the background above, it’s natural to see the motivation behind using Linear Attention:
- During training, for entire sequences, it reduces complexity to .
- Lightning Attention is essentially a parallelism-optimized version of Linear Attention.
- During inference, it reduces the per-token generation complexity to constant.
The specific structure of Linear Attention varies greatly. Except for the standard , it is also possible to choose different functions .
In earlier work, one of the authors of MiniMax-01 discussed that the scaling operation in
has a negative impact on Linear Attention. As a result, the scaling operation was replaced with a LayerNorm operation. MiniMax-01 follows this design, which is why the Linear Attention formulas in this technical report replace Softmax with Norm:
Where are the query, key, and value, respectively, with for sequence
length and for feature dimension. The equation can be transformed into its linear variant using right matrix multiplication:

Referring to the model structure code of MiniMax-01, the specific structural design of the Linear Attention used in MiniMax-01 can be viewed in the left diagrams (part of Figure 3 in the paper).
Note: The diagrams in the technical report suggest a computational complexity of . Because the diagrams do not easily represent the form of Linear Attention under a Causal Mask.
Lightning Attention
Causal Mask Makes Linear Attention Hard to Parallelize
To simplify, we ignore normalization. With a causal mask , standard (quadratic) attention is:
We call this the left-product form.
A causal mask is independent of the padding mask. The padding mask prevents attending to padded tokens in variable-length batches, whereas the causal mask enforces “no peeking ahead”:This mask is applied at every decoder self-attention layer, even with batch size = 1 (no padding).
By contrast, linear attention uses a right-product (prefix-sum) formulation:
- Without a causal mask, we can compute once and then
in time and fully in parallel.
- With a causal mask, each prefix sum must be built up sequentially, so it loses parallelism (though the overall cost stays ).
In summary:
- Right-product (linear) attention is but only parallelizable when no causal mask is applied.
- Left-product (quadratic) attention is and can be parallelized regardless of masking.
Lightning Attention splits the sequence of length into blocks of size , using:
- Left-product (quadratic) attention within each block
- Right-product (linear) attention across blocks.
Let be the current position, and define as the number of completed blocks before . Then
where
is the pre-summed block matrix for block .
- Inter-block: accumulate once per block (right-product).
- Intra-block: within the final (partial) block, apply the causal mask via
up to position (left-product).
As shown in this study, the pseudocode then reads:

Data Processing
- Used Byte Pair Encoding (BPE) for tokenization, with a vocabulary size of up to 200K tokens.
- Data Quality Enhancement: A filtering pipeline is applied, integrating rule-based cleaning with deduplication techniques in line with established methods. Document quality is evaluated at a fine-grained level using a previous-generation reward labeler, i.e., a MoE model featuring 5B active parameters and a total of 60B parameters.
- Data Formatting Optimization: while additional formatting like Markdown can aid human readability and understanding, this study observes that excessive formatting may reduce data diversity and quality. To preserve format generalization and align with human preferences, a nested document structure is introduced, utilizing flexible templates for both dialogue and QA data, which balance natural comprehension and consistent structure across diverse interaction styles.
- Data Mixture Investigation: This study finds that although high-scoring content leads to better overall performance, removing lower-scoring content entirely can decrease the downstream task performance. To address this, a balanced sampling strategy is employed.
- Employed a repetition-aware experimental framework to evaluate the impact of repeated data.
MoE Optimization
The main goal of optimizing the MoE (Mixture of Experts) architecture is to minimize communication overhead, especially when using all-to-all communication (a2a) in MoE models.
Comparison Between Tensor Parallel (TP), Pipeline Parallel (PP), and Expert Parallel (EP)
Parallelism | What it splits | Why you need it |
Tensor (TP) | Each large weight matrix inside a layer (e.g. the expert’s projection kernels) across n GPUs | Reduces the per-GPU weight memory so each device can hold huge matrices—but if it over-shard, mat-muls get too small and GPU utilization decrease. |
Pipeline (PP) | Consecutive layers (or groups of layers) across m GPUs | Reduces per-GPU activation memory by only storing activations for the consecutive subset of layers. But it still buffer all tokens in that stage, and it lead to pipeline bubbles. |
Expert (EP) | Batch of tokens routed to different experts across k GPUs | Solves MoE’s all-to-all scatter/gather by overlapping communication of one token-group with compute on another—but it doesn’t shard inside the experts. |
Expert Parallel (EP) Implementation
To reduce communication overhead, this study introduces a token-grouping-based overlap scheme. In this approach, a2a communication occurs within EP communication groups and overlaps with token processing across different expert groups. To maintain communication correctness, each ProcessGroup executes communication operators sequentially, which prevents a2a communications from overlapping across groups and results in some idle time. Despite this, the method brought significant performance improvements.

Understand the EP Overlap
- Both of the two devices are in the same network connection (for example, NVLink). Any a2a issued on one “sub-group” uses the same physical links as any other. So, even if we separate tokens into two devices of expert, the a2a-dispatch and -combine process can’t be performed at the same time.
- By splitting tokens into two groups of , each dispatch/combine only moves half as many elements, therefore spending half of the time.
Problems of Current TP and PP Solutions
After EP, the models are separated into 2 groups. However, in each group, the model weight is still too large to be placed in a single GPU. We need more parallelization methods, like TP or PP. Analysis showed that using TP or PP in the MiniMax-Text-01 model involves a critical trade-off:
- Tensor Parallelism (TP): When TP is employed to partition the expert parameters, the computational intensity becomes excessively low, which affects the efficiency of the computation.
When TP is used to split expert weight matrices across GPUs, each GPU handles a smaller chunk of the computation. Each TP layer needs a collective operation (e.g. all-reduce) to assemble the full activation for the next layer. As we increase the TP parallelism, the “extra communication inside each layer” (per-layer all-reduce) increases in both latency and effective bandwidth consumption.
- Pipeline Parallelism (PP): Without TP, the parameter count becomes too large, which requires the activation of a larger PP configuration. The challenge emerges because PP does not reduce the memory footprint required for storing activations.
Memory Footprint refers to the total GPU memory required to hold everything needed for both the forward and backward passes, including model parameters, optimizer state (momentum/Adam), intermediate activations, and gradient tensors. Even when it splits the model’s layers across different devices via pipeline parallelism, the total memory footprint does not shrink, because each stage still needs to buffer all the activations for back-propagation.
This limitation is harmful to training models with long contexts, as the associated increase in memory usage does not lead to corresponding improvements in computational efficiency or training speed.
ETP and EDP to Improve TP and PP Performance:
- ETP (Expert Tensor Parallel): A novel ProcessGroup, which is designed to manage the weight partitioning of experts.
- EDP (Expert Data Parallel): Also a novel ProcessGroup, to encapsulate the data parallelism of identical experts. In the system, it defines the total number of GPUs involved in training as
world_size. The system must satisfy two key conditions:
And
This configuration grants the MoE component the flexibility to define the expert distribution, manage expert weight partitioning, and independently configure the ZeRO (Zero Redundancy Optimizer) algorithm. With this setup, the study is able to fully decouple the parallel strategies of MoE components from those of non-MoE components.
Building on this modification, the ETP can be flexibly configured to achieve an optimal trade-off between memory consumption and computational intensity. Additionally, an EP-ETP overlap strategy is introduced to reduce communication overhead.

These optimization strategies enable a balanced configuration of memory usage and computational intensity. As a result, the pure communication overhead of the MoE component is reduced by 50% compared to the pre-optimization state, leading to a substantial boost in training efficiency.
Long Context Optimization
Varlen Ring Attention: Efficient Attention for Variable-Length Sequences
When training large LLM models, handling sequences of varying lengths efficiently is a real challenge. Standard ring attention can handle very long context window but doesn’t support data-packing formats. If one input is very long, then we need to pad all other sequences in the batch to the same length. Without data-packing, this is a waste of resources.
The Solution: Varlen Ring Attention
Varlen Ring Attention removes these limitations by applying ring attention directly to packed sequences, without requiring fixed-length alignment. It tracks each sequence’s offset within the packed batch and adjusts the attention mask accordingly. This allows the attention computation (both causal and non-causal) to work on each sample’s actual length.

Linear Attention Sequence Parallelism+
In lightning attention, the LASP (Linear Attention Sequence Parallelism) algorithm utilizes the CP (Communication Processing) communication group to support the extension of long sequences. As shown in Figure (a) below, the LASP algorithm requires all CP ranks to participate in send-receive operations for exchanging intermediate key-value (KV) block results. As a result, the overall training efficiency is reduced. Specifically, for the LASP algorithm:
- Initialize KV to zero
- The , , and matrices are first partitioned along the sequence dimension according to the CP size. Each resulting segment is then further divided into smaller blocks based on the specified BlockSize . For segments that cannot be evenly divided by (e.g., Q7, K7, V7), padding is applied to handle the remainder.
- Intra-block Computation
- Inter-block Computation and Communication
This study proposes an optimized approach LASP+ (Figure b), which restructures the computation and communication workflow to eliminate interdependencies during processing:
- Local Prefix Sum Calculation: Each computing node (i.e., CP rank) begins by independently computing its local prefix sum.
- Global Synchronization via AllGather: AllGather operation is executed to synchronize data across all nodes.
- Prefix Sum Computation: Based on the assigned computation order, each node selects the appropriate CP rank’s and performs the prefix sum operation accordingly.

Prev
CivitAI’s Payment Issue
Next
Awesome-AI-Tutorials
Loading...