MiniMax-01: Scaling Foundation Models with Lightning Attention

link：https://arxiv.org/pdf/2501.08313

Main Contributions

  • 包括文本模型MiniMax-Text-01（1 million tokens / train、4 million tokens / infer） 和多模态模型MiniMax-VL-01（512 billion tokens / train）;
  • hybrid architecture（lightning、linear、softmax attention）提高了处理长文本的能力;
  • 优化的并行策略和高效的计算-通信重叠技术，用于MoE和lightning attention;
  • 在MoE中使用allgather + all-to-all，减少通信开销;

img

Model Structure

img

  • Block：channel mixer (lightning、softmax attn) + feature mixer（MoE FFNs）
  • Softmax Attention: use Group Query Attention with a group size of 8
  • MoE Experts: 32 experts, k = 2
  • Total Layer：(7 * lightning blocks + 1 * softmax block) * 10 layers
  • Parameters: 456 billion (45.9 billion are activated)
  • Lightning Attention：来源于TransNormerLLM（可解释图中G的由来，具体在下图）

img

Normhere refers to RMSNorm.

Mixture of Experts

  • 计算公式

img

其中，E表示专家总数，W_g表示路由权重，FFN_i表示第i个expert，TopK表示保留top k scores，其余值被设置为负无穷。

  • 训练策略

    关于MoE训练中每一个expert可接受的token数量大小，当前的训练策略可以划分为token-drop和dropless两种，MiniMax采用token-drop。

    • token-drop指每个expert只能接受有限的token数，超出部分会直接丢弃。
    • dropless动态调整缓冲区，为了等待那个接收 Token 最多的专家算完而陷入长久的空转，大幅拉低 MFU（模型算力利用率）。

  • 辅助损失

动机：

上述策略在规模增大时会出现routing collapse，退化为稠密模型。

解决方案：

使用Gshard的辅助损失，如下：

假设共有E个expert， 对于第 t 个 token，路由会输出E个expert的Softmax概率：

\[p_t = Softmax(x_tW_g)=[p_{t,1}, p_{t,2}, ...,p_{t,E}]\]

这些概率值会通过 TOPK 来选择出哪些被过滤：

\[z_{t,i} = p_{t,i}*\operatorname{I}[i\in \operatorname{TOPK}]\]

其中$\operatorname{I}[.]$表示Iverson括号。

假设token一共有 T 个，即$t=[1, 2,…,T]$,

对于expert_i，由于每一个token都会输出分配到这个专家的概率值，

那么，对expert_i来说，路由到自己的平均概率定义为：

\[m_i= \frac{1}{T}\sum_{t=1}^{T}p_{t,i}\]

对于token来说，分配到第i个expert的token占比为：

\[f_i = \frac{1}{TK}\sum_{t=1}^{T}z_{t,i}\]

所以，Gshard的损失如下：

img

其中，$f_i$表示 batch 中，第 i 个 expert 的 token 占所有 token 的比例，m_i表示第 i 个 experts 的 softmax 平均得分，$\alpha_{aux}$表示系数。通过损失函数确保所有专家都能得到充分且平衡的训练。

  • Global Router

img

动机：

存在某一个EP过载，而其余EP空闲的情况。

解决方案：

使用全局分发策略。

在分发之前，使用allgather统计设备内部去各个专家的数量。

在全局分发时，某一个expert多余出来的数据不会被丢弃，而是跨组分发到别的expert中，从而降低drop率。

下面通过一个例子，给出具体的解释：

先考虑Expert Parallel（EP），观察Expert和device之间的排布：

  • 假设有8张卡
  • 共64个Expert
  • EP=8，表示每张卡放置8个Expert

在Expert Parallel的设定下，Experts的排布如下所示，

img

每个设备持有8个Expert, 且将所有设备的集合称作为一个Group。

接下来引入Data Parallel（DP），结合DP+EP，考虑Expert和Device怎么排布：

  • 假设DP=2，即将Batch内的token拆分成两份
  • 假设每个Expert能处理的capacity是 x 个token，超过就会drop
  • 其余条件和上述一致

这个时候总的GPU数量=DPEP=28=16，如下所示：

img

不同组均持有64个Expert，不同的是，由于DP的存在，不同Group处理的是不同的token。

由此会带来一些负载均衡的问题（即使已经通过 Gshard loss 优化过 Router，使其分配较为合理，但由于每个Group看到的data都是micro-batch，不同token的分布会使得同样的专家在Group #0内频繁选中，而在Group #1内无人问津）

假设每一组都传来 2x 个token，

对于Group #0，Expert #1可能被分配了y个token（y>x），而对应的Group #1中Expert #1只分配了1个token（DP导致不同组下的token分布波动大），

所以存在对Expert #1来说，Group #0爆满，Group #1空闲的状态。

并且即使Group #1有足够的空闲可以分担总的容量，也不能自动送过去（需要全局a2a）。

allgather 同步 token counts 是为了解决什么？

想跨组“借用容量”，得先知道全局每个 expert 在各组的负载情况：

例如，

  • Group #0 组 Expert #1 收了多少？
  • Group #1 组 Expert #1 收了多少？

有了这个负载Table，就能做负载的均衡，也就是跨组分发。

Linear Attention

img

Linear attention采用右乘的技巧，将计算复杂度从O(N^2d)变成O(Nd^2)

但是在有cumsum的情况下，Linear attention阻碍了并行实现。具体来说，

目标： 计算右乘$K^TV$ 。 在推理阶段， 假设$M_i = \sum_{t=1}^{i}K^T_tV_t$，这里K的维度是从（1,D）到（i,D）。 可以将M_i的计算进行拆分，$M_i = M_{i-1} + K^T_iV_i$(cumsum形式)从而降低计算复杂度。 但是在训练阶段（或prefill阶段）， 需要一次性把一个batch的输入（N,D）直接生成输出（N,D）,使用上述技巧就不能并行训练。

