Skip to content

Add MiniMax Text model support#4662

Open
yicycyc wants to merge 4 commits into
PaddlePaddle:developfrom
yicycyc:feat/minimax-text-01-clean
Open

Add MiniMax Text model support#4662
yicycyc wants to merge 4 commits into
PaddlePaddle:developfrom
yicycyc:feat/minimax-text-01-clean

Conversation

@yicycyc

@yicycyc yicycyc commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

PR 新增模型支持:MiniMax-Text-01

本 PR 新增 MiniMaxAI/MiniMax-Text-01 支持。复现对齐对象为 HuggingFace/Transformers 端 MiniMax 实现以及 ModelScope 上的原始 safetensors 权重。

主要改动

  • 新增 paddleformers/transformers/minimax/
    • MiniMaxConfig
    • MiniMaxModel
    • MiniMaxForCausalLM
    • MiniMaxForCausalLMPipe
    • RMSNorm、RoPE、full attention eager/cache
    • linear attention / lightning attention
    • sparse MoE router / expert MLP
    • Stage3 验证所需的固定 expert 遍历顺序和空 expert graph 连接开关
  • 注册 MiniMax-Text-01 到:
    • paddleformers.transformers
    • AutoConfig
    • AutoModelForCausalLM
  • 新增 MiniMax-Text-01 SFT 输入链路:
    • 新增 minimax template
    • 支持 MiniMax chat special token 格式
  • 权重加载:
    • MiniMaxPretrainedModel._gen_aoa_config 中实现 HF/ModelScope safetensors 到 PaddleFormers 权重的自动转换规则
    • MiniMaxPretrainedModel._gen_inv_aoa_config 中实现 PaddleFormers 到 safetensors 的反向映射规则
    • 用户可直接使用原始 safetensors 目录加载,无需手动运行离线权重转换脚本
  • 单测:
    • 新增 tests/transformers/minimax/test_modeling.py
    • 覆盖 tiny config、MiniMaxModel forward、MiniMaxForCausalLM loss/backward、AutoModelForCausalLM.from_config
  • 文档:
    • 更新 README 模型列表
    • 更新 docs/zh/model_capability.md 能力矩阵

Stage3 训练跑通方式在下方训练验证中说明,通用 Trainer / Optimizer 侧改动不纳入本 PR。

前向对齐验证

模型:MiniMaxAI/MiniMax-Text-01

验证模型为一层缩层模型:

  • num_hidden_layers = 1
  • 模型目录:/root/claude_workspace/minimax-text-01-layer0-ms
  • 当前目录包含一层缩层模型所需 shard,以及 embedding / lm-head 相关 shard
  • attention backend:两侧均使用 eager
  • embedding / input ids 路径已纳入最终验证

bf16 eager 最终 hidden states:

dtype shape max abs diff mean abs diff p99 abs diff cosine 结论
bf16 (1, 4, 6144) 0.0 0.0 0.0 1.00000083 对齐

fp32 eager :

dtype shape max abs diff mean abs diff p99 abs diff cosine 结论
fp32 (1, 4, 6144) 5.7220459e-06 2.5422841e-07 1.9073486e-06 1.00000012 对齐

说明:

  • bf16 eager 与之前 Transformers 4.57 对齐基准达到 0 diff。
  • fp32 eager 为 1e-6 量级误差,符合 fp32 算子执行差异预期。

训练验证

使用 GSM8K 做 BF16 full-SFT,Paddle 侧使用 sharding stage3,Torch 侧使用 ms-swift ZeRO-3。

共同设置:

  • max_length = 512
  • max_steps = 300
  • per_device_train_batch_size = 1
  • gradient_accumulation_steps = 1
  • learning_rate = 1e-5
  • lr_scheduler_type = constant
  • warmup_steps = 0
  • weight_decay = 0
  • shuffle 关闭
  • attention eager
  • full-parameter SFT
  • 可训练参数量:8,146,085,888

loss 曲线如下:

