Magistral-Small-2509 FSDP2 实践文档

本 issue 发布 Magistral-Small-2509 从 Megatron-FSDP2 迁移到纯 FSDP2 训练后端的两篇实践文档，前一篇说明实现方法，后一篇说明测试方法与结果。按维护者意见，这两篇经验文档不随合入请求以文档形式提交到仓库 docs/ 目录，全文以本 issue 为准。

对应代码见 PR #3133，当前提交为 d44a33d276ca64d197739d886372db4d35a728b4；关联任务为 issue #368；安装、数据准备和训练命令见 examples/magistral-2509/README_v1.md。

逐步数据、实验配置、汇总结果、图表和单元测试输出已整理为公开证据包，其 SHA-256 为 2ca862ecf16425bb368aa5a19b90d63854e58b94751454d545978742cfbccdf6。该证据包收录了 2000 步后端对照实验的全部逐步数据、10 个功能配置各 20 步的数据、最终 30 步回归数据、实验配置、汇总结果和单元测试输出；为避免公开训练样本文本，其中不包含原始训练日志全文。

实验在开发分支上分阶段执行，各阶段对应的代码版本如下：30 步训练、分布式检查点恢复和 2000 步后端对比使用提交 cdefc58b；Ulysses 与 Ring Attention 的对照实验使用提交 1134a0fd；其余优化功能的对照实验和此前的最终 30 步回归使用提交 7a151c16。当前合入请求包含任务实现、审查修订和文档移除三个提交。为解决与 master 的冲突，分支已 rebase 到 0029fcfe，提交内容不变，哈希分别为 c6853001、0ecf62e8 和 d44a33d2。其中审查修订提交另行通过了 70 项单元测试、八卡 HCCL all_reduce 与矩阵乘法检查、hybrid_cp_algo 上下文并行配置下的视觉注意力检查，以及 24B 真实数据 30 步训练；文档移除提交只删除文档及其引用，没有改动运行代码、配置或测试逻辑。

---

Magistral-Small-2509 FSDP2 模型迁移

本文记录把 Magistral-Small-2509 从 Megatron-FSDP2 后端迁移到 MindSpeed-MM 纯 FSDP2 训练后端的实现方法。该实现复用 Transformers 中的 Mistral3 模型，迁移工作的重点是保持多图像输入之间的视觉注意力相互独立、确定上下文并行的序列切分边界，以及适配 Pixtral 的数据处理流程。测试方法与性能结果见本文后半部分《Magistral-Small-2509 FSDP2 优化与测试》。

模型结构

该模型由 Pixtral 视觉编码器、多模态投影模块、Mistral 语言模型和输出层 lm_head 组成。适配类 Mistral3ForConditionalGenerationFSDP 继承自 Mistral3ForConditionalGeneration，并以 model_id: mistral3 注册到 ModelHub。训练默认冻结视觉编码器和多模态投影模块的参数。

当上下文并行和分块损失计算都关闭时，该模型直接调用 Transformers 的 forward；当其中任一功能启用时，该模型改为调用 _feature_forward。这两条前向计算路径都通过 register_forward_pre_hook 在视觉编码器执行之前设置注意力参数。

实现代码位于 mindspeed_mm/fsdp/ 目录下，主要文件的职责如下：

• models/mistral3/modeling_mistral3.py 实现模型适配。
• ops/flash_attn/flash_attn.py 把 Transformers 的注意力参数转换为 NPU 融合注意力算子所需的形式。
• data/data_utils/func_utils/ 目录下的 template.py、mm_plugin.py、model_args.py 和 convert.py 共同处理对话文本与图像输入。
• data/dataloader/data_collator.py 注册该模型使用的数据整理器。

多图像视觉注意力

Pixtral 把多张图像的图像块（patch）序列拼接成一条序列后再计算注意力。不同图像之间的注意力必须保持独立，因此需要向注意力算子传入每张图像各自的序列长度。

方法 _prepare_vision_attention 根据 image_sizes 和 patch_size 计算每张图像的图像块数量。当两张图像分别包含 64 个和 96 个图像块时，Transformers 使用的累计偏移量为 [0,64,160]，而 _convert_cu_seq_lens 会把它转换为 NPU 融合注意力算子所需的累计序列长度 [64,160]。这一长度必须在多模态投影模块执行空间合并之前计算，因为空间合并会改变标记数量。

