已开启
feat(torch): add Magistral-Small-2509 FSDP2 training #3133
chenchen创建于 9月11日
feat(torch): add Magistral-Small-2509 FSDP2 training #3133
已开启
合并受阻
9月11日 关联了issue：[Feature]: 基于FSDP2后端训练MM大模型
atomgit-bot
9月11日 评论：
9月11日 评论：
变更摘要
本 PR 为 Magistral-Small-2509（Mistral3）增加纯 FSDP2 训练支持（对应 issue #368）：新增 Mistral3ForConditionalGenerationFSDP 模型适配并注册为 model_id: mistral3，实现多图视觉注意力按每张图像序列长度分别计算且视觉编码器不参与上下文并行（CP）通信，语言序列支持 Ulysses、Ring 和 Hybrid 并行，并适配分块损失（ChunkLoss）；同时补充数据侧注册（mistral 模板、pixtral 插件与数据整理器、fix_mistral_regex 参数）及 8 卡/4 卡训练脚本与 YAML 配置、相关测试。训练默认冻结视觉编码器与投影模块，使用 BF16；PR 描述中说明还包含三份使用与迁移文档（无内嵌补丁）。
主要改动
- Mistral3 模型适配：新增
Mistral3ForConditionalGenerationFSDP（注册于model_id: mistral3），默认forward直接委托上游Mistral3ForConditionalGeneration；当启用分块损失或上下文并行时改走_feature_forward，分块损失路径下lm_head直接接收loss_function与labels并返回 loss，CP 路径通过_context_parallel_model_forward在图像合并与解码器之间切分序列，并将 CP 组内各分片损失求和。 - 视觉注意力多图隔离：
_prepare_vision_attention前置钩子按各图像尺寸生成cu_seq_lens/max_length并设置skip_context_parallel=True，视觉编码器在各 rank 上复制处理完整图像；_context_parallel_model_forward中通过skip_ulysses=True保证视觉塔不进入 Ulysses all-to-all，防止不同图像的视觉 token 被拼接在一起计算。 - flash attention 打包序列转换重构：
_convert_cu_seq_lens重写1TND/1NTD布局处理，校验 Q/K 长度成对出现、去掉 HF 偏移中的零前缀并转换为 NPU 累积结束位置，单序列或等长序列回退到稠密BSND/BNSD；flash_attention_forward新增skip_context_parallel参数，供视觉塔等序列切分之外的复制模块跳过 CP 分支。 - 数据侧注册：新增
mistral对话模板（[INST]…[/INST]、[SYSTEM_PROMPT]等格式，绑定pixtral插件与[IMG]图像 token）；PixtralPlugin将图像占位符展开为按行image_break_token、末位image_end_token的 patch 网格，并在 transformers < 4.49.0 时移除image_sizes入参；data_collator注册pixtral并复用DataCollatorForQwen2vl；ProcessorArguments新增fix_mistral_regex字段并透传给AutoTokenizer.from_pretrained。 - 训练示例与测试：新增
examples/magistral-2509/下 8 卡（finetune_magistral_2509_v1.sh）与 4 卡（optimizer: adamw_swap）微调脚本及两份 YAML 配置（FSDP2 分片计划、BF16、冻结model.vision_tower与model.multi_modal_projector）；新增 30 步合成小模型 smoke 测试train_smoke.py、NPU 视觉注意力探针vision_attention_probe.py，以及覆盖数据侧注册与视觉注意力行为的单元测试。
不准确？
ascend-robot
9月11日 评论：
9月11日 评论：
atomgit-bot
9月11日 评论：
9月11日 评论：
此处折叠了223条消息 查看更多
6 小时前 删除了label：docs-ci-pipeline-running
6 小时前 添加了label：docs-ci-pipeline-success
ascend-robot
6 小时前 评论：
6 小时前 评论：
6 小时前 删除了label：ci-pipeline-failed ci-pipeline-running
6 小时前 添加了label：ci-pipeline-failed
What this PR does / why we need it?
为 #368 增加 Magistral-Small-2509 纯 FSDP2 训练支持，使该模型不再依赖 Megatron-FSDP2 后端即可训练，并与现有 Megatron-FSDP2 配置进行精度和性能比较。
hybrid_cp_algo上下文并行。position_ids补齐，并过滤因最大长度截断而不完整的 Pixtral 图像标记序列。forward参数、缺少image_sizes时的处理和相应测试。Does this PR introduce any user-facing change?
新增以下训练入口和说明：
examples/magistral-2509/finetune_magistral_2509_v1.sh和magistral_2509_config_v1.yaml提供单机 8 卡 BF16 配置。examples/magistral-2509/finetune_magistral_2509_v1_4card.sh和对应 YAML 提供使用adamw_swap的单机 4 卡配置。examples/magistral-2509/README_v1.md说明安装、数据准备、训练和检查点恢复方法。模型迁移与性能优化实践文档不随本 PR 提交，全文见 issue #703。
How was this patch tested?
当前提交为
89fb8fe1930a9125000851d8c096c233600e77fa。2026-10-12 将原来的三个提交压缩为一个提交，Git 文件树与压缩前的d44a33d276ca64d197739d886372db4d35a728b4完全一致，运行代码、配置、测试和依赖均未改变。2026-10-13 在该单提交中补齐视觉注意力测试替身的is_kvallgather_enable()接口（返回False），修复与当前主线组合时的 UT 失败；另将训练说明中安装指南、数据构造指南的两处链接更新到当前主线目录。只修改测试与说明链接，未修改运行代码、训练配置或依赖。下文训练实验继续引用实际被测的历史版本，未重复执行训练实验。实测结果
3213496c（rebase 后为0ecf62e8，现已包含在单提交中）完成 30 步真实数据训练，损失和梯度范数均为有限值，没有 NPU 显存溢出；相关范围的 70 项单元测试通过。a6543278的 GitCode Actions #374：代码检查及 ST 通过，UT 因测试替身缺少is_kvallgather_enable()失败。测试修复提交34a89080已在昨日和当前主线的预合并版本上分别通过 93 项相关 CPU 测试（昇腾导入依赖使用模拟包）；其 Actions #386 尚在运行，文档检查发现上游目录调整导致两处相对链接失效。当前单提交89fb8fe1已同时包含测试接口修复与两处链接修复，运行代码未变；CLA、文档门禁和完整 CI 待重新核对。文档移除提交只删除两篇实践文档和三处文档链接，没有改动运行代码、配置、测试逻辑或依赖，因此没有重复执行上述实验。
压缩前，本分支为解决与
master的冲突 rebase 到0029fcfe，三个历史提交的内容不变，哈希相应变化：任务实现提交由74760a54变为c6853001，审查修订提交由3213496c变为0ecf62e8，文档移除提交由306dd193变为d44a33d2。冲突只出现在mindspeed_mm/fsdp/data/dataloader/data_collator.py的注册段落，上游新增的text_only与本 PR 新增的pixtral两项注册均已保留。rebase 后docs/目录不再有任何改动。下文引用流水线记录时沿用当时的提交哈希。优化功能
各功能分别与关闭该功能的配置比较。每个配置连续运行 20 步，损失差均低于 2%，且没有发生 NPU 显存溢出。
损失差按
sum(abs(loss_pure-loss_base)) / sum(abs(loss_base))计算，性能统计排除前 100 个预热步骤。20 步的功能对照实验只能支持功能正确性、损失差、峰值显存和短期单步耗时这几项结论，不能替代长期稳定性测试。全部性能结果来自单机 8 × Ascend 910B3，不能直接外推到多机或其他硬件。实验分阶段执行：30 步训练、DCP 恢复和 2000 步后端比较使用
cdefc58b；Ulysses 与 Ring Attention 使用1134a0fd；其余优化功能和此前最终回归使用7a151c16。上述任务实现、审查修订和文档移除原对应 rebase 后的三个历史提交c6853001、0ecf62e8和d44a33d2，先合并为单提交a6543278，后续测试兼容修复并入当前单提交89fb8fe1。审查修订提交另行通过 70 项单元测试、八卡 HCCLall_reduce与矩阵乘法检查、hybrid_cp_algo上下文并行配置下的视觉注意力检查和 24B 真实数据 30 步训练。2000 步实验使用非打包数据路径，后续数据处理修复不会在该实验中触发。提交
74760a54的流水线#4809和提交3213496c的流水线#4818首次运行时，Qwen3.5 系统测试都因为出现了单个耗时约 5 秒的异常步骤而未达到单步耗时阈值。流水线#4814和#4818的重试都已通过；这两次通过的运行同样出现过同类异常步骤，只是该步骤落在预热区间内。除该异常步骤以外，参与判定的步骤耗时均低于基准，因此没有证据表明存在持续的性能回归。原始失败和通过结果均已保留。Magistral-Small-2509 是稠密模型，没有专家层，因此无法执行有意义的专家并行对照。请维护者确认专家并行在本模型上不适用。
复现与材料
2ca862ecf16425bb368aa5a19b90d63854e58b94751454d545978742cfbccdf6。