minimax_gsm8k_len512_300_loss_compare

结论:

  • 两侧首步 loss 均为 27.x,没有出现异常 loss scale。
  • 300 steps 内两侧 loss 都稳定下降。

Stage3 / ZeRO3 说明

MiniMax sparse MoE 会根据 token 动态选择 expert。Stage3 / ZeRO3 下,如果不同 rank 只执行本 rank 命中的 expert,不同 rank 的 expert 执行顺序可能不一致,从而导致 collective 顺序不一致、训练卡住或梯度状态异常。

本地成功验证使用:

  • 固定 expert 遍历顺序
  • 对空 expert 加零贡献,使其保持在 autograd graph 中
  • 8 卡 Paddle sharding stage3
  • ms-swift ZeRO-3 中将 MiniMax sparse MoE block 作为 leaf module

生成对齐

使用真实文本输入做 greedy generation,对齐对象为 Transformers 4.57.6 MiniMax 实现。

输入:

用户:请用一句话介绍杭州。
助手:

输入 token:

[8526, 192831, 829, 49864, 10201, 30860, 487, 75827, 861]

bf16 生成前 10 个 token:

dtype HF tokens Paddle tokens 是否一致
bf16 [162551, 197066, 109245, 4698, 7169, 29897, 4791, 69130, 190900, 177743] [162551, 197066, 109245, 4698, 7169, 29897, 4791, 69130, 190900, 177743]

@paddle-bot

paddle-bot Bot commented Jul 13, 2026

Copy link
Copy Markdown

Thanks for your contribution!

@risemeup1111 risemeup1111 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已完成初次 review。当前还有需要先修复后再合入的问题,主要集中在 MiniMax 的 packed attention 处理和训练 loss 传参;详细意见已放在行级评论中。

本地未能运行新增测试:当前评审环境缺少 paddle,直接 import 仓库还缺少 colorlog;CI 页面目前显示通过。

Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.

Comment on lines +441 to +443
attention_interface: Callable = eager_attention_forward
if getattr(self.config, "_attn_implementation", "eager") != "eager":
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级:P1

问题: eager 模式下这里走的是本文件里的本地 eager_attention_forward,但该实现只处理 attention_mask,会忽略调用处传入的 attn_mask_startend_row_indices。SFT collator 默认支持生成 start/end row indices;当 full attention 层使用 eager + packing 时,样本边界不会被还原成 mask,token 会跨样本注意。

影响: PR 描述里的 eager SFT/对齐路径在打包样本上会产生错误 attention,训练 loss 和梯度会被跨样本上下文污染。

处理要求:请针对该评论修复并提交新的 commit。

建议: 直接复用共享 attention registry 的 eager 实现;它已经会把 attn_mask_startend_row_indices 转成 dense mask:

Suggested change
attention_interface: Callable = eager_attention_forward
if getattr(self.config, "_attn_implementation", "eager") != "eager":
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[getattr(self.config, "_attn_implementation", "eager")]

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Full Attention 改用共享的 ALL_ATTENTION_FUNCTIONS

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已确认当前 head 已改为共享 ALL_ATTENTION_FUNCTIONS,full attention 在 eager 下会走统一实现并处理 attn_mask_startend_row_indices。这个点已解决。

input_attention_mask = causal_mask
input_mask_startend = attn_mask_startend_row_indices
else:
input_attention_mask = attention_mask

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级:P1

问题: attn_mask_startend_row_indices 只传给 full_attention 分支;linear_attention 分支把 input_mask_startend 固定为 None,而 MiniMaxForCausalLM.forward 在同时收到 row indices 和 attention_mask 时还会清空 attention_mask。因此启用默认的 use_attn_mask_startend_row_indices=True 且 packing/greedy_intokens 把多个样本拼到同一行时,linear attention 会把整个拼接序列当成一个连续序列累计 attn_weights_inter,后一个样本可以读到前一个样本。