多图像输入使用 TND 张量布局，三个维度分别是总序列长度、注意力头数和每个注意力头的维度。即使多张图像的尺寸完全相同，也必须分别保留各自的序列长度，不能用统一长度代替。回归测试会检查前向输出以及输入和参数的梯度，并验证改变其中一张图像不会影响另一张图像对应的输出。

上下文并行与损失计算

每个上下文并行进程都处理完整的视觉输入，视觉注意力因此设置 skip_context_parallel=True，不参与 Ulysses 或 Ring Attention 的通信。只有在图像特征与文本嵌入合并之后，语言序列才按上下文并行组进行切分。

方法 _context_parallel_model_forward 在切分序列之前先生成全局位置索引，然后把嵌入、位置索引和缓存位置一并分配给各个进程。语言部分的注意力使用仓库已有的 Ulysses、Ring Attention 或混合上下文并行算法（hybrid_cp_algo）实现，损失函数取用对应的标签分片，并在上下文并行组内聚合各分片的损失。做对照实验时必须保持全局批大小与样本顺序一致，否则两次运行的损失不可比较。

分块损失计算改变了 lm_head.forward 的调用方式，使输出层按块计算损失，该模型为此使用一条独立的前向计算分支。激活重计算、异步激活卸载和优化器状态卸载这三项功能直接复用仓库已有的实现，本次迁移没有改动它们。专家并行适用于混合专家模型（MoE），而 Magistral-Small-2509 是稠密模型，不含专家层，因此无法在该模型上执行有意义的专家并行对照实验；该功能是否判定为不适用，需要任务维护者确认。

数据处理

训练配置使用 dataset_type: huggingface 指定数据集类型，使用 template: mistral 指定对话模板，并使用 collator_id: pixtral 指定数据整理器。对话模板负责生成 Mistral 格式的控制标记；PixtralPlugin 负责按图像尺寸把 <image> 占位符展开为对应数量的图像标记；数据整理器负责把一个批次内的文本补齐到相同长度，并传递 pixel_values 和 image_sizes。

图像占位符展开后的标记数量按 processor.patch_size * processor.spatial_merge_size 计算，并在每行末尾保留 [IMG_BREAK]、在整张图像末尾保留 [IMG_END]。这个数量是多模态投影模块执行空间合并之后的标记数量，与前面视觉注意力使用的图像块数量并不相同，两者不能互相代入。参数 fix_mistral_regex 会透传给 tokenizer，其实际行为取决于所安装的 Transformers 版本。

启用数据打包时，数据整理器会同时补齐 input_ids、labels、attention_mask 和 position_ids，这四个张量的序列长度必须保持一致。Mistral 并不使用多维旋转位置编码接口，但补齐后的 position_ids 仍然必须写回批次，否则打包后送入模型的各个张量长度不一致，训练会直接失败。

按最大长度截断时，截断位置可能落在某一张图像的标记序列内部。此时文本中保留的图像标记数量与图像处理器实际生成的特征数量不再一致，模型无法完成图像特征的合并。方法 PixtralPlugin.validate_input_ids 根据每张图像的尺寸、图像块尺寸和空间合并比例计算预期的图像标记数量，数据预处理器则在整理批次之前过滤掉数量不一致的样本。该校验只丢弃不完整的样本，既不修改完整样本，也不会用填充标记伪造缺失的图像内容。

在开始完整模型测试之前，应当先比较两种训练后端得到的标记序列、标签、图像尺寸和样本顺序，并记录模型版本、数据处理配置和数据集版本。训练日志还应当记录被过滤掉的样本数量，以便判断最大长度设置或数据质量筛选是否丢弃了过多样本。

---

Magistral-Small-2509 FSDP2 优化与测试

本文记录 Magistral-Small-2509 从 Megatron-FSDP2 迁移到纯 FSDP2 之后，在正确性、精度、性能和显存四个方面的测试方法与结果。实现说明见本文前半部分《Magistral-Small-2509 FSDP2 模型迁移》。

测试环境

