feat(torch): add Magistral-Small-2509 FSDP2 training and A3 optimizations-MindSpeed-MM-AtomGit

Aexamples/magistral-2509/README_v1.md+163-0

@@ -0,0 +1,163 @@
	1	+# Magistral-Small-2509 FSDP2 训练说明
	2	+
	3	+本示例使用纯 FSDP2 后端训练 Magistral-Small-2509。24B 模型的真实数据训练、检查点恢复、2000 步精度与性能比较以及优化功能测试结果见[实践文档 issue](https://gitcode.com/Ascend/MindSpeed-MM/issues/703)。
	4	+完整的两侧对照配置、启动入口和固定资产来源见[2000 步对照说明](baseline/README.md)。
	5	+
	6	+## 安装与模型文件
	7	+
	8	+按[安装指导](../../docs/zh/guides/installation/install_guide.md)准备环境。小模型 NPU 测试使用以下版本：
	9	+
	10	+\| 组件 \| 版本 \|
	11	+\| --- \| --- \|
	12	+\| 设备 \| 单机 910B3，使用 1 或 4 张卡。 \|
	13	+\| CANN / 驱动 \| 9.1.0 / 26.0.rc1。 \|
	14	+\| Python \| 3.12.13。 \|
	15	+\| PyTorch / torch_npu \| 2.10.0+cpu / 2.10.0.post4。 \|
	16	+\| Transformers \| 4.57.0。 \|
	17	+\| MindSpeed \| `c356513e331e55d53b799cfde51c7b23901d23dc`。 \|
	18	+
	19	+NPU 计算由 torch_npu 提供。MindSpeed 需包含 `mindspeed.fsdp`；安装和启动使用同一 Python 环境。
	20	+
	21	+```bash
	22	+python -m pip install -e .
	23	+python - <<'PYCODE'
	24	+import torch, torch_npu, transformers, mindspeed.fsdp
	25	+print(torch.__version__, torch_npu.__version__, transformers.__version__)
	26	+PYCODE
	27	+```
	28	+
	29	+从
	30	+[mistralai/Magistral-Small-2509](https://huggingface.co/mistralai/Magistral-Small-2509/tree/main)
	31	+下载 `config.json`、Hugging Face 格式的 safetensors 权重分片、索引和 `tekken.json`。本示例不使用
	32	+`consolidated.safetensors`。
	33	+
	34	+数据处理还需 `tokenizer.json`、`tokenizer_config.json`、`special_tokens_map.json`、
	35	+`chat_template.jinja`、`preprocessor_config.json` 和 `processor_config.json`，可从
	36	+[unsloth/Magistral-Small-2509](https://huggingface.co/unsloth/Magistral-Small-2509/tree/main)
	37	+取得并放入同一目录。这些是第三方转换的配置文件；分别记录两处来源的版本和文件校验和，保留官方权重与 `config.json`。
	38	+
	39	+训练前检查文件能否加载，再用代表性文本、特殊 token 和图像与原数据处理结果比较：
	40	+
	41	+```bash
	42	+python - <<'PYCODE'
	43	+from transformers import AutoConfig, AutoTokenizer, AutoProcessor
	44	+model_dir = '/home/data/Magistral-Small-2509' # 修改为实际路径。
	45	+config = AutoConfig.from_pretrained(model_dir, trust_remote_code=False)
	46	+tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=False)
	47	+processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=False)
	48	+assert config.model_type == 'mistral3'
	49	+assert processor.patch_size * processor.spatial_merge_size == 28
	50	+print(type(tokenizer).__name__, type(processor).__name__)
	51	+PYCODE
	52	+```
	53	+
	54	+## 数据与训练配置
	55	+
	56	+按[图文数据构造指南](../../docs/zh/features/data/building_data_for_VLModel.md#real-data)准备
	57	+COCO 图像和 LLaVA 对话。数据包含 `images` 和 `messages`，每条样本的 `<image>` 数量应与图像数量一致。
	58	+本示例处理文本和图像对话，未测试工具调用、视频或音频。
	59	+
	60	+编辑 [8 卡配置](magistral_2509_config_v1.yaml)：
	61	+
	62	+1. 设置 `data.dataset_param.preprocess_parameters.model_name_or_path`。
	63	+ `model.model_name_or_path` 和 `training.load` 通过 YAML 别名引用该路径，权重按默认
	64	+ `load_format: auto` 加载。
	65	+2. 设置 `basic_parameters` 下的 `dataset_dir`、`dataset` 和 `cache_dir`。
	66	+ 更换数据或处理配置后使用新缓存目录。
	67	+3. 按测试目的设置 `training.train_iters`：短测为 30 步，正式精度对照为 2000 步。改变总步数会改变学习率预热和调度，
	68	+ 应同步调整两侧配置。
	69	+4. 如需保存检查点或断点续训，设置 `training.save`、有效的 `save_interval`，并配置优化器和随机数状态的保存、恢复。
	70	+ 默认未设置保存路径，`no_save_optim/no_save_rng/no_load_optim/no_load_rng` 均为 true。
	71	+
	72	+默认使用 BF16、FSDP2 参数分片和激活重计算；冻结视觉编码器及多模态投影模块，仅训练语言模型和输出层。微批大小为 1，梯度累积次数为 1，8
	73	+卡全局批大小为 8；序列截断长度为 2048，学习率为 `1e-5`，使用余弦衰减、0.1 预热比例、随机种子 42 和 5000 个训练步。
	74	+上下文并行（CP）与分块损失计算（ChunkLoss）默认关闭。
	75	+
	76	+## 启动
	77	+
	78	+确认 8 张卡可用后，在仓库根目录运行：
	79	+
	80	+```bash
	81	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	82	+export NON_MEGATRON=true
	83	+bash examples/magistral-2509/finetune_magistral_2509_v1.sh
	84	+```
	85	+
	86	+脚本使用 8 个进程和端口 6000，日志位于 `logs/`；`pipefail` 会传递训练失败的退出状态。`torchrun` 必须属于上述 Python
	87	+环境。多机配置尚未测试。
	88	+
	89	+[4 卡配置](magistral_2509_config_v1_4card.yaml)使用 `adamw_swap` 卸载优化器状态，梯度累积次数为 2，无
	90	+CP 时全局批大小为 8。运行对应的 `finetune_magistral_2509_v1_4card.sh`。该配置需要额外的主机内存和数据传输，当前
	91	+24B 模型仍需重新测试。
	92	+
	93	+开启 CP 时，全局批大小应按数据并行组数计算，并调整梯度累积次数。
	94	+普通 `packing: true` 可与 CP 一起使用，样本之间允许因果注意力；它不提供样本间注意力隔离。
	95	+当前 Flash Attention 和 CP 路径不支持 `neat_packing`、左侧 padding 或自定义注意力掩码；
	96	+此类输入会在模型执行前报错。需要样本隔离时请关闭 packing，不能将普通 packing 当作等价替代。
	97	+各优化选项的实现见[实践文档 issue](https://gitcode.com/Ascend/MindSpeed-MM/issues/703)。
	98	+
	99	+### A3 性能配置
	100	+
	101	+[优化配置](magistral_2509_config_optimized.yaml)面向 PyTorch 2.10、torch_npu 2.10、Triton Ascend 和
	102	+Ascend A3。保留完整模型、全局批大小、训练数据与优化器，只启用以下执行优化：
	103	+
	104	+- `num_to_forward_prefetch` 和 `num_to_backward_prefetch` 显式预取后续 FSDP 单元，增加通信与计算重叠。
	105	+- `use_triton_split_copy` 使用 Triton 将 rank-major all-gather 缓冲区直接复制到参数。该开关默认关闭，
	106	+ 仅支持 PyTorch 2.10 的 NPU FSDP，启用后在当前进程内生效；不满足布局、类型和别名条件时调用原操作。
	107	+- `contiguous_linear_grad_modules` 只选择 MLP 的无偏置原生 Linear，直接生成连续的权重梯度，减少后续转置。
	108	+ 参数身份和 state_dict 名称保持不变；CPU、FP32 与 autocast 使用原生 Linear。
	109	+- ChunkLoss 过滤忽略标签后，只有一个 chunk 时直接使用普通 autograd，避免在前向保存完整的权重梯度。
	110	+ `infinite_sampler` 让 worker 跨 epoch 持续预取，保持有限 sampler 的样本顺序和恢复位置。
	111	+
	112	+A3 对照使用 8 个逻辑 NPU（Ascend910_9392），并非整机全部 16 个芯片。环境为 CANN 9.1、
	113	+PyTorch 2.10.0、torch_npu 2.10.0.post4、Triton Ascend 3.2.2、Python 3.12.13 和 Transformers 4.57.0。
	114	+配置保留全部 40 层语言模型重计算；部分重计算在 2048 长度压力检查中发生分配重试或 OOM，未采用。
	115	+ChunkLoss 在多块累加后及时释放临时梯度和 logits，降低后续块的瞬时显存；数值计算保持原样。
	116	+
	117	+对于每个 rank 仅训练一次的独立进程，可选 `training.freeze_gc_after_init: true` 冻结初始化对象至进程退出，
	118	+训练期间新建对象的自动循环 GC 继续启用，阈值保持不变。该选项全局默认关闭，本 A3 优化配置显式开启，与 `manual_gc_interval` 互斥；
	119	+再次训练需要新进程。初始化回收的耗时应计入完整运行，引用计数仍生效，但冻结对象构成的不可达循环可能保留至进程退出，因此不适用于嵌入式训练服务。
	120	+
	121	+先按本机修改模型、数据和 CANN 路径，再运行：
	122	+
	123	+```bash
	124	+bash examples/magistral-2509/finetune_magistral_2509_v1.sh \
	125	+ examples/magistral-2509/magistral_2509_config_optimized.yaml
	126	+```
	127	+
	128	+完整官方权重（revision `a31cc96ab10cf19bc42c628fedf1e359e0853c49`）和固定 256 条 LLaVA/COCO 样本的
	129	+100 步对照中，global batch 8、cutoff 2048、seed 42 和优化器保持一致。数据取
	130	+`liuhaotian/LLaVA-Instruct-150K` revision `9d451dc7629cfe0469f6ae4432b765cd603d5fcb` 的前 256 条，
	131	+保留原顺序和 COCO 图像。测量时将 `training.train_iters` 设为 100、`load_format` 设为 `hf`，
	132	+并启用 `load_strict`、`no_load_optim` 和 `no_load_rng`。每组使用空 Triton 缓存，
	133	+CANN 缓存保留，加载前逐片 SHA256 校验同时预热文件缓存。第 6–100 步均保留慢步：
	134	+
	135	+\| 指标 \| 普通配置 \| 优化、自动 GC \| 优化、初始化 GC 冻结 \|
	136	+\| --- \| --- \| --- \| --- \|
	137	+\| 平均步时 \| 1987.70 ms。 \| 1774.61 ms。 \| 1755.09 ms（减少 11.70%）。 \|
	138	+\| 中位数步时 \| 1956.8 ms。 \| 1759.8 ms。 \| 1757.3 ms。 \|
	139	+\| p95 步时 \| 2234.7 ms。 \| 1875.1 ms。 \| 1868.3 ms。 \|
	140	+\| 完整 torchrun 时间 \| 286.02 秒。 \| 277.54 秒。 \| 274.54 秒（减少 4.01%）。 \|
	141	+\| 最大 rank 的 peak allocated \| 50205.92 MiB。 \| 50079.92 MiB。 \| 50079.92 MiB。 \|
	142	+\| 最大 rank 的 peak reserved \| 54560 MiB。 \| 55600 MiB。 \| 55600 MiB。 \|
	143	+\| 八个 rank 的分配重试 \| 全部为 0。 \| 全部为 0。 \| 全部为 0。 \|
	144	+
	145	+完整运行包含初始化、权重加载、JIT、GC、训练与进程退出。两组相对基线的 100 步归一化平均绝对损失差分别为
	146	+0.841% 和 0.901%，逐步相对损失误差的平均值分别为 2.960% 和 3.348%，最大值分别为 35.884% 和 51.397%。
	147	+这些指标的分母不同；不表示每一步都在 2% 内，也不代替历史 2000 步后端比较。
	148	+此处是固定工作负载的单次对照，配置的 5000 步训练尚未在该优化版本重跑。详细来源见
	149	+[实践文档 issue](https://gitcode.com/Ascend/MindSpeed-MM/issues/703)。
	150	+
	151	+预取需要更多显存，Triton 首次编译有额外启动成本。比较时固定模型和输入，同时报告预热后的步时、
	152	+完整启动时间及全程峰值显存。不要用局部复制或矩阵乘法的耗时比例代替完整模型收益。
	153	+
	154	+## 测试
	155	+
	156	+安装完整依赖后运行 CPU 回归测试：
	157	+
	158	+```bash
	159	+NON_MEGATRON=true python -m pytest tests/ut_fsdp/data -q
	160	+```
	161	+
	162	+NPU 测试命令见[测试说明](../../tests/st/fsdp/mistral3/README.md)。这些测试使用小模型和合成输入，完整 24B
	163	+模型按上述训练脚本另行测试。

Aexamples/magistral-2509/baseline/README.md+74-0

@@ -0,0 +1,74 @@
	1	+# 2000 步后端对照的配置与入口
	2	+
	3	+本目录补充历史对照实验的完整 Megatron-FSDP2 模型、数据和分片配置，以及纯 FSDP2 YAML。
	4	+配置来自提交 `cdefc58b76fbfd876ece4361e04735f8e1f9ea4b` 的 2000 步实验，仅替换了机器相关路径。
	5	+历史环境为 Python 3.10.21、PyTorch 2.7.1、torch_npu 2.7.1.post8、Transformers 4.57.0 和 CANN 9.0.0。
	6	+Megatron 侧另外需要仓库安装指导要求的 Megatron/MindSpeed 依赖；原运行使用 MindSpeed `d76dbddd`。
	7	+本目录新增启动器经过参数及配置核对，尚未用它重跑 24B 的 2000 步实验。
	8	+
	9	+## 模型与数据准备
	10	+
	11	+模型权重及 `config.json` 使用官方仓库，tokenizer/processor 使用第三方转换配置。新的复现固定到：
	12	+
	13	+- 官方模型：`mistralai/Magistral-Small-2509`，revision `a31cc96ab10cf19bc42c628fedf1e359e0853c49`。
	14	+- tokenizer/processor：`unsloth/Magistral-Small-2509`，revision `dcfa8f675a2adaea5dad53b32c0c948b47c02f55`。
	15	+
	16	+使用已安装的 `huggingface_hub` 下载到同一目录：
	17	+
	18	+```bash
	19	+python - <<'PY'
	20	+from huggingface_hub import snapshot_download
	21	+model_dir = '/path/to/Magistral-Small-2509'
	22	+snapshot_download('mistralai/Magistral-Small-2509', revision='a31cc96ab10cf19bc42c628fedf1e359e0853c49',
	23	+ local_dir=model_dir, token=False,
	24	+ allow_patterns=['config.json', 'tekken.json', 'model*.safetensors', 'model.safetensors.index.json'])
	25	+snapshot_download('unsloth/Magistral-Small-2509', revision='dcfa8f675a2adaea5dad53b32c0c948b47c02f55',
	26	+ local_dir=model_dir, token=False,
	27	+ allow_patterns=['tokenizer.json', 'tokenizer_config.json', 'special_tokens_map.json',
	28	+ 'chat_template.jinja', 'preprocessor_config.json', 'processor_config.json'])
	29	+PY
	30	+```
	31	+
	32	+[asset_manifest.json](asset_manifest.json)提供这组配置与 tokenizer 文件的 SHA-256。
	33	+该清单经过文件下载、校验和 Transformers 4.57.0 的加载检查；它是新的复现来源记录，
	34	+不是对 9 月历史实验资产身份的补证。历史实验没有完整保存这些资产的 revision 与逐文件校验值。
	35	+
	36	+数据按[图文数据构造指南](../../../docs/zh/features/data/building_data_for_VLModel.md#real-data)准备
	37	+COCO 图像及 LLaVA 对话。两侧必须使用同一份图文 JSON、图像和样本顺序；
	38	+文件名中的 `16k` 不代表可以自行替换为任意同样大小的子集。
	39	+Megatron 检查点转换使用本目录的[转换工具](convert_megatron_checkpoint.py)。
	40	+官方权重需要先应用 Transformers 的 Mistral3 键映射，再补上 `TransformersModel.model` 包装层；
	41	+通用 HF→DCP 转换结果不能直接用于此 Megatron 基线。转换仅调整名称，保留张量值和数据类型。
	42	+启动器在正式执行前检查全部模型键，避免基线的部分加载逻辑留下随机初始化权重。
	43	+
	44	+```bash
	45	+NON_MEGATRON=true python examples/magistral-2509/baseline/convert_megatron_checkpoint.py \
	46	+ --model-dir /path/to/Magistral-Small-2509 --dcp-dir /path/to/Magistral-Small-2509-dcp
	47	+```
	48	+
	49	+## 顺序执行两侧训练
	50	+
	51	+在仓库根目录加载 CANN 环境，并激活相应后端的 Python 环境。
	52	+`--source-dir` 指定实际训练源码；重跑历史实验应使用上述 `cdefc58b`，验证新修订则记录新源码 SHA。
	53	+数据路径指向已构造的同一份 JSON。每次对照使用新的输出目录，以免复用旧预处理缓存。
	54	+
	55	+```bash
	56	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	57	+python examples/magistral-2509/baseline/run_comparison.py --backend megatron \
	58	+ --source-dir /path/to/MindSpeed-MM \
	59	+ --model-dir /path/to/Magistral-Small-2509 --dcp-dir /path/to/Magistral-Small-2509-dcp \
	60	+ --dataset /path/to/data/mllm_format_llava_instruct_data_16k.json --output-dir /path/to/output/megatron
	61	+python examples/magistral-2509/baseline/run_comparison.py --backend pure \
	62	+ --source-dir /path/to/MindSpeed-MM \
	63	+ --model-dir /path/to/Magistral-Small-2509 \
	64	+ --dataset /path/to/data/mllm_format_llava_instruct_data_16k.json --output-dir /path/to/output/pure
	65	+```
	66	+
	67	+两侧均使用八张卡、全局批大小 8、微批大小 1、2048-token 截断、2000 步和相同的余弦学习率调度。
	68	+启动器生成路径已解析的配置，打印完整命令，并传递训练失败的退出状态；增加 `--dry-run` 可仅核对配置和命令。
	69	+`run_identity.json` 记录实际源码 SHA、数据集与资产文件校验值、软件版本和相关环境变量。
	70	+日志输出到标准输出，可重定向到各自输出目录之外。该配置不保存检查点。
	71	+
	72	+精度按 `sum(abs(loss_pure-loss_base))/sum(abs(loss_base))` 计算，任务阈值为 2%。
	73	+性能统计第 101–2000 步平均步时，纯 FSDP2 不高于 Megatron-FSDP2 才满足该次配置的性能要求。
	74	+历史逐步数据与结果见[实践文章及证据包](https://gitcode.com/Ascend/MindSpeed-MM/issues/703)。

Aexamples/magistral-2509/baseline/asset_manifest.json+68-0

@@ -0,0 +1,68 @@
	1	+{
	2	+ "purpose": "Pinned asset sources for new reproductions; historical September asset revisions were not fully recorded.",
	3	+ "files": [
	4	+ {
	5	+ "name": "config.json",
	6	+ "repo": "mistralai/Magistral-Small-2509",
	7	+ "revision": "a31cc96ab10cf19bc42c628fedf1e359e0853c49",
	8	+ "sha256": "c1684445dc6f3d92415403e6790d909970da46059c6a90909466ad783d069182",
	9	+ "bytes": 1166
	10	+ },
	11	+ {
	12	+ "name": "model.safetensors.index.json",
	13	+ "repo": "mistralai/Magistral-Small-2509",
	14	+ "revision": "a31cc96ab10cf19bc42c628fedf1e359e0853c49",
	15	+ "sha256": "ddee03ee0188767bbd5b39bc64744b3441f1a8f8c503727980fe92e590917a51",
	16	+ "bytes": 57944
	17	+ },
	18	+ {
	19	+ "name": "tekken.json",
	20	+ "repo": "mistralai/Magistral-Small-2509",
	21	+ "revision": "a31cc96ab10cf19bc42c628fedf1e359e0853c49",
	22	+ "sha256": "05e051a32ae9d6333459da0083d36507116334fc85378321cfb32dfff7187531",
	23	+ "bytes": 19399765
	24	+ },
	25	+ {
	26	+ "name": "chat_template.jinja",
	27	+ "repo": "unsloth/Magistral-Small-2509",
	28	+ "revision": "dcfa8f675a2adaea5dad53b32c0c948b47c02f55",
	29	+ "sha256": "817cbd470e260e28e69015cc34b1abfb2a7ab76500839c743ed0fee3d1a32c9e",
	30	+ "bytes": 4690
	31	+ },
	32	+ {
	33	+ "name": "preprocessor_config.json",
	34	+ "repo": "unsloth/Magistral-Small-2509",
	35	+ "revision": "dcfa8f675a2adaea5dad53b32c0c948b47c02f55",
	36	+ "sha256": "e2281c57d4a9d3f5e43a1988a4b183792fa7172332ee33c7eba9583194b84cb4",
	37	+ "bytes": 662
	38	+ },
	39	+ {
	40	+ "name": "processor_config.json",
	41	+ "repo": "unsloth/Magistral-Small-2509",
	42	+ "revision": "dcfa8f675a2adaea5dad53b32c0c948b47c02f55",
	43	+ "sha256": "4ceef0d0cbf062ffd522d85fcbf5248c2db8659c589afcc47b2f26810e7d9a58",
	44	+ "bytes": 189
	45	+ },
	46	+ {
	47	+ "name": "special_tokens_map.json",
	48	+ "repo": "unsloth/Magistral-Small-2509",
	49	+ "revision": "dcfa8f675a2adaea5dad53b32c0c948b47c02f55",
	50	+ "sha256": "1864ac9a603857f894d6cb47ac270f50685b05359128b8b9e59edf5ad56cdb49",
	51	+ "bytes": 21431
	52	+ },
	53	+ {
	54	+ "name": "tokenizer.json",
	55	+ "repo": "unsloth/Magistral-Small-2509",
	56	+ "revision": "dcfa8f675a2adaea5dad53b32c0c948b47c02f55",
	57	+ "sha256": "158915c7c928301e0d32f1da363e45c294d907bd4c64a8e855f0bd1a47b9e870",
	58	+ "bytes": 17078001
	59	+ },
	60	+ {
	61	+ "name": "tokenizer_config.json",
	62	+ "repo": "unsloth/Magistral-Small-2509",
	63	+ "revision": "dcfa8f675a2adaea5dad53b32c0c948b47c02f55",
	64	+ "sha256": "bb0bbbbbe9fae009c439441a1c1a62c2759e930d97cf1a431a33d8ee81716881",
	65	+ "bytes": 202929
	66	+ }
	67	+ ]
	68	+}

Aexamples/magistral-2509/baseline/convert_megatron_checkpoint.py+62-0

@@ -0,0 +1,62 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Convert official Magistral HF weights for TransformersModel's Megatron wrapper."""
	3	+
	4	+import argparse
	5	+from pathlib import Path
	6	+
	7	+from transformers import AutoConfig, Mistral3ForConditionalGeneration
	8	+import torch
	9	+from torch.distributed.checkpoint import FileSystemReader
	10	+
	11	+
	12	+class MegatronMistral3WeightTransform:
	13	+ """Apply Transformers' HF key mapping, then add TransformersModel.model."""
	14	+
	15	+ def hf_to_dcp(self, key, tensor):
	16	+ from mindspeed_mm.fsdp.checkpoint.hf_utils import convert_weight_key
	17	+
	18	+ return "model." + convert_weight_key(key, Mistral3ForConditionalGeneration), tensor
	19	+
	20	+
	21	+def validate_megatron_checkpoint(model_dir, dcp_dir):
	22	+ """Check every model key before the baseline's partial-load planner can run."""
	23	+ config = AutoConfig.from_pretrained(model_dir, local_files_only=True)
	24	+ with torch.device("meta"):
	25	+ model = Mistral3ForConditionalGeneration(config)
	26	+ expected = {"model.model." + key: tuple(value.shape) for key, value in model.state_dict().items()}
	27	+ metadata = FileSystemReader(Path(dcp_dir) / "release").read_metadata()
	28	+ actual = {key for key in metadata.state_dict_metadata if key.startswith("model.")}
	29	+ missing, unexpected = sorted(expected.keys() - actual), sorted(actual - expected.keys())
	30	+ wrong_shapes = sorted(key for key in actual & expected.keys()
	31	+ if tuple(metadata.state_dict_metadata[key].size) != expected[key])
	32	+ if missing or unexpected or wrong_shapes:
	33	+ raise ValueError(
	34	+ "Megatron initial checkpoint model keys do not match: "
	35	+ f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}, "
	36	+ f"wrong_shapes={len(wrong_shapes)} {wrong_shapes[:5]}. "
	37	+ "Convert weights with convert_megatron_checkpoint.py."
	38	+ )
	39	+ return len(expected)
	40	+
	41	+
	42	+def convert_checkpoint(model_dir, dcp_dir):
	43	+ """Stream source shards through the production DCP writer without casting weights."""
	44	+ from mindspeed_mm.fsdp.tasks.checkpoint.utils import hf_to_dcp_sharded
	45	+
	46	+ hf_to_dcp_sharded(str(model_dir), str(dcp_dir), weight_transform=MegatronMistral3WeightTransform())
	47	+ return validate_megatron_checkpoint(model_dir, dcp_dir)
	48	+
	49	+
	50	+def main():
	51	+ parser = argparse.ArgumentParser(description=__doc__)
	52	+ parser.add_argument("--model-dir", type=Path, required=True)
	53	+ parser.add_argument("--dcp-dir", type=Path, required=True)
	54	+ args = parser.parse_args()
	55	+ if args.dcp_dir.exists():
	56	+ parser.error("--dcp-dir must be a new directory; existing checkpoints are preserved")
	57	+ count = convert_checkpoint(args.model_dir, args.dcp_dir)
	58	+ print(f"Validated {count} Megatron model keys.")
	59	+
	60	+
	61	+if __name__ == "__main__":
	62	+ main()

Aexamples/magistral-2509/baseline/data.json+58-0

@@ -0,0 +1,58 @@
	1	+{
	2	+ "dataset_param": {
	3	+ "dataset_type": "huggingface",
	4	+ "preprocess_parameters": {
	5	+ "model_name_or_path": "/path/to/Magistral-Small-2509",
	6	+ "use_fast_tokenizer": true,
	7	+ "split_special_tokens": false,
	8	+ "image_max_pixels": 262144,
	9	+ "image_min_pixels": 1024,
	10	+ "video_max_pixels": 16384,
	11	+ "video_min_pixels": 0,
	12	+ "video_fps": 2.0,
	13	+ "video_maxlen": 64,
	14	+ "fix_mistral_regex": true,
	15	+ "trust_remote_code": false
	16	+ },
	17	+ "basic_parameters": {
	18	+ "template": "mistral",
	19	+ "dataset_dir": "/path/to/data/",
	20	+ "dataset": "/path/to/data/mllm_format_llava_instruct_data_16k.json",
	21	+ "cache_dir": "/path/to/output/cache",
	22	+ "enable_thinking": false,
	23	+ "overwrite_cache": false,
	24	+ "train_on_prompt": false,
	25	+ "mask_history": false,
	26	+ "preprocessing_batch_size": 1000,
	27	+ "preprocessing_num_workers": 8,
	28	+ "max_samples": null,
	29	+ "tool_format": null
	30	+ },
	31	+ "attr": {
	32	+ "system": null,
	33	+ "images": "images",
	34	+ "videos": null,
	35	+ "audios": null,
	36	+ "messages": "messages",
	37	+ "role_tag": "role",
	38	+ "content_tag": "content",
	39	+ "user_tag": "user",
	40	+ "assistant_tag": "assistant",
	41	+ "observation_tag": null,
	42	+ "function_tag": null,
	43	+ "system_tag": null,
	44	+ "packing": false
	45	+ }
	46	+ },
	47	+ "dataloader_param": {
	48	+ "dataloader_mode": "sampler",
	49	+ "drop_last": true,
	50	+ "sampler_type": "BaseRandomBatchSampler",
	51	+ "collate_param": {
	52	+ "model_name": "qwen3vl",
	53	+ "ignore_pad_token_for_loss": true
	54	+ },
	55	+ "pin_memory": true,
	56	+ "shuffle": true
	57	+ }
	58	+}

Aexamples/magistral-2509/baseline/fsdp2_config.yaml+18-0

@@ -0,0 +1,18 @@
	1	+sharding_size: auto
	2	+reshard_after_forward: True
	3	+sub_modules_to_wrap:
	4	+ - language_model.embed_tokens
	5	+ - language_model.layers.{*}
	6	+ - vision_tower.transformer.layers.{*}
	7	+ - multi_modal_projector.linear_1
	8	+ - multi_modal_projector.linear_2
	9	+ - lm_head
	10	+recompute_modules:
	11	+ - language_model.layers.{*}
	12	+ - vision_tower.transformer.layers.{*}
	13	+param_dtype: "bf16"
	14	+reduce_dtype: "bf16"
	15	+cast_forward_inputs: True
	16	+num_to_forward_prefetch: 0
	17	+num_to_backward_prefetch: 0
	18	+offload_to_cpu: False

Aexamples/magistral-2509/baseline/model.json+27-0

@@ -0,0 +1,27 @@
	1	+{
	2	+ "model_id": "mistral3",
	3	+ "init_from_hf_path": "/path/to/Magistral-Small-2509",
	4	+ "image_token_index": 10,
	5	+ "text_decoder": {
	6	+ "model_id": "mistral",
	7	+ "num_layers": 40,
	8	+ "hidden_size": 5120,
	9	+ "num_attention_heads": 32,
	10	+ "max_position_embeddings": 131072,
	11	+ "freeze": false
	12	+ },
	13	+ "vision_encoder": {
	14	+ "model_id": "pixtral",
	15	+ "num_layers": 24,
	16	+ "hidden_size": 1024,
	17	+ "num_attention_heads": 16,
	18	+ "freeze": true
	19	+ },
	20	+ "vision_projector": {
	21	+ "freeze": true
	22	+ },
	23	+ "patch": {
	24	+ "clip_grad_async": true,
	25	+ "scale_grad": true
	26	+ }
	27	+}

Aexamples/magistral-2509/baseline/pure_fsdp2.yaml+159-0

@@ -0,0 +1,159 @@
	1	+# 并行策略
	2	+parallel:
	3	+ tensor_parallel_size: 1
	4	+ fully_shard_parallel_size: auto
	5	+ fsdp_plan:
	6	+ apply_modules:
	7	+ - model.vision_tower.transformer.layers.{*}
	8	+ - model.vision_tower
	9	+ - model.multi_modal_projector
	10	+ - model.language_model.embed_tokens
	11	+ - model.language_model.layers.{*}
	12	+ - model.language_model
	13	+ - lm_head
	14	+ param_dtype: bf16
	15	+ reduce_dtype: bf16
	16	+ ring_attention_size: 1
	17	+ ulysses_parallel_size: 1
	18	+
	19	+### 数据相关配置
	20	+data:
	21	+ dataset_param:
	22	+ dataset_type: huggingface
	23	+ #数据集属性
	24	+ attr:
	25	+ # system: null
	26	+ images: images
	27	+ # videos: null
	28	+ messages: messages
	29	+ role_tag: role
	30	+ content_tag: content
	31	+ user_tag: user
	32	+ assistant_tag: assistant
	33	+ # observation_tag: null
	34	+ # function_tag: null
	35	+ # system_tag: null
	36	+
	37	+ # 数据预处理
	38	+ preprocess_parameters:
	39	+ model_name_or_path: &HF_MODEL_LOAD_PATH /path/to/Magistral-Small-2509
	40	+ use_fast_tokenizer: true
	41	+ split_special_tokens: false
	42	+ image_max_pixels: 262144
	43	+ image_min_pixels: 1024
	44	+ video_max_pixels: 16384
	45	+ video_min_pixels: 0
	46	+ video_fps: 2.0
	47	+ video_maxlen: 64
	48	+ fix_mistral_regex: true
	49	+ trust_remote_code: false
	50	+
	51	+ basic_parameters:
	52	+ cutoff_len: 2048
	53	+ template: mistral
	54	+ enable_thinking: false
	55	+ train_on_prompt: false
	56	+ mask_history: false
	57	+ # tool_format: null
	58	+ dataset_dir: /path/to/data/
	59	+ dataset: &DATASET_PATH /path/to/data/mllm_format_llava_instruct_data_16k.json
	60	+ cache_dir: /path/to/output/cache
	61	+ overwrite_cache: false
	62	+ preprocessing_batch_size: 1000
	63	+ preprocessing_num_workers: 8
	64	+ max_samples: null
	65	+
	66	+ # 数据加载
	67	+ dataloader_param:
	68	+ pin_memory: true
	69	+ shuffle: true
	70	+ dataloader_mode: sampler
	71	+ drop_last: true
	72	+ sampler_type: BaseRandomBatchSampler
	73	+ num_workers: 8
	74	+ collate_param:
	75	+ collator_id: pixtral
	76	+ ignore_pad_token_for_loss: true
	77	+ enable_preload: false
	78	+
	79	+# 模型配置
	80	+model:
	81	+ model_id: mistral3
	82	+ model_name_or_path: *HF_MODEL_LOAD_PATH
	83	+ trust_remote_code: false
	84	+ attn_implementation: flash_attention_2
	85	+ freeze:
	86	+ - model.vision_tower
	87	+ - model.multi_modal_projector
	88	+
	89	+# 优化特性配置
	90	+features:
	91	+ loss_cfg:
	92	+ loss_type: default
	93	+ recompute: true
	94	+ recompute_plan:
	95	+ apply_modules:
	96	+ - model.vision_tower.transformer.layers.{*}
	97	+ - model.language_model.layers.{*}
	98	+ enable_chunk_loss: false
	99	+ chunkloss_plan:
	100	+ apply_module: lm_head
	101	+ chunk_size: 1024
	102	+ activation_offload_plan:
	103	+ apply_modules:
	104	+ - model.language_model.layers.{*}
	105	+
	106	+# 训练配置
	107	+training:
	108	+ micro_batch_size: 1
	109	+ gradient_accumulation_steps: 1
	110	+ seed: 42
	111	+ lr: 1.0e-5
	112	+ lr_decay_style: cosine
	113	+ lr_warmup_ratio: 0.1
	114	+ weight_decay: 0
	115	+ train_iters: 2000
	116	+ clip_grad: 0.0
	117	+ init_model_with_meta_device: true
	118	+ optimizer: adamw
	119	+ adam_fused: true
	120	+ log_interval: 1
	121	+ save_interval: 10000
	122	+ load: *HF_MODEL_LOAD_PATH
	123	+ load_format: hf
	124	+ # save: save_path
	125	+ no_load_optim: true
	126	+ no_load_rng: true
	127	+ no_save_optim: true
	128	+ no_save_rng: true
	129	+ use_deter_comp: false
	130	+ plugin:
	131	+ - mindspeed_mm/fsdp/models/mistral3
	132	+ - mindspeed_mm/fsdp/data/datasets/huggingface
	133	+
	134	+# 工具配置
	135	+tools:
	136	+ profile:
	137	+ enable: false
	138	+ profile_type: static
	139	+ ranks: [0]
	140	+ static_param:
	141	+ level: level1
	142	+ with_stack: false
	143	+ with_memory: false
	144	+ record_shapes: false
	145	+ with_cpu: true
	146	+ save_path: ./profiling
	147	+ start_step: 10
	148	+ end_step: 11
	149	+ data_simplification: false
	150	+ aic_metrics_type: PipeUtilization
	151	+ memory_profile:
	152	+ enable: false
	153	+ start_step: 1
	154	+ end_step: 2
	155	+ save_path: ./memory_snapshot
	156	+ dump_ranks: [0]
	157	+ stacks: all
	158	+ max_entries: null
	159	+ mem_info: false

Aexamples/magistral-2509/baseline/run_comparison.py+130-0

@@ -0,0 +1,130 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Launch one backend using the archived 2000-step comparison configuration."""
	3	+
	4	+import argparse
	5	+import hashlib
	6	+from importlib.metadata import PackageNotFoundError, version
	7	+import json
	8	+import os
	9	+from pathlib import Path
	10	+import shutil
	11	+import subprocess
	12	+import sys
	13	+
	14	+import yaml
	15	+
	16	+
	17	+ENVIRONMENT = {
	18	+ "CUDA_DEVICE_MAX_CONNECTIONS": "2", "ASCEND_SLOG_PRINT_TO_STDOUT": "0", "ASCEND_GLOBAL_LOG_LEVEL": "3",
	19	+ "TASK_QUEUE_ENABLE": "1", "COMBINED_ENABLE": "1", "CPU_AFFINITY_CONF": "1", "HCCL_CONNECT_TIMEOUT": "1200",
	20	+ "NPU_ASD_ENABLE": "0", "ASCEND_LAUNCH_BLOCKING": "0", "ACLNN_CACHE_LIMIT": "100000",
	21	+ "TOKENIZERS_PARALLELISM": "false", "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:True",
	22	+ "MULTI_STREAM_MEMORY_REUSE": "2", "HF_HUB_OFFLINE": "1", "HF_DATASETS_OFFLINE": "1",
	23	+}
	24	+
	25	+
	26	+def file_sha256(path):
	27	+ digest = hashlib.sha256()
	28	+ with path.open("rb") as stream:
	29	+ for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
	30	+ digest.update(block)
	31	+ return digest.hexdigest()
	32	+
	33	+
	34	+def main():
	35	+ parser = argparse.ArgumentParser(description=__doc__)
	36	+ parser.add_argument("--backend", choices=("megatron", "pure"), required=True)
	37	+ parser.add_argument("--model-dir", type=Path, required=True)
	38	+ parser.add_argument("--dataset", type=Path, required=True)
	39	+ parser.add_argument("--output-dir", type=Path, required=True)
	40	+ parser.add_argument("--dcp-dir", type=Path, help="Converted DCP weights, required for the Megatron backend.")
	41	+ parser.add_argument("--source-dir", type=Path, default=Path(__file__).resolve().parents[3])
	42	+ parser.add_argument("--master-port", type=int, default=6000)
	43	+ parser.add_argument("--dry-run", action="store_true", help="Write resolved configs and print the command without training.")
	44	+ args = parser.parse_args()
	45	+ if args.backend == "megatron" and args.dcp_dir is None:
	46	+ parser.error("--dcp-dir is required for the Megatron backend")
	47	+ source = args.source_dir.resolve()
	48	+ model_dir, dataset = args.model_dir.resolve(), args.dataset.resolve()
	49	+ output = args.output_dir.resolve()
	50	+ output.mkdir(parents=True, exist_ok=True)
	51	+ config_dir = Path(__file__).resolve().parent
	52	+ command = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node", "8",
	53	+ "--nnodes", "1", "--node_rank", "0", "--master_addr", "localhost",
	54	+ "--master_port", str(args.master_port)]
	55	+ env = os.environ.copy()
	56	+ for name, value in ENVIRONMENT.items():
	57	+ env.setdefault(name, value)
	58	+ if args.backend == "pure":
	59	+ config = yaml.safe_load((config_dir / "pure_fsdp2.yaml").read_text())
	60	+ config["model"]["model_name_or_path"] = str(model_dir)
	61	+ config["training"]["load"] = str(model_dir)
	62	+ preprocess = config["data"]["dataset_param"]["preprocess_parameters"]
	63	+ preprocess["model_name_or_path"] = str(model_dir)
	64	+ basic = config["data"]["dataset_param"]["basic_parameters"]
	65	+ basic.update(dataset_dir=str(dataset.parent), dataset=str(dataset), cache_dir=str(output / "cache"))
	66	+ config_path = output / "pure_fsdp2.yaml"
	67	+ config_path.write_text(yaml.safe_dump(config, sort_keys=False))
	68	+ command += ["mindspeed_mm/fsdp/train/trainer.py", str(config_path)]
	69	+ env["NON_MEGATRON"] = "true"
	70	+ else:
	71	+ model = json.loads((config_dir / "model.json").read_text())
	72	+ model["init_from_hf_path"] = str(model_dir)
	73	+ data = json.loads((config_dir / "data.json").read_text())
	74	+ data["dataset_param"]["preprocess_parameters"]["model_name_or_path"] = str(model_dir)
	75	+ data["dataset_param"]["basic_parameters"].update(
	76	+ dataset_dir=str(dataset.parent), dataset=str(dataset), cache_dir=str(output / "cache"))
	77	+ for name, value in (("model.json", model), ("data.json", data)):
	78	+ (output / name).write_text(json.dumps(value, indent=2) + "\n")
	79	+ command += [
	80	+ "pretrain_transformers.py", "--use-mcore-models", "--tensor-model-parallel-size", "1",
	81	+ "--pipeline-model-parallel-size", "1", "--context-parallel-size", "1",
	82	+ "--context-parallel-algo", "ulysses_cp_algo", "--micro-batch-size", "1", "--global-batch-size", "8",
	83	+ "--tokenizer-type", "NullTokenizer", "--vocab-size", "131072", "--seq-length", "2048",
	84	+ "--make-vocab-size-divisible-by", "1", "--normalization", "RMSNorm", "--use-fused-rmsnorm",
	85	+ "--swiglu", "--use-fused-swiglu", "--no-masked-softmax-fusion", "--lr", "1.0e-5",
	86	+ "--lr-decay-style", "cosine", "--weight-decay", "0", "--train-iters", "2000",
	87	+ "--lr-warmup-fraction", "0.1", "--clip-grad", "0.0", "--adam-beta1", "0.9", "--adam-beta2", "0.999",
	88	+ "--no-gradient-accumulation-fusion", "--seed", "42", "--load", str(args.dcp_dir.resolve()),
	89	+ "--use-flash-attn", "--no-load-optim", "--no-load-rng", "--no-save-optim", "--no-save-rng",
	90	+ "--num-workers", "8", "--use-torch-fsdp2", "--untie-embeddings-and-output-weights",
	91	+ "--ckpt-format", "torch_dcp", "--fsdp2-config-path", str(config_dir / "fsdp2_config.yaml"),
	92	+ "--optimizer-selection", "fused_torch_adamw", "--use-cpu-initialization", "--log-tps",
	93	+ "--init-model-with-meta-device", "--mm-data", str(output / "data.json"),
	94	+ "--mm-model", str(output / "model.json"), "--mm-tool", "./mindspeed_mm/tools/tools.json",
	95	+ "--log-interval", "1", "--distributed-backend", "nccl",
	96	+ ]
	97	+ env.pop("NON_MEGATRON", None)
	98	+ identity = {"cwd": str(source), "command": command, "backend": args.backend,
	99	+ "environment": {name: env[name] for name in ENVIRONMENT}}
	100	+ if (source / ".git").exists():
	101	+ git = shutil.which("git")
	102	+ if git is None:
	103	+ parser.error("git is required to record the training source commit")
	104	+ identity["source_head"] = subprocess.check_output([git, "rev-parse", "HEAD"], cwd=source, text=True).strip()
	105	+ identity["dataset_sha256"] = file_sha256(dataset) if dataset.is_file() else None
	106	+ identity["asset_sha256"] = {entry["name"]: file_sha256(model_dir / entry["name"])
	107	+ for entry in json.loads((config_dir / "asset_manifest.json").read_text())["files"]
	108	+ if (model_dir / entry["name"]).is_file()}
	109	+ identity["versions"] = {}
	110	+ for package in ("torch", "torch-npu", "transformers", "mindspeed", "megatron-core"):
	111	+ try:
	112	+ identity["versions"][package] = version(package)
	113	+ except PackageNotFoundError:
	114	+ identity["versions"][package] = None
	115	+ (output / "run_identity.json").write_text(json.dumps(identity, indent=2) + "\n")
	116	+ print(json.dumps(identity, indent=2), flush=True)
	117	+ if not args.dry_run:
	118	+ if not model_dir.is_dir() or not dataset.is_file():
	119	+ parser.error("--model-dir and --dataset must exist")
	120	+ if args.backend == "megatron" and not args.dcp_dir.is_dir():
	121	+ parser.error("--dcp-dir must exist")
	122	+ if args.backend == "megatron":
	123	+ from convert_megatron_checkpoint import validate_megatron_checkpoint
	124	+
	125	+ validate_megatron_checkpoint(model_dir, args.dcp_dir)
	126	+ subprocess.run(command, cwd=source, env=env, check=True)
	127	+
	128	+
	129	+if __name__ == "__main__":
	130	+ main()

Aexamples/magistral-2509/finetune_magistral_2509_v1.sh+33-0

@@ -0,0 +1,33 @@
	1	+# 让 torchrun 的失败码穿透 tee，不被管道尾部的成功码掩盖
	2	+set -o pipefail
	3	+
	4	+# 根据实际情况修改 ascend-toolkit 路径
	5	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	6	+export NON_MEGATRON=true
	7	+export HCCL_CONNECT_TIMEOUT=1200
	8	+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
	9	+export MULTI_STREAM_MEMORY_REUSE=2
	10	+export TASK_QUEUE_ENABLE=1
	11	+export CPU_AFFINITY_CONF=1
	12	+
	13	+NPUS_PER_NODE=8
	14	+MASTER_ADDR=localhost
	15	+MASTER_PORT=6000
	16	+NNODES=1
	17	+NODE_RANK=0
	18	+WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
	19	+
	20	+DISTRIBUTED_ARGS="
	21	+ --nproc_per_node $NPUS_PER_NODE \
	22	+ --nnodes $NNODES \
	23	+ --node_rank $NODE_RANK \
	24	+ --master_addr $MASTER_ADDR \
	25	+ --master_port $MASTER_PORT
	26	+"
	27	+
	28	+logfile=$(date +%Y%m%d)_$(date +%H%M%S)
	29	+config_path=${1:-examples/magistral-2509/magistral_2509_config_v1.yaml}
	30	+mkdir -p logs
	31	+torchrun $DISTRIBUTED_ARGS mindspeed_mm/fsdp/train/trainer.py \
	32	+ "$config_path" \
	33	+ 2>&1 \| tee logs/train_${logfile}.log

Aexamples/magistral-2509/finetune_magistral_2509_v1_4card.sh+34-0

@@ -0,0 +1,34 @@
	1	+# 4 卡（4×64 GB）全参微调：依赖 training.optimizer: adamw_swap，见同目录 magistral_2509_config_v1_4card.yaml。
	2	+# 8 卡请用 finetune_magistral_2509_v1.sh。
	3	+# 让 torchrun 的失败码穿透 tee，不被管道尾部的成功码掩盖
	4	+set -o pipefail
	5	+
	6	+# 根据实际情况修改 ascend-toolkit 路径
	7	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	8	+export NON_MEGATRON=true
	9	+export HCCL_CONNECT_TIMEOUT=1200
	10	+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
	11	+export MULTI_STREAM_MEMORY_REUSE=2
	12	+export TASK_QUEUE_ENABLE=1
	13	+export CPU_AFFINITY_CONF=1
	14	+
	15	+NPUS_PER_NODE=4
	16	+MASTER_ADDR=localhost
	17	+MASTER_PORT=6001
	18	+NNODES=1
	19	+NODE_RANK=0
	20	+WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
	21	+
	22	+DISTRIBUTED_ARGS="
	23	+ --nproc_per_node $NPUS_PER_NODE \
	24	+ --nnodes $NNODES \
	25	+ --node_rank $NODE_RANK \
	26	+ --master_addr $MASTER_ADDR \
	27	+ --master_port $MASTER_PORT
	28	+"
	29	+
	30	+logfile=$(date +%Y%m%d)_$(date +%H%M%S)
	31	+mkdir -p logs
	32	+torchrun $DISTRIBUTED_ARGS mindspeed_mm/fsdp/train/trainer.py \
	33	+ examples/magistral-2509/magistral_2509_config_v1_4card.yaml \
	34	+ 2>&1 \| tee logs/train_${logfile}.log

Aexamples/magistral-2509/magistral_2509_config_optimized.yaml+169-0

@@ -0,0 +1,169 @@
	1	+# NPU performance configuration for PyTorch 2.10, Triton and Ascend A3.
	2	+# 并行策略
	3	+parallel:
	4	+ tensor_parallel_size: 1
	5	+ fully_shard_parallel_size: auto
	6	+ fsdp_plan:
	7	+ apply_modules:
	8	+ - model.vision_tower.transformer.layers.{*}
	9	+ - model.vision_tower
	10	+ - model.multi_modal_projector
	11	+ - model.language_model.embed_tokens
	12	+ - model.language_model.layers.{*}
	13	+ - model.language_model
	14	+ - lm_head
	15	+ param_dtype: bf16
	16	+ reduce_dtype: bf16
	17	+ num_to_forward_prefetch: 1
	18	+ num_to_backward_prefetch: 1
	19	+ use_triton_split_copy: true
	20	+ ring_attention_size: 1
	21	+ ulysses_parallel_size: 1
	22	+
	23	+### 数据相关配置
	24	+data:
	25	+ dataset_param:
	26	+ dataset_type: huggingface
	27	+ #数据集属性
	28	+ attr:
	29	+ # system: null
	30	+ images: images
	31	+ # videos: null
	32	+ messages: messages
	33	+ role_tag: role
	34	+ content_tag: content
	35	+ user_tag: user
	36	+ assistant_tag: assistant
	37	+ # observation_tag: null
	38	+ # function_tag: null
	39	+ # system_tag: null
	40	+
	41	+ # 数据预处理
	42	+ preprocess_parameters:
	43	+ model_name_or_path: &HF_MODEL_LOAD_PATH /home/data/Magistral-Small-2509
	44	+ use_fast_tokenizer: true
	45	+ split_special_tokens: false
	46	+ image_max_pixels: 262144
	47	+ image_min_pixels: 1024
	48	+ video_max_pixels: 16384
	49	+ video_min_pixels: 0
	50	+ video_fps: 2.0
	51	+ video_maxlen: 64
	52	+ fix_mistral_regex: true
	53	+ trust_remote_code: false
	54	+
	55	+ basic_parameters:
	56	+ cutoff_len: 2048
	57	+ template: mistral
	58	+ enable_thinking: false
	59	+ train_on_prompt: false
	60	+ mask_history: false
	61	+ # tool_format: null
	62	+ dataset_dir: /home/usr/data/
	63	+ dataset: &DATASET_PATH /home/usr/data/mllm_format_llava_instruct_data.json
	64	+ cache_dir: ./cache_dir/
	65	+ overwrite_cache: false
	66	+ preprocessing_batch_size: 1000
	67	+ preprocessing_num_workers: 8
	68	+ max_samples: null
	69	+
	70	+ # 数据加载
	71	+ dataloader_param:
	72	+ pin_memory: true
	73	+ shuffle: true
	74	+ dataloader_mode: sampler
	75	+ drop_last: true
	76	+ sampler_type: BaseRandomBatchSampler
	77	+ infinite_sampler: true
	78	+ num_workers: 8
	79	+ collate_param:
	80	+ collator_id: pixtral
	81	+ ignore_pad_token_for_loss: true
	82	+ enable_preload: false
	83	+
	84	+# 模型配置
	85	+model:
	86	+ model_id: mistral3
	87	+ model_name_or_path: *HF_MODEL_LOAD_PATH
	88	+ trust_remote_code: false
	89	+ attn_implementation: flash_attention_2
	90	+ freeze:
	91	+ - model.vision_tower
	92	+ - model.multi_modal_projector
	93	+
	94	+# 优化特性配置
	95	+features:
	96	+ loss_cfg:
	97	+ loss_type: default
	98	+ recompute: true
	99	+ recompute_plan:
	100	+ apply_modules:
	101	+ - model.vision_tower.transformer.layers.{*}
	102	+ - model.language_model.layers.{*}
	103	+ enable_chunk_loss: true
	104	+ contiguous_linear_grad_modules:
	105	+ - model.language_model.layers.{*}.mlp.gate_proj
	106	+ - model.language_model.layers.{*}.mlp.up_proj
	107	+ - model.language_model.layers.{*}.mlp.down_proj
	108	+ chunkloss_plan:
	109	+ apply_module: lm_head
	110	+ chunk_size: 1024
	111	+ activation_offload_plan:
	112	+ apply_modules:
	113	+ - model.language_model.layers.{*}
	114	+
	115	+# 训练配置
	116	+training:
	117	+ micro_batch_size: 1
	118	+ gradient_accumulation_steps: 1
	119	+ seed: 42
	120	+ lr: 1.0e-5
	121	+ lr_decay_style: cosine
	122	+ lr_warmup_ratio: 0.1
	123	+ weight_decay: 0
	124	+ train_iters: 5000
	125	+ clip_grad: 0.0
	126	+ init_model_with_meta_device: true
	127	+ optimizer: adamw
	128	+ adam_fused: true
	129	+ log_interval: 1
	130	+ save_interval: 10000
	131	+ load: *HF_MODEL_LOAD_PATH
	132	+ # save: save_path
	133	+ load_strict: true
	134	+ freeze_gc_after_init: true
	135	+ no_load_optim: true
	136	+ no_load_rng: true
	137	+ no_save_optim: true
	138	+ no_save_rng: true
	139	+ use_deter_comp: false
	140	+ plugin:
	141	+ - mindspeed_mm/fsdp/models/mistral3
	142	+ - mindspeed_mm/fsdp/data/datasets/huggingface
	143	+
	144	+# 工具配置
	145	+tools:
	146	+ profile:
	147	+ enable: false
	148	+ profile_type: static
	149	+ ranks: [0]
	150	+ static_param:
	151	+ level: level1
	152	+ with_stack: false
	153	+ with_memory: false
	154	+ record_shapes: false
	155	+ with_cpu: true
	156	+ save_path: ./profiling
	157	+ start_step: 10
	158	+ end_step: 11
	159	+ data_simplification: false
	160	+ aic_metrics_type: PipeUtilization
	161	+ memory_profile:
	162	+ enable: false
	163	+ start_step: 1
	164	+ end_step: 2
	165	+ save_path: ./memory_snapshot
	166	+ dump_ranks: [0]
	167	+ stacks: all
	168	+ max_entries: null
	169	+ mem_info: false

Aexamples/magistral-2509/magistral_2509_config_v1.yaml+158-0

@@ -0,0 +1,158 @@
	1	+# 并行策略
	2	+parallel:
	3	+ tensor_parallel_size: 1
	4	+ fully_shard_parallel_size: auto
	5	+ fsdp_plan:
	6	+ apply_modules:
	7	+ - model.vision_tower.transformer.layers.{*}
	8	+ - model.vision_tower
	9	+ - model.multi_modal_projector
	10	+ - model.language_model.embed_tokens
	11	+ - model.language_model.layers.{*}
	12	+ - model.language_model
	13	+ - lm_head
	14	+ param_dtype: bf16
	15	+ reduce_dtype: bf16
	16	+ ring_attention_size: 1
	17	+ ulysses_parallel_size: 1
	18	+
	19	+### 数据相关配置
	20	+data:
	21	+ dataset_param:
	22	+ dataset_type: huggingface
	23	+ #数据集属性
	24	+ attr:
	25	+ # system: null
	26	+ images: images
	27	+ # videos: null
	28	+ messages: messages
	29	+ role_tag: role
	30	+ content_tag: content
	31	+ user_tag: user
	32	+ assistant_tag: assistant
	33	+ # observation_tag: null
	34	+ # function_tag: null
	35	+ # system_tag: null
	36	+
	37	+ # 数据预处理
	38	+ preprocess_parameters:
	39	+ model_name_or_path: &HF_MODEL_LOAD_PATH /home/data/Magistral-Small-2509
	40	+ use_fast_tokenizer: true
	41	+ split_special_tokens: false
	42	+ image_max_pixels: 262144
	43	+ image_min_pixels: 1024
	44	+ video_max_pixels: 16384
	45	+ video_min_pixels: 0
	46	+ video_fps: 2.0
	47	+ video_maxlen: 64
	48	+ fix_mistral_regex: true
	49	+ trust_remote_code: false
	50	+
	51	+ basic_parameters:
	52	+ cutoff_len: 2048
	53	+ template: mistral
	54	+ enable_thinking: false
	55	+ train_on_prompt: false
	56	+ mask_history: false
	57	+ # tool_format: null
	58	+ dataset_dir: /home/usr/data/
	59	+ dataset: &DATASET_PATH /home/usr/data/mllm_format_llava_instruct_data.json
	60	+ cache_dir: ./cache_dir/
	61	+ overwrite_cache: false
	62	+ preprocessing_batch_size: 1000
	63	+ preprocessing_num_workers: 8
	64	+ max_samples: null
	65	+
	66	+ # 数据加载
	67	+ dataloader_param:
	68	+ pin_memory: true
	69	+ shuffle: true
	70	+ dataloader_mode: sampler
	71	+ drop_last: true
	72	+ sampler_type: BaseRandomBatchSampler
	73	+ num_workers: 8
	74	+ collate_param:
	75	+ collator_id: pixtral
	76	+ ignore_pad_token_for_loss: true
	77	+ enable_preload: false
	78	+
	79	+# 模型配置
	80	+model:
	81	+ model_id: mistral3
	82	+ model_name_or_path: *HF_MODEL_LOAD_PATH
	83	+ trust_remote_code: false
	84	+ attn_implementation: flash_attention_2
	85	+ freeze:
	86	+ - model.vision_tower
	87	+ - model.multi_modal_projector
	88	+
	89	+# 优化特性配置
	90	+features:
	91	+ loss_cfg:
	92	+ loss_type: default
	93	+ recompute: true
	94	+ recompute_plan:
	95	+ apply_modules:
	96	+ - model.vision_tower.transformer.layers.{*}
	97	+ - model.language_model.layers.{*}
	98	+ enable_chunk_loss: false
	99	+ chunkloss_plan:
	100	+ apply_module: lm_head
	101	+ chunk_size: 1024
	102	+ activation_offload_plan:
	103	+ apply_modules:
	104	+ - model.language_model.layers.{*}
	105	+
	106	+# 训练配置
	107	+training:
	108	+ micro_batch_size: 1
	109	+ gradient_accumulation_steps: 1
	110	+ seed: 42
	111	+ lr: 1.0e-5
	112	+ lr_decay_style: cosine
	113	+ lr_warmup_ratio: 0.1
	114	+ weight_decay: 0
	115	+ train_iters: 5000
	116	+ clip_grad: 0.0
	117	+ init_model_with_meta_device: true
	118	+ optimizer: adamw
	119	+ adam_fused: true
	120	+ log_interval: 1
	121	+ save_interval: 10000
	122	+ load: *HF_MODEL_LOAD_PATH
	123	+ # save: save_path
	124	+ no_load_optim: true
	125	+ no_load_rng: true
	126	+ no_save_optim: true
	127	+ no_save_rng: true
	128	+ use_deter_comp: false
	129	+ plugin:
	130	+ - mindspeed_mm/fsdp/models/mistral3
	131	+ - mindspeed_mm/fsdp/data/datasets/huggingface
	132	+
	133	+# 工具配置
	134	+tools:
	135	+ profile:
	136	+ enable: false
	137	+ profile_type: static
	138	+ ranks: [0]
	139	+ static_param:
	140	+ level: level1
	141	+ with_stack: false
	142	+ with_memory: false
	143	+ record_shapes: false
	144	+ with_cpu: true
	145	+ save_path: ./profiling
	146	+ start_step: 10
	147	+ end_step: 11
	148	+ data_simplification: false
	149	+ aic_metrics_type: PipeUtilization
	150	+ memory_profile:
	151	+ enable: false
	152	+ start_step: 1
	153	+ end_step: 2
	154	+ save_path: ./memory_snapshot
	155	+ dump_ranks: [0]
	156	+ stacks: all
	157	+ max_entries: null
	158	+ mem_info: false

Aexamples/magistral-2509/magistral_2509_config_v1_4card.yaml+161-0

@@ -0,0 +1,161 @@
	1	+# 并行策略
	2	+parallel:
	3	+ tensor_parallel_size: 1
	4	+ fully_shard_parallel_size: auto
	5	+ fsdp_plan:
	6	+ apply_modules:
	7	+ - model.vision_tower.transformer.layers.{*}
	8	+ - model.vision_tower
	9	+ - model.multi_modal_projector
	10	+ - model.language_model.embed_tokens
	11	+ - model.language_model.layers.{*}
	12	+ - model.language_model
	13	+ - lm_head
	14	+ param_dtype: bf16
	15	+ reduce_dtype: bf16
	16	+ ring_attention_size: 1
	17	+ ulysses_parallel_size: 1
	18	+
	19	+### 数据相关配置
	20	+data:
	21	+ dataset_param:
	22	+ dataset_type: huggingface
	23	+ #数据集属性
	24	+ attr:
	25	+ # system: null
	26	+ images: images
	27	+ # videos: null
	28	+ messages: messages
	29	+ role_tag: role
	30	+ content_tag: content
	31	+ user_tag: user
	32	+ assistant_tag: assistant
	33	+ # observation_tag: null
	34	+ # function_tag: null
	35	+ # system_tag: null
	36	+
	37	+ # 数据预处理
	38	+ preprocess_parameters:
	39	+ model_name_or_path: &HF_MODEL_LOAD_PATH /home/data/Magistral-Small-2509
	40	+ use_fast_tokenizer: true
	41	+ split_special_tokens: false
	42	+ image_max_pixels: 262144
	43	+ image_min_pixels: 1024
	44	+ video_max_pixels: 16384
	45	+ video_min_pixels: 0
	46	+ video_fps: 2.0
	47	+ video_maxlen: 64
	48	+ fix_mistral_regex: true
	49	+ trust_remote_code: false
	50	+
	51	+ basic_parameters:
	52	+ cutoff_len: 2048
	53	+ template: mistral
	54	+ enable_thinking: false
	55	+ train_on_prompt: false
	56	+ mask_history: false
	57	+ # tool_format: null
	58	+ dataset_dir: /home/usr/data/
	59	+ dataset: &DATASET_PATH /home/usr/data/mllm_format_llava_instruct_data.json
	60	+ cache_dir: ./cache_dir/
	61	+ overwrite_cache: false
	62	+ preprocessing_batch_size: 1000
	63	+ preprocessing_num_workers: 8
	64	+ max_samples: null
	65	+
	66	+ # 数据加载
	67	+ dataloader_param:
	68	+ pin_memory: true
	69	+ shuffle: true
	70	+ dataloader_mode: sampler
	71	+ drop_last: true
	72	+ sampler_type: BaseRandomBatchSampler
	73	+ num_workers: 8
	74	+ collate_param:
	75	+ collator_id: pixtral
	76	+ ignore_pad_token_for_loss: true
	77	+ enable_preload: false
	78	+
	79	+# 模型配置
	80	+model:
	81	+ model_id: mistral3
	82	+ model_name_or_path: *HF_MODEL_LOAD_PATH
	83	+ trust_remote_code: false
	84	+ attn_implementation: flash_attention_2
	85	+ freeze:
	86	+ - model.vision_tower
	87	+ - model.multi_modal_projector
	88	+
	89	+# 优化特性配置
	90	+features:
	91	+ loss_cfg:
	92	+ loss_type: default
	93	+ recompute: true
	94	+ recompute_plan:
	95	+ apply_modules:
	96	+ - model.vision_tower.transformer.layers.{*}
	97	+ - model.language_model.layers.{*}
	98	+ enable_chunk_loss: false
	99	+ chunkloss_plan:
	100	+ apply_module: lm_head
	101	+ chunk_size: 1024
	102	+ activation_offload_plan:
	103	+ apply_modules:
	104	+ - model.language_model.layers.{*}
	105	+
	106	+# 训练配置
	107	+training:
	108	+ micro_batch_size: 1
	109	+ gradient_accumulation_steps: 2 # 4 卡下保持 GBS 8（1 × 2 × 4）
	110	+ seed: 42
	111	+ lr: 1.0e-5
	112	+ lr_decay_style: cosine
	113	+ lr_warmup_ratio: 0.1
	114	+ weight_decay: 0
	115	+ train_iters: 5000
	116	+ clip_grad: 0.0
	117	+ init_model_with_meta_device: true
	118	+ # 4 卡候选：通过换出优化器状态减少设备驻留；需要额外主机内存和传输。
	119	+ # 修复后的 24B 配方尚未重新验证，历史短测不能保证当前长图/长序列可运行。
	120	+ # 使用前核对峰值内存、loss、步时和 checkpoint，不作为默认 8 卡配方的性能替代。
	121	+ optimizer: adamw_swap
	122	+ adam_fused: true
	123	+ log_interval: 1
	124	+ save_interval: 10000
	125	+ load: *HF_MODEL_LOAD_PATH
	126	+ # save: save_path
	127	+ no_load_optim: true
	128	+ no_load_rng: true
	129	+ no_save_optim: true
	130	+ no_save_rng: true
	131	+ use_deter_comp: false
	132	+ plugin:
	133	+ - mindspeed_mm/fsdp/models/mistral3
	134	+ - mindspeed_mm/fsdp/data/datasets/huggingface
	135	+
	136	+# 工具配置
	137	+tools:
	138	+ profile:
	139	+ enable: false
	140	+ profile_type: static
	141	+ ranks: [0]
	142	+ static_param:
	143	+ level: level1
	144	+ with_stack: false
	145	+ with_memory: false
	146	+ record_shapes: false
	147	+ with_cpu: true
	148	+ save_path: ./profiling
	149	+ start_step: 10
	150	+ end_step: 11
	151	+ data_simplification: false
	152	+ aic_metrics_type: PipeUtilization
	153	+ memory_profile:
	154	+ enable: false
	155	+ start_step: 1
	156	+ end_step: 2
	157	+ save_path: ./memory_snapshot
	158	+ dump_ranks: [0]
	159	+ stacks: all
	160	+ max_entries: null
	161	+ mem_info: false

Mmindspeed_mm/fsdp/data/data_utils/func_utils/collator.py+2-0

	@@ -265,6 +265,8 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
265	# remove shift offset for postion ids	265	# remove shift offset for postion ids
266	if packed_postion_ids is not None:	266	if packed_postion_ids is not None:
267	features["position_ids"] = postprocess_position_ids(features["position_ids"], packed_postion_ids)	267	features["position_ids"] = postprocess_position_ids(features["position_ids"], packed_postion_ids)
		268	+ elif packed_postion_ids is not None:
		269	+ features["position_ids"] = packed_postion_ids
268		270
269	if (	271	if (
270	self.model is not None	272	self.model is not None

Mmindspeed_mm/fsdp/data/data_utils/func_utils/convert.py+32-6

	@@ -852,12 +852,19 @@ class SupervisedDatasetProcessor(DatasetProcessor):
852	videos: List["VideoInput"],	852	videos: List["VideoInput"],
853	audios: List["AudioInput"],	853	audios: List["AudioInput"],
854	tools: List[str],	854	tools: List[str],
855	- ) -> Tuple[List[int], List[int]]:	855	+ ) -> Tuple[List[int], List[int], Optional[str]]:
856	messages = self.template.mm_plugin.process_messages(prompt + response, images, videos, audios, self.processor)	856	messages = self.template.mm_plugin.process_messages(prompt + response, images, videos, audios, self.processor)
857	input_ids, labels = self.template.mm_plugin.process_token_ids(	857	input_ids, labels = self.template.mm_plugin.process_token_ids(
858	[], [], images, videos, audios, self.tokenizer, self.processor	858	[], [], images, videos, audios, self.tokenizer, self.processor
859	)	859	)
860	encoded_pairs = self.template.encode_multiturn(self.tokenizer, messages, system, tools)	860	encoded_pairs = self.template.encode_multiturn(self.tokenizer, messages, system, tools)
		861	+ # Snapshot the prefix before input_ids changes; validators consume the rest lazily.
		862	+ untruncated_input_ids = chain(
		863	+ tuple(input_ids),
		864	+ chain.from_iterable(chain.from_iterable(encoded_pairs)),
		865	+ (self.tokenizer.eos_token_id,) if self.template.efficient_eos else (),
		866	+ )
		867	+
861	total_length = len(input_ids) + (1 if self.template.efficient_eos else 0)	868	total_length = len(input_ids) + (1 if self.template.efficient_eos else 0)
862	if self.data_args.mask_history:	869	if self.data_args.mask_history:
863	encoded_pairs = encoded_pairs[::-1] # high priority for last turns	870	encoded_pairs = encoded_pairs[::-1] # high priority for last turns
	@@ -900,7 +907,15 @@ class SupervisedDatasetProcessor(DatasetProcessor):
900	input_ids += [self.tokenizer.eos_token_id]	907	input_ids += [self.tokenizer.eos_token_id]
901	labels += [self.tokenizer.eos_token_id]	908	labels += [self.tokenizer.eos_token_id]
902		909
903	- return input_ids, labels	910	+ validation_error = self.template.mm_plugin.validate_input_ids(
		911	+ input_ids=input_ids,
		912	+ untruncated_input_ids=untruncated_input_ids,
		913	+ images=images,
		914	+ videos=videos,
		915	+ audios=audios,
		916	+ tokenizer=self.tokenizer,
		917	+ )
		918	+ return input_ids, labels, validation_error
904		919
905	def preprocess_dataset(self, examples: Dict[str, List[Any]]) -> Dict[str, List[Any]]:	920	def preprocess_dataset(self, examples: Dict[str, List[Any]]) -> Dict[str, List[Any]]:
906	# build inputs with format `<bos> X Y <eos>` and labels with format `<ignore> ... <ignore> Y <eos>`	921	# build inputs with format `<bos> X Y <eos>` and labels with format `<ignore> ... <ignore> Y <eos>`
	@@ -926,7 +941,7 @@ class SupervisedDatasetProcessor(DatasetProcessor):
926	tool_schema = []	941	tool_schema = []
927	if '_tools' in examples:	942	if '_tools' in examples:
928	tool_schema = examples['_tools'][i]	943	tool_schema = examples['_tools'][i]
929	- input_ids, labels = self._encode_data_example(	944	+ input_ids, labels, validation_error = self._encode_data_example(
930	prompt=examples["_prompt"][i],	945	prompt=examples["_prompt"][i],
931	response=examples["_response"][i],	946	response=examples["_response"][i],
932	system=examples["_system"][i],	947	system=examples["_system"][i],
	@@ -935,6 +950,9 @@ class SupervisedDatasetProcessor(DatasetProcessor):
935	audios=examples["_audios"][i] or [],	950	audios=examples["_audios"][i] or [],
936	tools=tool_schema,	951	tools=tool_schema,
937	)	952	)
		953	+ if validation_error is not None:
		954	+ logger.warning_rank0(f"Dropped invalid example: {validation_error}")
		955	+ continue
938	model_inputs["input_ids"].append(input_ids)	956	model_inputs["input_ids"].append(input_ids)
939	model_inputs["attention_mask"].append([1] * len(input_ids))	957	model_inputs["attention_mask"].append([1] * len(input_ids))
940	model_inputs["labels"].append(labels)	958	model_inputs["labels"].append(labels)
	@@ -1224,7 +1242,7 @@ class PackedSupervisedDatasetProcessor(SupervisedDatasetProcessor):
1224	if "_tools" in examples:	1242	if "_tools" in examples:
1225	tool_schema = examples['_tools'][i]	1243	tool_schema = examples['_tools'][i]
1226		1244
1227	- input_ids, labels = self._encode_data_example(	1245	+ input_ids, labels, validation_error = self._encode_data_example(
1228	prompt=examples["_prompt"][i],	1246	prompt=examples["_prompt"][i],
1229	response=examples["_response"][i],	1247	response=examples["_response"][i],
1230	system=examples["_system"][i],	1248	system=examples["_system"][i],
	@@ -1233,6 +1251,9 @@ class PackedSupervisedDatasetProcessor(SupervisedDatasetProcessor):
1233	audios=examples["_audios"][i] or [],	1251	audios=examples["_audios"][i] or [],
1234	tools=tool_schema	1252	tools=tool_schema
1235	)	1253	)
		1254	+ if validation_error is not None:
		1255	+ logger.warning_rank0(f"Dropped invalid example: {validation_error}")
		1256	+ continue
1236	length = len(input_ids)	1257	length = len(input_ids)
1237	if length > self.data_args.cutoff_len:	1258	if length > self.data_args.cutoff_len:
1238	logger.warning_rank0(f"Dropped lengthy example with length {length} > {self.data_args.cutoff_len}.")	1259	logger.warning_rank0(f"Dropped lengthy example with length {length} > {self.data_args.cutoff_len}.")
	@@ -1425,14 +1446,19 @@ def load_tokenizer(model_args: "ProcessorArguments") -> "TokenizerModule":
1425	local_files_only=True,	1446	local_files_only=True,
1426	trust_remote_code=model_args.trust_remote_code,	1447	trust_remote_code=model_args.trust_remote_code,
1427	)	1448	)
1428	- tokenizer = AutoTokenizer.from_pretrained(	1449	+ tokenizer_kwargs = dict(
1429	- model_args.model_name_or_path,
1430	use_fast=model_args.use_fast_tokenizer,	1450	use_fast=model_args.use_fast_tokenizer,
1431	split_special_tokens=model_args.split_special_tokens,	1451	split_special_tokens=model_args.split_special_tokens,
1432	padding_side="right",	1452	padding_side="right",
1433	local_files_only=True,	1453	local_files_only=True,
1434	trust_remote_code=model_args.trust_remote_code,	1454	trust_remote_code=model_args.trust_remote_code,
1435	)	1455	)
		1456	+ if getattr(model_args, "fix_mistral_regex", False):
		1457	+ tokenizer_kwargs["fix_mistral_regex"] = True
		1458	+ tokenizer = AutoTokenizer.from_pretrained(
		1459	+ model_args.model_name_or_path,
		1460	+ **tokenizer_kwargs,
		1461	+ )
1436		1462
1437	try:	1463	try:
1438	processor = AutoProcessor.from_pretrained(	1464	processor = AutoProcessor.from_pretrained(

Mmindspeed_mm/fsdp/data/data_utils/func_utils/mm_plugin.py+82-1

 from copy import deepcopy
 from dataclasses import dataclass
 from io import BytesIO
++from typing import TYPE_CHECKING, Optional, TypedDict, Union, BinaryIO, Literal, List, Dict, Tuple, TypeVar, Any, Iterable
 import av
 import numpy as np
         if self.audio_token is not None and feature_extractor is None:
             raise ValueError("Audio feature extractor was not found, please check and update your model file.")
++    def validate_input_ids(
++            self,
++            input_ids: List[int],
++            untruncated_input_ids: Iterable[int],
++            images: List["ImageInput"],
++            videos: List["VideoInput"],
++            audios: List["AudioInput"],
++            tokenizer: "PreTrainedTokenizer",
++    ) -> Optional[str]:
++        r"""Return an error message when truncation makes multimodal inputs inconsistent."""
++        return None
++
     @staticmethod
     def _validate_messages(
             messages: List[Dict[str, str]],
         return messages
++@dataclass
++class PixtralPlugin(BasePlugin):
++    @override
++    def validate_input_ids(
++            self,
++            input_ids: List[int],
++            untruncated_input_ids: Iterable[int],
++            images: List["ImageInput"],
++            videos: List["VideoInput"],
++            audios: List["AudioInput"],
++            tokenizer: "PreTrainedTokenizer",
++    ) -> Optional[str]:
++        if not images:
++            return None
++
++        image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
++        expected = sum(token_id == image_token_id for token_id in untruncated_input_ids)
++        actual = sum(token_id == image_token_id for token_id in input_ids)
++        if actual != expected:
++            return (
++                "Pixtral image tokens were truncated: "
++                f"expected {expected} image tokens for {len(images)} image(s), got {actual}."
++            )
++        return None
++
++    @override
++    def process_messages(
++            self,
++            messages: list[dict[str, str]],
++            images: list["ImageInput"],
++            videos: list["VideoInput"],
++            audios: list["AudioInput"],
++            processor: Optional["MMProcessor"],
++    ) -> list[dict[str, str]]:
++        self._validate_input(processor, images, videos, audios)
++        self._validate_messages(messages, images, videos, audios)
++        messages = deepcopy(messages)
++        if not images:
++            return messages
++        if self.expand_mm_tokens:
++            mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
++            if "image_sizes" not in mm_inputs:
++                raise ValueError("Pixtral image processor did not return `image_sizes`.")
++            image_sizes = iter(mm_inputs["image_sizes"].tolist())
++            image_break_token: str = getattr(processor, "image_break_token")
++            image_end_token: str = getattr(processor, "image_end_token")
++        for message in messages:
++            content = message["content"]
++            while IMAGE_PLACEHOLDER in content:
++                if self.expand_mm_tokens:
++                    patch_size = processor.patch_size * getattr(processor, "spatial_merge_size", 1)
++                    height, width = next(image_sizes)
++                    num_height_tokens = height // patch_size
++                    num_width_tokens = width // patch_size
++                    replace_tokens = [
++                        [self.image_token] * num_width_tokens + [image_break_token]
++                        for _ in range(num_height_tokens)
++                    ]
++                    replace_tokens = [item for sublist in replace_tokens for item in sublist]  # flatten list
++                    replace_tokens[-1] = image_end_token
++                    replace_str = "".join(replace_tokens)
++                else:
++                    replace_str = self.image_token
++                content = content.replace(IMAGE_PLACEHOLDER, replace_str, 1)
++            message["content"] = content
++        return messages
++
++
 PLUGINS = {
     "base": BasePlugin,
     "qwen2_vl": Qwen2VLPlugin,
     "kimi_k3": KimiK3Plugin,
     "step3_vl": Step3VLPlugin,
     "minimax_m3_vl": MiniMaxM3VLPlugin,
++    "pixtral": PixtralPlugin,
 }

Mmindspeed_mm/fsdp/data/data_utils/func_utils/model_args.py+5-0

	@@ -91,3 +91,8 @@ class ProcessorArguments(BaseArguments):
91	default=False,	91	default=False,
92	metadata={"help": "Trust remote code for Huggingface."},	92	metadata={"help": "Trust remote code for Huggingface."},
93	)	93	)
		94	+ fix_mistral_regex: bool = field(
		95	+ default=False,
		96	+ metadata={
		97	+ "help": "Whether or not to fix the Mistral tokenizer's regex issue (specifically for Mistral/Mixtral models with incorrect spaces)."},
		98	+ )

Mmindspeed_mm/fsdp/data/data_utils/func_utils/template.py+14-0

	@@ -658,6 +658,20 @@ _register_template(
658	)	658	)
659		659
660		660
		661	+_register_template(
		662	+ name="mistral",
		663	+ params=RegisterParams(
		664	+ format_user=StringFormatter(slots=["[INST]{{content}}[/INST]"]),
		665	+ format_system=StringFormatter(slots=["[SYSTEM_PROMPT]{{content}}[/SYSTEM_PROMPT]"]),
		666	+ format_observation=StringFormatter(slots=["""[TOOL_RESULTS]{"content": {{content}}}[/TOOL_RESULTS]"""]),
		667	+ default_system="""First draft your thinking process (inner monologue) until you arrive at a response. Format your response using Markdown, and use LaTeX for any mathematical equations. Write both your thoughts and the response in the same language as the input.
		668	+Your thinking process must follow the template below:[THINK]Your thoughts or/and draft, like working through an exercise on scratch paper. Be as casual and as long as you want until you are confident to generate the response. Use the same language as the input.[/THINK]Here, provide a self-contained response. """,
		669	+ format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
		670	+ ),
		671	+ mm_plugin=get_mm_plugin(name="pixtral", image_token="[IMG]"),
		672	+)
		673	+
		674	+
661	_register_template(	675	_register_template(
662	name="default",	676	name="default",
663	params=RegisterParams(	677	params=RegisterParams(

Mmindspeed_mm/fsdp/data/dataloader/data_collator.py+4-0

	@@ -261,3 +261,7 @@ _register_data_collator(
261	"text_only",	261	"text_only",
262	DataCollatorForTextOnly,	262	DataCollatorForTextOnly,
263	)	263	)
		264	+_register_data_collator(
		265	+ "pixtral",
		266	+ DataCollatorForQwen2vl,
		267	+)

Mmindspeed_mm/fsdp/distributed/fully_shard_parallel.py+8-0

	@@ -78,6 +78,14 @@ def fully_shard_parallel_modules(model: torch.nn.Module, fsdp_mesh: DeviceMesh,
78	"DDP mode is enabled (fully_shard_parallel_size=1) instead of FSDP wrapping")	78	"DDP mode is enabled (fully_shard_parallel_size=1) instead of FSDP wrapping")
79	return model	79	return model
80		80
		81	+ if getattr(fsdp_plan, "use_triton_split_copy", False):
		82	+ if get_device_type() != "npu":
		83	+ raise RuntimeError("use_triton_split_copy requires NPU")
		84	+ from mindspeed_mm.fsdp.ops.fully_shard.npu_split_copy import enable_npu_split_copy
		85	+
		86	+ enable_npu_split_copy()
		87	+ print_rank(logger.info, "Enabled Triton NPU FSDP split copy for this process")
		88	+
81	is_container = hasattr(model, 'get_sub_models') and callable(getattr(model, 'get_sub_models'))	89	is_container = hasattr(model, 'get_sub_models') and callable(getattr(model, 'get_sub_models'))
82		90
83	if not is_container and hasattr(model, 'fully_shard') and callable(getattr(model, 'fully_shard')):	91	if not is_container and hasattr(model, 'fully_shard') and callable(getattr(model, 'fully_shard')):

Mmindspeed_mm/fsdp/features/apply_features.py+14-0

	@@ -176,6 +176,19 @@ class FeaturesApplier:
176	chunkloss_module = get_chunkloss_module(model, plan)	176	chunkloss_module = get_chunkloss_module(model, plan)
177	apply_chunkloss_module(chunkloss_module)	177	apply_chunkloss_module(chunkloss_module)
178		178
		179	+ def apply_contiguous_linear_grad(self, model):
		180	+ plan = getattr(self.config, "contiguous_linear_grad_modules", [])
		181	+ if not plan:
		182	+ return
		183	+ from mindspeed_mm.fsdp.ops.contiguous_linear import apply_contiguous_linear
		184	+
		185	+ modules = self.get_needed_modules(model, plan)
		186	+ if not modules:
		187	+ raise ValueError("contiguous_linear_grad_modules did not match any model modules")
		188	+ for _, module in modules:
		189	+ apply_contiguous_linear(module)
		190	+ logger.info("Applied contiguous Linear weight gradients to %d modules", len(modules))
		191	+
179	def apply_aux_loss_capture(self, model):	192	def apply_aux_loss_capture(self, model):
180	# This function is designed to automatically capture router logits from each MoE layer	193	# This function is designed to automatically capture router logits from each MoE layer
181	# when 'loss_cfg.router_aux_loss_coef' is configured and greater than 0.	194	# when 'loss_cfg.router_aux_loss_coef' is configured and greater than 0.
	@@ -211,6 +224,7 @@ class FeaturesApplier:
211	hook_optimizer_step(model, optimizer)	224	hook_optimizer_step(model, optimizer)
212		225
213	def pre_fully_shard_apply(self, model):	226	def pre_fully_shard_apply(self, model):
		227	+ self.apply_contiguous_linear_grad(model)
214	# The order of these operations is critical and must not be changed.	228	# The order of these operations is critical and must not be changed.
215	# 1. Prepare Chunk Layer: Replaces the decoder forward with the single-chunk forward so	229	# 1. Prepare Chunk Layer: Replaces the decoder forward with the single-chunk forward so
216	# recompute and activation offload are applied to each chunk independently.	230	# recompute and activation offload are applied to each chunk independently.

Mmindspeed_mm/fsdp/features/memory/chunkloss/chunkloss.py+12-1

	@@ -61,7 +61,7 @@ class ChunkLoss(torch.autograd.Function):
61	hidden_states_chunks, grad_inputs_chunks, loss_kwargs_chunks	61	hidden_states_chunks, grad_inputs_chunks, loss_kwargs_chunks
62	):	62	):
63	# Compute both gradients and loss value for this chunk	63	# Compute both gradients and loss value for this chunk
64	- (chunk_grad_input, chunk_grad_weight), (per_chunk_loss, _) = torch.func.grad_and_value(	64	+ (chunk_grad_input, chunk_grad_weight), (per_chunk_loss, chunk_aux) = torch.func.grad_and_value(
65	loss_forward, argnums=(0, 1), has_aux=True	65	loss_forward, argnums=(0, 1), has_aux=True
66	)(hidden_states_chunk, head_weight, None, **loss_kwargs)	66	)(hidden_states_chunk, head_weight, None, **loss_kwargs)
67		67
	@@ -69,6 +69,8 @@ class ChunkLoss(torch.autograd.Function):
69	accumulated_loss.add_(per_chunk_loss)	69	accumulated_loss.add_(per_chunk_loss)
70	grad_inputs_chunk.copy_(chunk_grad_input)	70	grad_inputs_chunk.copy_(chunk_grad_input)
71	grad_weight.add_(chunk_grad_weight)	71	grad_weight.add_(chunk_grad_weight)
		72	+ # Release these tensors before the next chunk allocates its results.
		73	+ del chunk_grad_input, chunk_grad_weight, per_chunk_loss, chunk_aux
72		74
73	# Save computed gradients for use in backward pass	75	# Save computed gradients for use in backward pass
74	ctx.save_for_backward(grad_inputs, grad_weight)	76	ctx.save_for_backward(grad_inputs, grad_weight)
	@@ -259,6 +261,8 @@ def chunk_loss(hidden_states, head_weight, head_bias, loss_forward, loss_kwargs_
259	Returns:	261	Returns:
260	The total accumulated loss as a scalar tensor.	262	The total accumulated loss as a scalar tensor.
261	"""	263	"""
		264	+ if head_bias is not None:
		265	+ raise NotImplementedError("head_bias is not supported in ChunkLoss")
262	kwargs = loss_kwargs_chunks[0]	266	kwargs = loss_kwargs_chunks[0]
263	alpha = kwargs.get("alpha")	267	alpha = kwargs.get("alpha")
264		268
	@@ -282,6 +286,13 @@ def chunk_loss(hidden_states, head_weight, head_bias, loss_forward, loss_kwargs_
282	for labels in torch.split(shift_labels, chunk_size, dim=1)	286	for labels in torch.split(shift_labels, chunk_size, dim=1)
283	]	287	]
284		288
		289	+ # With a single chunk, ordinary autograd already has the same logits bound.
		290	+ # Avoid computing and retaining a full weight gradient during the forward;
		291	+ # let the backward compute it once with its actual upstream gradient.
		292	+ if len(loss_kwargs_chunks) == 1:
		293	+ loss, _ = loss_forward(hidden_states, head_weight, None, **loss_kwargs_chunks[0])
		294	+ return loss
		295	+
285	return ChunkLoss.apply(	296	return ChunkLoss.apply(
286	hidden_states,	297	hidden_states,
287	head_weight,	298	head_weight,

Amindspeed_mm/fsdp/models/mistral3/__init__.py+0-0

The file is empty

Amindspeed_mm/fsdp/models/mistral3/modeling_mistral3.py+296-0

@@ -0,0 +1,296 @@
	1	+# Copyright 2025 HuggingFace Inc. team.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
	3	+
	4	+"""Mistral3 model adapter for FSDP2 training."""
	5	+
	6	+from itertools import accumulate
	7	+from typing import Optional, Union
	8	+
	9	+import torch
	10	+from transformers import Mistral3ForConditionalGeneration
	11	+from transformers.cache_utils import Cache
	12	+from transformers.models.mistral3.modeling_mistral3 import Mistral3CausalLMOutputWithPast
	13	+from transformers.utils import can_return_tuple
	14	+
	15	+from mindspeed_mm.fsdp.distributed.context_parallel.communication import (
	16	+ gather_forward_split_backward,
	17	+ split_forward_gather_backward_with_cp,
	18	+)
	19	+from mindspeed_mm.fsdp.distributed.parallel_state import (
	20	+ get_parallel_state,
	21	+ is_parallel_state_initialized,
	22	+)
	23	+from mindspeed_mm.fsdp.utils.register import model_register
	24	+
	25	+
	26	+def _prepare_vision_attention(module, args, kwargs):
	27	+ """Preserve image boundaries when Pixtral flattens patches into one sequence.
	28	+
	29	+ The complete vision tower is replicated across CP ranks. Its attention must
	30	+ stay local, even when the language decoder uses Ulysses, Ring or Hybrid CP.
	31	+ """
	32	+ if module.config._attn_implementation not in ("flash_attention_2", "flash_attention_3"):
	33	+ return args, kwargs
	34	+ pixel_values = args[0] if args else kwargs["pixel_values"]
	35	+ image_sizes = args[1] if len(args) > 1 else kwargs.get("image_sizes")
	36	+ if image_sizes is None:
	37	+ image_sizes = [pixel_values.shape[-2:]] * pixel_values.shape[0]
	38	+ if isinstance(image_sizes, torch.Tensor):
	39	+ image_sizes = image_sizes.tolist()
	40	+ lengths = [
	41	+ (int(height) // module.patch_size) * (int(width) // module.patch_size)
	42	+ for height, width in image_sizes
	43	+ ]
	44	+ cu_seqlens = [0, *accumulate(lengths)]
	45	+ max_length = max(lengths)
	46	+ kwargs = dict(kwargs)
	47	+ kwargs.update(
	48	+ cu_seq_lens_q=cu_seqlens,
	49	+ cu_seq_lens_k=cu_seqlens,
	50	+ max_length_q=max_length,
	51	+ max_length_k=max_length,
	52	+ skip_context_parallel=True,
	53	+ )
	54	+ return args, kwargs
	55	+
	56	+
	57	+def _context_parallel_enabled() -> bool:
	58	+ """Whether this step must split the sequence across context-parallel ranks."""
	59	+ if not is_parallel_state_initialized():
	60	+ return False
	61	+ return get_parallel_state().is_cp_enable()
	62	+
	63	+
	64	+def _validate_dense_attention_mask(attention_mask):
	65	+ """Reject masks that the NPU/CP dense causal path cannot preserve."""
	66	+ if attention_mask is None:
	67	+ return
	68	+ if attention_mask.ndim != 2:
	69	+ raise ValueError("Mistral3 FSDP2 Flash Attention and CP require a 2D right-padding attention mask.")
	70	+ # Read both validation results together, avoiding two device-to-host synchronizations.
	71	+ invalid_values, invalid_padding = torch.stack((
	72	+ ((attention_mask != 0) & (attention_mask != 1)).any(),
	73	+ ((attention_mask[:, 1:] != 0) & (attention_mask[:, :-1] == 0)).any(),
	74	+ )).cpu().tolist()
	75	+ if invalid_values:
	76	+ raise ValueError(
	77	+ "Mistral3 FSDP2 Flash Attention and CP do not support neat_packing. "
	78	+ "Use packing=False when sample isolation is required."
	79	+ )
	80	+ if invalid_padding:
	81	+ raise ValueError("Mistral3 FSDP2 Flash Attention and CP support right padding only.")
	82	+
	83	+
	84	+@model_register.register("mistral3")
	85	+class Mistral3ForConditionalGenerationFSDP(Mistral3ForConditionalGeneration):
	86	+ """Register the Transformers Mistral3 model for FSDP2 training."""
	87	+
	88	+ def __init__(self, config):
	89	+ super().__init__(config)
	90	+ self.model.vision_tower.register_forward_pre_hook(_prepare_vision_attention, with_kwargs=True)
	91	+
	92	+ @classmethod
	93	+ def _from_config(cls, config, **kwargs):
	94	+ config.use_cache = False
	95	+ return super()._from_config(config, **kwargs)
	96	+
	97	+ def forward(
	98	+ self,
	99	+ input_ids: Optional[torch.LongTensor] = None,
	100	+ pixel_values: Optional[torch.FloatTensor] = None,
	101	+ attention_mask: Optional[torch.Tensor] = None,
	102	+ position_ids: Optional[torch.LongTensor] = None,
	103	+ past_key_values: Optional[Cache] = None,
	104	+ inputs_embeds: Optional[torch.FloatTensor] = None,
	105	+ labels: Optional[torch.LongTensor] = None,
	106	+ use_cache: Optional[bool] = None,
	107	+ output_attentions: Optional[bool] = None,
	108	+ output_hidden_states: Optional[bool] = None,
	109	+ return_dict: Optional[bool] = None,
	110	+ cache_position: Optional[torch.LongTensor] = None,
	111	+ logits_to_keep: Union[int, torch.Tensor] = 0,
	112	+ image_sizes: Optional[torch.Tensor] = None,
	113	+ **kwargs,
	114	+ ) -> Union[tuple, Mistral3CausalLMOutputWithPast]:
	115	+ """Use the Transformers implementation unless an FSDP2 feature requires adaptation."""
	116	+ context_parallel = _context_parallel_enabled()
	117	+ text_config = getattr(getattr(self, "config", None), "text_config", None)
	118	+ if context_parallel or getattr(text_config, "_attn_implementation", None) in (
	119	+ "flash_attention_2", "flash_attention_3"
	120	+ ):
	121	+ _validate_dense_attention_mask(attention_mask)
	122	+ if not getattr(self, "enable_chunk_loss", False) and not context_parallel:
	123	+ forward = super().forward
	124	+ else:
	125	+ forward = self._feature_forward
	126	+
	127	+ return forward(
	128	+ input_ids=input_ids,
	129	+ pixel_values=pixel_values,
	130	+ attention_mask=attention_mask,
	131	+ position_ids=position_ids,
	132	+ past_key_values=past_key_values,
	133	+ inputs_embeds=inputs_embeds,
	134	+ labels=labels,
	135	+ use_cache=use_cache,
	136	+ output_attentions=output_attentions,
	137	+ output_hidden_states=output_hidden_states,
	138	+ return_dict=return_dict,
	139	+ cache_position=cache_position,
	140	+ logits_to_keep=logits_to_keep,
	141	+ image_sizes=image_sizes,
	142	+ **kwargs,
	143	+ )
	144	+
	145	+ @can_return_tuple
	146	+ def _feature_forward(
	147	+ self,
	148	+ input_ids: Optional[torch.LongTensor] = None,
	149	+ pixel_values: Optional[torch.FloatTensor] = None,
	150	+ attention_mask: Optional[torch.Tensor] = None,
	151	+ position_ids: Optional[torch.LongTensor] = None,
	152	+ past_key_values: Optional[Cache] = None,
	153	+ inputs_embeds: Optional[torch.FloatTensor] = None,
	154	+ labels: Optional[torch.LongTensor] = None,
	155	+ use_cache: Optional[bool] = None,
	156	+ output_attentions: Optional[bool] = None,
	157	+ output_hidden_states: Optional[bool] = None,
	158	+ return_dict: Optional[bool] = None,
	159	+ cache_position: Optional[torch.LongTensor] = None,
	160	+ logits_to_keep: Union[int, torch.Tensor] = 0,
	161	+ image_sizes: Optional[torch.Tensor] = None,
	162	+ **kwargs,
	163	+ ) -> Union[tuple, Mistral3CausalLMOutputWithPast]:
	164	+ """Run the adapted forward pass for chunked loss or context parallelism."""
	165	+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	166	+ output_hidden_states = (
	167	+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	168	+ )
	169	+
	170	+ if not _context_parallel_enabled():
	171	+ outputs = self.model(
	172	+ input_ids=input_ids,
	173	+ pixel_values=pixel_values,
	174	+ attention_mask=attention_mask,
	175	+ position_ids=position_ids,
	176	+ past_key_values=past_key_values,
	177	+ inputs_embeds=inputs_embeds,
	178	+ use_cache=use_cache,
	179	+ output_attentions=output_attentions,
	180	+ output_hidden_states=output_hidden_states,
	181	+ return_dict=True,
	182	+ cache_position=cache_position,
	183	+ image_sizes=image_sizes,
	184	+ **kwargs,
	185	+ )
	186	+ image_hidden_states = outputs.image_hidden_states
	187	+ else:
	188	+ outputs, image_hidden_states = self._context_parallel_model_forward(
	189	+ input_ids=input_ids,
	190	+ pixel_values=pixel_values,
	191	+ attention_mask=attention_mask,
	192	+ position_ids=position_ids,
	193	+ past_key_values=past_key_values,
	194	+ inputs_embeds=inputs_embeds,
	195	+ use_cache=use_cache,
	196	+ output_attentions=output_attentions,
	197	+ output_hidden_states=output_hidden_states,
	198	+ cache_position=cache_position,
	199	+ image_sizes=image_sizes,
	200	+ **kwargs,
	201	+ )
	202	+
	203	+ hidden_states = outputs[0]
	204	+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
	205	+ if getattr(self, "enable_chunk_loss", False):
	206	+ logits = None
	207	+ loss = self.lm_head(hidden_states[:, slice_indices, :], self.loss_function, labels)
	208	+ else:
	209	+ logits = self.lm_head(hidden_states[:, slice_indices, :])
	210	+ loss = None
	211	+ if labels is not None:
	212	+ loss = self.loss_function(
	213	+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
	214	+ )
	215	+
	216	+ if loss is not None and _context_parallel_enabled():
	217	+ ps = get_parallel_state()
	218	+ loss = gather_forward_split_backward(loss.unsqueeze(0), ps.get_cp_group(), dim=0).sum()
	219	+
	220	+ return Mistral3CausalLMOutputWithPast(
	221	+ loss=loss,
	222	+ logits=logits,
	223	+ past_key_values=outputs.past_key_values,
	224	+ hidden_states=outputs.hidden_states,
	225	+ attentions=outputs.attentions,
	226	+ image_hidden_states=image_hidden_states,
	227	+ )
	228	+
	229	+ def _context_parallel_model_forward(
	230	+ self,
	231	+ input_ids=None,
	232	+ pixel_values=None,
	233	+ attention_mask=None,
	234	+ position_ids=None,
	235	+ past_key_values=None,
	236	+ inputs_embeds=None,
	237	+ use_cache=None,
	238	+ output_attentions=None,
	239	+ output_hidden_states=None,
	240	+ cache_position=None,
	241	+ image_sizes=None,
	242	+ **kwargs,
	243	+ ):
	244	+ """``Mistral3Model.forward`` with the sequence split across CP ranks before the decoder.
	245	+
	246	+ The embedding and image-merge steps are upstream's own methods; only the split is added, in the
	247	+ one place ``Mistral3Model.forward`` gives no hook for. Splitting must happen after the image
	248	+ placeholders are filled in, because the merge indexes the full sequence.
	249	+ """
	250	+ model = self.model
	251	+ if (input_ids is None) ^ (inputs_embeds is not None):
	252	+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
	253	+
	254	+ if inputs_embeds is None:
	255	+ inputs_embeds = model.get_input_embeddings()(input_ids)
	256	+
	257	+ image_features = None
	258	+ if pixel_values is not None:
	259	+ image_features = model.get_image_features(
	260	+ pixel_values=pixel_values,
	261	+ vision_feature_layer=self.config.vision_feature_layer,
	262	+ image_sizes=image_sizes,
	263	+ )
	264	+ image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
	265	+ special_image_mask = model.get_placeholder_mask(
	266	+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
	267	+ )
	268	+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
	269	+
	270	+ total_seq_len = inputs_embeds.shape[1]
	271	+ if cache_position is None:
	272	+ cache_position = torch.arange(total_seq_len, device=inputs_embeds.device)
	273	+ if position_ids is None:
	274	+ position_ids = cache_position.unsqueeze(0)
	275	+
	276	+ inputs_embeds = split_forward_gather_backward_with_cp(inputs_embeds, dim=1)
	277	+ position_ids = split_forward_gather_backward_with_cp(position_ids, dim=position_ids.dim() - 1)
	278	+ cache_position = split_forward_gather_backward_with_cp(
	279	+ cache_position.unsqueeze(0), dim=1
	280	+ ).squeeze(0)
	281	+
	282	+ attention_mask = None
	283	+
	284	+ outputs = model.language_model(
	285	+ attention_mask=attention_mask,
	286	+ position_ids=position_ids,
	287	+ past_key_values=past_key_values,
	288	+ inputs_embeds=inputs_embeds,
	289	+ use_cache=use_cache,
	290	+ output_attentions=output_attentions,
	291	+ output_hidden_states=output_hidden_states,
	292	+ return_dict=True,
	293	+ cache_position=cache_position,
	294	+ **kwargs,
	295	+ )
	296	+ return outputs, image_features

Amindspeed_mm/fsdp/ops/contiguous_linear.py+49-0

@@ -0,0 +1,49 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Bias-free NPU Linear backward with a contiguous weight gradient."""
	3	+from types import MethodType
	4	+
	5	+import torch
	6	+import torch.nn.functional as F
	7	+
	8	+
	9	+class ContiguousLinear(torch.autograd.Function):
	10	+ generate_vmap_rule = True
	11	+
	12	+ @staticmethod
	13	+ def forward(inputs, weight):
	14	+ return F.linear(inputs, weight)
	15	+
	16	+ @staticmethod
	17	+ def setup_context(ctx, inputs, output):
	18	+ ctx.save_for_backward(*inputs)
	19	+
	20	+ @staticmethod
	21	+ def backward(ctx, grad_output):
	22	+ inputs, weight = ctx.saved_tensors
	23	+ dy = grad_output.reshape(-1, weight.shape[0])
	24	+ x = inputs.reshape(-1, weight.shape[1])
	25	+ dx = (dy @ weight).reshape_as(inputs) if ctx.needs_input_grad[0] else None
	26	+ dw = dy.t() @ x if ctx.needs_input_grad[1] else None
	27	+ return dx, dw
	28	+
	29	+
	30	+def _linear_forward(module, input):
	31	+ inputs = input
	32	+ supported = (
	33	+ inputs.device.type == "npu"
	34	+ and inputs.dtype in (torch.bfloat16, torch.float16)
	35	+ and module.weight.dtype == inputs.dtype
	36	+ and inputs.ndim >= 2
	37	+ and torch.is_grad_enabled()
	38	+ and not torch.is_autocast_enabled("npu")
	39	+ )
	40	+ if supported:
	41	+ return ContiguousLinear.apply(inputs, module.weight)
	42	+ return F.linear(inputs, module.weight, module.bias)
	43	+
	44	+
	45	+def apply_contiguous_linear(module):
	46	+ """Retain parameter identities and names; reject custom Linear subclasses."""
	47	+ if type(module) is not torch.nn.Linear or module.bias is not None:
	48	+ raise ValueError("contiguous_linear_grad_modules requires bias-free torch.nn.Linear modules")
	49	+ module.forward = MethodType(_linear_forward, module)

Mmindspeed_mm/fsdp/ops/flash_attn/flash_attn.py+36-32

	@@ -48,16 +48,16 @@ def _convert_cu_seq_lens(
48	cu_seq_lens_q: Optional[Union[torch.Tensor, List[int]]],	48	cu_seq_lens_q: Optional[Union[torch.Tensor, List[int]]],
49	cu_seq_lens_k: Optional[Union[torch.Tensor, List[int]]],	49	cu_seq_lens_k: Optional[Union[torch.Tensor, List[int]]],
50	input_layout: Optional[str] = None	50	input_layout: Optional[str] = None
51	-) -> Tuple[Optional[List[int]], Optional[str]]:	51	+) -> Tuple[Optional[List[int]], Optional[List[int]], Optional[str]]:
52	- """	52	+ """Convert query and key sequence offsets and validate the input layout.
53	- Converts cu_seq_lens to a list and validates input_layout constraints.
54		53
55	Args:	54	Args:
56	- cu_seq_lens: Cumulative sequence lengths, can be a Tensor or List.	55	+ cu_seq_lens_q: Query cumulative sequence lengths.
57	- input_layout: The layout of the input tensor (e.g., 'BNSD', 'BSND').	56	+ cu_seq_lens_k: Key cumulative sequence lengths.
		57	+ input_layout: Input tensor layout.
58		58
59	Returns:	59	Returns:
60	- A tuple containing the processed cu_seq_lens (as list) and the validated layout.	60	+ The converted query lengths, key lengths, and normalized input layout.
61	"""	61	"""
62		62
63	# 1. Convert Tensor to List if necessary	63	# 1. Convert Tensor to List if necessary
	@@ -72,36 +72,35 @@ def _convert_cu_seq_lens(
72		72
73	# Check constraints for packed layouts (BNSD/BSND)	73	# Check constraints for packed layouts (BNSD/BSND)
74	if input_layout in ["BNSD", "BSND"]:	74	if input_layout in ["BNSD", "BSND"]:
75	- if (cu_seq_lens_q is not None or cu_seq_lens_k is not None) and (len(cu_seq_lens_q) > 2 or len(cu_seq_lens_k) > 2):	75	+ query_is_packed = cu_seq_lens_q is not None and len(cu_seq_lens_q) > 2
		76	+ key_is_packed = cu_seq_lens_k is not None and len(cu_seq_lens_k) > 2
		77	+ if query_is_packed or key_is_packed:
76	raise RuntimeError(	78	raise RuntimeError(
77	f"NPU flash attention layout {input_layout} does not support packing data "	79	f"NPU flash attention layout {input_layout} does not support packing data "
78	f"when cu_seq_lens length > 2"	80	f"when cu_seq_lens length > 2"
79	)	81	)
80		82
81	- # Auto-correct layout if cu_seq_lens is missing	83	+ elif input_layout in ["1TND", "1NTD"]:
82	- elif input_layout == "1TND" and (cu_seq_lens_q is None or len(cu_seq_lens_q) == 2) and (cu_seq_lens_k is None or len(cu_seq_lens_k) == 2):	84	+ if (cu_seq_lens_q is None) != (cu_seq_lens_k is None):
83	- # Fallback to standard layout if packing info is missing	85	+ raise ValueError("Packed attention requires both query and key sequence lengths")
84	- input_layout = "BSND"	86	+ if cu_seq_lens_q is not None:
85	- cu_seq_lens_q = None	87	+ # HF offsets start at zero; npu_fusion_attention expects only
86	- cu_seq_lens_k = None	88	+ # cumulative end positions. Also accept callers already using
87	-	89	+ # the NPU convention.
88	- elif input_layout == "1NTD":	90	+ cu_seq_lens_q = list(cu_seq_lens_q)
89	- if (cu_seq_lens_q is None or len(cu_seq_lens_q) == 2) and (cu_seq_lens_k is None or len(cu_seq_lens_k) == 2):	91	+ cu_seq_lens_k = list(cu_seq_lens_k)
90	- # Fallback to standard layout if packing info is missing	92	+ if cu_seq_lens_q and cu_seq_lens_q[0] == 0:
91	- input_layout = "BNSD"	93	+ cu_seq_lens_q = cu_seq_lens_q[1:]
		94	+ if cu_seq_lens_k and cu_seq_lens_k[0] == 0:
		95	+ cu_seq_lens_k = cu_seq_lens_k[1:]
		96	+ if not cu_seq_lens_q or not cu_seq_lens_k:
		97	+ raise ValueError("Packed attention requires nonempty sequence lengths")
		98	+ # Equal-length packed sequences are still separate sequences. A
		99	+ # layout change alone cannot turn their token dimension into batch.
		100	+ if cu_seq_lens_q is None or (len(cu_seq_lens_q) == 1 and len(cu_seq_lens_k) == 1):
		101	+ input_layout = "BSND" if input_layout == "1TND" else "BNSD"
92	cu_seq_lens_q = None	102	cu_seq_lens_q = None
93	cu_seq_lens_k = None	103	cu_seq_lens_k = None
94	- elif len(cu_seq_lens_q) > 2 and len(cu_seq_lens_k) > 2:
95	- # Judge whether cu_seq_lens is an arithmetic sequence (i.e., all sequences have equal length).
96	- # If true, switch to BNSD computation
97	- cu_seq_q = torch.tensor(cu_seq_lens_q)
98	- diffs_q = torch.diff(cu_seq_q)
99	- cu_seq_k = torch.tensor(cu_seq_lens_k)
100	- diffs_k = torch.diff(cu_seq_k)
101	- if torch.all(diffs_q == diffs_q[0]).item() and torch.all(diffs_k == diffs_k[0]).item():
102	- input_layout = "BNSD"
103	- cu_seq_lens_q = None
104	- cu_seq_lens_k = None
105		104
106	return cu_seq_lens_q, cu_seq_lens_k, input_layout	105	return cu_seq_lens_q, cu_seq_lens_k, input_layout
107		106
	@@ -276,6 +275,7 @@ def flash_attention_forward(
276	total_seq_len: int = None, # unaligned cp need this	275	total_seq_len: int = None, # unaligned cp need this
277	seq_split_lens: torch.Tensor = None, # unaligned cp need this	276	seq_split_lens: torch.Tensor = None, # unaligned cp need this
278	skip_flash_attn_recompute: bool = False,	277	skip_flash_attn_recompute: bool = False,
		278	+ skip_context_parallel: bool = False, # Replicated modules outside the sequence split, e.g. a vision tower.
279	**kwargs,	279	**kwargs,
280	) -> tuple[torch.Tensor, None]:	280	) -> tuple[torch.Tensor, None]:
281	if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None:	281	if kwargs.get("output_attentions", False) or kwargs.get("head_mask") is not None:
	@@ -302,7 +302,11 @@ def flash_attention_forward(
302		302
303	# Convert cumulative sequence lengths to list format for NPU API.	303	# Convert cumulative sequence lengths to list format for NPU API.
304	# (For ring cp, cu_seqlens info in seq_split_lens)	304	# (For ring cp, cu_seqlens info in seq_split_lens)
305	- if not dist.is_initialized() or not (get_parallel_state().is_ring_enable() and seq_split_lens is not None):	305	+ if (
		306	+ skip_context_parallel
		307	+ or not dist.is_initialized()
		308	+ or not (get_parallel_state().is_ring_enable() and seq_split_lens is not None)
		309	+ ):
306	cu_seq_lens_q, cu_seq_lens_k, input_layout = _convert_cu_seq_lens(cu_seq_lens_q, cu_seq_lens_k, input_layout=input_layout)	310	cu_seq_lens_q, cu_seq_lens_k, input_layout = _convert_cu_seq_lens(cu_seq_lens_q, cu_seq_lens_k, input_layout=input_layout)
307		311
308	# Determine sequence length based on layout.	312	# Determine sequence length based on layout.
	@@ -349,8 +353,8 @@ def flash_attention_forward(
349		353
350	# Modification: ============= CONTEXT PARALLEL (CP) =============	354	# Modification: ============= CONTEXT PARALLEL (CP) =============
351	ps = get_parallel_state()	355	ps = get_parallel_state()
352	- is_ulysses_enabled = ps.is_ulysses_enable() if dist.is_initialized() else False	356	+ is_ulysses_enabled = ps.is_ulysses_enable() if dist.is_initialized() and not skip_context_parallel else False
353	- is_ring_enabled = ps.is_ring_enable() if dist.is_initialized() else False	357	+ is_ring_enabled = ps.is_ring_enable() if dist.is_initialized() and not skip_context_parallel else False
354		358
355	if use_npu_fusion_fa:	359	if use_npu_fusion_fa:
356	if input_layout in ["BSND", "1TND"]:	360	if input_layout in ["BSND", "1TND"]:

Amindspeed_mm/fsdp/ops/fully_shard/npu_split_copy.py+66-0

@@ -0,0 +1,66 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Opt-in NPU FSDP copy from rank-major buffers to contiguous parameters."""
	3	+import torch
	4	+import torch_npu
	5	+import triton
	6	+import triton.language as tl
	7	+
	8	+
	9	+@triton.jit
	10	+def _copy_parameter(source, destination, N, STRIDE, OFFSET, BLOCK: tl.constexpr):
	11	+ index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
	12	+ rank = tl.program_id(1)
	13	+ values = tl.load(source + rank * STRIDE + OFFSET + index, index < N, other=0)
	14	+ tl.store(destination + rank * N + index, values, index < N)
	15	+
	16	+
	17	+def split_with_sizes_copy(source, sizes, dim, out):
	18	+ supported = (
	19	+ source.device.type == 'npu'
	20	+ and source.layout == torch.strided
	21	+ and source.ndim == 2 and dim == 1 and source.is_contiguous()
	22	+ and not source.is_neg() and not source.is_conj()
	23	+ and source.dtype in (torch.bfloat16, torch.float16, torch.float32, torch.uint8, torch.int8)
	24	+ and source.numel() <= 2**31 - 1 and len(sizes) == len(out)
	25	+ and not source.requires_grad and not any(t.requires_grad for t in out)
	26	+ and all(n >= 0 for n in sizes) and sum(sizes) == source.shape[1]
	27	+ and all(t.device == source.device and t.dtype == source.dtype
	28	+ and t.layout == torch.strided and not t.is_neg() and not t.is_conj()
	29	+ and t.shape == (source.shape[0], n) and t.is_contiguous()
	30	+ for t, n in zip(out, sizes))
	31	+ )
	32	+ if supported:
	33	+ # Logical strides do not describe blocked NPU storage such as FRACTAL_NZ.
	34	+ # Only the plain NCHW/ND formats have the row layout used by this kernel.
	35	+ supported = all(torch_npu.get_npu_format(t) in (0, 2) for t in (source, *out))
	36	+ if supported:
	37	+ pointers = [source.untyped_storage().data_ptr()]
	38	+ pointers.extend(t.untyped_storage().data_ptr() for t in out)
	39	+ supported = len(set(pointers)) == len(pointers)
	40	+ if not supported:
	41	+ torch.split_with_sizes_copy(source, sizes, dim=dim, out=out)
	42	+ return
	43	+ offset = 0
	44	+ for n, output in zip(sizes, out):
	45	+ if n:
	46	+ _copy_parameter[(triton.cdiv(n, 32768), source.shape[0])](
	47	+ source, output, n, source.shape[1], offset, 32768)
	48	+ offset += n
	49	+ torch.autograd.graph.increment_version(out)
	50	+
	51	+
	52	+_registration = None
	53	+
	54	+
	55	+def enable_npu_split_copy():
	56	+ """Install once per process; unsupported layouts use the native operation."""
	57	+ global _registration
	58	+ if _registration is not None:
	59	+ return
	60	+ if not torch.__version__.split("+")[0].startswith("2.10."):
	61	+ raise RuntimeError("use_triton_split_copy requires PyTorch 2.10")
	62	+ from torch.distributed.fsdp._fully_shard import _fsdp_collectives # noqa: F401
	63	+
	64	+ library = torch.library.Library("fsdp", "IMPL", "PrivateUse1")
	65	+ library.impl("split_with_sizes_copy", split_with_sizes_copy, allow_override=True)
	66	+ _registration = library

Mmindspeed_mm/fsdp/params/feature_args.py+6-0

	@@ -253,6 +253,12 @@ class FeatureArguments(BaseArguments):
253	)	253	)
254	chunkloss_plan: ChunkLossPlanConfig = field(default_factory=ChunkLossPlanConfig)	254	chunkloss_plan: ChunkLossPlanConfig = field(default_factory=ChunkLossPlanConfig)
255		255
		256	+ contiguous_linear_grad_modules: List[str] = field(
		257	+ default_factory=list,
		258	+ metadata={"help": "Bias-free Linear modules computing contiguous weight gradients on NPU BF16/FP16. "
		259	+ "Empty disables the optimization; other dtypes and autocast retain native Linear."},
		260	+ )
		261	+
256	enable_activation_offload: bool = field(	262	enable_activation_offload: bool = field(
257	default=False,	263	default=False,
258	metadata={"help": "Whether apply activation offload"}	264	metadata={"help": "Whether apply activation offload"}

Mmindspeed_mm/fsdp/params/parallel_args.py+3-0

	@@ -26,6 +26,9 @@ class FSDPPlanConfig(BaseArguments):
26	num_to_forward_prefetch: Optional[int] = 0	26	num_to_forward_prefetch: Optional[int] = 0
27	num_to_backward_prefetch: Optional[int] = 0	27	num_to_backward_prefetch: Optional[int] = 0
28		28
		29	+ # Explicit opt-in: the operator registration applies to this process.
		30	+ use_triton_split_copy: bool = False
		31	+
29	# pregather settings	32	# pregather settings
30	pregather: bool = False	33	pregather: bool = False
31		34

Mmindspeed_mm/fsdp/params/training_args.py+8-0

	@@ -322,6 +322,12 @@ class TrainingArguments(BaseArguments):
322	default=0,	322	default=0,
323	metadata={"help": "Training step interval to trigger manual garbage collection. When the value is set to 0, garbage collection is not triggered between training steps. When the value is greater than 0, default python garbage collection is disabled and triggered every `manual_gc_interval` steps."},	323	metadata={"help": "Training step interval to trigger manual garbage collection. When the value is set to 0, garbage collection is not triggered between training steps. When the value is greater than 0, default python garbage collection is disabled and triggered every `manual_gc_interval` steps."},
324	)	324	)
		325	+ freeze_gc_after_init: bool = field(
		326	+ default=False,
		327	+ metadata={"help": "Freeze setup objects for the lifetime of this standalone rank process. "
		328	+ "Automatic GC still collects new cycles. Supports one training invocation per process; "
		329	+ "cannot be combined with manual_gc_interval."},
		330	+ )
325	use_deter_comp: bool = field(	331	use_deter_comp: bool = field(
326	default=False,	332	default=False,
327	metadata={"help": "Whether to use deterministic computation for reproducibility."},	333	metadata={"help": "Whether to use deterministic computation for reproducibility."},
	@@ -380,6 +386,8 @@ class TrainingArguments(BaseArguments):
380		386
381	if self.manual_gc_interval < 0:	387	if self.manual_gc_interval < 0:
382	raise ValueError(f"`manual_gc_interval` must be a non-negative integer, got {self.manual_gc_interval}.")	388	raise ValueError(f"`manual_gc_interval` must be a non-negative integer, got {self.manual_gc_interval}.")
		389	+ if self.freeze_gc_after_init and self.manual_gc_interval > 0:
		390	+ raise ValueError("freeze_gc_after_init cannot be combined with manual_gc_interval")
383		391
384	def compute_distributed_training(	392	def compute_distributed_training(
385	self, parallel_args	393	self, parallel_args

Mmindspeed_mm/fsdp/train/train_engine.py+6-0

	@@ -322,6 +322,12 @@ class TrainEngine:
322	"Please provide validation dataset or set val_interval to 0."	322	"Please provide validation dataset or set val_interval to 0."
323	)	323	)
324		324
		325	+ if getattr(args.training, "freeze_gc_after_init", False):
		326	+ from mindspeed_mm.fsdp.utils.gc_utils import freeze_training_initialization
		327	+
		328	+ count = freeze_training_initialization()
		329	+ print_rank(logger.info, f"Froze {count} setup objects; automatic GC remains enabled for new objects")
		330	+
325	# Get data iterator	331	# Get data iterator
326	train_dataloader_iter, _, _ = build_iterations(self.train_dataloader)	332	train_dataloader_iter, _, _ = build_iterations(self.train_dataloader)
327	param_dtype = get_dtype(args.parallel.fsdp_plan.param_dtype) if args.parallel.fsdp_plan.param_dtype else None	333	param_dtype = get_dtype(args.parallel.fsdp_plan.param_dtype) if args.parallel.fsdp_plan.param_dtype else None

Amindspeed_mm/fsdp/utils/gc_utils.py+24-0

@@ -0,0 +1,24 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Process lifetime GC policy for a standalone training invocation."""
	3	+import gc
	4	+
	5	+
	6	+_training_initialization_frozen = False
	7	+
	8	+
	9	+def freeze_training_initialization():
	10	+ """Freeze setup objects once; automatic GC still collects new cycles.
	11	+
	12	+ This is an explicit process lifetime option. Existing permanent-generation
	13	+ objects are retained, and this function never unfreezes the caller's objects.
	14	+ Use a fresh rank process for another training invocation.
	15	+ """
	16	+ global _training_initialization_frozen
	17	+ if _training_initialization_frozen:
	18	+ raise RuntimeError("freeze_gc_after_init supports one training invocation per process")
	19	+ if not gc.isenabled():
	20	+ raise RuntimeError("freeze_gc_after_init requires automatic GC to remain enabled")
	21	+ _training_initialization_frozen = True
	22	+ gc.collect()
	23	+ gc.freeze()
	24	+ return gc.get_freeze_count()

Atests/st/fsdp/mistral3/README.md+54-0

@@ -0,0 +1,54 @@
	1	+# Mistral3 NPU 回归测试
	2	+
	3	+脚本使用真实 NPU 计算，测试视觉注意力和合成小模型训练。
	4	+依赖版本见[训练说明](../../../../examples/magistral-2509/README_v1.md)，从仓库根目录运行：
	5	+
	6	+```bash
	7	+source /usr/local/Ascend/ascend-toolkit/set_env.sh
	8	+export NON_MEGATRON=true
	9	+export OMP_NUM_THREADS=1
	10	+export MAX_JOBS=4
	11	+python -m torch.distributed.run --standalone --nproc_per_node=1 \
	12	+ tests/st/fsdp/mistral3/vision_attention_probe.py --mode default
	13	+python -m torch.distributed.run --standalone --nproc_per_node=4 \
	14	+ tests/st/fsdp/mistral3/vision_attention_probe.py --mode ulysses
	15	+python -m torch.distributed.run --standalone --nproc_per_node=4 \
	16	+ tests/st/fsdp/mistral3/vision_attention_probe.py --mode ring
	17	+python -m torch.distributed.run --standalone --nproc_per_node=4 \
	18	+ tests/st/fsdp/mistral3/vision_attention_probe.py --mode hybrid
	19	+python -m torch.distributed.run --standalone --nproc_per_node=1 \
	20	+ tests/st/fsdp/mistral3/train_smoke.py --mode default
	21	+python -m torch.distributed.run --standalone --nproc_per_node=1 \
	22	+ tests/st/fsdp/mistral3/train_smoke.py --mode multi
	23	+python -m torch.distributed.run --standalone --nproc_per_node=4 \
	24	+ tests/st/fsdp/mistral3/train_smoke.py --mode ring
	25	+python -m torch.distributed.run --standalone --nproc_per_node=4 \
	26	+ tests/st/fsdp/mistral3/train_smoke.py --mode hybrid
	27	+```
	28	+
	29	+`vision_attention_probe.py` 与逐图 eager 注意力比较前向输出、输入及参数梯度，检查图像间独立性。它在已初始化的 CP
	30	+进程组中对通信函数设置断言：视觉编码器若调用 Ulysses 或 Ring 通信，测试失败。该检查只用于视觉测试进程。
	31	+
	32	+`train_smoke.py` 使用正常的分布式通信。每个进程必须完成连续的 1–30 步，损失和梯度均为有限值，并打印最终 `PASS`。
	33	+多进程输出可能拼接，应逐个解析 JSON 对象并核对各进程记录。
	34	+
	35	+测试结果见[实践文档 issue](https://gitcode.com/Ascend/MindSpeed-MM/issues/703)。
	36	+这些脚本不加载 24B 权重，也不使用真实数据和默认 trainer；运行时应另存日志。
	37	+
	38	+`compare_training.py` 进一步比较完整小模型的一步训练：固定相同权重、双图输入、标签和位置编号，
	39	+先保存 CP=1 的 FSDP2 参考，再比较 Ulysses 或 Ring 的损失、全部可训练参数梯度和 SGD 参数更新。
	40	+BF16 对照要求损失相对误差小于 2%，梯度和参数更新的相对 L2 误差均小于 0.05。
	41	+使用独立的空目录保存每组参考，两个命令必须使用相同进程数和相同 `--chunk-loss` 设置。
	42	+
	43	+```bash
	44	+python -m torch.distributed.run --standalone --nproc_per_node=2 \
	45	+ tests/st/fsdp/mistral3/compare_training.py --mode reference --reference-dir /tmp/mistral3-reference
	46	+python -m torch.distributed.run --standalone --nproc_per_node=2 \
	47	+ tests/st/fsdp/mistral3/compare_training.py --mode ulysses --reference-dir /tmp/mistral3-reference
	48	+python -m torch.distributed.run --standalone --nproc_per_node=2 \
	49	+ tests/st/fsdp/mistral3/compare_training.py --mode ring --reference-dir /tmp/mistral3-reference
	50	+```
	51	+
	52	+如需检查 ChunkLoss，在参考和被测命令中都增加 `--chunk-loss`，并使用新的参考目录。
	53	+Hybrid 使用四个进程和 `--mode hybrid`，其 CP=1 参考也应使用四个进程。
	54	+该对照不测量吞吐或峰值显存，不能替代 24B 真实数据的 2000 步后端比较。

Atests/st/fsdp/mistral3/check_npu_split_copy.py+152-0

@@ -0,0 +1,152 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Check device values, mutation versions, streams, and inference-mode outputs."""
	3	+import argparse
	4	+import json
	5	+from pathlib import Path
	6	+
	7	+import torch
	8	+import torch_npu
	9	+from torch.distributed.fsdp._fully_shard import _fsdp_collectives
	10	+
	11	+from mindspeed_mm.fsdp.ops.fully_shard.npu_split_copy import enable_npu_split_copy
	12	+
	13	+
	14	+def main():
	15	+ parser = argparse.ArgumentParser(description=__doc__)
	16	+ parser.add_argument('--output', type=Path, required=True)
	17	+ args = parser.parse_args()
	18	+ torch.npu.set_device(0)
	19	+ torch.manual_seed(43)
	20	+ source = torch.randn(2, 16, device='npu')
	21	+ native = [torch.empty(2, n, device='npu') for n in [3, 13]]
	22	+ before = [t._version for t in native]
	23	+ torch.ops.fsdp.split_with_sizes_copy(source, [3, 13], dim=1, out=native)
	24	+ native_delta = [t._version - old for t, old in zip(native, before)]
	25	+ # Negative dispatch can reject mutable TensorLists before either kernel.
	26	+ # Capture the same FSDP operation before installing the NPU implementation.
	27	+ native_negative_errors = {}
	28	+ for negate_source, negate_output in [(True, False), (False, True)]:
	29	+ original = torch.randn(2, 15, device='npu')
	30	+ if negate_source:
	31	+ original = original._neg_view()
	32	+ outputs = [torch.empty(2, n, device='npu') for n in [3, 5, 7]]
	33	+ if negate_output:
	34	+ outputs = [t._neg_view() for t in outputs]
	35	+ error_message = None
	36	+ try:
	37	+ torch.ops.fsdp.split_with_sizes_copy(original, [3, 5, 7], dim=1, out=outputs)
	38	+ except RuntimeError as error:
	39	+ error_message = str(error).splitlines()[0]
	40	+ native_negative_errors[negate_source, negate_output] = error_message
	41	+ enable_npu_split_copy()
	42	+ enable_npu_split_copy()
	43	+ cases = []
	44	+ for dtype in (torch.bfloat16, torch.float16, torch.float32, torch.uint8, torch.int8):
	45	+ for world in (1, 2, 8):
	46	+ for sizes in ([1, 7, 33, 65539], [0, 5, 3]):
	47	+ source = torch.randint(-63, 64, (world, sum(sizes)), device='npu').to(dtype)
	48	+ expected = list(torch.split_with_sizes_copy(source, sizes, dim=1))
	49	+ actual = [torch.empty_like(t) for t in expected]
	50	+ versions = [t._version for t in actual]
	51	+ torch.ops.fsdp.split_with_sizes_copy(source, sizes, dim=1, out=actual)
	52	+ delta = [t._version - old for t, old in zip(actual, versions)]
	53	+ if delta != [native_delta[0]] * len(sizes):
	54	+ raise AssertionError(f'Mutation versions differ: {delta} vs {native_delta}')
	55	+ versions = [t._version for t in actual]
	56	+ with torch.autograd._unsafe_preserve_version_counter(tuple(actual)):
	57	+ torch.ops.fsdp.split_with_sizes_copy(source, sizes, dim=1, out=actual)
	58	+ for target, reference, old in zip(actual, expected, versions):
	59	+ torch.testing.assert_close(target, reference, rtol=0, atol=0)
	60	+ if target._version != old:
	61	+ raise AssertionError('FSDP version preservation failed')
	62	+ cases.append({'world': world, 'dtype': str(dtype), 'sizes': sizes})
	63	+ source = torch.randn(8, 30, device='npu')[:, ::2]
	64	+ sizes = [3, 5, 7]
	65	+ expected = list(torch.split_with_sizes_copy(source, sizes, dim=1))
	66	+ actual = [torch.empty(n, 8, device='npu').t() for n in sizes]
	67	+ torch.ops.fsdp.split_with_sizes_copy(source, sizes, dim=1, out=actual)
	68	+ for a, b in zip(actual, expected):
	69	+ torch.testing.assert_close(a, b, rtol=0, atol=0)
	70	+ cases.append({'fallback': 'noncontiguous'})
	71	+ source = torch.empty(8, 101, device='npu')
	72	+ actual = [torch.empty(8, n, device='npu') for n in [1, 33, 67]]
	73	+ stream = torch.npu.Stream()
	74	+ stream.wait_stream(torch.npu.current_stream())
	75	+ with torch.npu.stream(stream):
	76	+ source.fill_(3.25)
	77	+ torch.ops.fsdp.split_with_sizes_copy(source, [1, 33, 67], dim=1, out=actual)
	78	+ torch.npu.current_stream().wait_stream(stream)
	79	+ for a in actual:
	80	+ torch.testing.assert_close(a, torch.full_like(a, 3.25), rtol=0, atol=0)
	81	+ cases.append({'stream': 'nondefault producer and consumer dependency'})
	82	+ torch.npu.config.allow_internal_format = True
	83	+ for cast_source, cast_output in [(True, False), (False, True)]:
	84	+ source = torch.randn(8, 96, device='npu', dtype=torch.bfloat16)
	85	+ if cast_source:
	86	+ source = torch_npu.npu_format_cast(source, 29)
	87	+ expected = [torch.empty(8, n, device='npu', dtype=source.dtype) for n in [32, 64]]
	88	+ actual = [torch.empty_like(t) for t in expected]
	89	+ if cast_output:
	90	+ expected = [torch_npu.npu_format_cast(t, 29) for t in expected]
	91	+ actual = [torch_npu.npu_format_cast(t, 29) for t in actual]
	92	+ formats = [torch_npu.get_npu_format(t) for t in [source, *actual]]
	93	+ if 29 not in formats:
	94	+ raise AssertionError(f'Blocked format was not constructed: {formats}')
	95	+ native_error = None
	96	+ try:
	97	+ torch.split_with_sizes_copy(source, [32, 64], dim=1, out=expected)
	98	+ except RuntimeError as error:
	99	+ native_error = str(error).splitlines()[0]
	100	+ if native_error is None:
	101	+ torch.ops.fsdp.split_with_sizes_copy(source, [32, 64], dim=1, out=actual)
	102	+ for a, b in zip(actual, expected):
	103	+ torch.testing.assert_close(a, b, rtol=0, atol=0)
	104	+ else:
	105	+ try:
	106	+ torch.ops.fsdp.split_with_sizes_copy(source, [32, 64], dim=1, out=actual)
	107	+ except RuntimeError as error:
	108	+ if str(error).splitlines()[0] != native_error:
	109	+ raise AssertionError('Fallback changed the native exception') from error
	110	+ else:
	111	+ raise AssertionError('Expected the same native exception')
	112	+ cases.append({'fallback': 'blocked NPU format', 'formats': formats,
	113	+ 'cast_source': cast_source, 'cast_output': cast_output,
	114	+ 'native_error': native_error})
	115	+ source = torch.randn(2, 15, device='npu')
	116	+ for negate_source, negate_output in [(True, False), (False, True)]:
	117	+ original = source._neg_view() if negate_source else source
	118	+ expected = [torch.empty(2, n, device='npu') for n in [3, 5, 7]]
	119	+ actual = [torch.empty_like(t) for t in expected]
	120	+ if negate_output:
	121	+ expected = [t._neg_view() for t in expected]
	122	+ actual = [t._neg_view() for t in actual]
	123	+ native_error = native_negative_errors[negate_source, negate_output]
	124	+ if native_error is None:
	125	+ torch.split_with_sizes_copy(original, [3, 5, 7], dim=1, out=expected)
	126	+ torch.ops.fsdp.split_with_sizes_copy(original, [3, 5, 7], dim=1, out=actual)
	127	+ for a, b in zip(actual, expected):
	128	+ torch.testing.assert_close(a, b, rtol=0, atol=0)
	129	+ else:
	130	+ try:
	131	+ torch.ops.fsdp.split_with_sizes_copy(original, [3, 5, 7], dim=1, out=actual)
	132	+ except RuntimeError as error:
	133	+ if str(error).splitlines()[0] != native_error:
	134	+ raise AssertionError('Negative-view fallback changed the native exception') from error
	135	+ else:
	136	+ raise AssertionError('Expected the same native negative-view exception')
	137	+ cases.append({'fallback': 'logical negative view', 'negate_source': negate_source,
	138	+ 'negate_output': negate_output, 'native_error': native_error})
	139	+ with torch.inference_mode():
	140	+ source = torch.randn(2, 15, device='npu')
	141	+ actual = [torch.empty(2, n, device='npu') for n in [3, 5, 7]]
	142	+ torch.ops.fsdp.split_with_sizes_copy(source, [3, 5, 7], dim=1, out=actual)
	143	+ for a, b in zip(actual, source.split([3, 5, 7], dim=1)):
	144	+ torch.testing.assert_close(a, b, rtol=0, atol=0)
	145	+ cases.append({'inference_mode': 'outputs with no version counter'})
	146	+ result = {'passed': len(cases), 'cases': cases, 'native_version_delta': native_delta}
	147	+ args.output.write_text(json.dumps(result, indent=2) + '\n')
	148	+ print(json.dumps({'passed': len(cases), 'native_version_delta': native_delta}), flush=True)
	149	+
	150	+
	151	+if __name__ == '__main__':
	152	+ main()

Atests/st/fsdp/mistral3/compare_training.py+114-0

@@ -0,0 +1,114 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Compare a synthetic full Mistral3 training step with a CP=1 reference."""
	3	+
	4	+import argparse
	5	+import json
	6	+import os
	7	+from pathlib import Path
	8	+
	9	+import torch
	10	+import torch.distributed as dist
	11	+import torch_npu # noqa: F401
	12	+
	13	+from train_smoke import build_model, get_parallel_sizes
	14	+from mindspeed_mm.fsdp.distributed.parallel_state import init_parallel_state
	15	+from mindspeed_mm.fsdp.features.memory.chunkloss.chunkloss_lm_head import apply_chunkloss_module
	16	+from mindspeed_mm.fsdp.loss.loss_func import build_loss_func
	17	+
	18	+
	19	+def full_tensor(value):
	20	+ return value.full_tensor() if hasattr(value, "full_tensor") else value
	21	+
	22	+
	23	+def relative_l2(actual, expected):
	24	+ return ((actual - expected).float().norm() / expected.float().norm().clamp_min(1e-8)).item()
	25	+
	26	+
	27	+def main():
	28	+ parser = argparse.ArgumentParser(description=__doc__)
	29	+ parser.add_argument("--mode", choices=("reference", "ulysses", "ring", "hybrid"), required=True)
	30	+ parser.add_argument("--reference-dir", type=Path, required=True)
	31	+ parser.add_argument("--chunk-loss", action="store_true")
	32	+ parser.add_argument("--chunk-size", type=int, default=8)
	33	+ parser.add_argument("--contiguous-linear-grad", action="store_true")
	34	+ parser.add_argument("--triton-split-copy", action="store_true")
	35	+ args = parser.parse_args()
	36	+ if args.chunk_size <= 0:
	37	+ parser.error("--chunk-size must be positive")
	38	+ torch.npu.set_device(int(os.environ["LOCAL_RANK"]))
	39	+ dist.init_process_group("hccl")
	40	+ rank, world = dist.get_rank(), dist.get_world_size()
	41	+ ulysses, ring = (1, 1) if args.mode == "reference" else get_parallel_sizes(args.mode, world)
	42	+ state = init_parallel_state(data_parallel_size=world // (ulysses * ring),
	43	+ fully_shard_parallel_size=world, ulysses_parallel_size=ulysses,
	44	+ ring_attention_size=ring)
	45	+ torch.manual_seed(42)
	46	+ if args.triton_split_copy:
	47	+ from mindspeed_mm.fsdp.ops.fully_shard.npu_split_copy import enable_npu_split_copy
	48	+
	49	+ enable_npu_split_copy()
	50	+ model = build_model(state)
	51	+ if args.contiguous_linear_grad:
	52	+ from mindspeed_mm.fsdp.features.apply_features import FeaturesApplier
	53	+ from mindspeed_mm.fsdp.params.feature_args import FeatureArguments
	54	+
	55	+ plan = [f"model.language_model.layers.{{*}}.mlp.{name}"
	56	+ for name in ("gate_proj", "up_proj", "down_proj")]
	57	+ FeaturesApplier(FeatureArguments(contiguous_linear_grad_modules=plan)).apply_contiguous_linear_grad(model)
	58	+ if args.chunk_loss:
	59	+ model.enable_chunk_loss = True
	60	+ apply_chunkloss_module(model.lm_head)
	61	+ # Identical full inputs are replicated in both runs. CP alone changes how the
	62	+ # sequence is partitioned; the reference DP replicas contain identical data.
	63	+ generator = torch.Generator().manual_seed(123)
	64	+ ids = torch.randint(3, 254, (1, 64), generator=generator).to("npu")
	65	+ ids[:, 1:33] = 255
	66	+ labels = ids.clone()
	67	+ labels[:, :33] = -100
	68	+ pixels = torch.randn(2, 3, 16, 16, generator=generator).to("npu")
	69	+ sizes = torch.tensor([[16, 16], [16, 16]], device="npu")
	70	+ positions = torch.arange(64, device="npu").unsqueeze(0)
	71	+ model.loss_function = build_loss_func("default", labels=labels, chunk_size=args.chunk_size if args.chunk_loss else None)
	72	+ initial = {name: full_tensor(p.detach()).cpu() for name, p in model.named_parameters() if p.requires_grad}
	73	+ optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=0.01)
	74	+ output = model(input_ids=ids, labels=labels, pixel_values=pixels, image_sizes=sizes,
	75	+ position_ids=positions, attention_mask=torch.ones_like(ids), use_cache=False)
	76	+ assert torch.isfinite(output.loss), "Non-finite loss"
	77	+ output.loss.backward()
	78	+ gradients = {name: full_tensor(p.grad).detach().cpu() for name, p in model.named_parameters() if p.requires_grad}
	79	+ assert all(torch.isfinite(g).all() for g in gradients.values()), "Non-finite gradient"
	80	+ optimizer.step()
	81	+ updates = {name: full_tensor(p.detach()).cpu() - initial[name]
	82	+ for name, p in model.named_parameters() if p.requires_grad}
	83	+ record = {"loss": output.loss.detach().cpu(), "gradients": gradients, "updates": updates,
	84	+ "world_size": world, "chunk_loss": args.chunk_loss}
	85	+ args.reference_dir.mkdir(parents=True, exist_ok=True)
	86	+ reference = args.reference_dir / f"rank{rank}.pt"
	87	+ metrics = {"mode": args.mode, "rank": rank, "world_size": world, "chunk_loss": args.chunk_loss,
	88	+ "contiguous_linear_grad": args.contiguous_linear_grad, "triton_split_copy": args.triton_split_copy,
	89	+ "loss": record["loss"].item(), "model": "synthetic full Mistral3", "steps": 1}
	90	+ if args.mode == "reference":
	91	+ torch.save(record, reference)
	92	+ else:
	93	+ expected = torch.load(reference, map_location="cpu", weights_only=True)
	94	+ assert expected["world_size"] == world and expected["chunk_loss"] == args.chunk_loss
	95	+ loss_error = abs(record["loss"].item() - expected["loss"].item()) / abs(expected["loss"].item())
	96	+ actual_grad = torch.cat([gradients[name].float().flatten() for name in gradients])
	97	+ expected_grad = torch.cat([expected["gradients"][name].float().flatten() for name in gradients])
	98	+ actual_update = torch.cat([updates[name].float().flatten() for name in updates])
	99	+ expected_update = torch.cat([expected["updates"][name].float().flatten() for name in updates])
	100	+ metrics.update(loss_relative_error=loss_error, gradient_relative_l2=relative_l2(actual_grad, expected_grad),
	101	+ update_relative_l2=relative_l2(actual_update, expected_update))
	102	+ # BF16 attention/reductions may differ with parallel partitioning. Check
	103	+ # updates themselves, not final parameter values dominated by initialization.
	104	+ assert loss_error < 0.02, metrics
	105	+ assert metrics["gradient_relative_l2"] < 0.05, metrics
	106	+ assert metrics["update_relative_l2"] < 0.05, metrics
	107	+ metrics["result"] = "PASS"
	108	+ print(json.dumps(metrics), flush=True)
	109	+ dist.barrier()
	110	+ dist.destroy_process_group()
	111	+
	112	+
	113	+if __name__ == "__main__":
	114	+ main()

Atests/st/fsdp/mistral3/train_smoke.py+164-0

@@ -0,0 +1,164 @@
	1	+"""Run a 30-step synthetic Mistral3 training smoke test on NPUs."""
	2	+
	3	+import argparse
	4	+import json
	5	+import os
	6	+import time
	7	+
	8	+import torch
	9	+import torch.distributed as dist
	10	+import torch_npu # noqa: F401
	11	+from torch.distributed.fsdp import MixedPrecisionPolicy
	12	+from transformers import Mistral3Config
	13	+
	14	+from mindspeed_mm.fsdp.distributed.parallel_state import init_parallel_state
	15	+from mindspeed_mm.fsdp.loss.loss_func import build_loss_func
	16	+from mindspeed_mm.fsdp.models.mistral3.modeling_mistral3 import Mistral3ForConditionalGenerationFSDP
	17	+from mindspeed_mm.fsdp.ops.fully_shard.fully_shard import fully_shard
	18	+
	19	+
	20	+PARALLEL_SIZES = {
	21	+ "default": (1, 1),
	22	+ "multi": (1, 1),
	23	+ "ulysses": None,
	24	+ "ring": None,
	25	+ "hybrid": (2, 2),
	26	+}
	27	+
	28	+
	29	+def parse_args():
	30	+ parser = argparse.ArgumentParser()
	31	+ parser.add_argument("--mode", choices=PARALLEL_SIZES, required=True)
	32	+ return parser.parse_args()
	33	+
	34	+
	35	+def get_parallel_sizes(mode, world_size):
	36	+ if mode == "ulysses":
	37	+ return world_size, 1
	38	+ if mode == "ring":
	39	+ return 1, world_size
	40	+ return PARALLEL_SIZES[mode]
	41	+
	42	+
	43	+def build_config():
	44	+ config = Mistral3Config(
	45	+ vision_config={
	46	+ "hidden_size": 128,
	47	+ "intermediate_size": 256,
	48	+ "num_hidden_layers": 2,
	49	+ "num_attention_heads": 2,
	50	+ "image_size": 32,
	51	+ "patch_size": 2,
	52	+ },
	53	+ text_config={
	54	+ "hidden_size": 128,
	55	+ "intermediate_size": 256,
	56	+ "num_hidden_layers": 2,
	57	+ "num_attention_heads": 2,
	58	+ "num_key_value_heads": 2,
	59	+ "vocab_size": 256,
	60	+ "max_position_embeddings": 1024,
	61	+ },
	62	+ image_token_index=255,
	63	+ spatial_merge_size=2,
	64	+ vision_feature_layer=-1,
	65	+ )
	66	+ config._attn_implementation = "eager"
	67	+ config.use_cache = False
	68	+ return config
	69	+
	70	+
	71	+def build_model(parallel_state):
	72	+ model = Mistral3ForConditionalGenerationFSDP(build_config()).to("npu")
	73	+ for module in model.modules():
	74	+ if hasattr(module, "config"):
	75	+ module.config._attn_implementation = "flash_attention_2"
	76	+
	77	+ model.loss_function = build_loss_func("default")
	78	+ model.train()
	79	+ model.model.vision_tower.requires_grad_(False)
	80	+ model.model.multi_modal_projector.requires_grad_(False)
	81	+
	82	+ policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16)
	83	+ for layer in model.model.language_model.layers:
	84	+ fully_shard(layer, mesh=parallel_state.get_hsdp_device_mesh(), mp_policy=policy)
	85	+ fully_shard(model, mesh=parallel_state.get_hsdp_device_mesh(), mp_policy=policy)
	86	+ return model
	87	+
	88	+
	89	+def assert_finite_gradients(model, step):
	90	+ for name, parameter in model.named_parameters():
	91	+ if parameter.grad is None:
	92	+ continue
	93	+ gradient = parameter.grad.to_local() if hasattr(parameter.grad, "to_local") else parameter.grad
	94	+ assert torch.isfinite(gradient).all(), (step, name)
	95	+
	96	+
	97	+def main():
	98	+ args = parse_args()
	99	+ rank = int(os.environ["RANK"])
	100	+ world_size = int(os.environ["WORLD_SIZE"])
	101	+ torch.npu.set_device(int(os.environ["LOCAL_RANK"]))
	102	+ dist.init_process_group("hccl")
	103	+
	104	+ ulysses_size, ring_size = get_parallel_sizes(args.mode, world_size)
	105	+ parallel_state = init_parallel_state(
	106	+ data_parallel_size=world_size // (ulysses_size * ring_size),
	107	+ fully_shard_parallel_size=world_size,
	108	+ ulysses_parallel_size=ulysses_size,
	109	+ ring_attention_size=ring_size,
	110	+ )
	111	+
	112	+ torch.manual_seed(42)
	113	+ model = build_model(parallel_state)
	114	+ optimizer = torch.optim.AdamW(
	115	+ [parameter for parameter in model.parameters() if parameter.requires_grad],
	116	+ lr=1e-4,
	117	+ )
	118	+
	119	+ image_count = 1 if args.mode == "default" else 2
	120	+ input_ids = torch.randint(3, 254, (1, 256), device="npu")
	121	+ input_ids[:, 1:1 + 16 * image_count] = 255
	122	+ labels = input_ids.clone()
	123	+ labels[labels == 255] = -100
	124	+ pixel_values = torch.randn(image_count, 3, 16, 16, device="npu")
	125	+ image_sizes = torch.tensor([[16, 16]] * image_count, device="npu")
	126	+
	127	+ for step in range(1, 31):
	128	+ start_time = time.monotonic()
	129	+ optimizer.zero_grad(set_to_none=True)
	130	+ output = model(
	131	+ input_ids=input_ids,
	132	+ pixel_values=pixel_values,
	133	+ image_sizes=image_sizes,
	134	+ labels=labels,
	135	+ use_cache=False,
	136	+ )
	137	+ loss = output.loss
	138	+ assert torch.isfinite(loss).all(), (step, loss)
	139	+ loss.backward()
	140	+ assert_finite_gradients(model, step)
	141	+ optimizer.step()
	142	+ torch.npu.synchronize()
	143	+ print(json.dumps({
	144	+ "mode": args.mode,
	145	+ "rank": rank,
	146	+ "step": step,
	147	+ "loss": loss.item(),
	148	+ "seconds": time.monotonic() - start_time,
	149	+ }), flush=True)
	150	+
	151	+ print(json.dumps({
	152	+ "mode": args.mode,
	153	+ "rank": rank,
	154	+ "result": "PASS",
	155	+ "steps": 30,
	156	+ "model": "synthetic small Mistral3",
	157	+ "parameters": sum(parameter.numel() for parameter in model.parameters()),
	158	+ }), flush=True)
	159	+ dist.barrier()
	160	+ dist.destroy_process_group()
	161	+
	162	+
	163	+if __name__ == "__main__":
	164	+ main()

Atests/st/fsdp/mistral3/vision_attention_probe.py+209-0

@@ -0,0 +1,209 @@
	1	+"""Compare Pixtral eager attention with NPU fused attention."""
	2	+
	3	+import argparse
	4	+import importlib
	5	+import json
	6	+import os
	7	+
	8	+import torch
	9	+import torch.distributed as dist
	10	+import torch_npu
	11	+from transformers import Mistral3Config
	12	+
	13	+from mindspeed_mm.fsdp.distributed.parallel_state import init_parallel_state
	14	+from mindspeed_mm.fsdp.models.mistral3.modeling_mistral3 import Mistral3ForConditionalGenerationFSDP
	15	+
	16	+
	17	+PARALLEL_SIZES = {
	18	+ "default": None,
	19	+ "ulysses": None,
	20	+ "ring": None,
	21	+ "hybrid": (2, 2),
	22	+}
	23	+
	24	+
	25	+def parse_args():
	26	+ parser = argparse.ArgumentParser()
	27	+ parser.add_argument("--mode", choices=PARALLEL_SIZES, default="default")
	28	+ return parser.parse_args()
	29	+
	30	+
	31	+def get_parallel_sizes(mode, world_size):
	32	+ if mode == "ulysses":
	33	+ return world_size, 1
	34	+ if mode == "ring":
	35	+ return 1, world_size
	36	+ if mode == "hybrid":
	37	+ return PARALLEL_SIZES[mode]
	38	+ return 1, 1
	39	+
	40	+
	41	+def relative_error(actual, expected):
	42	+ actual = actual.float()
	43	+ expected = expected.float()
	44	+ difference = (actual - expected).abs()
	45	+ return {
	46	+ "max_abs": difference.max().item(),
	47	+ "relative_l2": (
	48	+ (actual - expected).norm() / expected.norm().clamp_min(1e-10)
	49	+ ).item(),
	50	+ }
	51	+
	52	+
	53	+def check_close(actual, expected, relative_l2_limit):
	54	+ result = relative_error(actual, expected)
	55	+ assert torch.isfinite(actual).all()
	56	+ assert torch.isfinite(expected).all()
	57	+ assert result["relative_l2"] <= relative_l2_limit, result
	58	+ return result
	59	+
	60	+
	61	+def build_vision_tower():
	62	+ config = Mistral3Config(
	63	+ vision_config={
	64	+ "hidden_size": 128,
	65	+ "intermediate_size": 256,
	66	+ "num_hidden_layers": 2,
	67	+ "num_attention_heads": 2,
	68	+ "image_size": 32,
	69	+ "patch_size": 2,
	70	+ },
	71	+ text_config={
	72	+ "hidden_size": 128,
	73	+ "intermediate_size": 256,
	74	+ "num_hidden_layers": 2,
	75	+ "num_attention_heads": 2,
	76	+ "num_key_value_heads": 2,
	77	+ "vocab_size": 256,
	78	+ },
	79	+ image_token_index=255,
	80	+ spatial_merge_size=2,
	81	+ )
	82	+ config._attn_implementation = "eager"
	83	+ return Mistral3ForConditionalGenerationFSDP(config).model.vision_tower.to(
	84	+ "npu", dtype=torch.bfloat16
	85	+ ).eval()
	86	+
	87	+
	88	+def install_attention_checks():
	89	+ flash_attention = importlib.import_module("mindspeed_mm.fsdp.ops.flash_attn.flash_attn")
	90	+ fusion_attention = torch_npu.npu_fusion_attention
	91	+ calls = []
	92	+
	93	+ def traced_fusion_attention(args, *kwargs):
	94	+ calls.append((args[4], kwargs.get("actual_seq_qlen"), kwargs.get("actual_seq_kvlen")))
	95	+ return fusion_attention(args, *kwargs)
	96	+
	97	+ def reject_context_parallel(args, *kwargs):
	98	+ raise AssertionError("Replicated vision tower entered context-parallel communication")
	99	+
	100	+ torch_npu.npu_fusion_attention = traced_fusion_attention
	101	+ flash_attention.all_to_all = reject_context_parallel
	102	+ flash_attention.do_ring_attention = reject_context_parallel
	103	+ return calls
	104	+
	105	+
	106	+def run_case(vision_tower, parameters, calls, sizes):
	107	+ pixel_values = torch.randn(
	108	+ len(sizes), 3, 16, 24, device="npu", dtype=torch.bfloat16, requires_grad=True
	109	+ )
	110	+ image_sizes = torch.tensor(sizes, device="npu")
	111	+
	112	+ vision_tower.config._attn_implementation = "eager"
	113	+ expected = torch.cat([
	114	+ vision_tower(
	115	+ pixel_values[index:index + 1],
	116	+ image_sizes=image_sizes[index:index + 1],
	117	+ ).last_hidden_state
	118	+ for index in range(len(sizes))
	119	+ ], dim=1)
	120	+ weights = torch.randn_like(expected)
	121	+ expected_gradients = torch.autograd.grad(
	122	+ (expected.float() * weights.float()).sum(),
	123	+ (pixel_values, *parameters),
	124	+ )
	125	+
	126	+ vision_tower.config._attn_implementation = "flash_attention_2"
	127	+ calls.clear()
	128	+ actual = vision_tower(pixel_values, image_sizes=image_sizes).last_hidden_state
	129	+ actual_gradients = torch.autograd.grad(
	130	+ (actual.float() * weights.float()).sum(),
	131	+ (pixel_values, *parameters),
	132	+ )
	133	+
	134	+ forward_result = check_close(actual, expected, 0.015)
	135	+ gradient_results = [
	136	+ check_close(actual_gradient, expected_gradient, 0.04)
	137	+ for actual_gradient, expected_gradient in zip(actual_gradients, expected_gradients)
	138	+ ]
	139	+
	140	+ if len(sizes) > 1:
	141	+ sequence_ends = []
	142	+ for height, width in sizes:
	143	+ length = (height // 2) * (width // 2)
	144	+ sequence_ends.append((sequence_ends[-1] if sequence_ends else 0) + length)
	145	+ assert all(call == ("TND", sequence_ends, sequence_ends) for call in calls), calls
	146	+
	147	+ first_image_end = sequence_ends[0]
	148	+ repeated_output = vision_tower(pixel_values, image_sizes=image_sizes).last_hidden_state
	149	+ cross_image_gradient = torch.autograd.grad(
	150	+ repeated_output[:, :first_image_end].float().square().sum(), pixel_values
	151	+ )[0][1:]
	152	+ assert cross_image_gradient.count_nonzero().item() == 0
	153	+
	154	+ changed_pixels = pixel_values.detach().clone()
	155	+ changed_pixels[1:].mul_(4).add_(2)
	156	+ with torch.no_grad():
	157	+ changed_output = vision_tower(changed_pixels, image_sizes=image_sizes).last_hidden_state
	158	+ assert torch.equal(
	159	+ repeated_output[:, :first_image_end],
	160	+ changed_output[:, :first_image_end],
	161	+ )
	162	+
	163	+ return {
	164	+ "sizes": sizes,
	165	+ "forward": forward_result,
	166	+ "max_gradient_relative_l2": max(result["relative_l2"] for result in gradient_results),
	167	+ "calls": calls.copy(),
	168	+ "isolation": "PASS",
	169	+ }
	170	+
	171	+
	172	+def main():
	173	+ args = parse_args()
	174	+ rank = int(os.environ.get("RANK", 0))
	175	+ world_size = int(os.environ.get("WORLD_SIZE", 1))
	176	+ torch.npu.set_device(int(os.environ.get("LOCAL_RANK", 0)))
	177	+ dist.init_process_group("hccl")
	178	+
	179	+ ulysses_size, ring_size = get_parallel_sizes(args.mode, world_size)
	180	+ assert world_size % (ulysses_size * ring_size) == 0
	181	+ init_parallel_state(
	182	+ data_parallel_size=world_size // (ulysses_size * ring_size),
	183	+ fully_shard_parallel_size=world_size,
	184	+ ulysses_parallel_size=ulysses_size,
	185	+ ring_attention_size=ring_size,
	186	+ )
	187	+
	188	+ torch.manual_seed(42)
	189	+ vision_tower = build_vision_tower()
	190	+ parameters = tuple(vision_tower.parameters())
	191	+ calls = install_attention_checks()
	192	+
	193	+ for sizes in ([[16, 16]], [[16, 16], [16, 16]], [[16, 16], [16, 24]]):
	194	+ result = run_case(vision_tower, parameters, calls, sizes)
	195	+ result.update({"rank": rank, "mode": args.mode})
	196	+ print(json.dumps(result), flush=True)
	197	+
	198	+ print(json.dumps({
	199	+ "rank": rank,
	200	+ "mode": args.mode,
	201	+ "result": "PASS",
	202	+ "world_size": world_size,
	203	+ }), flush=True)
	204	+ dist.barrier()
	205	+ dist.destroy_process_group()
	206	+
	207	+
	208	+if __name__ == "__main__":
	209	+ main()

Mtests/st/run_configs/finetune_qwen3_5_35B/qwen3_5_35B_config.yaml+2-0

	@@ -79,6 +79,8 @@ data:
79	dataloader_mode: sampler	79	dataloader_mode: sampler
80	drop_last: true	80	drop_last: true
81	sampler_type: BaseRandomBatchSampler	81	sampler_type: BaseRandomBatchSampler
		82	+ # Keep workers prefetching when 50 samples yield only one full global batch.
		83	+ infinite_sampler: true
82	num_workers: 8	84	num_workers: 8
83	collate_param:	85	collate_param:
84	model_name: qwen3vl	86	model_name: qwen3vl

Atests/ut_fsdp/checkpoint/test_magistral_megatron_conversion.py+79-0

@@ -0,0 +1,79 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
	2	+"""Exercise real BF16 safetensors -> DCP -> Megatron wrapper state loading."""
	3	+
	4	+import importlib.util
	5	+from pathlib import Path
	6	+
	7	+import pytest
	8	+from safetensors.torch import save_file
	9	+import torch
	10	+import torch.distributed.checkpoint as dcp
	11	+from transformers import Mistral3Config, Mistral3ForConditionalGeneration, MistralConfig, PixtralVisionConfig
	12	+
	13	+
	14	+@pytest.fixture
	15	+def converter():
	16	+ path = Path(__file__).resolve().parents[3] / "examples/magistral-2509/baseline/convert_megatron_checkpoint.py"
	17	+ spec = importlib.util.spec_from_file_location("magistral_megatron_converter", path)
	18	+ module = importlib.util.module_from_spec(spec)
	19	+ spec.loader.exec_module(module)
	20	+ return module
	21	+
	22	+
	23	+@pytest.fixture
	24	+def checkpoint(tmp_path):
	25	+ config = Mistral3Config(
	26	+ text_config=MistralConfig(vocab_size=128, hidden_size=32, intermediate_size=64,
	27	+ num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=2,
	28	+ head_dim=8, max_position_embeddings=64),
	29	+ vision_config=PixtralVisionConfig(hidden_size=32, intermediate_size=64, num_hidden_layers=1,
	30	+ num_attention_heads=4, image_size=32, patch_size=16),
	31	+ image_token_index=10,
	32	+ )
	33	+ model = Mistral3ForConditionalGeneration(config)
	34	+ original = {key: value.to(torch.bfloat16).contiguous() for key, value in model.state_dict().items()}
	35	+ official = {}
	36	+ for key, tensor in original.items():
	37	+ if key.startswith("model.language_model."):
	38	+ name = key.replace("model.language_model.", "language_model.model.", 1)
	39	+ elif key.startswith("model."):
	40	+ name = key.removeprefix("model.")
	41	+ else:
	42	+ name = "language_model." + key
	43	+ official[name] = tensor
	44	+ hf = tmp_path / "hf"
	45	+ hf.mkdir()
	46	+ config.save_pretrained(hf)
	47	+ save_file(official, hf / "model.safetensors")
	48	+ return hf, original
	49	+
	50	+
	51	+def test_official_bf16_weights_round_trip_into_wrapper(converter, checkpoint, tmp_path):
	52	+ hf, original = checkpoint
	53	+ target = tmp_path / "converted"
	54	+ converter.convert_checkpoint(hf, target)
	55	+ assert converter.validate_megatron_checkpoint(hf, target) == len(original)
	56	+ state = {"model": {"model." + key: torch.zeros_like(value) for key, value in original.items()}}
	57	+ dcp.load(state, checkpoint_id=target / "release")
	58	+ for key, value in original.items():
	59	+ assert state["model"]["model." + key].dtype == torch.bfloat16
	60	+ assert torch.equal(state["model"]["model." + key], value), key
	61	+
	62	+
	63	+def test_generic_conversion_rejected_before_training(converter, checkpoint, tmp_path):
	64	+ hf, _ = checkpoint
	65	+ target = tmp_path / "wrong_namespace"
	66	+ from mindspeed_mm.fsdp.tasks.checkpoint.utils import hf_to_dcp_sharded
	67	+
	68	+ hf_to_dcp_sharded(str(hf), str(target))
	69	+ with pytest.raises(ValueError, match="Megatron initial checkpoint model keys do not match"):
	70	+ converter.validate_megatron_checkpoint(hf, target)
	71	+
	72	+
	73	+def test_wrong_tensor_shapes_rejected(converter, checkpoint, tmp_path):
	74	+ hf, original = checkpoint
	75	+ target = tmp_path / "wrong_shapes" / "release"
	76	+ state = {"model": {"model." + key: value[:1] for key, value in original.items()}}
	77	+ dcp.save(state, checkpoint_id=target)
	78	+ with pytest.raises(ValueError, match="wrong_shapes="):
	79	+ converter.validate_megatron_checkpoint(hf, target.parent)

Atests/ut_fsdp/data/test_infinite_sampler.py+71-0

@@ -0,0 +1,71 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
	2	+
	3	+"""Preserve sample order and checkpoint position across infinite sampling."""
	4	+from itertools import islice
	5	+
	6	+import pytest
	7	+import torch
	8	+from torchdata.stateful_dataloader import StatefulDataLoader
	9	+
	10	+from mindspeed_mm.fsdp.data.dataloader.sampler import BaseRandomBatchSampler, SeedRandomBatchSampler
	11	+
	12	+
	13	+def finite_batches(sampler):
	14	+ while True:
	15	+ yield from sampler
	16	+
	17	+
	18	+@pytest.mark.parametrize("sampler_type", [BaseRandomBatchSampler, SeedRandomBatchSampler])
	19	+@pytest.mark.parametrize("total", [50, 96])
	20	+@pytest.mark.parametrize("shuffle", [False, True])
	21	+@pytest.mark.parametrize("rank", [0, 7])
	22	+def test_infinite_order_matches_finite_epochs(sampler_type, total, shuffle, rank):
	23	+ options = dict(batch_size=4, num_replicas=8, rank=rank, shuffle=shuffle, seed=17)
	24	+ finite = sampler_type(list(range(total)), **options)
	25	+ infinite = sampler_type(list(range(total)), infinite=True, **options)
	26	+ expected = list(islice(finite_batches(finite), 20))
	27	+ actual = list(islice(infinite, 20))
	28	+ torch.testing.assert_close(torch.tensor(actual), torch.tensor(expected), rtol=0, atol=0)
	29	+
	30	+
	31	+@pytest.mark.parametrize("sampler_type", [BaseRandomBatchSampler, SeedRandomBatchSampler])
	32	+@pytest.mark.parametrize("rank", [0, 7])
	33	+def test_infinite_sampler_resumes_with_dropped_tail(sampler_type, rank):
	34	+ options = dict(batch_size=4, num_replicas=8, rank=rank, shuffle=True, seed=17, infinite=True)
	35	+ sampler = sampler_type(list(range(50)), **options)
	36	+ iterator = iter(sampler)
	37	+ list(islice(iterator, 5))
	38	+ state = sampler.state_dict()
	39	+ expected = list(islice(iterator, 8))
	40	+ restored = sampler_type(list(range(50)), **options)
	41	+ restored.load_state_dict(state)
	42	+ actual = list(islice(restored, 8))
	43	+ torch.testing.assert_close(torch.tensor(actual), torch.tensor(expected), rtol=0, atol=0)
	44	+
	45	+
	46	+@pytest.mark.parametrize("sampler_type", [BaseRandomBatchSampler, SeedRandomBatchSampler])
	47	+@pytest.mark.parametrize("workers", [0, 2])
	48	+@pytest.mark.parametrize("prefix", [11, 13])
	49	+def test_stateful_loader_resumes_across_epoch_boundary(sampler_type, workers, prefix):
	50	+ def build():
	51	+ dataset = list(range(96))
	52	+ sampler = sampler_type(dataset, batch_size=4, num_replicas=2, rank=1,
	53	+ shuffle=True, seed=17, infinite=True)
	54	+ worker_options = {"multiprocessing_context": "spawn"} if workers else {}
	55	+ return StatefulDataLoader(dataset, batch_sampler=sampler, num_workers=workers,
	56	+ persistent_workers=workers > 0, **worker_options)
	57	+
	58	+ loader, restored = build(), build()
	59	+ try:
	60	+ iterator = iter(loader)
	61	+ list(islice(iterator, prefix))
	62	+ state = loader.state_dict()
	63	+ expected = torch.stack(list(islice(iterator, 8)))
	64	+ restored.load_state_dict(state)
	65	+ actual = torch.stack(list(islice(iter(restored), 8)))
	66	+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
	67	+ finally:
	68	+ for item in (loader, restored):
	69	+ iterator = getattr(item, "_iterator", None)
	70	+ if iterator is not None and hasattr(iterator, "_shutdown_workers"):
	71	+ iterator._shutdown_workers()

Atests/ut_fsdp/data/test_mistral3_pixtral.py+415-0

Atests/ut_fsdp/data/test_mistral3_vision_attention.py+159-0

Atests/ut_fsdp/loss/test_chunkloss_buffer_lifetime.py+40-0

Atests/ut_fsdp/loss/test_chunkloss_single_chunk.py+42-0

Atests/ut_fsdp/models/mistral3/test_modeling_mistral3.py+304-0

Atests/ut_fsdp/ops/test_contiguous_linear.py+86-0

Atests/ut_fsdp/utils/test_gc_utils.py+86-0