Add MiniMax Text model support#4662
Conversation
|
Thanks for your contribution! |
risemeup1111
left a comment
There was a problem hiding this comment.
已完成初次 review。当前还有需要先修复后再合入的问题,主要集中在 MiniMax 的 packed attention 处理和训练 loss 传参;详细意见已放在行级评论中。
本地未能运行新增测试:当前评审环境缺少 paddle,直接 import 仓库还缺少 colorlog;CI 页面目前显示通过。
| attention_interface: Callable = eager_attention_forward | ||
| if getattr(self.config, "_attn_implementation", "eager") != "eager": | ||
| attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] |
There was a problem hiding this comment.
问题: 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:
| 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")] |
There was a problem hiding this comment.
Full Attention 改用共享的 ALL_ATTENTION_FUNCTIONS
There was a problem hiding this comment.
已确认当前 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 |
There was a problem hiding this comment.
问题: 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."
)There was a problem hiding this comment.
Linear Attention 暂不支持 packing,遇到 row indices 或 4D packing mask 时明确报错,并补充文档和测试。
|
|
||
| loss = None | ||
| if labels is not None: | ||
| loss, _ = self.criterion(logits, labels) |
There was a problem hiding this comment.
问题: 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。
| loss, _ = self.criterion(logits, labels) | |
| loss, _ = self.criterion(logits, labels, loss_mask) |
There was a problem hiding this comment.
loss_mask 已传入 Criterion,并增加回归测试
|
|
||
| aux_loss = None | ||
| if output_router_logits: | ||
| aux_loss = None |
There was a problem hiding this comment.
问题: 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=None、router_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_lossThere was a problem hiding this comment.
禁用未完整实现的 router logits/aux loss,删除无效代码
…t-01-clean # Conflicts: # README.md # docs/zh/model_capability.md
risemeup1111
left a comment
There was a problem hiding this comment.
已复查新提交。此前 3 个 P1 问题均已通过代码修改处理:full attention 改用共享 attention 实现,linear attention 对当前不支持的 packing mask 明确报错,loss_mask 已传入 criterion;P2 的 router logits/aux loss 也已通过显式禁用并补充说明处理,回复可接受。
未发现新的阻塞问题。CI 当前通过;本地测试仍因评审环境缺少 paddle 未能运行。
PR 新增模型支持:MiniMax-Text-01
本 PR 新增
MiniMaxAI/MiniMax-Text-01支持。复现对齐对象为 HuggingFace/Transformers 端 MiniMax 实现以及 ModelScope 上的原始 safetensors 权重。主要改动
paddleformers/transformers/minimax/MiniMaxConfigMiniMaxModelMiniMaxForCausalLMMiniMaxForCausalLMPipepaddleformers.transformersAutoConfigAutoModelForCausalLMminimaxtemplateMiniMaxPretrainedModel._gen_aoa_config中实现 HF/ModelScope safetensors 到 PaddleFormers 权重的自动转换规则MiniMaxPretrainedModel._gen_inv_aoa_config中实现 PaddleFormers 到 safetensors 的反向映射规则tests/transformers/minimax/test_modeling.pyMiniMaxModelforward、MiniMaxForCausalLMloss/backward、AutoModelForCausalLM.from_configdocs/zh/model_capability.md能力矩阵Stage3 训练跑通方式在下方训练验证中说明,通用 Trainer / Optimizer 侧改动不纳入本 PR。
前向对齐验证
模型:
MiniMaxAI/MiniMax-Text-01验证模型为一层缩层模型:
num_hidden_layers = 1/root/claude_workspace/minimax-text-01-layer0-msbf16 eager 最终 hidden states:
(1, 4, 6144)fp32 eager :
(1, 4, 6144)5.7220459e-062.5422841e-071.9073486e-06说明:
1e-6量级误差,符合 fp32 算子执行差异预期。训练验证
使用 GSM8K 做 BF16 full-SFT,Paddle 侧使用 sharding stage3,Torch 侧使用 ms-swift ZeRO-3。
共同设置:
max_length = 512max_steps = 300per_device_train_batch_size = 1gradient_accumulation_steps = 1learning_rate = 1e-5lr_scheduler_type = constantwarmup_steps = 0weight_decay = 08,146,085,888loss 曲线如下:
结论:
27.x,没有出现异常 loss scale。Stage3 / ZeRO3 说明
MiniMax sparse MoE 会根据 token 动态选择 expert。Stage3 / ZeRO3 下,如果不同 rank 只执行本 rank 命中的 expert,不同 rank 的 expert 执行顺序可能不一致,从而导致 collective 顺序不一致、训练卡住或梯度状态异常。
本地成功验证使用:
生成对齐
使用真实文本输入做 greedy generation,对齐对象为 Transformers 4.57.6 MiniMax 实现。
输入:
输入 token:
bf16 生成前 10 个 token:
[162551, 197066, 109245, 4698, 7169, 29897, 4791, 69130, 190900, 177743][162551, 197066, 109245, 4698, 7169, 29897, 4791, 69130, 190900, 177743]