变更说明

变更摘要

该 PR 主要围绕 Issue #369，为 Magistral-LoRA 训练路径新增纯 FSDP2 入口，并保留原有 Megatron-FSDP2 路径作为对比基线。核心改动包括：新增纯 FSDP2 的模型/处理器/插件与 LoRA 优化器集成、Launcher 与配置、DCP/key 校验、数值评测和性能工具，以及配套的开发与性能实践文档；同时引入确定性的 LoRA 初始化与优化器同步机制，用于在两条路径间对齐训练完整性与数值等价性。

主要改动

• 新增纯 FSDP2 Magistral-LoRA 入口：新增 mindspeed_mm/fsdp/models/mistral3/__init__.py 中的 PureMistral3ForConditionalGeneration，以及 magistral3_fsdp2_lora_config.yaml 和 finetune_magistral_2509_pure_fsdp2_lora.sh，提供独立的模型适配器、训练配置和启动脚本，将 LoRA 训练切到纯 FSDP2 路径。
• 接入 Pixtral 多模态插件与 tokenizer/processor 兼容路径：在 mm_plugin.py 中用 PixtralPlugin 替换 MiniMaxM3VLPlugin，在 template.py 注册带 pixtral 插件的 mistral 模板并在 data_collator.py 注册 pixtral 收集器，同时在 convert.py 中重写 load_tokenizer/update_tokenizer_with_chat_template，以支持 Mistral3/Tekken tokenizer 与 Pixtral 图像处理。
• 新增确定性 LoRA 初始化与优化器同步：新增 task3_deterministic_lora.py 提供 apply_once、sync_optimizer_lora 和 audit_optimizer，并在 lora_utils.py、lora_patch.py、train_engine.py、training.py 中接入，使 LoRA 初始化不依赖 wrapper 遍历顺序，并在优化器首次更新前同步模型与主参数。
• 调整训练与 checkpoint 逻辑：train_engine.py 在 clip_grad <= 0 时跳过冗余的全局 FSDP2 梯度范数计算，LoRA 保存后提前返回避免重复保存基座模型；torch_dcp_patch.py 将 DCP key 检查改为严格模式，存在无法映射到模型的 checkpoint key 时抛出异常。
• 新增评测与诊断工具：新增 scripts/task3_final_eval.py、scripts/measure_task3_performance.py 以及 task3_diagnostics.py、task3_runtime_audit.py、task3_layer0_audit.py、task3_stage_timing.py，用于分阶段门控地检查 checkpoint 完整性、数值误差和运行时性能。

代码审查

Closing Summary

已逐一审查全部 14 个变更文件，结论如下：

• checkpoint/common/hf_to_dcp.py：已审查，新增 model.safetensors.index.json 索引解析逻辑，未发现需要上报的问题。
• examples/magistral-2509/data.json：已审查，仅路径/预处理参数值调整，JSON 键无拼写问题。
• examples/magistral-2509/finetune_magistral_2509_lora.sh：已审查，仅新增两个默认开启的环境变量导出，无问题。
• examples/magistral-2509/finetune_magistral_2509_pure_fsdp2_lora.sh：已审查（新增），set -euo pipefail、sed 生成运行时配置、trap 清理等，未发现确定性问题。
• examples/magistral-2509/fsdp2_lora_config.yaml：已审查，reduce_dtype 由 bf16 改为 fp32，属有意配置变更，键名正确。
• examples/magistral-2509/magistral3_fsdp2_lora_config.yaml：已审查（新增），YAML 键名无拼写问题。
• examples/magistral-2509/model.json：已审查，仅路径值变更，键名正确。
• mindspeed_mm/fsdp/data/data_utils/func_utils/convert.py：已审查，load_tokenizer 重构及 Mistral3 处理器/Tokenizer 兼容逻辑，未发现确定性问题。
• mindspeed_mm/fsdp/data/data_utils/func_utils/mm_plugin.py：已审查，发现 2 个问题（1 个 P2、1 个 P3），见上方报告。
• mindspeed_mm/fsdp/data/data_utils/func_utils/template.py：已审查，删除 minimax_m3_vl 模板、新增 mistral 模板，未发现确定性问题。
• mindspeed_mm/fsdp/data/dataloader/data_collator.py：已审查，新增 pixtral -> DataCollatorForQwen2vl 注册，复用已有 collator，未发现确定性问题。
• mindspeed_mm/fsdp/models/mistral3/__init__.py：已审查（新增），模型适配类，未发现确定性问题。
• mindspeed_mm/fsdp/ops/fully_shard/fully_shard.py：已审查，为 2.10.0 复用 2.9.0 补丁路径，未发现确定性问题。
• mindspeed_mm/fsdp/tools/task3_diagnostics.py：已审查（新增），惰性导入的诊断门面，默认关闭，未发现确定性问题。