Ligtning Attention

Lightning Attention提出了一种新的tiling技术，有效地规避了cumsum操作。

关键：将注意力计算划分成了两部分：

  1. intra-block (left product，可并行)
  2. inter-block (right product，block级cumsum，非序列级)

以下是Lightning Attention的算法：

详细推导请查看 回顾 Linear Attention 的并行化计算 - Chunkwise Parallel。

img

img

Minimax-Hybrid-lightning

lightning attenrion在大规模下性能不如softmax attention，因此提出了7层lightning + 1层softmax的架构。

Computation Optimization

MoE Optimization

问题：

MoE 模型的最大瓶颈在于 Expert Parallel (EP) 带来的 GPU 间“a2a”通信（如下图w/o EP overlap，三步串行，网卡和算力交替空闲）。

img

概念说明：

  • a2a-dispatch（全对全分发）： GPU 必须把手里的 Token 发送到对应专家所在的 GPU。
  • expert（专家计算）： GPU 运行 FFN（前馈网络）计算。
  • a2a-combine（全对全合并）： 计算完后，必须把结果传回原 GPU 以维持序列顺序。

解决方案：

将一个 Batch 的 Token 进一步拆分成两组（或多组）。当第一组 Token 正在进行 GPU 间的 All-to-All 分发（a2a-dispatch）时，GPU 已经开始计算本地专家的任务。

详细解释如下（@蛋蛋）：

w/o EP overlap vs EP overlap 梳理

以论文中的设置为例，

一共有2个设备，处在同一个EP Group中，假设每个设备中有2个expert。

做如下定义：

img

先来看没有EP overlap时的流程：

img

再来看EP overlap时的流程：

首先，将token 划分成多组，如下：

img

接下来是主要的流程：

注意，原文中说底层通信算子被设计为串行，因此上一个通信操作没处理完的时候，后续的通信必须等待。

img

可知，

Y的时间通过overlap节省掉了。

EP overlap 结合 TP 时出现的问题

全局TP并行时，MoE的计算强度低；然而，选择不使用TP会导致过大的参数数量，需要激活更大的Pipeline Parallelism ( PP )配置。

解决方案：

将MoE的并行配置从全局中解耦。

  1. ETP (Expert Tensor Parallel - 专家张量并行 - 一个专家拆到multi-devices - 增加计算密度)

  2. EDP (Expert Data Parallel - 专家数据并行- 一个专家复制到不同的 GPU 组 - 提高吞吐)

  3. EP-ETP重叠策略

img

Varlen Ring Attention

回顾：RingAttention

ring attention + flash attention：超长上下文之路 - 知乎

Ring Attention Paper

问题：

传统 Transformer 的自注意力机制（Self-Attention）需要计算并存储全部 Query-Key 交互的中间结果，显存占用与序列长度呈平方级增长）。

解决方案：

RingAttention 将分块计算（Blockwise Computation）与环形通信（Ring Communication）相结合，将长序列分散到多个设备（GPU/TPU）上处理。

分块与分布： 将输入序列切分成多个块，每个设备只负责计算其中一小块 Query 对应的注意力。

环形传递： 设备之间组成一个环。每个设备在计算完本地 Key-Value (KV) 块的注意力后，将 KV 块传给下一个设备，同时接收上一个设备的 KV 块。

img

问题：

  • ring attention对data-packing没有优化
  • flash attention提供了varlen的接口，但是没有ring attention的实现
  • 现有实现对序列存在巨大的padding浪费

MiniMax的创新：Varlen + ring attention + data-packing format

直接对拼接后的序列之间应用ring attention，规避padding。具体而言，该实现涉及在环形注意力计算中区分每个序列对应的注意力掩码的偏移量。关键的修改是将原始的因果计算转化为varlen因果计算，类似地将非因果计算转化为varlen非因果计算，如图11所示。（带颜色的部分都是不同序列的拼接，不是padding）

img

Improved Linear Attention Sequence Parallelism

问题：

CP等级之间具有顺序依赖性，从而迫使计算以串行方式进行。因此，这种顺序依赖严重阻碍了训练过程的整体效率，因为系统的内在并行性没有得到充分的利用。

解决方案：

通过计算local prefix sum -> global prefix sum的方法实现并行。

img

虽然在增加总通信量和临时内存使用方面产生了额外的成本，但它们所赋予的实质性性能好处是明确的。这些增强大大超过了通信和内存消耗的相关开销。通过全面的测试和验证，LASP +的计算速度可以达到原LASP算法的1 / Npcn，其中Npcn为并行计算节点数。

Lightning Attention

MiniMax实现了四种针对lightning attention的优化策略：批处理内核融合、分离预填充和解码执行、多级填充和跨步批处理矩阵扩展。

  • Batched Kernel Fusion read paper

  • Separated Prefill and Decoding Execution

    未优化前（串行）：

    GPU 先处理那 18 个 Decoding 的请求，再处理那 2 个 Prefill 的请求。

    1. 处理 Prefill 时：计算核心跑满。
    2. 处理 Decoding 时：显存带宽跑满，但大部分计算核心（SMs）是闲置的。

总耗时 = Decoding时间 + Prefill时间。

优化后（CUDA Streams 并行）：

利用 CUDA Streams（GPU 的多流技术），让这两个任务在同一块 GPU 上同时跑。

总耗时 = Max(Decoding时间, Prefill时间)。

  • Multi-level Padding

推理阶段，将256的block size替换成34，68，128三级。

  • StridedBatchedMatmul Extension

read paper