24B 模型的实验运行在单机 8 × Ascend 910B3 上，软件环境为 Python 3.10.21、PyTorch 2.7.1、torch_npu 2.7.1.post8、Transformers 4.57.0、CANN 9.0.0 和驱动 26.0.rc1。两种后端的实验使用相同的 Magistral-Small-2509 权重和 LLaVA-Instruct 图文数据，训练精度为 BF16，并冻结视觉编码器和多模态投影模块。

视觉注意力的专项测试使用两层 Pixtral 编码器，隐藏维度为 128，注意力头数为 2，图像块尺寸为 2。测试输入覆盖单张图像、两张相同尺寸的图像和两张不同尺寸的图像这三种情况。

实验分阶段执行，各阶段对应的代码版本如下：30 步训练、分布式检查点恢复和 2000 步后端对比使用提交 cdefc58b；Ulysses 与 Ring Attention 的对照实验使用提交 1134a0fd；其余优化功能的对照实验和最终的 30 步回归使用提交 7a151c16。

视觉注意力正确性

视觉注意力的正确性测试以 Transformers 的 eager 注意力实现为参考，该参考实现逐张图像单独计算注意力；被测实现则把多张图像拼接为一个批次，交给 npu_fusion_attention 一次算完。两者的输出和梯度按下表比较。

指标	通过条件	实测结果
前向输出的相对 L2 误差	不大于 0.015	最大值为 0.0004810383
输入和参数梯度的相对 L2 误差	不大于 0.04	最大值为 0.0052183196
图像之间的独立性	某张图像的输出对另一张图像输入的梯度为零，且改变另一张图像后该图像的输出保持不变	全部两张图像的测试用例通过
视觉编码器的上下文并行通信	不调用 Ulysses 或 Ring Attention 的通信	Ulysses、Ring Attention 和混合上下文并行算法三种四卡配置均通过

相对 L2 误差按 ||actual-reference||₂ / max(||reference||₂, 1e-10) 计算。单卡测试共 3 项，上下文并行配置下的测试共 36 项，合计 39 项全部通过。

真实数据训练与检查点恢复

纯 FSDP2 后端使用 24B 权重和真实图文数据连续训练了 30 步。这 30 个损失值和梯度范数都是有限值，第 6 至 30 步的平均单步耗时为 2445.04 ms。保存下来的分布式检查点包含模型参数、AdamW 优化器状态、学习率调度器状态、随机数状态和数据加载状态，共 138 个文件、284,644,466,498 字节。从第 30 步恢复训练后，第 31 步的损失为 1.193712，梯度范数为 13.338。

最终代码版本另外完成了一轮 30 步真实数据回归。全部训练步骤连续执行，损失值和梯度范数都是有限值，也没有发生 NPU 显存溢出。数据预处理在 4096 个候选样本中识别并过滤掉了 132 个因最大长度截断而不完整的 Pixtral 样本。

2000 步精度与性能对比

Megatron-FSDP2 和纯 FSDP2 两种后端在同一台主机上先后串行运行，各自完成 2000 步训练。两种后端使用相同的权重、数据顺序、全局批大小、冻结参数范围、优化器和学习率调度。损失差按 sum(abs(loss_pure-loss_base)) / sum(abs(loss_base)) 计算，性能统计使用第 101 至 2000 步的数据，以排除前 100 个预热步骤的影响。

指标	Megatron-FSDP2	纯 FSDP2	结论
平均损失	1.1729174753	1.1720026138	归一化平均绝对损失差为 0.708614%，小于 2% 的验收阈值。
平均单步耗时	2397.8312 ms	2374.4979 ms	纯 FSDP2 的平均单步耗时降低了 0.9731%。
单步耗时中位数	2385.35 ms	2364.70 ms	中位数给出的结论与平均值一致。
最大已分配显存	48945.97 MB	49161.47 MB	该指标不属于本组实验的验收条件。

两种后端的 2000 个训练步骤都连续执行，损失值和梯度范数都是有限值。该结果满足任务要求的损失差小于 2% 且性能不低于基线这两项条件。

优化功能对照实验

每项优化功能都与关闭该功能的配置比较，同一组对照中的其余配置项保持一致。每个配置连续运行 20 步，并从第 6 步开始统计平均单步耗时。所有配置的损失值都是有限值，也都没有发生 NPU 显存溢出。

功能	归一化平均绝对损失差	最大已分配显存变化	最大保留显存变化	平均单步耗时变化
Ulysses 上下文并行，组大小 2	0.272835%	-3377.44 MB	-3318 MB	+25.45%
Ring Attention，组大小 2	0.160284%	-2030.08 MB	-1938 MB	+39.78%
激活重计算	0.295354%	-7120.50 MB	-4520 MB	-50.07%
分块损失计算，块大小 256	1.785513%	-372.00 MB	+1680 MB	+0.95%
异步激活卸载	1.389279%	0 MB	-520 MB	-3.25%

表中显存变化为负值表示开启该功能后显存占用减少。判断短期单步耗时差异时必须先确定对照组自身的波动幅度：上下文并行对照组重复运行的平均单步耗时变化为 -1.14%，微批大小为 2 的对照组重复运行的变化为 -3.89%。异步激活卸载的单步耗时差异落在这一波动范围之内，因此不能认定为性能提升，它在本次实验中站得住的结果是最大保留显存减少了 520 MB。分块损失计算只使最大已分配显存减少 372 MB，同时使最大保留显存增加 1680 MB，因此本次实验没有证明该功能带来整体的显存或性能收益。

激活重计算使最大已分配显存减少约 6.95 GiB，是本次单卡 64 GiB 环境中效果最明显的显存优化手段。该组实验同时观察到平均单步耗时下降，但 20 步的测试长度不足以证明这一性能差异能够稳定复现，因此不作为性能结论。当显存接近容量上限时，建议首先启用激活重计算；如果显存仍然不足，再根据序列长度和可以接受的通信开销，在 Ulysses 和 Ring Attention 之间选择。

Ulysses 和 Ring Attention 都是语言序列的可选上下文并行实现，它们并没有替换默认的注意力实现：默认情况下语言部分的注意力仍然由 Transformers 的 Mistral3 调用 torch_npu.npu_fusion_attention 完成，视觉编码器也继续使用 NPU 融合注意力，并跳过上下文并行通信。

各项方法的作用与代价

方法	作用与代价
NPU 融合注意力	该算子避免显式保存完整的注意力矩阵，从而减少显存占用；代价是多图像输入必须额外提供每张图像的序列长度。
FSDP2	该策略把参数、梯度和优化器状态分片到各个进程，从而减少单卡显存占用；代价是增加集合通信开销。
激活重计算	该方法在反向传播时重新计算前向激活值，从而减少显存占用；代价是增加计算量。
优化器状态卸载	adamw_swap 在主机内存和设备显存之间传输优化器状态，从而减少设备显存占用；代价是增加主机内存占用和数据传输开销。
上下文并行	该方法在进程之间切分语言序列，使每个进程只保存一部分序列数据；代价是增加与序列相关的通信开销。
分块损失计算	该方法分块计算输出层损失，减少同时驻留的 logits 数量；但在本次实验中，它在使最大已分配显存减少 372 MB 的同时使最大保留显存增加了 1680 MB。
异步激活卸载	该方法把激活值转存到主机内存，并尝试让数据传输与计算重叠；本次实验只观测到最大保留显存减少 520 MB，单步耗时差异落在对照组波动范围内。

专家并行适用于混合专家模型。Magistral-Small-2509 是稠密模型，不含专家层，因此无法在该模型上执行有意义的专家并行对照实验；该功能是否判定为不适用，需要任务维护者确认。

结果的适用边界

2000 步对比实验支持的是当前配置下的精度和性能结论。20 步的功能对照实验只能用于确认功能是否正确、峰值显存如何变化以及短期单步耗时的量级，不能代替长期稳定性测试。全部性能结果都来自一台 8 × Ascend 910B3 主机，不能直接外推到多机环境或其他硬件。本地自动化测试全部通过，也不代表平台持续集成、人工代码审查或任务维护者已经完成验收。

结果图

2000 步损失对比

第 101 至 2000 步单步耗时对比

各项优化功能的显存与单步耗时对比

最终代码版本的 30 步回归