问题计数：P0 0 个，P1 0 个，P2 1 个，P3 1 个。

总体风险判断：本 PR 以新增 pure-FSDP2 Magistral-LoRA 路径为主，核心风险集中在 PixtralPlugin.process_messages 的图像 token 展开逻辑上——其中 transformers < 4.49 多图样本会因只取首图尺寸而崩溃（P2），极端长宽比图像存在空列表越界（P3，低置信度）。其余配置、脚本、注册与版本分支改动未见明确缺陷。建议重点验证多图样本与旧版 transformers 组合下的数据预处理路径。

我已完成对全部 16 个变更文件的逐一审查。最终报告如下。

审查结论

已报告的问题（按优先级）

各变更文件审查确认

1. mindspeed_mm/fsdp/train/train_engine.py — 已审，无确凿问题（clip_grad<=0 跳过 norm 计算与代码库另一引擎语义一致；LoRA save 提前 return 属文档化意图）。
2. mindspeed_mm/fsdp/utils/lora_utils.py — 已审，无确凿问题（新增的确定性 LoRA 初始化与其他副本算法一致）。
3. mindspeed_mm/models/transformers/cp_utils.py — 已审，无问题（新增 get_safe_context_parallel_world_size 合理）。
4. mindspeed_mm/models/transformers/mistral3/mistral3.py — 已审；存在用全新 MMMistralModel/MMMistralAttention 覆盖既有模块的写法，但权重是否被丢弃取决于加载顺序，证据不足，未报告。
5. mindspeed_mm/models/transformers/mistral3/modeling_mistral.py — 已审，无问题。
6. mindspeed_mm/models/transformers/mistral3/modules.py — 已审，无问题。
7. mindspeed_mm/patchs/torch_dcp_patch.py — 已审，报告 1 项（P2）。
8. mindspeed_mm/task3_deterministic_lora.py — 已审，无确凿问题（诊断/同步逻辑由 env 开关控制，默认开启但语义自洽）。
9. mindspeed_mm/task3_layer0_audit.py — 已审，无问题（默认关闭的诊断代码）。
10. mindspeed_mm/task3_runtime_audit.py — 已审，无问题。
11. mindspeed_mm/task3_stage_timing.py — 已审，无问题。
12. mindspeed_mm/tasks/finetune/lora/lora_patch.py — 已审，无确凿问题（新增确定性初始化与 lora_utils.py 副本一致）。
13. mindspeed_mm/training.py — 已审，报告 1 项（P0）。
14. pretrain_transformers.py — 已审，无问题（Task3Diagnostics 在 import 之后使用，正常）。
15. scripts/measure_task3_performance.py — 已审，无问题。
16. scripts/task3_final_eval.py — 已审，无问题。

整体风险判断

高风险：mindspeed_mm/training.py 第 11 行的 NameError 会让整个 Megatron-FSDP2 桥接训练入口在导入阶段直接崩溃，属于必须修复的启动级缺陷；torch_dcp_patch.py 的严格校验则存在破坏断点续训的回归风险。建议优先修复 P0 项后再合入。

⛔ 需要修改

✅ 文档门禁通过！