影响: MiniMax 的 SFT 打包训练 attention 语义错误,loss/梯度会跨样本污染;即使 full attention 层正确处理了 row indices,linear attention 层仍会泄漏上下文。

处理要求:请针对该评论修复并提交新的 commit。

期望:MiniMaxLightningAttention 增加基于 start/end row indices 的 segment reset,并补一个 packed 两段样本互不影响的测试;如果暂不支持 packing,则至少在含 linear_attention 时显式拒绝 attn_mask_startend_row_indices/4D packing mask,并在 MiniMax 训练配置或文档中要求关闭 packing。修复形态可以类似:

if attn_mask_startend_row_indices is not None and "linear_attention" in self.config.layer_types:
    raise ValueError(
        "MiniMax linear_attention does not support packed attn_mask_startend_row_indices yet; "
        "disable packing/use_attn_mask_startend_row_indices or implement segment-reset support."
    )

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Linear Attention 暂不支持 packing,遇到 row indices 或 4D packing mask 时明确报错,并补充文档和测试。


loss = None
if labels is not None:
loss, _ = self.criterion(logits, labels)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级:P1

问题: forward 签名接收了 loss_mask,但计算 loss 时没有传给 CriterionLayer。共享 CriterionLayer.forward(logits, labels, loss_mask) 会用这个 mask 控制哪些 token 参与 SFT loss,Phi3/Phi4 等同类 CausalLM 也都把 loss_mask 传下去。

影响: 如果训练链路通过 loss_mask 屏蔽 prompt、padding 或特殊 token,而不是把对应 label 全部改成 -100,MiniMax 会把这些本应忽略的 token 计入 loss,训练目标错误。

处理要求:请针对该评论修复并提交新的 commit。

Suggested change
loss, _ = self.criterion(logits, labels)
loss, _ = self.criterion(logits, labels, loss_mask)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

loss_mask 已传入 Criterion,并增加回归测试


aux_loss = None
if output_router_logits:
aux_loss = None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 优先级:P2

问题: output_router_logits=True 时这里仍然把 aux_loss 设为 None,而 MiniMaxModel 也没有收集或返回任何 router logits;MiniMaxSparseMoeBlock 当前还在调用 gate 后直接丢掉了第一个返回值。这样配置里的 output_router_logits / router_aux_loss_coef 和本文件里的 load_balancing_loss_func 都不会生效。

影响: 调用方打开 MoE router 输出时会得到 aux_loss=Nonerouter_logits=None,不符合 MoECausalLMOutputWithPast 的输出契约,也无法做 router load-balancing loss 或诊断。

处理要求:请针对该评论进行回复(同意并已修改请回复 Done,不同意请说明理由)。

期望: 请二选一明确处理:要么把 raw router logits 从 block/layer/model 逐层传出并在这里计算 aux loss;要么移除/禁用这些配置入口,避免用户以为 MiniMax 已支持 router aux loss。实现形态可以参考:

router_logits, routing_weights, selected_experts = self.gate(hidden_states_flat)
...
return final_hidden_states.reshape([batch_size, sequence_length, hidden_dim]), router_logits

# model.forward: collect all_router_logits when output_router_logits=True
aux_loss = load_balancing_loss_func(outputs.router_logits, self.config.num_local_experts, self.config.num_experts_per_tok, attention_mask)
if labels is not None:
    loss = loss + self.config.router_aux_loss_coef * aux_loss

@yicycyc yicycyc Jul 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

禁用未完整实现的 router logits/aux loss,删除无效代码

@risemeup1111 risemeup1111 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已复查新提交。此前 3 个 P1 问题均已通过代码修改处理:full attention 改用共享 attention 实现,linear attention 对当前不支持的 packing mask 明确报错,loss_mask 已传入 criterion;P2 的 router logits/aux loss 也已通过显式禁用并补充说明处理,回复可接受。

未发现新的阻塞问题。CI 当前通过;本地测试仍因评审环境缺少 paddle 未能运行。

Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants