feat(torch): migrate Magistral-LoRA to pure FSDP2-MindSpeed-MM-AtomGit

Mcheckpoint/common/hf_to_dcp.py+17-3

	@@ -1,5 +1,6 @@
1	-import argparse	1	+import argparse
2	-from typing import Optional	2	+import json
		3	+from typing import Optional
3	from pathlib import Path	4	from pathlib import Path
4	from concurrent.futures import ThreadPoolExecutor, as_completed	5	from concurrent.futures import ThreadPoolExecutor, as_completed
5	from tqdm import tqdm	6	from tqdm import tqdm
	@@ -76,7 +77,20 @@ def hf_to_dcp_sharded(
76	save_root_dir.joinpath(LATEST_TXT).write_text("release")	77	save_root_dir.joinpath(LATEST_TXT).write_text("release")
77		78
78	storage_writer = FileSystemWriter(save_path)	79	storage_writer = FileSystemWriter(save_path)
79	- files = sorted(list(Path(hf_dir).glob("*.safetensors")))	80	+ hf_root = Path(hf_dir)
		81	+ index_path = hf_root / "model.safetensors.index.json"
		82	+ if index_path.is_file():
		83	+ index = json.loads(index_path.read_text())
		84	+ shard_names = sorted(set(index.get("weight_map", {}).values()))
		85	+ files = [hf_root / name for name in shard_names]
		86	+ missing = [path for path in files if not path.is_file()]
		87	+ if missing:
		88	+ missing_names = ", ".join(path.name for path in missing)
		89	+ raise FileNotFoundError(f"Missing safetensors shards listed by {index_path}: {missing_names}")
		90	+ else:
		91	+ files = sorted(
		92	+ path for path in hf_root.glob("*.safetensors") if path.name != "consolidated.safetensors"
		93	+ )
80		94
81	# Prepare task arguments	95	# Prepare task arguments
82	tasks = []	96	tasks = []

Aexamples/magistral-2509/README.md+395-0

@@ -0,0 +1,395 @@
	1	+# Magistral-Small-2509 使用指南
	2	+
	3	+<p align="left">
	4	+</p>
	5	+
	6	+## 目录
	7	+
	8	+- [版本说明](#版本说明)
	9	+ - [参考实现](#参考实现)
	10	+ - [变更记录](#变更记录)
	11	+- [环境安装](#环境安装)
	12	+ - [仓库拉取](#1-环境准备)
	13	+ - [环境搭建](#2-环境搭建)
	14	+- [权重下载及转换](#权重下载及转换)
	15	+ - [权重下载](#权重下载)
	16	+ - [权重转换](#权重转换)
	17	+- [数据集准备及处理](#数据集准备及处理)
	18	+- [微调](#微调)
	19	+ - [准备工作](#1-准备工作)
	20	+ - [配置参数](#2-配置参数)
	21	+ - [启动微调](#3-启动微调)
	22	+ - [启动推理](#4-启动推理)
	23	+- [lora微调](#lora微调)
	24	+- [环境变量声明](#环境变量声明)
	25	+
	26	+## 版本说明
	27	+
	28	+### 参考实现
	29	+
	30	+```shell
	31	+url=https://github.com/hiyouga/LLaMAFactory
	32	+commit_id=68119e5
	33	+```
	34	+
	35	+### 变更记录
	36	+
	37	+2026.1.19: 首次支持Magistral-Small-2509模型
	38	+
	39	+---
	40	+<a id="jump1"></a>
	41	+
	42	+## 环境安装
	43	+
	44	+<a id="jump1.1"></a>
	45	+
	46	+### 1. 环境准备
	47	+
	48	+【模型开发时推荐使用配套的环境版本】
	49	+
	50	+请参考[安装指南](../../docs/zh/guides/installation/install_guide.md)，完成昇腾软件安装。
	51	+
	52	+<a id="jump1.2"></a>
	53	+
	54	+### 2. 环境搭建
	55	+
	56	+```bash
	57	+git clone https://gitcode.com/Ascend/MindSpeed-MM.git
	58	+git clone https://github.com/NVIDIA/Megatron-LM.git
	59	+cd Megatron-LM
	60	+git checkout core_v0.12.1
	61	+cp -r megatron ../MindSpeed-MM/
	62	+cd ..
	63	+cd MindSpeed-MM
	64	+mkdir logs data ckpt
	65	+
	66	+# 安装加速库。纯 FSDP2 入口需要 MindSpeed 中提供 mindspeed.fsdp 的版本。
	67	+git clone https://gitcode.com/Ascend/MindSpeed.git
	68	+cd MindSpeed
	69	+# 不要切到仅支持旧 Megatron-FSDP2 路线的 core_r0.12.1 commit。
	70	+# 当前候选已用下面这个可复现 revision 完成模型入口 smoke。
	71	+git checkout 4fb7dc00614aa195e0f3a0fa4336f72110fa1555
	72	+
	73	+# 安装mindspeed及依赖
	74	+pip install -e .
	75	+cd ..
	76	+# 安装mindspeed mm及依赖
	77	+pip install -e .
	78	+
	79	+```
	80	+
	81	+<a id="jump2"></a>
	82	+
	83	+## 权重下载及转换
	84	+
	85	+<a id="jump2.1"></a>
	86	+
	87	+### 权重下载
	88	+
	89	+从Hugging Face等网站下载开源模型权重
	90	+
	91	+> [!NOTE]
	92	+>
	93	+> 如无法顺利访问HuggingFace社区下载资源，推荐前往ModelScope下载，需关注待下载文件的正确性与安全性。
	94	+
	95	+- [Magistral-Small-2509](https://huggingface.co/mistralai/Magistral-Small-2509)
	96	+
	97	+将模型权重保存在`ckpt/hf_path/`目录下，例如`ckpt/hf_path/Magistral-Small-2509`。
	98	+
	99	+<a id="jump2.2"></a>
	100	+
	101	+### 权重转换
	102	+
	103	+```shell
	104	+mm-convert Mistral3Converter hf_to_dcp --hf_dir "ckpt/hf_path/Magistral-Small-2509" --dcp_dir "ckpt/convert_path/Magistral-Small-2509"
	105	+```
	106	+
	107	+---
	108	+
	109	+<a id="jump3"></a>
	110	+
	111	+## 数据集准备及处理
	112	+
	113	+- 使用真实数据集训练：参考[针对VL模型的数据构造 · 使用真实数据集](../../docs/zh/features/data/building_data_for_VLModel.md#real-data)（下载COCO2017 → 下载LLaVA-Instruct-150K标注 → 运行转换脚本生成`mllm_format_llava_instruct_data.json`）。
	114	+- 使用虚构数据做功能/性能测试：参考[针对VL模型的数据构造 · 使用虚构数据](../../docs/zh/features/data/building_data_for_VLModel.md#mock-data)。
	115	+
	116	+## 微调
	117	+
	118	+<a id="jump4.1"></a>
	119	+
	120	+### 1. 准备工作
	121	+
	122	+配置脚本前需要完成前置准备工作，包括：环境安装、权重下载、数据集准备及处理，详情可查看对应章节。
	123	+
	124	+<a id="jump4.2"></a>
	125	+
	126	+### 2. 配置参数
	127	+
	128	+【数据目录配置】
	129	+
	130	+根据实际情况修改`data.json`中的数据集路径，包括`model_name_or_path`、`dataset_dir`、`dataset`等字段。
	131	+model_name_or_path为原始模型路径
	132	+
	133	+【模型路径配置】
	134	+
	135	+根据实际情况修改`model.json`中的权重路径，包括`init_from_hf_path`
	136	+init_from_hf_path为原始模型路径
	137	+
	138	+【模型保存加载及日志信息配置】
	139	+
	140	+根据实际情况配置`examples/magistral-2509/finetune_magistral_2509.sh`的参数，包括加载、保存路径以及保存间隔`--save-interval`（注意：分布式优化器保存文件较大耗时较长，请谨慎设置保存间隔）, 以Magistral-Small-2509为例：
	141	+
	142	+```shell
	143	+...
	144	+
	145	+# 加载路径：权重转换后路径
	146	+LOAD_PATH="ckpt/convert_path/Magistral-Small-2509"
	147	+# 保存路径
	148	+SAVE_PATH="Magistral-Small-2509_finetune_result"
	149	+...
	150	+GPT_ARGS="
	151	+ ...
	152	+ --no-load-optim \ # 不加载优化器状态，若需加载请移除
	153	+ --no-load-rng \ # 不加载随机数状态，若需加载请移除
	154	+ --no-save-optim \ # 不保存优化器状态，若需保存请移除
	155	+ --no-save-rng \ # 不保存随机数状态，若需保存请移除
	156	+ ...
	157	+"
	158	+...
	159	+OUTPUT_ARGS="
	160	+ --log-interval 1 \ # 日志间隔
	161	+ --save-interval 5000 \ # 保存间隔
	162	+ ...
	163	+ --log-tps \ # 增加此参数可使能在训练中打印每步语言模块的平均序列长度，并在训练结束后计算每秒吞吐tokens量。
	164	+"
	165	+```
	166	+
	167	+```shell
	168	+$save_dir
	169	+ ├── latest_checkpointed_iteration.txt
	170	+ ├── ...
	171	+```
	172	+
	173	+【单机运行配置】
	174	+
	175	+配置`examples/magistral-2509/finetune_magistral_2509.sh`参数如下
	176	+
	177	+```shell
	178	+ # 根据实际情况修改 ascend-toolkit 路径
	179	+ source /usr/local/Ascend/cann/set_env.sh
	180	+ NPUS_PER_NODE=8 # A2单机可跑
	181	+ MASTER_ADDR=localhost
	182	+ MASTER_PORT=6000
	183	+ NNODES=1
	184	+ NODE_RANK=0
	185	+ WORLD_SIZE=$(($NPUS_PER_NODE * $NNODES))
	186	+```
	187	+
	188	+【多机运行配置】
	189	+
	190	+配置`examples/magistral-2509/finetune_magistral_2509.sh`参数如下
	191	+
	192	+```shell
	193	+ # 根据实际情况修改 ascend-toolkit 路径
	194	+ source /usr/local/Ascend/cann/set_env.sh
	195	+ # 根据分布式集群实际情况配置分布式参数
	196	+ NPUS_PER_NODE=8 # 每个节点的卡数，以实际情况填写
	197	+ MASTER_ADDR="your master node IP" # 都需要修改为主节点的IP地址（不能为localhost）
	198	+ MASTER_PORT=6000
	199	+ NNODES=2 # 集群里的节点数，以实际情况填写
	200	+ NODE_RANK="current node id" # 当前节点的RANK，多个节点不能重复，主节点为0, 其他节点可以是1,2..
	201	+ WORLD_SIZE=$(($NPUS_PER_NODE * $NNODES))
	202	+```
	203	+
	204	+【Data Packing】（可选）
	205	+如需开启Data Packing，在`data.json`中将`packing`设置为`true`，并根据需要调整`finetune_magistral_2509.sh`中的SEQ_LEN
	206	+
	207	+<a id="jump4.3"></a>
	208	+
	209	+### 3. 启动微调
	210	+
	211	+以Magistral-Small-2509为例，启动微调训练任务。
	212	+
	213	+```shell
	214	+bash examples/magistral-2509/finetune_magistral_2509.sh
	215	+```
	216	+
	217	+<a id="jump4.4"></a>
	218	+
	219	+### 4. 启动推理
	220	+
	221	+训练完成之后，将保存在`SAVE_PATH`目录下的权重转换成huggingface格式
	222	+
	223	+```shell
	224	+mm-convert Mistral3Converter dcp_to_hf \
	225	+ --load_dir "Magistral-Small-2509_finetune_result/iter_000xx" \
	226	+ --save_dir "ckpt/dcp_to_hf/Magistral-Small-2509" \
	227	+ --model_assets_dir "ckpt/hf_path/Magistral-Small-2509"
	228	+```
	229	+
	230	+其中，`iter_000xx`表示保存的第xx步的权重，`--save_dir`表示转换后的权重保存路径，`--model_assets_dir`原始huggingface权重的路径。
	231	+
	232	+完成权重转换之后，即可使用transformers库进行推理。
	233	+
	234	+<a id="jump5"></a>
	235	+
	236	+## lora微调
	237	+
	238	+与模型微调一样，需进行权重转换与脚本配置
	239	+
	240	+### lora权重转换
	241	+
	242	+```shell
	243	+mm-convert Mistral3Converter hf_to_dcp --hf_dir "ckpt/hf_path/Magistral-Small-2509" --dcp_dir "ckpt/convert_path/Magistral-Small-2509-lora-base" --is_lora_base true
	244	+```
	245	+
	246	+### 配置参数
	247	+
	248	+lora微调使用finetune_magistral_2509_lora.sh脚本，数据、模型路径等配置与微调一致
	249	+
	250	+首次lora微调，需将finetune_magistral_2509_lora.sh中的LOAD_PATH设置为转换后的权重路径，如ckpt/convert_path/Magistral-Small-2509-lora-base
	251	+若要进行lora续训，则在finetune_magistral_2509_lora.sh中的GPT_ARGS中添加--load-base-model 参数
	252	+--load-base-model配置为lora权重转换部分得到的权重，LOAD_PATH配置为上一次lora保存的权重
	253	+
	254	+### 启动lora微调
	255	+
	256	+```shell
	257	+bash examples/magistral-2509/finetune_magistral_2509_lora.sh
	258	+```
	259	+
	260	+### lora微调后权重合并
	261	+
	262	+```shell
	263	+mm-convert Mistral3Converter merge_mm_lora_dcp_weight_to_base_hf \
	264	+ --base_hf_dir ckpt/hf_path/Magistral-Small-2509 \ # 原始权重路径
	265	+ --lora_dcp_dir /path/to/Magistral-Small-2509_lora_finetune_result \ # lora权重保存路径
	266	+ --lora_target_modules q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj \ # lora计算模块
	267	+ --save_merged_hf_dir /path/to/save # 合并后保存路径
	268	+```
	269	+
	270	+<a id="jump6"></a>
	271	+
	272	+## 环境变量声明
	273	+
	274	+以下列出常见的环境变量，完整说明请见[环境变量参考](../../docs/zh/reference/environment_variables.md)，训练命令参数请见[训练参数参考](../../docs/zh/reference/training_arguments.md)。
	275	+
	276	+\| 环境变量 \| 描述 \| 取值说明 \|
	277	+\|-------------------------------\|--------------------------------------------------------------------\|----------------------------------------------------------------------------------------------\|
	278	+\| `ASCEND_SLOG_PRINT_TO_STDOUT` \| 是否开启日志打印 \| `0`: 关闭日志打屏<br>`1`: 开启日志打屏 \|
	279	+\| `ASCEND_GLOBAL_LOG_LEVEL` \| 设置应用类日志的日志级别及各模块日志级别，仅支持调试日志 \| `0`: 对应DEBUG级别<br>`1`: 对应INFO级别<br>`2`: 对应WARNING级别<br>`3`: 对应ERROR级别<br>`4`: 对应NULL级别，不输出日志 \|
	280	+\| `TASK_QUEUE_ENABLE` \| 用于控制开启task_queue算子下发队列优化的等级 \| `0`: 关闭<br>`1`: 开启Level 1优化<br>`2`: 开启Level 2优化 \|
	281	+\| `COMBINED_ENABLE` \| 设置combined标志。设置为0表示关闭此功能；设置为1表示开启，用于优化非连续两个算子组合类场景 \| `0`: 关闭<br>`1`: 开启 \|
	282	+\| `CPU_AFFINITY_CONF` \| 控制CPU端算子任务的处理器亲和性，即设定任务绑核 \| 设置`0`或未设置: 表示不启用绑核功能<br>`1`: 表示开启粗粒度绑核<br>`2`: 表示开启细粒度绑核 \|
	283	+\| `HCCL_CONNECT_TIMEOUT` \| 用于限制不同设备之间socket建链过程的超时等待时间 \| 需要配置为整数，取值范围`[120,7200]`，默认值为`120`，单位`s` \|
	284	+\| `PYTORCH_NPU_ALLOC_CONF` \| 控制缓存分配器行为 \| `expandable_segments:<value>`: 使能内存池扩展段功能，即虚拟内存特征 \|
	285	+\| `HCCL_EXEC_TIMEOUT` \| 控制设备间执行时同步等待的时间，在该配置时间内各设备进程等待其他设备执行通信同步 \| 需要配置为整数，取值范围`[68,17340]`，默认值为`1800`，单位`s` \|
	286	+\| `ACLNN_CACHE_LIMIT` \| 配置单算子执行API在Host侧缓存的算子信息条目个数 \| 需要配置为整数，取值范围`[1, 10,000,000]`，默认值为`10000` \|
	287	+\| `TOKENIZERS_PARALLELISM` \| 用于控制Hugging Face的transformers库中的分词器（tokenizer）在多线程环境下的行为 \| `False`: 禁用并行分词<br>`True`: 开启并行分词 \|
	288	+\| `MULTI_STREAM_MEMORY_REUSE` \| 配置多流内存复用是否开启 \| `0`: 关闭多流内存复用<br>`1`: 开启多流内存复用 \|
	289	+\| `NPU_ASD_ENABLE` \| 控制是否开启TorchNPU的特征值检测功能 \| 设置`0`或未设置: 关闭特征值检测<br>`1`: 表示开启特征值检测，只打印异常日志，不告警<br>`2`:开启特征值检测，并告警<br>`3`:开启特征值检测，并告警，同时会在device侧info级别日志中记录过程数据 \|
	290	+\| `ASCEND_LAUNCH_BLOCKING` \| 控制算子执行时是否启动同步模式 \| `0`: 采用异步方式执行<br>`1`: 强制算子采用同步模式运行 \|
	291	+\| `NPUS_PER_NODE` \| 配置一个计算节点上使用的NPU数量 \| 整数值（如 `1`, `8` 等） \|
	292	+
	293	+### LoRA 后端对照复现
	294	+
	295	+普通训练分别使用以下入口，模型和优化器参数仍由原有脚本与 YAML 提供：
	296	+
	297	+- Megatron-FSDP2：`finetune_magistral_2509_lora.sh` 和 `fsdp2_lora_config.yaml`。
	298	+- 纯 FSDP2：`finetune_magistral_2509_pure_fsdp2_lora.sh` 和
	299	+ `magistral3_fsdp2_lora_config.yaml`。
	300	+
	301	+普通数据配置不限制样本数量，默认启用 shuffle。根据实际数据集填写模型、
	302	+标注与媒体路径后，即可用各自的普通训练脚本启动。
	303	+
	304	+#### 固定输入的对照入口
	305	+
	306	+需要比较两条后端时，使用 `reproduce_lora.py`。它只检查启动条件、锁定对照参数，
	307	+再调用上述训练入口，不复制模型实现或生成另一份 YAML。
	308	+
	309	+先激活安装好依赖的 Python 环境，在仓库根目录执行：
	310	+
	311	+```bash
	312	+# 默认四卡；仅检查语法，不加载模型或启动训练。
	313	+python examples/magistral-2509/reproduce_lora.py old
	314	+python examples/magistral-2509/reproduce_lora.py pure
	315	+
	316	+# 先分别完成短步检查，前一路退出后再运行后一路。
	317	+python examples/magistral-2509/reproduce_lora.py old --steps 5 --run
	318	+python examples/magistral-2509/reproduce_lora.py pure --steps 5 --run
	319	+
	320	+# 资源允许且短测通过后，分别运行完整对照。
	321	+python examples/magistral-2509/reproduce_lora.py old --steps 2000 --run
	322	+python examples/magistral-2509/reproduce_lora.py pure --steps 2000 --run
	323	+```
	324	+
	325	+默认检查不要求模型文件；真正运行时，会在 torchrun 前检查模型、数据、图片及
	326	+旧路径的 DCP metadata。只读到文件存在不证明两个格式的权重张量数值等价，
	327	+加载期间还会执行严格模型键检查。
	328	+
	329	+\| 项目 \| 对照设置 \|
	330	+\| --- \| --- \|
	331	+\| 卡数与 batch \| 默认 4 卡，MBS=1、GBS=4、累积=1；TP=PP=CP=1 \|
	332	+\| 数据 \| 同一份 4 条固定样本；shuffle=false；worker=1；seq/cutoff=512 \|
	333	+\| 随机源与精度 \| seed=42；确定性 LoRA 初始化；BF16 参数、FP32 归约 \|
	334	+\| 优化器 \| lr=1e-5；cosine；weight decay=0；clip grad=0；pure 使用 fused AdamW \|
	335	+\| LoRA \| rank=8、alpha=16、dropout=0 \|
	336	+\| FSDP2 \| recompute=true；reshard_after_forward=true；prefetch=2/2 \|
	337	+\| 状态保存 \| 默认不保存 checkpoint，不开启 profiler \|
	338	+
	339	+2000 步使用 1% warmup，即 20 步。少于 100 步时关闭 warmup，避免小数步数取整差异；
	340	+100 步及以上要求步数为 100 的倍数。短测学习率计划因此不是 2000 步运行的前缀。
	341	+
	342	+固定小样本用于可重复的后端比较，不代表完整数据集覆盖或模型泛化能力。
	343	+它只在显式复现入口生效：old 使用 `task3_repro_data.json`；
	344	+pure 在解析配置后选择同一数据文件并设置样本限制与顺序。
	345	+普通 `data.json` 和 YAML 的训练默认值不受影响。
	346	+
	347	+#### 资产与配置约束
	348	+
	349	+复用已有资产，不必再次复制大模型：
	350	+
	351	+- HF：`./.magistral_run/hf`，包含配置、Tekken tokenizer、权重索引及全部分片。
	352	+- 样本：`./.magistral_run/data/task3_repro_mllm_format.json`，含四条固定样本及其图片。
	353	+- old DCP：`./ckpt/convert_path/Magistral-Small-2509-lora-base`，
	354	+ 或显式设置 `LOAD_PATH`；支持直接 metadata、release 与 iteration 目录布局。
	355	+
	356	+环境中与对照值冲突的变量会被拒绝。`CONFIG_PATH` 和自定义训练入口不能用于
	357	+这个固定模式；`CANN_ENV_SCRIPT`、`MASTER_PORT`、`RUN_ROOT` 可按安装环境设置。
	358	+两卡仅用于诊断，可加 `--cards 2`，累积步数会变为 2 以保持 GBS=4；
	359	+两卡结果不能替代四卡性能结论。
	360	+
	361	+`task3_repro.sh` 负责资产检查、日志、manifest 与退出码；
	362	+`task3_repro_train.py` 负责对照运行的参数衔接。
	363	+确定性 LoRA 初始化、old optimizer 主参数同步和 DCP 键检查由模型目录中的
	364	+`repro_hooks.py` 提供，仅由显式复现入口启用。
	365	+普通训练不导入这些钩子，也不修改公共 `train_engine.py` 或 `training.py`。
	366	+
	367	+范数优化只在 clip grad=0 且没有要求非有限值检查时跳过范数统计；
	368	+正阈值裁剪及显式异常检查仍走原函数。当前复现入口保留该开关。
	369	+完整阶段计时和 profiler 实验不属于此入口；如需 profiling，
	370	+使用仓库原生 `tools.profile` 配置，结果单独统计。
	371	+
	372	+#### 验证与结果
	373	+
	374	+运行目录保留 `manifest.txt`、`launcher.log` 和退出码。检查两路的权重加载、
	375	+LoRA 初始化与 optimizer 同步记录、实际输入、loss 和原生 iteration 耗时。
	376	+日志中的参数快照便于复查，但不能单独证明模型已经应用了所有设置。
	377	+
	378	+本地回归无需模型权重或 NPU：
	379	+
	380	+```bash
	381	+python -m unittest discover -s tests/ut_fsdp/models/mistral3 -p 'test_*.py'
	382	+python examples/magistral-2509/reproduce_lora.py old
	383	+python examples/magistral-2509/reproduce_lora.py pure
	384	+```
	385	+
	386	+测试覆盖配置传递、退出码、资产错误、严格 DCP 键检查、CP 初始化边界和图像处理。
	387	+未安装 torch 时，真实张量存储别名检查会标记跳过；不能算作该检查通过。
	388	+
	389	+精度曲线可用 [TrainingLogParser](https://curryrice233.github.io/TrainingLogParser/)
	390	+查看。性能统计另用同一计时范围、预热区间和有效步数计算，不能把启动脚本的
	391	+简单输出当成完整性能验收。历史四卡 2000 步结果对应当时记录的代码与环境；
	392	+当前修改的本地测试不替代 NPU 运行或新的性能测量。
	393	+
	394	+迁移和性能实践报告提交到 [Issue #369](https://gitcode.com/Ascend/MindSpeed-MM/issues/369)，
	395	+本 README 说明安装、使用与验证方式。

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

@@ -0,0 +1,58 @@
	1	+{
	2	+ "dataset_param": {
	3	+ "dataset_type": "huggingface",
	4	+ "preprocess_parameters": {
	5	+ "model_name_or_path": "./.magistral_run/hf",
	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": "./.magistral_run/data",
	20	+ "dataset": "./.magistral_run/data/mllm_format_llava_instruct_data.json",
	21	+ "cache_dir": null,
	22	+ "enable_thinking": false,
	23	+ "overwrite_cache": false,
	24	+ "train_on_prompt": false,
	25	+ "mask_history": false,
	26	+ "preprocessing_batch_size": 2,
	27	+ "preprocessing_num_workers": 1,
	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/finetune_magistral_2509_lora.sh+164-0

@@ -0,0 +1,164 @@
	1	+#!/usr/bin/env bash
	cxiaolong8 天前脚本改的太复杂，建议保留原始的风格即可
	2	+set -eo pipefail
	3	+
	4	+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
	5	+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
	6	+CANN_ENV_SCRIPT="${CANN_ENV_SCRIPT:-/usr/local/Ascend/cann/set_env.sh}"
	7	+if [[ -f "$CANN_ENV_SCRIPT" ]]; then
	8	+ source "$CANN_ENV_SCRIPT"
	9	+elif [[ "${SKIP_CANN_ENV:-0}" != "1" ]]; then
	10	+ echo "CANN environment script does not exist: $CANN_ENV_SCRIPT" >&2
	11	+ exit 2
	12	+fi
	13	+if [[ -n "${MMSPEED_MM_VENV:-}" ]]; then
	14	+ source "$MMSPEED_MM_VENV/bin/activate"
	15	+fi
	16	+
	17	+# 普通训练入口：环境、参数和 torchrun。固定输入的对照由独立复现入口负责。
	18	+PYTHON_BIN="${PYTHON_BIN:-python}"
	19	+MASTER_ADDR="${MASTER_ADDR:-localhost}"
	20	+MASTER_PORT="${MASTER_PORT:-6000}"
	21	+NNODES="${NNODES:-1}"
	22	+NODE_RANK="${NODE_RANK:-0}"
	23	+NPUS_PER_NODE="${NPUS_PER_NODE:-4}"
	24	+TRAIN_ITERS="${TRAIN_ITERS:-2000}"
	25	+LR_WARMUP_FRACTION="${LR_WARMUP_FRACTION:-0.01}"
	26	+MBS="${MBS:-1}"
	27	+GRAD_ACC_STEP="${GRAD_ACC_STEP:-1}"
	28	+SEQ_LEN="${SEQ_LEN:-512}"
	29	+LOG_INTERVAL="${LOG_INTERVAL:-1}"
	30	+SAVE_INTERVAL="${SAVE_INTERVAL:-10000}"
	31	+EVAL_INTERVAL="${EVAL_INTERVAL:-2000}"
	32	+EVAL_ITERS="${EVAL_ITERS:-2000}"
	33	+TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-pretrain_transformers.py}"
	34	+
	35	+MM_DATA="${MM_DATA:-$REPO_ROOT/examples/magistral-2509/data.json}"
	36	+MM_MODEL="${MM_MODEL:-$REPO_ROOT/examples/magistral-2509/model.json}"
	37	+MM_TOOL="${MM_TOOL:-$REPO_ROOT/mindspeed_mm/tools/tools.json}"
	38	+LOAD_PATH="${LOAD_PATH:-$REPO_ROOT/ckpt/convert_path/Magistral-Small-2509-lora-base}"
	39	+# An empty SAVE_PATH is intentional for comparison/smoke runs: no checkpoint
	40	+# is written unless the caller explicitly supplies a destination.
	41	+SAVE_PATH="${SAVE_PATH-$REPO_ROOT/Magistral-Small-2509_lora_finetune_result}"
	42	+FSDP2_PATH="${FSDP2_PATH:-$REPO_ROOT/examples/magistral-2509/fsdp2_lora_config.yaml}"
	43	+LOG_DIR="${LOG_DIR:-$REPO_ROOT/logs}"
	44	+
	45	+export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-2}"
	46	+export ASCEND_SLOG_PRINT_TO_STDOUT="${ASCEND_SLOG_PRINT_TO_STDOUT:-0}"
	47	+export ASCEND_GLOBAL_LOG_LEVEL="${ASCEND_GLOBAL_LOG_LEVEL:-3}"
	48	+export TASK_QUEUE_ENABLE="${TASK_QUEUE_ENABLE:-2}"
	49	+export COMBINED_ENABLE="${COMBINED_ENABLE:-1}"
	50	+export CPU_AFFINITY_CONF="${CPU_AFFINITY_CONF:-1}"
	51	+export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-1200}"
	52	+export NPU_ASD_ENABLE="${NPU_ASD_ENABLE:-0}"
	53	+export ASCEND_LAUNCH_BLOCKING="${ASCEND_LAUNCH_BLOCKING:-0}"
	54	+export ACLNN_CACHE_LIMIT="${ACLNN_CACHE_LIMIT:-100000}"
	55	+export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
	56	+export PYTORCH_NPU_ALLOC_CONF="${PYTORCH_NPU_ALLOC_CONF:-expandable_segments:True}"
	57	+export MULTI_STREAM_MEMORY_REUSE="${MULTI_STREAM_MEMORY_REUSE:-1}"
	58	+export PYTHONPATH="$REPO_ROOT/MindSpeed:$REPO_ROOT:$REPO_ROOT/megatron:${PYTHONPATH:-}"
	59	+
	60	+TP="${TP:-1}"
	61	+PP="${PP:-1}"
	62	+CP="${CP:-1}"
	63	+for name in NPUS_PER_NODE NNODES TP PP CP MBS GRAD_ACC_STEP; do
	64	+ if ! [[ "${!name}" =~ ^[1-9][0-9]*$ ]]; then
	65	+ echo "$name must be a positive decimal integer" >&2
	66	+ exit 2
	67	+ fi
	68	+done
	69	+WORLD_SIZE=$((NPUS_PER_NODE * NNODES))
	70	+if (( WORLD_SIZE % (TP * PP * CP) != 0 )); then
	71	+ echo "WORLD_SIZE must be divisible by TPPPCP" >&2
	72	+ exit 2
	73	+fi
	74	+DP=$((WORLD_SIZE / TP / PP / CP))
	75	+GBS=$((MBS * GRAD_ACC_STEP * DP))
	76	+
	77	+for required_path in "$MM_DATA" "$MM_MODEL" "$MM_TOOL" "$FSDP2_PATH" "$LOAD_PATH"; do
	78	+ if [[ ! -e "$required_path" ]]; then
	79	+ echo "required path does not exist: $required_path" >&2
	80	+ exit 2
	81	+ fi
	82	+done
	83	+command -v "$PYTHON_BIN" >/dev/null 2>&1 \|\| {
	84	+ echo "Python interpreter not found: $PYTHON_BIN" >&2
	85	+ exit 2
	86	+}
	87	+
	88	+cd "$REPO_ROOT"
	89	+DISTRIBUTED_ARGS=(
	90	+ --nproc_per_node "$NPUS_PER_NODE"
	91	+ --nnodes "$NNODES"
	92	+ --node_rank "$NODE_RANK"
	93	+ --master_addr "$MASTER_ADDR"
	94	+ --master_port "$MASTER_PORT"
	95	+)
	96	+GPT_ARGS=(
	97	+ --use-mcore-models
	98	+ --tensor-model-parallel-size "$TP"
	99	+ --pipeline-model-parallel-size "$PP"
	100	+ --context-parallel-size "$CP"
	101	+ --context-parallel-algo ulysses_cp_algo
	102	+ --micro-batch-size "$MBS"
	103	+ --global-batch-size "$GBS"
	104	+ --tokenizer-type NullTokenizer
	105	+ --vocab-size 131072
	106	+ --seq-length "$SEQ_LEN"
	107	+ --make-vocab-size-divisible-by 1
	108	+ --normalization RMSNorm
	109	+ --use-fused-rmsnorm
	110	+ --swiglu
	111	+ --use-fused-swiglu
	112	+ --no-masked-softmax-fusion
	113	+ --lr 1.0e-5
	114	+ --lr-decay-style cosine
	115	+ --weight-decay 0
	116	+ --train-iters "$TRAIN_ITERS"
	117	+ --lr-warmup-fraction "$LR_WARMUP_FRACTION"
	118	+ --clip-grad 0.0
	119	+ --adam-beta1 0.9
	120	+ --adam-beta2 0.999
	121	+ --no-gradient-accumulation-fusion
	122	+ --seed 42
	123	+ --load "$LOAD_PATH"
	124	+ --use-flash-attn
	125	+ --no-load-optim
	126	+ --no-load-rng
	127	+ --no-save-optim
	128	+ --no-save-rng
	129	+ --num-workers 1
	130	+ --use-torch-fsdp2
	131	+ --untie-embeddings-and-output-weights
	132	+ --ckpt-format torch_dcp
	133	+ --fsdp2-config-path "$FSDP2_PATH"
	134	+ --optimizer-selection fused_torch_adamw
	135	+ --use-cpu-initialization
	136	+ --log-tps
	137	+ --init-model-with-meta-device
	138	+)
	139	+LORA_ARGS=(
	140	+ --lora-r 8
	141	+ --lora-alpha 16
	142	+ --lora-dropout 0
	143	+ --lora-target-modules q_proj k_proj v_proj o_proj gate_proj up_proj down_proj
	144	+)
	145	+MM_ARGS=(--mm-data "$MM_DATA" --mm-model "$MM_MODEL" --mm-tool "$MM_TOOL")
	146	+OUTPUT_ARGS=(
	147	+ --log-interval "$LOG_INTERVAL"
	148	+ --save-interval "$SAVE_INTERVAL"
	149	+ --eval-interval "$EVAL_INTERVAL"
	150	+ --eval-iters "$EVAL_ITERS"
	151	+ --save "$SAVE_PATH"
	152	+)
	153	+
	154	+mkdir -p "$LOG_DIR"
	155	+LOGFILE="$LOG_DIR/train_$(date -u +%Y%m%dT%H%M%SZ).log"
	156	+echo "ENTRY=old-megatron-fsdp2 CONFIG=$FSDP2_PATH WORLD=$WORLD_SIZE GBS=$GBS TRAIN_ITERS=$TRAIN_ITERS SEQ_LEN=$SEQ_LEN SEED=42" \| tee "$LOGFILE"
	157	+set +e
	158	+"$PYTHON_BIN" -m torch.distributed.run "${DISTRIBUTED_ARGS[@]}" "$TRAIN_ENTRYPOINT" \
	159	+ "${GPT_ARGS[@]}" "${MM_ARGS[@]}" "${OUTPUT_ARGS[@]}" "${LORA_ARGS[@]}" \
	160	+ --distributed-backend nccl 2>&1 \| tee -a "$LOGFILE"
	161	+TRAIN_STATUS=${PIPESTATUS[0]}
	162	+set -e
	163	+echo "TRAIN_EXIT_CODE=$TRAIN_STATUS" \| tee -a "$LOGFILE"
	164	+exit "$TRAIN_STATUS"

Aexamples/magistral-2509/finetune_magistral_2509_pure_fsdp2_lora.sh+46-0

@@ -0,0 +1,46 @@
	1	+#!/bin/bash
	2	+set -euo pipefail
	3	+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
	4	+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
	5	+CANN_ENV_SCRIPT="${CANN_ENV_SCRIPT:-/usr/local/Ascend/cann/set_env.sh}"
	6	+# CANN's setup script appends to these variables and is not nounset-safe.
	7	+export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}"
	8	+export PYTHONPATH="${PYTHONPATH:-}"
	9	+export CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH:-}"
	10	+if [[ -f "$CANN_ENV_SCRIPT" ]]; then
	11	+ # The Ascend runtime is site-specific; callers may override this path.
	12	+ source "$CANN_ENV_SCRIPT"
	13	+fi
	14	+if [[ -n "${MMSPEED_MM_VENV:-}" ]]; then
	15	+ # Optional virtual environment; do not assume a private machine path.
	16	+ source "$MMSPEED_MM_VENV/bin/activate"
	17	+fi
	18	+export NON_MEGATRON=true
	19	+export MULTI_STREAM_MEMORY_REUSE="${MULTI_STREAM_MEMORY_REUSE:-1}"
	20	+export TASK_QUEUE_ENABLE="${TASK_QUEUE_ENABLE:-2}"
	21	+export ASCEND_LAUNCH_BLOCKING=0
	22	+export ACLNN_CACHE_LIMIT=100000
	23	+export CPU_AFFINITY_CONF=1
	24	+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
	25	+export TOKENIZERS_PARALLELISM=false
	26	+export PYTHONPATH="$REPO_ROOT/MindSpeed:$REPO_ROOT:$REPO_ROOT/megatron:${PYTHONPATH:-}"
	27	+cd "$REPO_ROOT"
	28	+NPUS_PER_NODE=${NPUS_PER_NODE:-4}
	29	+NNODES=${NNODES:-1}
	30	+NODE_RANK=${NODE_RANK:-0}
	31	+MASTER_ADDR=${MASTER_ADDR:-localhost}
	32	+MASTER_PORT=${MASTER_PORT:-6027}
	33	+PYTHON_BIN=${PYTHON_BIN:-python}
	34	+TRAIN_ENTRYPOINT="${TRAIN_ENTRYPOINT:-mindspeed_mm/fsdp/train/trainer.py}"
	35	+LOG_DIR="${LOG_DIR:-./.magistral_run/logs}"
	36	+mkdir -p "$LOG_DIR"
	37	+LOGFILE="$LOG_DIR/pure-fsdp2-mistral3-$(date -u +%Y%m%dT%H%M%SZ).log"
	38	+CONFIG="${CONFIG_PATH:-examples/magistral-2509/magistral3_fsdp2_lora_config.yaml}"
	39	+if [[ ! -f "$CONFIG" ]]; then
	40	+ echo "configuration file does not exist: $CONFIG" >&2
	41	+ exit 2
	42	+fi
	43	+
	44	+echo "ENTRY=pure-fsdp2 CONFIG=$CONFIG WORLD=$NPUS_PER_NODE NNODES=$NNODES MASTER_ADDR=$MASTER_ADDR MASTER_PORT=$MASTER_PORT" \| tee "$LOGFILE"
	45	+"$PYTHON_BIN" -m torch.distributed.run --nproc_per_node "$NPUS_PER_NODE" --nnodes "$NNODES" --node_rank "$NODE_RANK" --master_addr "$MASTER_ADDR" --master_port "$MASTER_PORT" \
	46	+ "$TRAIN_ENTRYPOINT" "$CONFIG" 2>&1 \| tee -a "$LOGFILE"

Aexamples/magistral-2509/fsdp2_lora_config.yaml+19-0

@@ -0,0 +1,19 @@
	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	+use_reentrant: False
	14	+param_dtype: "bf16"
	15	+reduce_dtype: "fp32"
	16	+cast_forward_inputs: True
	17	+num_to_forward_prefetch: 2
	18	+num_to_backward_prefetch: 2
	19	+offload_to_cpu: False

Aexamples/magistral-2509/magistral3_fsdp2_lora_config.yaml+132-0

@@ -0,0 +1,132 @@
	1	+parallel:
	2	+ tensor_parallel_size: 1
	3	+ fully_shard_parallel_size: auto
	4	+ fsdp_plan:
	5	+ apply_modules:
	6	+ - model.vision_tower.transformer.layers.{*}
	7	+ - model.vision_tower
	8	+ - model.multi_modal_projector
	9	+ - model.language_model.embed_tokens
	10	+ - model.language_model.layers.{*}
	11	+ - model.language_model
	12	+ - lm_head
	13	+ param_dtype: bf16
	14	+ reduce_dtype: fp32
	15	+ cast_forward_inputs: true
	16	+ reshard_after_forward: true
	17	+ num_to_forward_prefetch: 2
	18	+ num_to_backward_prefetch: 2
	19	+ ring_attention_size: 1
	20	+ ulysses_parallel_size: 1
	21	+
	22	+data:
	23	+ dataset_param:
	24	+ dataset_type: huggingface
	25	+ attr:
	26	+ images: images
	27	+ messages: messages
	28	+ role_tag: role
	29	+ content_tag: content
	30	+ user_tag: user
	31	+ assistant_tag: assistant
	32	+ preprocess_parameters:
	33	+ model_name_or_path: ./.magistral_run/hf
	34	+ use_fast_tokenizer: true
	35	+ split_special_tokens: false
	36	+ image_max_pixels: 262144
	37	+ image_min_pixels: 1024
	38	+ video_max_pixels: 16384
	39	+ video_min_pixels: 0
	40	+ video_fps: 2.0
	41	+ video_maxlen: 64
	42	+ fix_mistral_regex: true
	43	+ trust_remote_code: false
	44	+ basic_parameters:
	45	+ cutoff_len: 512
	46	+ template: mistral
	47	+ enable_thinking: false
	48	+ train_on_prompt: false
	49	+ mask_history: false
	50	+ dataset_dir: ./.magistral_run/data
	51	+ dataset: ./.magistral_run/data/mllm_format_llava_instruct_data.json
	52	+ cache_dir: ./.magistral_run/fsdp_cache
	53	+ overwrite_cache: false
	54	+ preprocessing_batch_size: 2
	55	+ preprocessing_num_workers: 1
	56	+ max_samples: null
	57	+ dataloader_param:
	58	+ pin_memory: true
	59	+ shuffle: true
	60	+ dataloader_mode: sampler
	61	+ drop_last: true
	62	+ sampler_type: BaseRandomBatchSampler
	63	+ num_workers: 1
	64	+ collate_param:
	65	+ collator_id: pixtral
	66	+ ignore_pad_token_for_loss: true
	67	+ enable_preload: false
	68	+
	69	+model:
	70	+ model_id: mistral3
	71	+ model_name_or_path: ./.magistral_run/hf
	72	+ trust_remote_code: false
	73	+ attn_implementation: flash_attention_2
	74	+ dtype: bf16
	75	+ freeze:
	76	+ - model.vision_tower
	77	+ - model.multi_modal_projector
	78	+
	79	+features:
	80	+ loss_cfg:
	81	+ loss_type: raw
	82	+ recompute: true
	83	+ recompute_plan:
	84	+ apply_modules:
	85	+ - model.vision_tower.transformer.layers.{*}
	86	+ - model.language_model.layers.{*}
	87	+
	88	+training:
	89	+ micro_batch_size: 1
	90	+ gradient_accumulation_steps: 1
	91	+ seed: 42
	92	+ lr: 1.0e-5
	93	+ lr_decay_style: cosine
	94	+ lr_warmup_ratio: 0.01
	95	+ weight_decay: 0.0
	96	+ train_iters: 2000
	97	+ clip_grad: 0.0
	98	+ init_model_with_meta_device: true
	99	+ optimizer: adamw
	100	+ adam_fused: true
	101	+ save_interval: 2000
	102	+ no_load_optim: true
	103	+ no_load_rng: true
	104	+ no_save_optim: true
	105	+ no_save_rng: true
	106	+ load: ./.magistral_run/hf
	107	+ load_format: hf
	108	+ save: ./.magistral_run/ckpt/pure_fsdp2_lora_2000
	109	+ save_format: dcp
	110	+ use_deter_comp: false
	111	+ lora:
	112	+ enable: true
	113	+ rank: 8
	114	+ alpha: 16
	115	+ dropout: 0.0
	116	+ target_modules:
	117	+ - model.language_model.layers.{*}.self_attn.q_proj
	118	+ - model.language_model.layers.{*}.self_attn.k_proj
	119	+ - model.language_model.layers.{*}.self_attn.v_proj
	120	+ - model.language_model.layers.{*}.self_attn.o_proj
	121	+ - model.language_model.layers.{*}.mlp.gate_proj
	122	+ - model.language_model.layers.{*}.mlp.up_proj
	123	+ - model.language_model.layers.{*}.mlp.down_proj
	124	+ plugin:
	125	+ - mindspeed_mm/fsdp/models/mistral3
	126	+ - mindspeed_mm/fsdp/data/datasets/huggingface
	127	+
	128	+tools:
	129	+ profile:
	130	+ enable: false
	131	+ memory_profile:
	132	+ enable: false

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

@@ -0,0 +1,27 @@
	1	+{
	2	+ "model_id": "mistral3",
	3	+ "init_from_hf_path": "./.magistral_run/hf",
	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/reproduce_lora.py+72-0

@@ -0,0 +1,72 @@
	1	+#!/usr/bin/env python3
	2	+"""Reproduce an old/pure LoRA comparison with the checked-in training entries.
	3	+
	4	+The default checks syntax without loading weights. Pass --run to train.
	5	+Model, optimizer and FSDP parameters remain in the existing YAML/launchers.
	6	+"""
	7	+import argparse
	8	+import json
	9	+import os
	10	+from pathlib import Path
	11	+import subprocess
	12	+import sys
	13	+
	14	+
	15	+def locked_environment(cards, steps, inherited):
	16	+ if cards not in (2, 4) or not isinstance(steps, int) or steps < 1:
	17	+ raise ValueError("reference requires 2/4 cards and a positive step count")
	18	+ if steps >= 100 and steps % 100:
	19	+ raise ValueError("For >=100 steps, use a multiple of 100 so 1% warmup is an integer")
	20	+ locked = {
	21	+ "REFERENCE_WORLD_SIZE": "4", "NPUS_PER_NODE": str(cards), "NNODES": "1",
	22	+ "NODE_RANK": "0", "TP": "1", "PP": "1", "CP": "1",
	23	+ "MBS": "1", "GLOBAL_BATCH_SIZE": "4", "GRAD_ACC_STEP": str(4 // cards),
	24	+ "SEQ_LEN": "512", "TRAIN_ITERS": str(steps),
	25	+ "LR_WARMUP_FRACTION": "0" if steps < 100 else "0.01",
	26	+ "TASK3_SEED": "42", "TASK3_DETERMINISTIC_LORA": "1",
	27	+ "TASK3_REPRO_RECOMPUTE": "1", "TASK3_REPRO_PREFETCH": "2",
	28	+ "TASK3_REPRO_ADAM_FUSED": "1", "TASK3_REPRO_RESHARD_AFTER_FORWARD": "1",
	29	+ "TASK3_REPRO_PROFILE": "0", "TASK3_REPRO_SAVE_CHECKPOINT": "0",
	30	+ "TASK3_REPRO_SKIP_ZERO_NORM": "1", "TASK3_REPRO_STRICT_DCP": "1",
	31	+ "ALLOW_NON_REFERENCE_TOPOLOGY": "1" if cards == 2 else "0",
	32	+ }
	33	+ blocked = ("CONFIG_PATH", "FSDP2_PATH", "MM_MODEL", "MM_DATA", "MM_TOOL",
	34	+ "TRAIN_ENTRYPOINT", "TASK3_PURE_TRAINER_ENTRYPOINT",
	35	+ "TASK3_REPRO_ATTN_IMPLEMENTATION", "TASK3_REPRO_DATA", "HCCL_ALGO",
	36	+ "HCCL_BUFFSIZE", "SAVE_PATH")
	37	+ for key in blocked:
	38	+ if inherited.get(key):
	39	+ raise ValueError(f"Unset {key}: reference mode uses the checked-in entry/config")
	40	+ for key, value in locked.items():
	41	+ if key in inherited and inherited[key] != value:
	42	+ raise ValueError(f"{key} conflicts with reference value {value}; unset it first")
	43	+ return {inherited, locked}
	44	+
	45	+
	46	+def main():
	47	+ parser = argparse.ArgumentParser(description=__doc__)
	48	+ parser.add_argument("mode", choices=("old", "pure"))
	49	+ parser.add_argument("--cards", type=int, choices=(2, 4), default=4)
	50	+ parser.add_argument("--steps", type=int, default=2000,
	51	+ help="<100: no warmup; >=100: multiple of 100, 1%% warmup")
	52	+ parser.add_argument("--run", action="store_true", help="launch training after asset checks")
	53	+ args = parser.parse_args()
	54	+ try:
	55	+ env = locked_environment(args.cards, args.steps, os.environ)
	56	+ except ValueError as exc:
	57	+ parser.error(str(exc))
	58	+ env["PYTHON_BIN"] = sys.executable
	59	+ print("TASK3_REFERENCE=" + json.dumps({
	60	+ "cards": args.cards, "steps": args.steps, "global_batch": 4,
	61	+ "gradient_accumulation": 4 // args.cards,
	62	+ "effective_warmup_ratio": float(env["LR_WARMUP_FRACTION"]),
	63	+ "checkpoint": False, "strict_dcp_model_keys": True,
	64	+ }, sort_keys=True), flush=True)
	65	+ command = ["bash", str(Path(__file__).with_name("task3_repro.sh"))]
	66	+ if not args.run:
	67	+ command.append("--check")
	68	+ return subprocess.call([*command, args.mode], env=env)
	69	+
	70	+
	71	+if __name__ == "__main__":
	72	+ sys.exit(main())

Aexamples/magistral-2509/requirements.txt+2-0

@@ -0,0 +1,2 @@
	1	+# Mistral3 Tekken tokenizer and Pixtral processor support for Transformers 4.57.0.
	2	+transformers[mistral-common]==4.57.0

Aexamples/magistral-2509/task3_repro.sh+436-0

@@ -0,0 +1,436 @@
	1	+#!/usr/bin/env bash
	2	+set -euo pipefail
	3	+
	4	+# Reproducibility wrapper for the two checked-in Task3 training paths.
	5	+# It records launcher inputs and source identity, but never rewrites YAML,
	6	+# downloads weights, converts checkpoints, or contains site-specific access data.
	7	+
	8	+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
	9	+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
	10	+PYTHON_BIN="${PYTHON_BIN:-}"
	11	+MODE="pure"
	12	+ACTION="run"
	13	+REFERENCE_WORLD_SIZE="${REFERENCE_WORLD_SIZE:-4}"
	14	+NPUS_PER_NODE="${NPUS_PER_NODE:-$REFERENCE_WORLD_SIZE}"
	15	+NNODES="${NNODES:-1}"
	16	+NODE_RANK="${NODE_RANK:-0}"
	17	+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-4}"
	18	+MBS="${MBS:-1}"
	19	+GRAD_ACC_STEP="${GRAD_ACC_STEP:-}"
	20	+MASTER_ADDR="${MASTER_ADDR:-localhost}"
	21	+MASTER_PORT="${MASTER_PORT:-6027}"
	22	+RUN_ROOT="${RUN_ROOT:-$REPO_ROOT/.magistral_run/repro-$(date -u +%Y%m%dT%H%M%SZ)}"
	23	+REFERENCE_SOURCE_COMMIT="${REFERENCE_SOURCE_COMMIT:-2247e1c383e5fb6c6ede01225a4d36ac2d5edb2}"
	24	+ALLOW_NON_REFERENCE_TOPOLOGY="${ALLOW_NON_REFERENCE_TOPOLOGY:-0}"
	25	+TASK3_REPRO_DATA="${TASK3_REPRO_DATA:-$REPO_ROOT/.magistral_run/data/task3_repro_mllm_format.json}"
	26	+
	27	+# Prefer the environment shipped beside the reproduction workspace. The
	28	+# checked-in launchers accept an absolute PYTHON_BIN, but leaving the choice
	29	+# to the host's `python3` makes a fresh machine fail before argument parsing
	30	+# when dependencies are installed only in the workspace venv.
	31	+if [[ -z "$PYTHON_BIN" ]]; then
	32	+ if [[ -n "${MMSPEED_MM_VENV:-}" && -x "$MMSPEED_MM_VENV/bin/python" ]]; then
	33	+ PYTHON_BIN="$MMSPEED_MM_VENV/bin/python"
	34	+ elif [[ -x "$REPO_ROOT/venv/bin/python" ]]; then
	35	+ PYTHON_BIN="$REPO_ROOT/venv/bin/python"
	36	+ elif [[ -x "$REPO_ROOT/../venv/bin/python" ]]; then
	37	+ PYTHON_BIN="$REPO_ROOT/../venv/bin/python"
	38	+ else
	39	+ PYTHON_BIN="python3"
	40	+ fi
	41	+fi
	42	+
	43	+usage() {
	44	+ cat <<'EOF'
	45	+Usage:
	46	+ task3_repro.sh [--check\|--dry-run\|--print] [old\|pure]
	47	+
	48	+Launcher-level environment variables:
	49	+ NPUS_PER_NODE, NNODES, NODE_RANK, MASTER_ADDR, MASTER_PORT
	50	+ GLOBAL_BATCH_SIZE, MBS, GRAD_ACC_STEP, PYTHON_BIN, MMSPEED_MM_VENV
	51	+ CANN_ENV_SCRIPT, CONFIG_PATH, LOG_DIR, SAVE_PATH, RUN_ROOT
	52	+ TASK3_REPRO_DATA (pure path; defaults to the shared mllm-format fixture)
	53	+ TRAIN_ITERS, LR_WARMUP_FRACTION, TASK3_REPRO_RECOMPUTE, TASK3_REPRO_PREFETCH
	54	+ TASK3_REPRO_ADAM_FUSED, TASK3_REPRO_RESHARD_AFTER_FORWARD
	55	+ (configuration overrides; no YAML is copied)
	56	+ TASK3_REPRO_SAVE_CHECKPOINT (default: 0; set 1 to reuse the configured path)
	57	+ TASK3_DETERMINISTIC_LORA (default: 1 for this reproduction wrapper)
	58	+
	59	+Examples:
	60	+ bash examples/magistral-2509/task3_repro.sh --check old
	61	+ TRAIN_ITERS=5 NPUS_PER_NODE=4 bash examples/magistral-2509/task3_repro.sh old
	62	+ TRAIN_ITERS=5 LR_WARMUP_FRACTION=0 NPUS_PER_NODE=2 \\
	63	+ ALLOW_NON_REFERENCE_TOPOLOGY=1 bash examples/magistral-2509/task3_repro.sh pure
	64	+ NPUS_PER_NODE=2 GRAD_ACC_STEP=2 ALLOW_NON_REFERENCE_TOPOLOGY=1 \
	65	+ bash examples/magistral-2509/task3_repro.sh old
	66	+
	67	+The default reference topology is 4 NPU and global batch 4. A non-reference
	68	+topology must be explicit and is recorded as diagnostic-only evidence.
	69	+The selected checked-in YAML or launcher remains the source of training
	70	+parameters; this wrapper does not create a modified parameter copy.
	71	+EOF
	72	+}
	73	+
	74	+while [[ $# -gt 0 ]]; do
	75	+ case "$1" in
	76	+ --check\|--dry-run\|--print)
	77	+ ACTION="${1#--}"
	78	+ ;;
	79	+ --help\|-h)
	80	+ usage
	81	+ exit 0
	82	+ ;;
	83	+ old\|pure)
	84	+ MODE="$1"
	85	+ ;;
	86	+ *)
	87	+ echo "unsupported argument: $1" >&2
	88	+ usage >&2
	89	+ exit 2
	90	+ ;;
	91	+ esac
	92	+ shift
	93	+done
	94	+
	95	+if [[ "$MODE" == "pure" ]]; then
	96	+ ENTRYPOINT="$SCRIPT_DIR/finetune_magistral_2509_pure_fsdp2_lora.sh"
	97	+ CONFIG_PATH="${CONFIG_PATH:-$SCRIPT_DIR/magistral3_fsdp2_lora_config.yaml}"
	98	+else
	99	+ ENTRYPOINT="$SCRIPT_DIR/finetune_magistral_2509_lora.sh"
	100	+ CONFIG_PATH="${CONFIG_PATH:-$SCRIPT_DIR/fsdp2_lora_config.yaml}"
	101	+fi
	102	+
	103	+[[ -f "$ENTRYPOINT" ]] \|\| { echo "missing entry point: $ENTRYPOINT" >&2; exit 2; }
	104	+[[ -f "$CONFIG_PATH" ]] \|\| { echo "missing canonical config: $CONFIG_PATH" >&2; exit 2; }
	105	+command -v "$PYTHON_BIN" >/dev/null 2>&1 \|\| {
	106	+ echo "Python interpreter not found: $PYTHON_BIN" >&2
	107	+ exit 2
	108	+}
	109	+
	110	+for name in NPUS_PER_NODE NNODES MBS GLOBAL_BATCH_SIZE REFERENCE_WORLD_SIZE; do
	111	+ value="${!name}"
	112	+ if ! [[ "$value" =~ ^[1-9][0-9]*$ ]]; then
	113	+ echo "$name must be a positive decimal integer" >&2
	114	+ exit 2
	115	+ fi
	116	+done
	117	+for name in TP PP CP; do
	118	+ if [[ "${!name:-1}" != "1" ]]; then
	119	+ echo "$name must be 1 for this data-parallel reproduction wrapper" >&2
	120	+ exit 2
	121	+ fi
	122	+done
	123	+if ! [[ "$NODE_RANK" =~ ^(0\|[1-9][0-9]*)$ ]] \|\| (( NODE_RANK >= NNODES )); then
	124	+ echo "NODE_RANK must be between 0 and NNODES-1" >&2
	125	+ exit 2
	126	+fi
	127	+
	128	+if [[ -z "$GRAD_ACC_STEP" ]]; then
	129	+ denominator=$((NPUS_PER_NODE * NNODES * MBS))
	130	+ if (( denominator <= 0 \|\| GLOBAL_BATCH_SIZE % denominator != 0 )); then
	131	+ echo "GLOBAL_BATCH_SIZE must be divisible by NPUS_PER_NODENNODESMBS" >&2
	132	+ exit 2
	133	+ fi
	134	+ GRAD_ACC_STEP=$((GLOBAL_BATCH_SIZE / denominator))
	135	+fi
	136	+
	137	+if ! [[ "$GRAD_ACC_STEP" =~ ^[1-9][0-9]*$ ]]; then
	138	+ echo "GRAD_ACC_STEP must be a positive decimal integer" >&2
	139	+ exit 2
	140	+fi
	141	+if (( NPUS_PER_NODE * NNODES * MBS * GRAD_ACC_STEP != GLOBAL_BATCH_SIZE )); then
	142	+ echo "batch mismatch: GLOBAL_BATCH_SIZE must equal NPUS_PER_NODENNODESMBS*GRAD_ACC_STEP" >&2
	143	+ exit 2
	144	+fi
	145	+if (( NPUS_PER_NODE < 1 \|\| NNODES < 1 \|\| MBS < 1 \|\| GRAD_ACC_STEP < 1 )); then
	146	+ echo "NPUS_PER_NODE, NNODES, MBS and GRAD_ACC_STEP must be positive" >&2
	147	+ exit 2
	148	+fi
	149	+
	150	+if (( NPUS_PER_NODE * NNODES != REFERENCE_WORLD_SIZE )) && [[ "$ALLOW_NON_REFERENCE_TOPOLOGY" != "1" ]]; then
	151	+ echo "topology is $((NPUS_PER_NODE * NNODES)) NPU, but the reference is $REFERENCE_WORLD_SIZE;" >&2
	152	+ echo "set ALLOW_NON_REFERENCE_TOPOLOGY=1 only for diagnostic evidence" >&2
	153	+ exit 2
	154	+fi
	155	+
	156	+if [[ "$ACTION" == "check" ]]; then
	157	+ bash -n "$ENTRYPOINT"
	158	+ "$PYTHON_BIN" -m py_compile "$REPO_ROOT/examples/magistral-2509/task3_repro_train.py"
	159	+ "$PYTHON_BIN" -m py_compile "$REPO_ROOT/mindspeed_mm/fsdp/models/mistral3/repro_hooks.py"
	160	+ "$PYTHON_BIN" - "$CONFIG_PATH" <<'PY'
	161	+from pathlib import Path
	162	+import sys
	163	+
	164	+try:
	165	+ import yaml
	166	+except ImportError as exc:
	167	+ raise SystemExit(f"PyYAML is required for config validation: {exc}")
	168	+
	169	+path = Path(sys.argv[1])
	170	+with path.open(encoding="utf-8") as handle:
	171	+ value = yaml.safe_load(handle)
	172	+if not isinstance(value, dict) or not value:
	173	+ raise SystemExit(f"configuration must be a non-empty mapping: {path}")
	174	+print(f"CONFIG_SYNTAX_CHECK=ok PATH={path}")
	175	+PY
	176	+ echo "ENTRYPOINT_SYNTAX_CHECK=ok PATH=$ENTRYPOINT"
	177	+ echo "REPRO_MODE=$MODE REFERENCE_WORLD_SIZE=$REFERENCE_WORLD_SIZE REQUESTED_WORLD_SIZE=$((NPUS_PER_NODE * NNODES))"
	178	+ echo "NO_RUNTIME_STARTED=true NO_MODEL_DOWNLOAD=true NO_CONFIG_REWRITE=true"
	179	+ exit 0
	180	+fi
	181	+
	182	+SOURCE_COMMIT="unknown"
	183	+if git -C "$REPO_ROOT" rev-parse HEAD >/dev/null 2>&1; then
	184	+ SOURCE_COMMIT="$(git -C "$REPO_ROOT" rev-parse HEAD)"
	185	+fi
	186	+
	187	+if [[ "$ACTION" == "print" \|\| "$ACTION" == "dry-run" ]]; then
	188	+ echo "REPRO_MODE=$MODE"
	189	+ echo "ENTRYPOINT=$ENTRYPOINT"
	190	+ echo "CONFIG_PATH=$CONFIG_PATH"
	191	+ echo "SOURCE_COMMIT=$SOURCE_COMMIT"
	192	+ echo "REFERENCE_SOURCE_COMMIT=$REFERENCE_SOURCE_COMMIT"
	193	+ echo "REFERENCE_WORLD_SIZE=$REFERENCE_WORLD_SIZE REQUESTED_WORLD_SIZE=$((NPUS_PER_NODE * NNODES))"
	194	+ echo "GLOBAL_BATCH_SIZE=$GLOBAL_BATCH_SIZE MBS=$MBS GRAD_ACC_STEP=$GRAD_ACC_STEP"
	195	+ echo "NO_MODEL_DOWNLOAD=true NO_CHECKPOINT_CONVERSION=true NO_CONFIG_REWRITE=true"
	196	+ exit 0
	197	+fi
	198	+
	199	+if [[ "$MODE" == "old" ]]; then
	200	+ MM_DATA="${MM_DATA:-$SCRIPT_DIR/task3_repro_data.json}"
	201	+ MM_MODEL="${MM_MODEL:-$REPO_ROOT/examples/magistral-2509/model.json}"
	202	+ MM_TOOL="${MM_TOOL:-$REPO_ROOT/mindspeed_mm/tools/tools.json}"
	203	+ LOAD_PATH="${LOAD_PATH:-$REPO_ROOT/ckpt/convert_path/Magistral-Small-2509-lora-base}"
	204	+ TRAIN_ITERS="${TRAIN_ITERS:-2000}"
	205	+ LR_WARMUP_FRACTION="${LR_WARMUP_FRACTION:-0.01}"
	206	+ for required_path in "$MM_DATA" "$MM_MODEL" "$MM_TOOL" "$LOAD_PATH"; do
	207	+ [[ -e "$required_path" ]] \|\| { echo "required old-path asset does not exist: $required_path" >&2; exit 2; }
	208	+ done
	209	+ "$PYTHON_BIN" - "$SCRIPT_DIR" "$LOAD_PATH" <<'PY'
	210	+from pathlib import Path
	211	+import runpy
	212	+import sys
	213	+
	214	+entry = runpy.run_path(str(Path(sys.argv[1]) / "task3_repro_train.py"))
	215	+entry["dcp_metadata"](Path(sys.argv[2]))
	216	+PY
	217	+fi
	218	+
	219	+if [[ "$MODE" == "pure" ]]; then
	220	+ # Fail before torchrun if a locally prepared HF/data fixture is incomplete.
	221	+ # A late rank-local FileNotFoundError otherwise hides the preparation error.
	222	+ if ! "$PYTHON_BIN" - "$CONFIG_PATH" "$REPO_ROOT" "$TASK3_REPRO_DATA" <<'PY'
	223	+from __future__ import annotations
	224	+
	225	+from pathlib import Path
	226	+import json
	227	+import runpy
	228	+import sys
	229	+
	230	+try:
	231	+ import yaml
	232	+except ImportError as exc:
	233	+ raise SystemExit(f"PyYAML is required for runtime asset validation: {exc}")
	234	+
	235	+config_path = Path(sys.argv[1]).resolve()
	236	+repo_root = Path(sys.argv[2]).resolve()
	237	+dataset_override = Path(sys.argv[3]).expanduser()
	238	+if not dataset_override.is_absolute():
	239	+ dataset_override = (repo_root / dataset_override).resolve()
	240	+with config_path.open(encoding="utf-8") as handle:
	241	+ config = yaml.safe_load(handle) or {}
	242	+
	243	+def resolve(value: object) -> Path \| None:
	244	+ if not isinstance(value, str) or not value:
	245	+ return None
	246	+ path = Path(value)
	247	+ return path if path.is_absolute() else repo_root / path
	248	+
	249	+model_path = resolve((config.get("model") or {}).get("model_name_or_path"))
	250	+data_param = (config.get("data") or {}).get("dataset_param") or {}
	251	+basic = data_param.get("basic_parameters") or {}
	252	+dataset_path = dataset_override
	253	+dataset_dir = dataset_override.parent
	254	+
	255	+missing = []
	256	+if model_path is None or not model_path.exists():
	257	+ missing.append(f"model directory: {model_path or '<not configured>'}")
	258	+if dataset_path is None or not dataset_path.is_file():
	259	+ missing.append(f"dataset JSON: {dataset_path or '<not configured>'}")
	260	+if missing:
	261	+ raise SystemExit(
	262	+ "PURE_ASSET_PREFLIGHT=fail\n"
	263	+ + "Prepare the local HF/data fixture before launching distributed training:\n"
	264	+ + "- " + "\n- ".join(missing)
	265	+ )
	266	+
	267	+entry = runpy.run_path(str(repo_root / "examples/magistral-2509/task3_repro_train.py"))
	268	+entry["validate_hf_assets"](model_path)
	269	+old_data = json.loads((repo_root / "examples/magistral-2509/task3_repro_data.json").read_text())
	270	+for section, fields in {
	271	+ "preprocess_parameters": ("model_name_or_path", "image_max_pixels", "image_min_pixels",
	272	+ "use_fast_tokenizer", "split_special_tokens", "fix_mistral_regex"),
	273	+ "basic_parameters": ("template", "enable_thinking", "train_on_prompt", "mask_history",
	274	+ "preprocessing_batch_size", "preprocessing_num_workers"),
	275	+}.items():
	276	+ for field in fields:
	277	+ if old_data["dataset_param"][section][field] != data_param[section][field]:
	278	+ raise SystemExit(f"Old/pure data configuration differs: {section}.{field}")
	279	+
	280	+try:
	281	+ records = json.loads(dataset_path.read_text(encoding="utf-8"))
	282	+except (OSError, json.JSONDecodeError) as exc:
	283	+ raise SystemExit(f"PURE_ASSET_PREFLIGHT=fail\nCannot read dataset JSON {dataset_path}: {exc}")
	284	+
	285	+if isinstance(records, dict):
	286	+ raise SystemExit(
	287	+ "PURE_ASSET_PREFLIGHT=fail\n"
	288	+ f"{dataset_path} is a configuration mapping, not a sample list; "
	289	+ "use the generated mllm-format sample JSON."
	290	+ )
	291	+if not isinstance(records, list) or not records:
	292	+ raise SystemExit(f"PURE_ASSET_PREFLIGHT=fail\nDataset must be a non-empty JSON list: {dataset_path}")
	293	+if len(records) != 4:
	294	+ raise SystemExit("Reference comparison requires exactly four fixed samples")
	295	+
	296	+missing_images = []
	297	+for index, record in enumerate(records):
	298	+ if not isinstance(record, dict):
	299	+ raise SystemExit(f"sample[{index}] must be a mapping")
	300	+ if not isinstance(record.get("messages"), list) or not record["messages"]:
	301	+ raise SystemExit(f"sample[{index}].messages must be a non-empty list")
	302	+ images = record.get("images", [])
	303	+ if isinstance(images, str):
	304	+ images = [images]
	305	+ if not isinstance(images, list) or not all(isinstance(image, str) and image for image in images):
	306	+ raise SystemExit(f"sample[{index}].images must contain non-empty filenames")
	307	+ for image in images:
	308	+ image_path = Path(image)
	309	+ if not image_path.is_absolute():
	310	+ image_path = dataset_dir / image_path
	311	+ if not image_path.exists():
	312	+ missing_images.append(f"sample[{index}]: {image_path}")
	313	+if missing_images:
	314	+ shown = missing_images[:8]
	315	+ suffix = "" if len(missing_images) <= len(shown) else f"\n- ... and {len(missing_images) - len(shown)} more"
	316	+ raise SystemExit(
	317	+ "PURE_ASSET_PREFLIGHT=fail\nMissing image assets:\n- "
	318	+ + "\n- ".join(shown)
	319	+ + suffix
	320	+ )
	321	+
	322	+print(
	323	+ "PURE_ASSET_PREFLIGHT=pass "
	324	+ f"model={model_path} dataset={dataset_path} samples={len(records)}"
	325	+)
	326	+PY
	327	+ then
	328	+ exit 2
	329	+ fi
	330	+fi
	331	+
	332	+mkdir -p "$RUN_ROOT"
	333	+LOG_DIR="${LOG_DIR:-$RUN_ROOT/logs}"
	334	+if [[ "$MODE" == "old" ]]; then
	335	+ # The comparison gate consumes logs, not a checkpoint. Preserve an explicit
	336	+ # caller path, but default to no save so repeated smoke runs do not create
	337	+ # large artifacts that must be deleted afterwards.
	338	+ SAVE_PATH="${SAVE_PATH-}"
	339	+else
	340	+ SAVE_PATH="configuration-defined"
	341	+fi
	342	+if command -v sha256sum >/dev/null 2>&1; then
	343	+ CONFIG_SHA256="$(sha256sum "$CONFIG_PATH" \| awk '{print $1}')"
	344	+elif command -v shasum >/dev/null 2>&1; then
	345	+ CONFIG_SHA256="$(shasum -a 256 "$CONFIG_PATH" \| awk '{print $1}')"
	346	+else
	347	+ CONFIG_SHA256="unavailable"
	348	+fi
	349	+MANIFEST="$RUN_ROOT/manifest.txt"
	350	+{
	351	+ echo "repro_mode=$MODE"
	352	+ echo "source_commit=$SOURCE_COMMIT"
	353	+ echo "reference_source_commit=$REFERENCE_SOURCE_COMMIT"
	354	+ echo "entrypoint=$ENTRYPOINT"
	355	+ echo "config_path=$CONFIG_PATH"
	356	+ echo "config_sha256=$CONFIG_SHA256"
	357	+ echo "reference_world_size=$REFERENCE_WORLD_SIZE"
	358	+ echo "requested_world_size=$((NPUS_PER_NODE * NNODES))"
	359	+ echo "npu_per_node=$NPUS_PER_NODE"
	360	+ echo "nnodes=$NNODES"
	361	+ echo "global_batch_size=$GLOBAL_BATCH_SIZE"
	362	+ echo "micro_batch_size=$MBS"
	363	+ echo "gradient_accumulation_steps=$GRAD_ACC_STEP"
	364	+ echo "master_addr=$MASTER_ADDR"
	365	+ echo "master_port=$MASTER_PORT"
	366	+ echo "python_bin=$PYTHON_BIN"
	367	+ echo "run_root=$RUN_ROOT"
	368	+ echo "log_dir=$LOG_DIR"
	369	+ echo "configured_save_path=$SAVE_PATH"
	370	+ if [[ "${TASK3_REPRO_SAVE_CHECKPOINT:-0}" == "1" ]]; then
	371	+ echo "save_path=$SAVE_PATH"
	372	+ else
	373	+ echo "save_path="
	374	+ fi
	375	+ echo "checkpoint_policy=${TASK3_REPRO_SAVE_CHECKPOINT:-0}"
	376	+ echo "skip_zero_norm=${TASK3_REPRO_SKIP_ZERO_NORM:-0}"
	377	+ echo "strict_dcp_model_keys=${TASK3_REPRO_STRICT_DCP:-0}"
	378	+ if [[ "$MODE" == "pure" ]]; then
	379	+ echo "task3_repro_data=$TASK3_REPRO_DATA"
	380	+ fi
	381	+ echo "asset_preflight=$([[ "$MODE" == "pure" ]] && echo pass \|\| echo old-path-checked)"
	382	+ if [[ "$MODE" == "old" ]]; then
	383	+ echo "train_iters=$TRAIN_ITERS"
	384	+ echo "lr_warmup_fraction=$LR_WARMUP_FRACTION"
	385	+ echo "seed=42"
	386	+ else
	387	+ echo "train_iters_override=${TRAIN_ITERS:-config-defined}"
	388	+ echo "lr_warmup_ratio_override=${LR_WARMUP_FRACTION:-config-defined}"
	389	+ echo "recompute_override=${TASK3_REPRO_RECOMPUTE:-config-defined}"
	390	+ echo "prefetch_override=${TASK3_REPRO_PREFETCH:-config-defined}"
	391	+ echo "adam_fused_override=${TASK3_REPRO_ADAM_FUSED:-config-defined}"
	392	+ echo "attn_implementation_override=${TASK3_REPRO_ATTN_IMPLEMENTATION:-config-defined}"
	393	+ echo "reshard_after_forward_override=${TASK3_REPRO_RESHARD_AFTER_FORWARD:-config-defined}"
	394	+ echo "seed=config-defined"
	395	+ fi
	396	+ if (( NPUS_PER_NODE * NNODES == REFERENCE_WORLD_SIZE )); then
	397	+ echo "topology_class=reference"
	398	+ else
	399	+ echo "topology_class=diagnostic-only"
	400	+ fi
	401	+ echo "created_at_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
	402	+} > "$MANIFEST"
	403	+
	404	+export CONFIG_PATH LOG_DIR NPUS_PER_NODE NNODES NODE_RANK MASTER_ADDR MASTER_PORT PYTHON_BIN
	405	+export MBS GRAD_ACC_STEP GLOBAL_BATCH_SIZE
	406	+export TASK3_REPRO_MODE="$MODE"
	407	+export TASK3_REPRO_SAVE_CHECKPOINT="${TASK3_REPRO_SAVE_CHECKPOINT:-0}"
	408	+export TASK3_REPRO_GRAD_ACC_STEP="$GRAD_ACC_STEP"
	409	+export TASK3_REPRO_TRAIN_ITERS="${TRAIN_ITERS:-}"
	410	+export TASK3_REPRO_LR_WARMUP_FRACTION="${LR_WARMUP_FRACTION:-}"
	411	+export TASK3_REPRO_RECOMPUTE="${TASK3_REPRO_RECOMPUTE:-}"
	412	+export TASK3_REPRO_PREFETCH="${TASK3_REPRO_PREFETCH:-}"
	413	+export TASK3_REPRO_ADAM_FUSED="${TASK3_REPRO_ADAM_FUSED:-}"
	414	+export TASK3_REPRO_ATTN_IMPLEMENTATION="${TASK3_REPRO_ATTN_IMPLEMENTATION:-}"
	415	+export TASK3_REPRO_RESHARD_AFTER_FORWARD="${TASK3_REPRO_RESHARD_AFTER_FORWARD:-}"
	416	+if [[ "$MODE" == "pure" ]]; then
	417	+ export TASK3_REPRO_DATA
	418	+fi
	419	+export TASK3_ORIGINAL_ENTRYPOINT="$REPO_ROOT/pretrain_transformers.py"
	420	+export TASK3_DETERMINISTIC_LORA="${TASK3_DETERMINISTIC_LORA:-1}"
	421	+export TRAIN_ENTRYPOINT="$REPO_ROOT/examples/magistral-2509/task3_repro_train.py"
	422	+if [[ "$MODE" == "old" ]]; then
	423	+ export MM_DATA MM_MODEL MM_TOOL LOAD_PATH FSDP2_PATH="$CONFIG_PATH"
	424	+ export SAVE_PATH
	425	+ export TRAIN_ITERS LR_WARMUP_FRACTION
	426	+fi
	427	+
	428	+cd "$REPO_ROOT"
	429	+set +e
	430	+bash "$ENTRYPOINT" 2>&1 \| tee "$RUN_ROOT/launcher.log"
	431	+STATUS=${PIPESTATUS[0]}
	432	+set -e
	433	+echo "exit_code=$STATUS" >> "$MANIFEST"
	434	+echo "manifest=$MANIFEST"
	435	+echo "log=$RUN_ROOT/launcher.log"
	436	+exit "$STATUS"

Aexamples/magistral-2509/task3_repro_data.json+58-0

@@ -0,0 +1,58 @@
	1	+{
	2	+ "dataset_param": {
	3	+ "dataset_type": "huggingface",
	4	+ "preprocess_parameters": {
	5	+ "model_name_or_path": "./.magistral_run/hf",
	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": "./.magistral_run/data",
	20	+ "dataset": "./.magistral_run/data/task3_repro_mllm_format.json",
	21	+ "cache_dir": null,
	22	+ "enable_thinking": false,
	23	+ "overwrite_cache": false,
	24	+ "train_on_prompt": false,
	25	+ "mask_history": false,
	26	+ "preprocessing_batch_size": 2,
	27	+ "preprocessing_num_workers": 1,
	28	+ "max_samples": 4,
	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": false
	57	+ }
	58	+}

Aexamples/magistral-2509/task3_repro_train.py+302-0

@@ -0,0 +1,302 @@
	1	+#!/usr/bin/env python3
	2	+"""Run one Task3 backend through the opt-in post-load alignment hook."""
	3	+
	4	+from __future__ import annotations
	5	+
	6	+import os
	7	+import json
	8	+import math
	9	+import runpy
	10	+from pathlib import Path
	11	+
	12	+
	13	+def dcp_metadata(directory):
	14	+ """Resolve a direct DCP or the repository's release/iteration layout."""
	15	+ direct = directory / ".metadata"
	16	+ if direct.is_file():
	17	+ return direct
	18	+ tracker = directory / "latest_checkpointed_iteration.txt"
	19	+ if tracker.is_file():
	20	+ tag = tracker.read_text().strip()
	21	+ if tag == "release":
	22	+ direct = directory / "release/.metadata"
	23	+ elif tag.isdecimal():
	24	+ direct = directory / f"iter_{int(tag):07d}/.metadata"
	25	+ else:
	26	+ raise ValueError("Invalid DCP tracker; expected release or an iteration number")
	27	+ if direct.is_file():
	28	+ return direct
	29	+ raise ValueError(f"Missing DCP metadata: {directory}")
	30	+
	31	+
	32	+def validate_hf_assets(directory):
	33	+ """Reject incomplete local model files before starting distributed ranks."""
	34	+ index = json.loads((directory / "model.safetensors.index.json").read_text())
	35	+ mapping = index.get("weight_map") if isinstance(index, dict) else None
	36	+ if not isinstance(mapping, dict) or not mapping:
	37	+ raise ValueError("HF weight_map must be a non-empty mapping")
	38	+ for name in mapping.values():
	39	+ if not isinstance(name, str) or not name or Path(name).is_absolute() or ".." in Path(name).parts:
	40	+ raise ValueError("HF shard must be located inside the model directory")
	41	+ for name in set(mapping.values()):
	42	+ if not (directory / name).is_file():
	43	+ raise ValueError(f"Missing HF shard: {name}")
	44	+ for name in ("config.json", "tekken.json"):
	45	+ if not (directory / name).is_file():
	46	+ raise ValueError(f"Missing model asset: {name}")
	47	+
	48	+
	49	+def _positive_int_from_env(name: str) -> int \| None:
	50	+ value = os.environ.get(name, "")
	51	+ if not value:
	52	+ return None
	53	+ try:
	54	+ parsed = int(value)
	55	+ except ValueError as exc:
	56	+ raise SystemExit(f"{name} must be a positive integer, got {value!r}") from exc
	57	+ if parsed < 1:
	58	+ raise SystemExit(f"{name} must be a positive integer, got {value!r}")
	59	+ return parsed
	60	+
	61	+
	62	+def _optional_bool_from_env(name: str) -> bool \| None:
	63	+ value = os.environ.get(name, "").strip().lower()
	64	+ if not value:
	65	+ return None
	66	+ if value in {"1", "true", "yes", "on"}:
	67	+ return True
	68	+ if value in {"0", "false", "no", "off"}:
	69	+ return False
	70	+ raise SystemExit(f"{name} must be a boolean, got {value!r}")
	71	+
	72	+
	73	+def _optional_nonnegative_int_from_env(name: str) -> int \| None:
	74	+ value = os.environ.get(name, "").strip()
	75	+ if not value:
	76	+ return None
	77	+ try:
	78	+ parsed = int(value)
	79	+ except ValueError as exc:
	80	+ raise SystemExit(f"{name} must be a non-negative integer, got {value!r}") from exc
	81	+ if parsed < 0:
	82	+ raise SystemExit(f"{name} must be a non-negative integer, got {value!r}")
	83	+ return parsed
	84	+
	85	+
	86	+def apply_pure_runtime_overrides(arguments: object) -> None:
	87	+ """Apply explicit smoke overrides without creating a second YAML config.
	88	+
	89	+ The canonical pure-FSDP2 YAML remains the source of model and optimizer
	90	+ semantics. These launcher-level values only make a requested smoke use
	91	+ the same step count, warmup, and effective global batch as the old-path
	92	+ launcher on a non-reference topology.
	93	+ """
	94	+
	95	+ training = getattr(arguments, "training", None)
	96	+ if training is None:
	97	+ raise SystemExit("pure-FSDP2 arguments do not contain a training section")
	98	+
	99	+ # The upstream pure YAML keeps its general example dataset. A Task3
	100	+ # comparison must use the same generated four-sample fixture as the old
	101	+ # launcher, so the wrapper changes only this runtime data pointer. Model,
	102	+ # optimizer and FSDP settings still come from the canonical YAML.
	103	+ repro_data = os.environ.get("TASK3_REPRO_DATA", "").strip()
	104	+ if repro_data:
	105	+ data = getattr(arguments, "data", None)
	106	+ dataset_param = getattr(data, "dataset_param", None) if data is not None else None
	107	+ basic_parameters = (
	108	+ getattr(dataset_param, "basic_parameters", None)
	109	+ if dataset_param is not None
	110	+ else None
	111	+ )
	112	+ if basic_parameters is None:
	113	+ raise SystemExit("pure-FSDP2 arguments do not contain data.dataset_param.basic_parameters")
	114	+ data_path = Path(repro_data).expanduser().resolve()
	115	+ basic_parameters.dataset = str(data_path)
	116	+ basic_parameters.dataset_dir = str(data_path.parent)
	117	+ # The ordinary YAML trains on the configured dataset with shuffling.
	118	+ # Only the explicit comparison entry uses the fixed four-sample order.
	119	+ basic_parameters.max_samples = 4
	120	+ data.dataloader_param.shuffle = False
	121	+ print(f"TASK3_REPRO_DATA_OVERRIDE path={data_path}", flush=True)
	122	+
	123	+ train_iters = _positive_int_from_env("TASK3_REPRO_TRAIN_ITERS")
	124	+ if train_iters is not None:
	125	+ training.train_iters = train_iters
	126	+
	127	+ grad_acc = _positive_int_from_env("TASK3_REPRO_GRAD_ACC_STEP")
	128	+ micro_batch = _positive_int_from_env("MBS")
	129	+ if micro_batch is not None:
	130	+ training.micro_batch_size = micro_batch
	131	+ if grad_acc is not None:
	132	+ training.gradient_accumulation_steps = grad_acc
	133	+ data_parallel_size = int(getattr(arguments.parallel, "data_parallel_size", 1))
	134	+ training.global_batch_size = training.micro_batch_size * data_parallel_size * grad_acc
	135	+ requested_batch = _positive_int_from_env("GLOBAL_BATCH_SIZE")
	136	+ if requested_batch is not None and training.global_batch_size != requested_batch:
	137	+ raise SystemExit(
	138	+ f"effective global batch {training.global_batch_size} differs from requested {requested_batch}"
	139	+ )
	140	+
	141	+ warmup = os.environ.get("TASK3_REPRO_LR_WARMUP_FRACTION", "")
	142	+ if warmup:
	143	+ try:
	144	+ training.lr_warmup_ratio = float(warmup)
	145	+ except ValueError as exc:
	146	+ raise SystemExit(
	147	+ f"TASK3_REPRO_LR_WARMUP_FRACTION must be numeric, got {warmup!r}"
	148	+ ) from exc
	149	+ if not math.isfinite(training.lr_warmup_ratio) or not 0 <= training.lr_warmup_ratio <= 1:
	150	+ raise SystemExit("TASK3_REPRO_LR_WARMUP_FRACTION must be between 0 and 1")
	151	+
	152	+ recompute = _optional_bool_from_env("TASK3_REPRO_RECOMPUTE")
	153	+ if recompute is not None:
	154	+ features = getattr(arguments, "features", None)
	155	+ if features is None:
	156	+ raise SystemExit("pure-FSDP2 arguments do not contain a features section")
	157	+ features.recompute = recompute
	158	+
	159	+ prefetch = _optional_nonnegative_int_from_env("TASK3_REPRO_PREFETCH")
	160	+ if prefetch is not None:
	161	+ fsdp_plan = getattr(getattr(arguments, "parallel", None), "fsdp_plan", None)
	162	+ if fsdp_plan is None:
	163	+ raise SystemExit("pure-FSDP2 arguments do not contain a parallel.fsdp_plan section")
	164	+ fsdp_plan.num_to_forward_prefetch = prefetch
	165	+ fsdp_plan.num_to_backward_prefetch = prefetch
	166	+
	167	+ adam_fused = _optional_bool_from_env("TASK3_REPRO_ADAM_FUSED")
	168	+ if adam_fused is not None:
	169	+ training.adam_fused = adam_fused
	170	+
	171	+ attention_impl = os.environ.get("TASK3_REPRO_ATTN_IMPLEMENTATION", "").strip()
	172	+ if attention_impl:
	173	+ if attention_impl not in {"eager", "flash_attention_2"}:
	174	+ raise SystemExit(
	175	+ "TASK3_REPRO_ATTN_IMPLEMENTATION must be eager or flash_attention_2, "
	176	+ f"got {attention_impl!r}"
	177	+ )
	178	+ model = getattr(arguments, "model", None)
	179	+ if model is None or not hasattr(model, "attn_implementation"):
	180	+ raise SystemExit("pure-FSDP2 arguments do not contain model.attn_implementation")
	181	+ model.attn_implementation = attention_impl
	182	+
	183	+ reshard_after_forward = _optional_bool_from_env("TASK3_REPRO_RESHARD_AFTER_FORWARD")
	184	+ if reshard_after_forward is not None:
	185	+ fsdp_plan = getattr(getattr(arguments, "parallel", None), "fsdp_plan", None)
	186	+ if fsdp_plan is None:
	187	+ raise SystemExit("pure-FSDP2 arguments do not contain a parallel.fsdp_plan section")
	188	+ fsdp_plan.reshard_after_forward = reshard_after_forward
	189	+
	190	+ if (
	191	+ train_iters is not None
	192	+ or grad_acc is not None
	193	+ or warmup
	194	+ or recompute is not None
	195	+ or prefetch is not None
	196	+ or adam_fused is not None
	197	+ or attention_impl
	198	+ or reshard_after_forward is not None
	199	+ ):
	200	+ print(
	201	+ "TASK3_REPRO_OVERRIDES "
	202	+ f"train_iters={training.train_iters} "
	203	+ f"gradient_accumulation_steps={training.gradient_accumulation_steps} "
	204	+ f"global_batch_size={training.global_batch_size} "
	205	+ f"lr_warmup_ratio={training.lr_warmup_ratio} "
	206	+ f"recompute={getattr(getattr(arguments, 'features', None), 'recompute', 'config-defined')} "
	207	+ f"prefetch={prefetch if prefetch is not None else 'config-defined'} "
	208	+ f"adam_fused={training.adam_fused} "
	209	+ f"attn_implementation={attention_impl or 'config-defined'} "
	210	+ f"reshard_after_forward={getattr(getattr(arguments.parallel, 'fsdp_plan', None), 'reshard_after_forward', 'config-defined')}",
	211	+ flush=True,
	212	+ )
	213	+
	214	+def main() -> None:
	215	+ """Import accelerator dependencies only when launching training."""
	216	+ mode = os.environ.get("TASK3_REPRO_MODE")
	217	+ if mode not in {"old", "pure"}:
	218	+ raise SystemExit("TASK3_REPRO_MODE must be old or pure")
	219	+ if os.environ.get("TASK3_REPRO_PROFILE", "0") not in {"", "0"}:
	220	+ raise SystemExit("Use the native tools.profile configuration for profiling, outside comparison runs")
	221	+
	222	+ if mode == "old":
	223	+ # Apply MindSpeed's pre-import patches before importing the Megatron
	224	+ # training module. In particular, requirements-basic replaces
	225	+ # torch.compile/torch.jit.script on Ascend; importing the hook first would
	226	+ # decorate fused functions with the unpatched Inductor backend.
	227	+ import mindspeed.megatron_adaptor # noqa: F401
	228	+
	229	+ from mindspeed_mm.fsdp.models.mistral3.repro_hooks import (
	230	+ patch_legacy_mistral3_attention,
	231	+ patch_legacy_mistral3_data,
	232	+ patch_megatron_training,
	233	+ )
	234	+
	235	+ patch_legacy_mistral3_attention()
	236	+ patch_legacy_mistral3_data()
	237	+ patch_megatron_training()
	238	+ target = Path(
	239	+ os.environ.get(
	240	+ "TASK3_ORIGINAL_ENTRYPOINT",
	241	+ str(Path(__file__).resolve().parents[2] / "pretrain_transformers.py"),
	242	+ )
	243	+ )
	244	+ runpy.run_path(str(target), run_name="__main__")
	245	+ else:
	246	+ # Keep the native FSDP2 trainer entry point. Constructing Trainer here
	247	+ # directly looks equivalent, but it changes the module/startup context and
	248	+ # was observed to raise the A3 initialization peak. Patch the already
	249	+ # imported TrainEngine and let the native __main__ block parse the config.
	250	+ from mindspeed_mm.fsdp.config.config_manager import ConfigManager
	251	+ from mindspeed_mm.fsdp.models.mistral3.repro_hooks import patch_pure_trainer
	252	+
	253	+ original_load_and_parse = ConfigManager.load_and_parse
	254	+
	255	+ def load_repro_config(self, args, *kwargs):
	256	+ arguments = original_load_and_parse(self, args, *kwargs)
	257	+ apply_pure_runtime_overrides(arguments)
	258	+ if os.environ.get("TASK3_REPRO_SAVE_CHECKPOINT", "0") != "1":
	259	+ training = getattr(arguments, "training", None)
	260	+ if training is not None:
	261	+ # TrainingArguments validates this field as a string. An
	262	+ # empty path is the project's no-checkpoint value and keeps
	263	+ # repeated comparison runs from writing large artifacts.
	264	+ training.save = ""
	265	+ # The engine guards both periodic and final saves with
	266	+ # training.save. Keep the interval intact: train_iters may
	267	+ # still be None until the trainer resolves train_epochs.
	268	+ # This is the parsed configuration, not proof of applied module hooks.
	269	+ # Record it after overrides and checkpoint policy, before model loading.
	270	+ def section_values(section, names):
	271	+ obj = getattr(arguments, section, None)
	272	+ return {name: getattr(obj, name, None) for name in names}
	273	+
	274	+ plan = getattr(getattr(arguments, "parallel", None), "fsdp_plan", None)
	275	+ print("TASK3_PARSED_CONFIG=" + json.dumps({
	276	+ "rank": os.environ.get("RANK", "0"),
	277	+ "training": section_values("training", (
	278	+ "seed", "train_iters", "micro_batch_size", "global_batch_size",
	279	+ "gradient_accumulation_steps", "lr_warmup_ratio", "adam_fused",
	280	+ "save", "save_interval")),
	281	+ "model": section_values("model", ("dtype", "attn_implementation")),
	282	+ "features": section_values("features", ("recompute",)),
	283	+ "fsdp_plan": {name: getattr(plan, name, None) for name in (
	284	+ "num_to_forward_prefetch", "num_to_backward_prefetch",
	285	+ "reshard_after_forward")},
	286	+ }, sort_keys=True, default=str), flush=True)
	287	+ return arguments
	288	+
	289	+ ConfigManager.load_and_parse = load_repro_config
	290	+
	291	+ patch_pure_trainer()
	292	+ target = Path(
	293	+ os.environ.get(
	294	+ "TASK3_PURE_TRAINER_ENTRYPOINT",
	295	+ str(Path(__file__).resolve().parents[2] / "mindspeed_mm/fsdp/train/trainer.py"),
	296	+ )
	297	+ )
	298	+ runpy.run_path(str(target), run_name="__main__")
	299	+
	300	+
	301	+if __name__ == "__main__":
	302	+ main()

Mmindspeed_mm/fsdp/data/data_utils/func_utils/convert.py+25-14

	@@ -21,7 +21,7 @@ from collections import defaultdict
21	from dataclasses import dataclass, field, asdict	21	from dataclasses import dataclass, field, asdict
22	from itertools import chain	22	from itertools import chain
23	from enum import Enum, unique	23	from enum import Enum, unique
24	-from typing import TYPE_CHECKING, Any, Optional, Union, Tuple, Literal, List, Dict, Type, TypedDict	24	+from typing import TYPE_CHECKING, Any, Optional, Union, Tuple, Literal, List, Dict, Type, TypedDict, Callable
25		25
26	import torch	26	import torch
27	from transformers import PreTrainedTokenizer, ProcessorMixin, AutoProcessor, AutoConfig, AutoTokenizer, PretrainedConfig	27	from transformers import PreTrainedTokenizer, ProcessorMixin, AutoProcessor, AutoConfig, AutoTokenizer, PretrainedConfig
	@@ -1414,17 +1414,34 @@ def get_vision_feature_select_strategy(config: "PretrainedConfig") -> int:
1414	return vision_feature_select_strategy	1414	return vision_feature_select_strategy
1415		1415
1416		1416
		1417	+
		1418	+TOKENIZER_LOADERS: Dict[str, Callable[["ProcessorArguments", "PretrainedConfig"], "TokenizerModule"]] = {}
		1419	+
		1420	+
		1421	+def register_tokenizer_loader(model_type: str, loader: Callable) -> Callable:
		1422	+ """Register an optional model-specific tokenizer/processor loader."""
		1423	+ if model_type in TOKENIZER_LOADERS:
		1424	+ raise ValueError(f"Tokenizer loader for {model_type!r} is already registered.")
		1425	+ TOKENIZER_LOADERS[model_type] = loader
		1426	+ return loader
		1427	+
		1428	+
1417	def load_tokenizer(model_args: "ProcessorArguments") -> "TokenizerModule":	1429	def load_tokenizer(model_args: "ProcessorArguments") -> "TokenizerModule":
1418	r"""	1430	r"""
1419	- Loads pretrained tokenizer and optionally loads processor.	1431	+ Loads the shared tokenizer and optionally loads a processor.
1420		1432
1421	- Note: including inplace operation of model_args.	1433	+ Model-specific tokenizers may register a loader from their FSDP2 model
		1434	+ package. The common path itself remains model agnostic.
1422	"""	1435	"""
1423	config = AutoConfig.from_pretrained(	1436	config = AutoConfig.from_pretrained(
1424	model_args.model_name_or_path,	1437	model_args.model_name_or_path,
1425	local_files_only=True,	1438	local_files_only=True,
1426	trust_remote_code=model_args.trust_remote_code,	1439	trust_remote_code=model_args.trust_remote_code,
1427	)	1440	)
		1441	+ model_loader = TOKENIZER_LOADERS.get(getattr(config, "model_type", None))
		1442	+ if model_loader is not None:
		1443	+ return model_loader(model_args, config)
		1444	+
1428	tokenizer = AutoTokenizer.from_pretrained(	1445	tokenizer = AutoTokenizer.from_pretrained(
1429	model_args.model_name_or_path,	1446	model_args.model_name_or_path,
1430	use_fast=model_args.use_fast_tokenizer,	1447	use_fast=model_args.use_fast_tokenizer,
	@@ -1452,12 +1469,10 @@ def load_tokenizer(model_args: "ProcessorArguments") -> "TokenizerModule":
1452	setattr(processor, "video_maxlen", model_args.video_maxlen)	1469	setattr(processor, "video_maxlen", model_args.video_maxlen)
1453	setattr(processor, "audio_sampling_rate", model_args.audio_sampling_rate)	1470	setattr(processor, "audio_sampling_rate", model_args.audio_sampling_rate)
1454	setattr(processor, "use_audio_in_video", model_args.use_audio_in_video)	1471	setattr(processor, "use_audio_in_video", model_args.use_audio_in_video)
1455	- except Exception as e:	1472	+ except Exception as exc:
1456	- logger.warning("Processor was not found: %s.", e)	1473	+ logger.warning("Processor was not found: %s.", exc)
1457	processor = None	1474	processor = None
1458		1475
1459	- # Avoid load tokenizer, see:
1460	- # https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/auto/processing_auto.py#L324
1461	if processor is not None and "Processor" not in processor.__class__.__name__:	1476	if processor is not None and "Processor" not in processor.__class__.__name__:
1462	processor = None	1477	processor = None
1463		1478
	@@ -1465,16 +1480,12 @@ def load_tokenizer(model_args: "ProcessorArguments") -> "TokenizerModule":
1465		1480
1466		1481
1467	def update_tokenizer_with_chat_template(tokenizer: "PreTrainedTokenizer", chat_template: str) -> "PreTrainedTokenizer":	1482	def update_tokenizer_with_chat_template(tokenizer: "PreTrainedTokenizer", chat_template: str) -> "PreTrainedTokenizer":
1468	- r"""	1483	+ r"""Update a tokenizer with a custom chat template file."""
1469	- Update tokenizer with custom chat_template file.
1470	- """
1471	-
1472	if not os.path.isfile(chat_template):	1484	if not os.path.isfile(chat_template):
1473	raise FileNotFoundError(f"The chat_template path {chat_template} does not exist or is not a file")	1485	raise FileNotFoundError(f"The chat_template path {chat_template} does not exist or is not a file")
1474		1486
1475	- with open(chat_template, 'r', encoding='utf-8') as f:	1487	+ with open(chat_template, "r", encoding="utf-8") as template_file:
1476	- chat_template_content = f.read()	1488	+ tokenizer.chat_template = template_file.read()
1477		1489
1478	logger.info("Apply custom chat_template %s to tokenizer.", chat_template)	1490	logger.info("Apply custom chat_template %s to tokenizer.", chat_template)
1479	- tokenizer.chat_template = chat_template_content
1480	return tokenizer	1491	return tokenizer

Mmindspeed_mm/fsdp/data/data_utils/func_utils/mm_plugin.py+128-5

	@@ -1,6 +1,7 @@
1	# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.	1	# Copyright 2025 HuggingFace Inc. and the LlamaFactory team.
2		2
3	import inspect	3	import inspect
		4	+import importlib.metadata
4	import math	5	import math
5	import os	6	import os
6	from copy import deepcopy	7	from copy import deepcopy
	@@ -14,6 +15,7 @@ import torch
14	from PIL import Image	15	from PIL import Image
15	from PIL.Image import Image as ImageObject	16	from PIL.Image import Image as ImageObject
16	from typing_extensions import override	17	from typing_extensions import override
		18	+from packaging import version
17		19
18	VIDEO_PLACEHOLDER = os.getenv("VIDEO_PLACEHOLDER", "<video>")	20	VIDEO_PLACEHOLDER = os.getenv("VIDEO_PLACEHOLDER", "<video>")
19	IMAGE_PLACEHOLDER = os.getenv("IMAGE_PLACEHOLDER", "<image>")	21	IMAGE_PLACEHOLDER = os.getenv("IMAGE_PLACEHOLDER", "<image>")
	@@ -1561,6 +1563,8 @@ class Step3VLPlugin(BasePlugin):
1561		1563
1562	@dataclass	1564	@dataclass
1563	class MiniMaxM3VLPlugin(BasePlugin):	1565	class MiniMaxM3VLPlugin(BasePlugin):
		1566	+ """MiniMax M3-VL processor retained for backward-compatible templates."""
		1567	+
1564	vision_bos_token: str = "]<]start of image[>["	1568	vision_bos_token: str = "]<]start of image[>["
1565	vision_eos_token: str = "]<]end of image[>["	1569	vision_eos_token: str = "]<]end of image[>["
1566		1570
	@@ -1582,7 +1586,7 @@ class MiniMaxM3VLPlugin(BasePlugin):
1582	image_processor: BaseImageProcessor = getattr(processor, "image_processor", None)	1586	image_processor: BaseImageProcessor = getattr(processor, "image_processor", None)
1583	video_processor = get_video_processor(processor)	1587	video_processor = get_video_processor(processor)
1584	mm_inputs = {}	1588	mm_inputs = {}
1585	- if len(images) != 0:	1589	+ if images:
1586	images = self._regularize_images(	1590	images = self._regularize_images(
1587	images,	1591	images,
1588	image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768),	1592	image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768),
	@@ -1590,8 +1594,7 @@ class MiniMaxM3VLPlugin(BasePlugin):
1590	resample=False,	1594	resample=False,
1591	)["images"]	1595	)["images"]
1592	mm_inputs.update(self._process_images(image_processor, images))	1596	mm_inputs.update(self._process_images(image_processor, images))
1593	-	1597	+ if videos:
1594	- if len(videos) != 0:
1595	videos = self._regularize_videos(	1598	videos = self._regularize_videos(
1596	videos,	1599	videos,
1597	image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256),	1600	image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256),
	@@ -1620,7 +1623,6 @@ class MiniMaxM3VLPlugin(BasePlugin):
1620	image_merge_length = getattr(processor.image_processor, "merge_size", 2) ** 2	1623	image_merge_length = getattr(processor.image_processor, "merge_size", 2) ** 2
1621	video_processor = get_video_processor(processor)	1624	video_processor = get_video_processor(processor)
1622	video_merge_length = getattr(video_processor, "merge_size", 2) ** 2	1625	video_merge_length = getattr(video_processor, "merge_size", 2) ** 2
1623	-
1624	image_idx, video_idx = 0, 0	1626	image_idx, video_idx = 0, 0
1625	for message in messages:	1627	for message in messages:
1626	content = message["content"]	1628	content = message["content"]
	@@ -1632,7 +1634,6 @@ class MiniMaxM3VLPlugin(BasePlugin):
1632	1,	1634	1,
1633	)	1635	)
1634	image_idx += 1	1636	image_idx += 1
1635	-
1636	while VIDEO_PLACEHOLDER in content:	1637	while VIDEO_PLACEHOLDER in content:
1637	video_seqlen = int(video_grid_thw[video_idx].prod() // video_merge_length) if self.expand_mm_tokens else 1	1638	video_seqlen = int(video_grid_thw[video_idx].prod() // video_merge_length) if self.expand_mm_tokens else 1
1638	content = content.replace(	1639	content = content.replace(
	@@ -1645,6 +1646,125 @@ class MiniMaxM3VLPlugin(BasePlugin):
1645	return messages	1646	return messages
1646		1647
1647		1648
		1649	+def _get_package_version(name: str) -> "Version":
		1650	+ try:
		1651	+ return version.parse(importlib.metadata.version(name))
		1652	+ except importlib.metadata.PackageNotFoundError:
		1653	+ return version.parse("0.0.0")
		1654	+ except Exception as e:
		1655	+ print(f"Warning: Failed to get version for package {name!r}: {e}")
		1656	+ return version.parse("0.0.0")
		1657	+
		1658	+
		1659	+def is_transformers_version_greater_than(content: str):
		1660	+ return _get_package_version("transformers") >= version.parse(content)
		1661	+
		1662	+
		1663	+@dataclass
		1664	+class PixtralPlugin(BasePlugin):
		1665	+ """Reusable Pixtral multimodal adapter for image-token expansion.
		1666	+
		1667	+ The model-specific tokenizer and processor construction live under the
		1668	+ Mistral3 model package. This common plugin only translates the generic
		1669	+ processor outputs into the multimodal message contract used by templates.
		1670	+ """
		1671	+
		1672	+ @override
		1673	+ def process_messages(
		1674	+ self,
		1675	+ messages: list[dict[str, str]],
		1676	+ images: list["ImageInput"],
		1677	+ videos: list["VideoInput"],
		1678	+ audios: list["AudioInput"],
		1679	+ processor: Optional["MMProcessor"],
		1680	+ ) -> list[dict[str, str]]:
		1681	+ self._validate_input(processor, images, videos, audios)
		1682	+ self._validate_messages(messages, images, videos, audios)
		1683	+ messages = deepcopy(messages)
		1684	+ image_sizes = iter(())
		1685	+ image_break_token = None
		1686	+ image_end_token = None
		1687	+ if self.expand_mm_tokens:
		1688	+ mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
		1689	+ # BC for transformers < 4.49.0. Some processors return image
		1690	+ # metadata without pixel_values, so do not use pixel_values as the
		1691	+ # guard for values consumed by the text expansion below.
		1692	+ raw_image_sizes = mm_inputs.get("image_sizes")
		1693	+ if raw_image_sizes is None:
		1694	+ if images:
		1695	+ raise ValueError("Pixtral processor did not return image_sizes for the provided images")
		1696	+ raw_image_sizes = []
		1697	+ if isinstance(raw_image_sizes, torch.Tensor):
		1698	+ normalized_image_sizes = raw_image_sizes.reshape(-1, 2).tolist()
		1699	+ else:
		1700	+ normalized_image_sizes = raw_image_sizes
		1701	+ # Older transformers versions may return one extra batch
		1702	+ # dimension. Remove only nested dimensions; retain all image
		1703	+ # pairs so multi-image messages are expanded fully.
		1704	+ while (
		1705	+ normalized_image_sizes
		1706	+ and isinstance(normalized_image_sizes[0], (list, tuple))
		1707	+ and normalized_image_sizes[0]
		1708	+ and isinstance(normalized_image_sizes[0][0], (list, tuple))
		1709	+ ):
		1710	+ normalized_image_sizes = normalized_image_sizes[0]
		1711	+ if (
		1712	+ normalized_image_sizes
		1713	+ and len(normalized_image_sizes) == 2
		1714	+ and isinstance(normalized_image_sizes[0], (int, float))
		1715	+ ):
		1716	+ normalized_image_sizes = [normalized_image_sizes]
		1717	+ image_sizes = iter((int(height), int(width)) for height, width in normalized_image_sizes)
		1718	+ image_break_token = getattr(processor, "image_break_token", None)
		1719	+ image_end_token = getattr(processor, "image_end_token", None)
		1720	+ if images and (image_break_token is None or image_end_token is None):
		1721	+ raise ValueError("Pixtral processor must provide image_break_token and image_end_token")
		1722	+ for message in messages:
		1723	+ content = message["content"]
		1724	+ while IMAGE_PLACEHOLDER in content:
		1725	+ if self.expand_mm_tokens:
		1726	+ patch_size = processor.patch_size * getattr(processor, "spatial_merge_size", 1)
		1727	+ height, width = next(image_sizes)
		1728	+ num_height_tokens = height // patch_size
		1729	+ num_width_tokens = width // patch_size
		1730	+ if num_height_tokens <= 0 or num_width_tokens <= 0:
		1731	+ raise ValueError(
		1732	+ f"Pixtral image size {(height, width)} is smaller than patch size {patch_size}"
		1733	+ )
		1734	+ replace_tokens = [[self.image_token] * num_width_tokens + [image_break_token]] * num_height_tokens
		1735	+ replace_tokens = [item for sublist in replace_tokens for item in sublist] # flatten list
		1736	+ replace_tokens[-1] = image_end_token
		1737	+ replace_str = "".join(replace_tokens)
		1738	+ else:
		1739	+ replace_str = self.image_token
		1740	+ content = content.replace(IMAGE_PLACEHOLDER, replace_str, 1)
		1741	+ message["content"] = content
		1742	+ return messages
		1743	+
		1744	+
		1745	+ @override
		1746	+ def get_mm_inputs(
		1747	+ self,
		1748	+ images: list["ImageInput"],
		1749	+ videos: list["VideoInput"],
		1750	+ audios: list["AudioInput"],
		1751	+ imglens: list[int],
		1752	+ vidlens: list[int],
		1753	+ audlens: list[int],
		1754	+ batch_ids: list[list[int]],
		1755	+ processor: Optional["MMProcessor"],
		1756	+ ) -> dict[str, Union[list[int], "torch.Tensor"]]:
		1757	+ self._validate_input(processor, images, videos, audios)
		1758	+ mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
		1759	+ # ref to this commit https://github.com/huggingface/transformers/pull/35122
		1760	+ # after transformers 4.49.0, the `image_sizes` is mandatory as an input parameter for Pixtral VisionEncoder forwarding.
		1761	+ # it can be passed into `LlavaConditionalGeneration` as a parameter.
		1762	+ if not is_transformers_version_greater_than("4.49.0"):
		1763	+ mm_inputs.pop("image_sizes", None)
		1764	+ return mm_inputs
		1765	+
		1766	+
		1767	+
1648	PLUGINS = {	1768	PLUGINS = {
1649	"base": BasePlugin,	1769	"base": BasePlugin,
1650	"qwen2_vl": Qwen2VLPlugin,	1770	"qwen2_vl": Qwen2VLPlugin,
	@@ -1657,6 +1777,9 @@ PLUGINS = {
1657	"kimi_k3": KimiK3Plugin,	1777	"kimi_k3": KimiK3Plugin,
1658	"step3_vl": Step3VLPlugin,	1778	"step3_vl": Step3VLPlugin,
1659	"minimax_m3_vl": MiniMaxM3VLPlugin,	1779	"minimax_m3_vl": MiniMaxM3VLPlugin,
		1780	+ # Pixtral message expansion is shared by any model using this processor;
		1781	+ # Mistral3-specific tokenizer construction stays in fsdp/models/mistral3.
		1782	+ "pixtral": PixtralPlugin,
1660	}	1783	}
1661		1784
1662		1785

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	+)
		atomgit-bot8月16日 🟠 High Priority 建议：恢复被删除的 minimax_m3_vl 模板注册（可与 mistral 模板并存），确保现有 examples/minimax_m3_vl 配置仍能正常加载。不准确？
		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+5-1

	@@ -20,7 +20,7 @@ class DataCollatorForQwen2vl:
20	tokenizer_module = load_tokenizer(process_args)	20	tokenizer_module = load_tokenizer(process_args)
21	tokenizer = tokenizer_module.get('tokenizer')	21	tokenizer = tokenizer_module.get('tokenizer')
22		22
23	- # chat_template的优先级高于template。如果chat_template和template同时为None，使用模型自带的chat_template	23	+ # Prefer chat_template over template. If both are unset, use the model default.
24	chat_template_path = dataset_param.basic_parameters.chat_template	24	chat_template_path = dataset_param.basic_parameters.chat_template
25	if chat_template_path is not None:	25	if chat_template_path is not None:
26	tokenizer = update_tokenizer_with_chat_template(tokenizer, chat_template_path)	26	tokenizer = update_tokenizer_with_chat_template(tokenizer, chat_template_path)
	@@ -257,6 +257,10 @@ _register_data_collator(
257	"step3_vl",	257	"step3_vl",
258	DataCollatorForStep3VL,	258	DataCollatorForStep3VL,
259	)	259	)
		260	+_register_data_collator(
		261	+ "pixtral",
		262	+ DataCollatorForQwen2vl,
		263	+)
260	_register_data_collator(	264	_register_data_collator(
261	"text_only",	265	"text_only",
262	DataCollatorForTextOnly,	266	DataCollatorForTextOnly,

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

@@ -0,0 +1,29 @@
	1	+from transformers import Mistral3ForConditionalGeneration
	2	+from .modeling_mistral import MMMistralModel
	3	+from .modules import MMMistralAttention
	4	+
	5	+from mindspeed_mm.fsdp.utils.register import model_register
	6	+from . import processing_mistral3 # noqa: F401 # registers the model-specific tokenizer loader
	7	+
	8	+
	9	+@model_register.register("mistral3")
	10	+class PureMistral3ForConditionalGeneration(Mistral3ForConditionalGeneration):
	11	+ """HF-native Mistral3 adapter for the plugin-style FSDP2 trainer."""
	12	+ def __init__(self, config):
	13	+ text_config = getattr(config, "text_config", None)
	14	+ if text_config is None:
	15	+ raise ValueError("Mistral3 FSDP2 requires config.text_config")
	16	+ config._attn_implementation = "eager"
	17	+ text_config._attn_implementation = "flash_attention_2"
	18	+ if getattr(config, "vision_config", None) is not None:
	19	+ config.vision_config._attn_implementation = "eager"
	20	+ super().__init__(config)
	21	+ self.model.language_model = MMMistralModel(config=text_config)
	22	+ for idx, layer in enumerate(self.model.language_model.layers):
	23	+ layer.self_attn = MMMistralAttention(config=text_config, layer_idx=idx)
	24	+
	25	+
	26	+ @classmethod
	27	+ def _from_config(cls, config, **kwargs):
	28	+ config.use_cache = False
	29	+ return cls(config)

Amindspeed_mm/fsdp/models/mistral3/context_parallel.py+75-0

@@ -0,0 +1,75 @@
	1	+"""Context-parallel helpers used by the pure-FSDP2 Mistral3 adapter."""
	2	+
	3	+from typing import Optional
	4	+
	5	+import torch
	6	+from torch import Tensor
	7	+from torch.distributed import ProcessGroup
	8	+
	9	+from megatron.core import mpu
	10	+from mindspeed.core.context_parallel.ulysses_context_parallel.unaligned_cp.mapping import all_to_all
	11	+
	12	+
	13	+_TOTAL_SEQ_LEN: Optional[int] = None
	14	+
	15	+
	16	+def get_safe_context_parallel_world_size() -> int:
	17	+ """Return CP size, treating an uninitialized Megatron group as CP=1."""
	18	+ try:
	19	+ return mpu.get_context_parallel_world_size()
	20	+ except (AssertionError, RuntimeError) as exc:
	21	+ message = str(exc).lower()
	22	+ if any(marker in message for marker in ("not initialized", "not init", "initialize model parallel")):
	23	+ return 1
	24	+ raise
	25	+
	26	+
	27	+def set_seq_len(seq_len: int) -> None:
	28	+ global _TOTAL_SEQ_LEN
	29	+ _TOTAL_SEQ_LEN = seq_len
	30	+
	31	+
	32	+def get_seq_len() -> Optional[int]:
	33	+ return _TOTAL_SEQ_LEN
	34	+
	35	+
	36	+def gather_seq_scatter_heads(
	37	+ input_tensor: Tensor,
	38	+ seq_dim: int,
	39	+ head_dim: int,
	40	+ gather_size: int,
	41	+ group: Optional[ProcessGroup] = None,
	42	+) -> Tensor:
	43	+ group = mpu.get_context_parallel_group() if group is None else group
	44	+ if not group:
	45	+ return input_tensor
	46	+ return all_to_all(input_tensor, group, scatter_dim=head_dim, gather_dim=seq_dim, gather_size=gather_size)
	47	+
	48	+
	49	+def gather_heads_scatter_seq(
	50	+ input_tensor: Tensor,
	51	+ head_dim: int,
	52	+ seq_dim: int,
	53	+ gather_size: int,
	54	+ group: Optional[ProcessGroup] = None,
	55	+) -> Tensor:
	56	+ group = mpu.get_context_parallel_group() if group is None else group
	57	+ if not group:
	58	+ return input_tensor
	59	+ return all_to_all(input_tensor, group, scatter_dim=seq_dim, gather_dim=head_dim, gather_size=gather_size)
	60	+
	61	+
	62	+def gather_seq_scatter_heads_qkv(
	63	+ query: Tensor,
	64	+ key: Tensor,
	65	+ value: Tensor,
	66	+ seq_dim: int,
	67	+ head_dim: int,
	68	+ gather_size: int,
	69	+ group: Optional[ProcessGroup] = None,
	70	+):
	71	+ return (
	72	+ gather_seq_scatter_heads(query, seq_dim, head_dim, gather_size, group),
	73	+ gather_seq_scatter_heads(key, seq_dim, head_dim, gather_size, group),
	74	+ gather_seq_scatter_heads(value, seq_dim, head_dim, gather_size, group),
	75	+ )

Amindspeed_mm/fsdp/models/mistral3/modeling_mistral.py+84-0

@@ -0,0 +1,84 @@
	1	+"""Mistral3 language-model adapter for pure FSDP2."""
	2	+
	3	+from typing import Optional
	4	+
	5	+import torch
	6	+from transformers import Cache, DynamicCache
	7	+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
	8	+from transformers.modeling_outputs import BaseModelOutputWithPast
	9	+from transformers.models.mistral.modeling_mistral import MistralModel
	10	+from transformers.processing_utils import Unpack
	11	+from transformers.utils import TransformersKwargs, auto_docstring
	12	+from transformers.utils.generic import check_model_inputs
	13	+
	14	+from mindspeed_mm.models.common.communications import split_forward_gather_backward_with_cp
	15	+
	16	+from .context_parallel import get_safe_context_parallel_world_size, set_seq_len
	17	+
	18	+
	19	+@auto_docstring
	20	+class MMMistralModel(MistralModel):
	21	+ @check_model_inputs
	22	+ @auto_docstring
	23	+ def forward(
	24	+ self,
	25	+ input_ids: Optional[torch.LongTensor] = None,
	26	+ attention_mask: Optional[torch.Tensor] = None,
	27	+ position_ids: Optional[torch.LongTensor] = None,
	28	+ past_key_values: Optional[Cache] = None,
	29	+ inputs_embeds: Optional[torch.FloatTensor] = None,
	30	+ use_cache: Optional[bool] = None,
	31	+ cache_position: Optional[torch.LongTensor] = None,
	32	+ **kwargs: Unpack[TransformersKwargs],
	33	+ ) -> BaseModelOutputWithPast:
	34	+ if (input_ids is None) ^ (inputs_embeds is not None):
	35	+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
	36	+
	37	+ if inputs_embeds is None:
	38	+ inputs_embeds = self.embed_tokens(input_ids)
	39	+
	40	+ if use_cache and past_key_values is None:
	41	+ past_key_values = DynamicCache(config=self.config)
	42	+
	43	+ if cache_position is None:
	44	+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
	45	+ cache_position = torch.arange(
	46	+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
	47	+ )
	48	+
	49	+ if position_ids is None:
	50	+ position_ids = cache_position.unsqueeze(0)
	51	+
	52	+ mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
	53	+ causal_mask = mask_function(
	54	+ config=self.config,
	55	+ input_embeds=inputs_embeds,
	56	+ attention_mask=attention_mask,
	57	+ cache_position=cache_position,
	58	+ past_key_values=past_key_values,
	59	+ position_ids=position_ids,
	60	+ )
	61	+
	62	+ set_seq_len(inputs_embeds.shape[1])
	63	+ if get_safe_context_parallel_world_size() > 1:
	64	+ inputs_embeds = split_forward_gather_backward_with_cp(inputs_embeds, dim=1)
	65	+ position_ids = split_forward_gather_backward_with_cp(position_ids, dim=1)
	66	+
	67	+ hidden_states = inputs_embeds
	68	+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
	69	+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
	70	+ hidden_states = decoder_layer(
	71	+ hidden_states,
	72	+ attention_mask=causal_mask,
	73	+ position_ids=position_ids,
	74	+ past_key_values=past_key_values,
	75	+ use_cache=use_cache,
	76	+ cache_position=cache_position,
	77	+ position_embeddings=position_embeddings,
	78	+ **kwargs,
	79	+ )
	80	+ hidden_states = self.norm(hidden_states)
	81	+ return BaseModelOutputWithPast(
	82	+ last_hidden_state=hidden_states,
	83	+ past_key_values=past_key_values if use_cache else None,
	84	+ )

Amindspeed_mm/fsdp/models/mistral3/modules.py+80-0

@@ -0,0 +1,80 @@
	1	+"""Mistral3 attention adapter for pure FSDP2."""
	2	+
	3	+from typing import Callable, Optional
	4	+
	5	+import torch
	6	+from transformers import Cache
	7	+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
	8	+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
	9	+from transformers.models.mistral.modeling_mistral import MistralAttention, apply_rotary_pos_emb, eager_attention_forward
	10	+from transformers.processing_utils import Unpack
	11	+from transformers.utils.deprecation import deprecate_kwarg
	12	+
	13	+from .context_parallel import (
	14	+ gather_heads_scatter_seq,
	15	+ gather_seq_scatter_heads_qkv,
	16	+ get_safe_context_parallel_world_size,
	17	+ get_seq_len,
	18	+)
	19	+
	20	+
	21	+class MMMistralAttention(MistralAttention):
	22	+ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
	23	+ def forward(
	24	+ self,
	25	+ hidden_states: torch.Tensor,
	26	+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
	27	+ attention_mask: Optional[torch.Tensor],
	28	+ past_key_values: Optional[Cache] = None,
	29	+ cache_position: Optional[torch.LongTensor] = None,
	30	+ **kwargs: Unpack[FlashAttentionKwargs],
	31	+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
	32	+ input_shape = hidden_states.shape[:-1]
	33	+ hidden_shape = (*input_shape, -1, self.head_dim)
	34	+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
	35	+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
	36	+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
	37	+ cos, sin = position_embeddings
	38	+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
	39	+ if past_key_values is not None:
	40	+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
	41	+ key_states, value_states = past_key_values.update(
	42	+ key_states, value_states, self.layer_idx, cache_kwargs
	43	+ )
	44	+
	45	+ attention_interface: Callable = eager_attention_forward
	46	+ if self.config._attn_implementation != "eager":
	47	+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
	48	+
	49	+ if get_safe_context_parallel_world_size() > 1:
	50	+ query_states, key_states, value_states = gather_seq_scatter_heads_qkv(
	51	+ query_states,
	52	+ key_states,
	53	+ value_states,
	54	+ seq_dim=2,
	55	+ head_dim=1,
	56	+ gather_size=get_seq_len(),
	57	+ )
	58	+
	59	+ attn_output, attn_weights = attention_interface(
	60	+ self,
	61	+ query_states,
	62	+ key_states,
	63	+ value_states,
	64	+ attention_mask,
	65	+ dropout=0.0 if not self.training else self.attention_dropout,
	66	+ scaling=self.scaling,
	67	+ sliding_window=getattr(self.config, "sliding_window", None),
	68	+ **kwargs,
	69	+ )
	70	+
	71	+ if get_safe_context_parallel_world_size() > 1:
	72	+ attn_output = gather_heads_scatter_seq(
	73	+ attn_output,
	74	+ seq_dim=1,
	75	+ head_dim=2,
	76	+ gather_size=self.config.num_attention_heads,
	77	+ )
	78	+
	79	+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
	80	+ return self.o_proj(attn_output), attn_weights

Amindspeed_mm/fsdp/models/mistral3/processing_mistral3.py+107-0

@@ -0,0 +1,107 @@
	1	+"""Mistral3 tokenizer and processor support for the pure FSDP2 model."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import os
	6	+import re
	7	+
	8	+from transformers import AutoProcessor, PixtralImageProcessor
	9	+from transformers.tokenization_mistral_common import MistralCommonTokenizer
	10	+
	11	+from mindspeed_mm.fsdp.data.data_utils.func_utils.convert import (
	12	+ TokenizerModule,
	13	+ register_tokenizer_loader,
	14	+)
	15	+
	16	+
	17	+def load_tokenizer(model_args, model_config) -> TokenizerModule:
	18	+ """Build the Tekken tokenizer and Pixtral processor used by Mistral3."""
	19	+ tokenizer_path = os.path.join(model_args.model_name_or_path, "tekken.json")
	20	+
	21	+ class Mistral3Tokenizer(MistralCommonTokenizer):
	22	+ """Preserve Tekken control tokens embedded in multimodal prompts."""
	23	+
	24	+ def _text_to_ids(self, text, add_special_tokens):
	25	+ control_tokens = self.tokenizer.instruct_tokenizer.tokenizer._special_tokens_reverse_vocab
	26	+ if not control_tokens:
	27	+ return super()._text_to_ids(text, add_special_tokens)
	28	+ pattern = re.compile("(" + "\|".join(
	29	+ re.escape(token) for token in sorted(control_tokens, key=len, reverse=True)
	30	+ ) + ")")
	31	+ ids = []
	32	+ for part in pattern.split(text):
	33	+ if not part:
	34	+ continue
	35	+ if part in control_tokens:
	36	+ ids.append(control_tokens[part])
	37	+ else:
	38	+ ids.extend(self.tokenizer.instruct_tokenizer.tokenizer.encode(part, bos=False, eos=False))
	39	+ if add_special_tokens:
	40	+ ids = [self.bos_token_id, *ids, self.eos_token_id]
	41	+ return ids
	42	+
	43	+ def convert_tokens_to_ids(self, tokens):
	44	+ control_tokens = self.tokenizer.instruct_tokenizer.tokenizer._special_tokens_reverse_vocab
	45	+ if isinstance(tokens, str) and tokens in control_tokens:
	46	+ return control_tokens[tokens]
	47	+ if isinstance(tokens, (list, tuple)):
	48	+ return [
	49	+ control_tokens[token] if token in control_tokens else self._piece_to_id(token)
	50	+ for token in tokens
	51	+ ]
	52	+ return super().convert_tokens_to_ids(tokens)
	53	+
	54	+ if not os.path.isfile(tokenizer_path):
	55	+ raise FileNotFoundError(f"Mistral3 tokenizer file not found: {tokenizer_path}")
	56	+
	57	+ # Native Tekken loads its tokenization rules from tekken.json. The shared
	58	+ # fix_mistral_regex option targets the HF fast-tokenizer implementation;
	59	+ # it does not apply here. MistralCommonTokenizer rejects unsupported kwargs,
	60	+ # including fix_mistral_regex/use_fast/split_special_tokens. Do not forward
	61	+ # them or silently build then filter an apparently effective option set.
	62	+ tokenizer = Mistral3Tokenizer(tokenizer_path, padding_side="right")
	63	+
	64	+ try:
	65	+ processor = AutoProcessor.from_pretrained(
	66	+ model_args.model_name_or_path,
	67	+ use_fast=model_args.use_fast_tokenizer,
	68	+ local_files_only=True,
	69	+ trust_remote_code=getattr(model_args, "trust_remote_code", False),
	70	+ )
	71	+ setattr(processor, "tokenizer", tokenizer)
	72	+ except (ImportError, OSError, ValueError):
	73	+ vision_config = model_config.vision_config.to_dict()
	74	+ patch_size = vision_config.get("patch_size", 14)
	75	+ image_size = vision_config.get("image_size", vision_config.get("max_image_size", 1540))
	76	+ image_processor = PixtralImageProcessor(
	77	+ size={"longest_edge": image_size},
	78	+ patch_size={"height": patch_size, "width": patch_size},
	79	+ )
	80	+
	81	+ class Mistral3Processor:
	82	+ def __init__(self):
	83	+ self.image_processor = image_processor
	84	+ self.tokenizer = tokenizer
	85	+ self.patch_size = patch_size
	86	+ self.spatial_merge_size = getattr(model_config, "spatial_merge_size", 2)
	87	+ self.image_token = "[IMG]"
	88	+ self.image_break_token = "[IMG_BREAK]"
	89	+ self.image_end_token = "[IMG_END]"
	90	+
	91	+ processor = Mistral3Processor()
	92	+
	93	+ setattr(processor, "image_max_pixels", model_args.image_max_pixels)
	94	+ vision_config = model_config.vision_config
	95	+ spatial_merge_size = getattr(model_config, "spatial_merge_size", 2)
	96	+ safe_min_pixels = (getattr(vision_config, "patch_size", 14) * spatial_merge_size * 2) ** 2
	97	+ setattr(processor, "image_min_pixels", max(model_args.image_min_pixels, safe_min_pixels))
	98	+ for name in (
	99	+ "image_do_pan_and_scan", "crop_to_patches", "video_max_pixels", "video_min_pixels",
	100	+ "video_fps", "video_maxlen", "audio_sampling_rate", "use_audio_in_video",
	101	+ ):
	102	+ setattr(processor, name, getattr(model_args, name))
	103	+
	104	+ return {"tokenizer": tokenizer, "processor": processor}
	105	+
	106	+
	107	+register_tokenizer_loader("mistral3", load_tokenizer)

Amindspeed_mm/fsdp/models/mistral3/repro_hooks.py+315-0

@@ -0,0 +1,315 @@
	1	+"""Opt-in Task3 alignment hooks for the checked-in reproduction scripts.
	2	+
	3	+This module is scoped to the FSDP2 Mistral3 model package. It is only imported
	4	+by the explicit Task3 reproduction entry point and is never imported by a
	5	+normal training run.
	6	+"""
	7	+
	8	+from __future__ import annotations
	9	+
	10	+import hashlib
	11	+import math
	12	+import os
	13	+import sys
	14	+from typing import Any
	15	+
	16	+import torch
	17	+
	18	+
	19	+_PURE_PATCHED = False
	20	+_MEGATRON_PATCHED = False
	21	+_LEGACY_MISTRAL3_PATCHED = False
	22	+_LEGACY_MISTRAL3_DATA_PATCHED = False
	23	+_ZERO_NORM_PATCHED = False
	24	+_DCP_CHECK_PATCHED = False
	25	+
	26	+
	27	+def tensor_region(value):
	28	+ """Identity of a live tensor region; never treat meta/empty tensors as aliases."""
	29	+ value = _local_tensor(value)
	30	+ if not isinstance(value, torch.Tensor) or value.is_meta or value.numel() == 0:
	31	+ return None
	32	+ return (str(value.device), str(value.dtype), tuple(value.shape), tuple(value.stride()), value.data_ptr())
	33	+
	34	+
	35	+def check_dcp_model_keys(planner, allow_adapter_initialization=False):
	36	+ """Check the actual flattened load plan, after the backend maps base keys."""
	37	+ expected = {key for key in planner.state_dict if key.startswith("model.")}
	38	+ stored = {key for key in planner.metadata.state_dict_metadata if key.startswith("model.")}
	39	+ if not expected:
	40	+ raise ValueError("Reference DCP load did not expose any model keys")
	41	+ missing, unexpected = sorted(expected - stored), sorted(stored - expected)
	42	+ loaded_regions = {tensor_region(planner.state_dict[key]) for key in expected & stored}
	43	+ loaded_regions.discard(None)
	44	+ aliases = [key for key in missing if tensor_region(planner.state_dict[key]) in loaded_regions]
	45	+ missing = [key for key in missing if key not in aliases]
	46	+ print(f"TASK3_REFERENCE_DCP_STORAGE_ALIASES count={len(aliases)}", flush=True)
	47	+ if allow_adapter_initialization:
	48	+ initialized = [key for key in missing if key.endswith((
	49	+ ".lora_A.default.weight", ".lora_B.default.weight"))]
	50	+ missing = [key for key in missing if key not in initialized]
	51	+ print(f"TASK3_REFERENCE_DCP_INIT_ADAPTERS count={len(initialized)}", flush=True)
	52	+ if missing or unexpected:
	53	+ raise ValueError(f"Reference DCP model keys mismatch: missing={missing[:10]}, unexpected={unexpected[:10]}")
	54	+ return len(expected & stored)
	55	+
	56	+
	57	+def patch_reference_dcp_check():
	58	+ """Only explicit comparison processes opt into strict post-load validation."""
	59	+ global _DCP_CHECK_PATCHED
	60	+ if _DCP_CHECK_PATCHED or os.getenv("TASK3_REPRO_STRICT_DCP", "0") != "1":
	61	+ return
	62	+ import torch.distributed.checkpoint as dcp
	63	+ original = dcp.load
	64	+
	65	+ def load(args, *kwargs):
	66	+ result = original(args, *kwargs)
	67	+ state_dict = kwargs.get("state_dict", args[0] if args else {})
	68	+ if "model" in state_dict:
	69	+ planner = kwargs.get("planner")
	70	+ if planner is None:
	71	+ raise ValueError("Reference model DCP load requires an inspectable planner")
	72	+ count = check_dcp_model_keys(planner, allow_adapter_initialization=(
	73	+ os.getenv("TASK3_DETERMINISTIC_LORA", "0") == "1"))
	74	+ print(f"TASK3_REFERENCE_DCP_KEYS matched={count}", flush=True)
	75	+ return result
	76	+
	77	+ dcp.load = load
	78	+ _DCP_CHECK_PATCHED = True
	79	+
	80	+
	81	+def zero_norm_wrapper(original):
	82	+ """Historical opt-in: omit norm reporting only when clipping is disabled."""
	83	+ def wrapped(model, max_norm, norm_type=2.0, error_if_nonfinite=False, foreach=None, **kwargs):
	84	+ # Upstream may request per-layer norm statistics even without clipping.
	85	+ # Preserve that work, and forward optional arguments to the native API.
	86	+ if max_norm == 0 and not error_if_nonfinite and not any(value is not None for value in kwargs.values()):
	87	+ return None
	88	+ return original(model, max_norm=max_norm, norm_type=norm_type,
	89	+ error_if_nonfinite=error_if_nonfinite, foreach=foreach, **kwargs)
	90	+ return wrapped
	91	+
	92	+
	93	+def _local_tensor(value: Any) -> Any:
	94	+ return value.to_local() if hasattr(value, "to_local") else value
	95	+
	96	+
	97	+def initialize_deterministic_lora(
	98	+ model: Any,
	99	+ label: str,
	100	+ *,
	101	+ normalize_dtype: bool = False,
	102	+) -> int:
	103	+ """Apply the historical Task3 LoRA starting-point contract."""
	104	+
	105	+ seed = int(os.getenv("TASK3_SEED", "42"))
	106	+ modules = model if isinstance(model, (list, tuple)) else [model]
	107	+ changed = 0
	108	+
	109	+ with torch.no_grad():
	110	+ for module in modules:
	111	+ for name, parameter in module.named_parameters():
	112	+ if ".lora_" not in name:
	113	+ continue
	114	+ if normalize_dtype and parameter.dtype != torch.float32:
	115	+ parameter.data = parameter.data.to(dtype=torch.float32)
	116	+ target = _local_tensor(parameter)
	117	+ if ".lora_A." in name:
	118	+ canonical = name.replace("._checkpoint_wrapped_module", "")
	119	+ marker = "language_model.layers."
	120	+ if marker in canonical:
	121	+ canonical = canonical[canonical.index(marker) :]
	122	+ digest = hashlib.sha256(f"{seed}:{canonical}".encode()).digest()
	123	+ local_seed = int.from_bytes(digest[:8], "little") % (2**63 - 1)
	124	+ generator = torch.Generator(device="cpu")
	125	+ generator.manual_seed(local_seed)
	126	+ fan_in = max(1, int(target.shape[1]))
	127	+ bound = 1.0 / math.sqrt(fan_in)
	128	+ value = torch.empty(tuple(target.shape), dtype=torch.float32, device="cpu")
	129	+ value.uniform_(-bound, bound, generator=generator)
	130	+ target.data.copy_(value.to(device=target.device, dtype=target.dtype))
	131	+ changed += 1
	132	+ elif ".lora_B." in name:
	133	+ target.data.zero_()
	134	+ changed += 1
	135	+
	136	+ if torch.distributed.is_available() and torch.distributed.is_initialized():
	137	+ if torch.distributed.get_rank() == 0:
	138	+ print(f"TASK3_DETERMINISTIC_LORA label={label} applied={changed}", flush=True)
	139	+ else:
	140	+ print(f"TASK3_DETERMINISTIC_LORA label={label} applied={changed}", flush=True)
	141	+ return changed
	142	+
	143	+
	144	+def _unwrap_optimizers(optimizer: Any, seen: set[int] \| None = None) -> list[Any]:
	145	+ if optimizer is None:
	146	+ return []
	147	+ seen = set() if seen is None else seen
	148	+ if id(optimizer) in seen:
	149	+ return []
	150	+ seen.add(id(optimizer))
	151	+ result = [optimizer] if hasattr(optimizer, "param_groups") else []
	152	+ mapping = getattr(optimizer, "optimizers_dict", None)
	153	+ if mapping is not None:
	154	+ for child in mapping.values():
	155	+ result.extend(_unwrap_optimizers(child, seen))
	156	+ for attr in ("optimizer", "chained_optimizers"):
	157	+ child = getattr(optimizer, attr, None)
	158	+ if child is None:
	159	+ continue
	160	+ children = child if isinstance(child, (list, tuple)) else [child]
	161	+ for item in children:
	162	+ result.extend(_unwrap_optimizers(item, seen))
	163	+ return result
	164	+
	165	+
	166	+def sync_megatron_lora_master(model: Any, optimizer: Any) -> tuple[int, int]:
	167	+ """Copy deterministic LoRA values into Megatron FP32 master groups."""
	168	+
	169	+ copied = skipped = 0
	170	+ for candidate in _unwrap_optimizers(optimizer):
	171	+ for model_attr, master_attr in (
	172	+ ("model_float16_groups", "shard_fp32_from_float16_groups"),
	173	+ ("float16_groups", "fp32_from_float16_groups"),
	174	+ ):
	175	+ model_groups = getattr(candidate, model_attr, None)
	176	+ master_groups = getattr(candidate, master_attr, None)
	177	+ if model_groups is None or master_groups is None:
	178	+ continue
	179	+ for model_group, master_group in zip(model_groups, master_groups):
	180	+ for model_param, master_param in zip(model_group, master_group):
	181	+ model_value = _local_tensor(model_param).detach()
	182	+ master_value = _local_tensor(master_param)
	183	+ if tuple(model_value.shape) != tuple(master_value.shape):
	184	+ skipped += 1
	185	+ continue
	186	+ master_value.copy_(model_value.to(dtype=master_value.dtype))
	187	+ copied += 1
	188	+ if torch.distributed.is_available() and torch.distributed.is_initialized():
	189	+ if torch.distributed.get_rank() == 0:
	190	+ print(f"TASK3_SYNC_OPTIMIZER_LORA copied={copied} skipped={skipped}", flush=True)
	191	+ return copied, skipped
	192	+
	193	+
	194	+def patch_legacy_mistral3_attention() -> None:
	195	+ """Keep the old bridge's historical attention contract in this wrapper.
	196	+
	197	+ The public bridge model is intentionally unchanged. The comparison
	198	+ launcher applies the historical nested Transformers settings before the
	199	+ bridge constructs the model, so the reproducibility aid cannot affect
	200	+ normal users or the pure-FSDP2 model registry.
	201	+ """
	202	+
	203	+ global _LEGACY_MISTRAL3_PATCHED
	204	+ if _LEGACY_MISTRAL3_PATCHED:
	205	+ return
	206	+
	207	+ from mindspeed_mm.models.transformers.mistral3 import mistral3
	208	+
	209	+ model_cls = mistral3.MultiModelMistral3ForConditionalGeneration
	210	+ original_init = model_cls.__init__
	211	+
	212	+ def wrapped_init(self: Any, config: Any) -> None:
	213	+ config._attn_implementation = "eager"
	214	+ text_config = getattr(config, "text_config", None)
	215	+ if text_config is None:
	216	+ raise ValueError("Mistral3 config must define a non-null text_config")
	217	+ text_config._attn_implementation = "flash_attention_2"
	218	+ vision_config = getattr(config, "vision_config", None)
	219	+ if vision_config is not None:
	220	+ vision_config._attn_implementation = "eager"
	221	+ original_init(self, config)
	222	+
	223	+ model_cls.__init__ = wrapped_init
	224	+ _LEGACY_MISTRAL3_PATCHED = True
	225	+
	226	+
	227	+def patch_legacy_mistral3_data() -> None:
	228	+ """Route only the repro old path to the FSDP2 Mistral3 processor.
	229	+
	230	+ The bridge data module keeps its generic tokenizer loader unchanged. The
	231	+ old-vs-pure comparison needs the same Tekken/Pixtral preprocessing, so the
	232	+ example wrapper installs this dispatch before the bridge imports its data
	233	+ modules. Other model types continue using the original loader.
	234	+ """
	235	+
	236	+ global _LEGACY_MISTRAL3_DATA_PATCHED
	237	+ if _LEGACY_MISTRAL3_DATA_PATCHED:
	238	+ return
	239	+
	240	+ from transformers import AutoConfig
	241	+ from mindspeed_mm.data.data_utils.func_utils import convert as bridge_convert
	242	+ from mindspeed_mm.fsdp.models.mistral3.processing_mistral3 import load_tokenizer as load_mistral3
	243	+
	244	+ original_load_tokenizer = bridge_convert.load_tokenizer
	245	+
	246	+ def load_tokenizer(model_args: Any) -> Any:
	247	+ config = AutoConfig.from_pretrained(
	248	+ model_args.model_name_or_path,
	249	+ trust_remote_code=getattr(model_args, "trust_remote_code", False),
	250	+ local_files_only=True,
	251	+ )
	252	+ if getattr(config, "model_type", None) == "mistral3":
	253	+ return load_mistral3(model_args, config)
	254	+ return original_load_tokenizer(model_args)
	255	+
	256	+ bridge_convert.load_tokenizer = load_tokenizer
	257	+ # Some bridge modules import the function directly during framework
	258	+ # bootstrap, before this wrapper gets control. Refresh those already
	259	+ # imported bindings as well; future imports see the patched public symbol.
	260	+ for module_name in (
	261	+ "mindspeed_mm.data.datasets.qwen2vl_dataset",
	262	+ "mindspeed_mm.data.dataloader.data_collator",
	263	+ ):
	264	+ module = sys.modules.get(module_name)
	265	+ if module is not None:
	266	+ module.load_tokenizer = load_tokenizer
	267	+ _LEGACY_MISTRAL3_DATA_PATCHED = True
	268	+
	269	+
	270	+def patch_pure_trainer() -> None:
	271	+ """Apply the opt-in reset after the native pure trainer loads weights."""
	272	+
	273	+ global _PURE_PATCHED, _ZERO_NORM_PATCHED
	274	+ if os.getenv("TASK3_REPRO_SKIP_ZERO_NORM", "0") == "1" and not _ZERO_NORM_PATCHED:
	275	+ from mindspeed_mm.fsdp.train import train_engine
	276	+ train_engine.clip_grad_norm = zero_norm_wrapper(train_engine.clip_grad_norm)
	277	+ _ZERO_NORM_PATCHED = True
	278	+ print("TASK3_REPRO_SKIP_ZERO_NORM enabled; zero-threshold norm reporting omitted", flush=True)
	279	+ if _PURE_PATCHED or os.getenv("TASK3_DETERMINISTIC_LORA", "0") != "1":
	280	+ return
	281	+ from mindspeed_mm.fsdp.train.train_engine import TrainEngine
	282	+
	283	+ original_init = TrainEngine.__init__
	284	+
	285	+ def wrapped_init(self, args: Any, *kwargs: Any) -> None:
	286	+ original_init(self, args, *kwargs)
	287	+ training = getattr(self.args, "training", None)
	288	+ lora = getattr(training, "lora", None)
	289	+ if getattr(lora, "enable", False):
	290	+ initialize_deterministic_lora(self.model, "pure-post-load")
	291	+
	292	+ TrainEngine.__init__ = wrapped_init
	293	+ _PURE_PATCHED = True
	294	+
	295	+
	296	+def patch_megatron_training() -> None:
	297	+ """Apply the opt-in reset after Megatron model/optimizer setup."""
	298	+
	299	+ global _MEGATRON_PATCHED
	300	+ if _MEGATRON_PATCHED or os.getenv("TASK3_DETERMINISTIC_LORA", "0") != "1":
	301	+ return
	302	+ patch_reference_dcp_check()
	303	+ import mindspeed_mm.training as training
	304	+
	305	+ original_setup = training.setup_model_and_optimizer
	306	+
	307	+ def wrapped_setup(args: Any, *kwargs: Any) -> Any:
	308	+ result = original_setup(args, *kwargs)
	309	+ model, optimizer = result[0], result[1]
	310	+ initialize_deterministic_lora(model, "old-post-load", normalize_dtype=True)
	311	+ sync_megatron_lora_master(model, optimizer)
	312	+ return result
	313	+
	314	+ training.setup_model_and_optimizer = wrapped_setup
	315	+ _MEGATRON_PATCHED = True

Atests/ut_fsdp/models/mistral3/test_reference_entry.py+147-0

@@ -0,0 +1,147 @@
	1	+"""Import and exercise the reference entry, including malformed local assets."""
	2	+import importlib.util
	3	+import json
	4	+import os
	5	+import re
	6	+from pathlib import Path
	7	+import shutil
	8	+import subprocess
	9	+import sys
	10	+import tempfile
	11	+import unittest
	12	+from unittest.mock import patch
	13	+
	14	+ROOT = Path(__file__).resolve().parents[4]
	15	+DIRECTORY = ROOT / "examples/magistral-2509"
	16	+
	17	+def load_reference_entry(directory):
	18	+ """Load the checked-in example at test setup, not during collection."""
	19	+ path = directory / "reproduce_lora.py"
	20	+ if not path.is_file():
	21	+ raise FileNotFoundError(f"Required checked-in reference entry is missing: {path}")
	22	+ spec = importlib.util.spec_from_file_location("reference_entry", path)
	23	+ reference = importlib.util.module_from_spec(spec)
	24	+ spec.loader.exec_module(reference)
	25	+ return reference
	26	+
	27	+
	28	+class ReferenceEntryTests(unittest.TestCase):
	29	+ @classmethod
	30	+ def setUpClass(cls):
	31	+ cls.reference = load_reference_entry(DIRECTORY)
	32	+ spec = importlib.util.spec_from_file_location("repro_assets", DIRECTORY / "task3_repro_train.py")
	33	+ cls.assets = importlib.util.module_from_spec(spec)
	34	+ spec.loader.exec_module(cls.assets)
	35	+
	36	+ def test_real_tensor_alias_is_not_equal_value_clone(self):
	37	+ try:
	38	+ import torch
	39	+ except ImportError:
	40	+ self.skipTest("torch is required for real tensor storage check")
	41	+ hook_spec = importlib.util.spec_from_file_location("reference_hooks_test",
	42	+ ROOT / "mindspeed_mm/fsdp/models/mistral3/repro_hooks.py")
	43	+ hooks = importlib.util.module_from_spec(hook_spec)
	44	+ hook_spec.loader.exec_module(hooks)
	45	+ tensor = torch.ones(4)
	46	+ self.assertEqual(hooks.tensor_region(tensor), hooks.tensor_region(tensor.detach()))
	47	+ self.assertNotEqual(hooks.tensor_region(tensor), hooks.tensor_region(tensor.clone()))
	48	+ self.assertIsNone(hooks.tensor_region(torch.empty(0)))
	49	+
	50	+ def test_short_warmup_is_integral(self):
	51	+ reference = self.reference
	52	+ self.assertEqual(reference.locked_environment(2, 5, {})["LR_WARMUP_FRACTION"], "0")
	53	+ self.assertEqual(reference.locked_environment(4, 2000, {})["LR_WARMUP_FRACTION"], "0.01")
	54	+ with self.assertRaises(ValueError):
	55	+ reference.locked_environment(2, 101, {})
	56	+
	57	+ def test_missing_checked_in_entry_has_explicit_setup_error(self):
	58	+ with tempfile.TemporaryDirectory() as temp:
	59	+ with self.assertRaisesRegex(FileNotFoundError, "Required checked-in reference entry"):
	60	+ load_reference_entry(Path(temp))
	61	+
	62	+ def test_actual_cli_defaults_to_no_training(self):
	63	+ result = subprocess.run([sys.executable, str(DIRECTORY / "reproduce_lora.py"),
	64	+ "pure", "--cards", "2"], cwd="/",
	65	+ env={"PATH": os.environ["PATH"]},
	66	+ text=True, capture_output=True, timeout=20)
	67	+ self.assertEqual(result.returncode, 0, result.stderr)
	68	+ self.assertIn("NO_RUNTIME_STARTED=true", result.stdout)
	69	+ self.assertIn('"gradient_accumulation": 2', result.stdout)
	70	+
	71	+ def test_preprocess_drift_is_rejected(self):
	72	+ with tempfile.TemporaryDirectory() as temp:
	73	+ root = Path(temp)
	74	+ directory = root / "examples/magistral-2509"
	75	+ directory.mkdir(parents=True)
	76	+ for name in ("magistral3_fsdp2_lora_config.yaml", "task3_repro_data.json", "task3_repro_train.py"):
	77	+ shutil.copyfile(DIRECTORY / name, directory / name)
	78	+ hf = root / ".magistral_run/hf"
	79	+ hf.mkdir(parents=True)
	80	+ (hf / "model.safetensors.index.json").write_text('{"weight_map":{"w":"model.safetensors"}}')
	81	+ for name in ("model.safetensors", "config.json", "tekken.json"):
	82	+ (hf / name).write_text("{}")
	83	+ data_path = root / "samples.json"
	84	+ data_path.write_text(json.dumps([{"messages": [{"role": "user", "content": "hi"}]}] * 4))
	85	+ blocks = re.findall(r"<<'PY'\n(.*?)\nPY", (DIRECTORY / "task3_repro.sh").read_text(), flags=re.S)
	86	+ preflight = next(block for block in blocks if "def resolve(" in block)
	87	+ argv = ["preflight", str(directory / "magistral3_fsdp2_lora_config.yaml"), str(root), str(data_path)]
	88	+ with patch.object(sys, "argv", argv):
	89	+ exec(compile(preflight, "asset-preflight", "exec"), {})
	90	+ path = directory / "task3_repro_data.json"
	91	+ original = path.read_text()
	92	+ data = json.loads(path.read_text())
	93	+ data["dataset_param"]["preprocess_parameters"]["image_max_pixels"] = 123
	94	+ path.write_text(json.dumps(data))
	95	+ with patch.object(sys, "argv", argv), self.assertRaisesRegex(SystemExit, "image_max_pixels"):
	96	+ exec(compile(preflight, "asset-preflight", "exec"), {})
	97	+ path.write_text(original)
	98	+ for sample, message in (({}, "messages"), ({"messages": ["hi"], "images": [None]}, "images")):
	99	+ data_path.write_text(json.dumps([sample] * 4))
	100	+ with self.subTest(sample=sample), patch.object(sys, "argv", argv), \
	101	+ self.assertRaisesRegex(SystemExit, message):
	102	+ exec(compile(preflight, "asset-preflight", "exec"), {})
	103	+
	104	+ def test_empty_weight_index_is_rejected(self):
	105	+ with tempfile.TemporaryDirectory() as temp:
	106	+ root = Path(temp)
	107	+ directory = root / "examples/magistral-2509"
	108	+ directory.mkdir(parents=True)
	109	+ hf = root / ".magistral_run/hf"
	110	+ hf.mkdir(parents=True)
	111	+ (hf / "model.safetensors.index.json").write_text('{"weight_map": {}}')
	112	+ with self.assertRaisesRegex(ValueError, "non-empty mapping"):
	113	+ self.assets.validate_hf_assets(hf)
	114	+
	115	+ def test_weight_index_cannot_escape_model_directory(self):
	116	+ with tempfile.TemporaryDirectory() as temp:
	117	+ root = Path(temp)
	118	+ directory = root / "examples/magistral-2509"
	119	+ directory.mkdir(parents=True)
	120	+ hf = root / ".magistral_run/hf"
	121	+ hf.mkdir(parents=True)
	122	+ (hf / "model.safetensors.index.json").write_text(
	123	+ json.dumps({"weight_map": {"weight": "../outside.safetensors"}}))
	124	+ with self.assertRaisesRegex(ValueError, "inside the model directory"):
	125	+ self.assets.validate_hf_assets(hf)
	126	+
	127	+ def test_run_requires_assets_before_subprocess(self):
	128	+ with tempfile.TemporaryDirectory() as temp:
	129	+ result = subprocess.run([sys.executable, str(DIRECTORY / "reproduce_lora.py"),
	130	+ "old", "--run"], cwd="/", text=True, capture_output=True,
	131	+ env={"PATH": os.environ["PATH"], "LOAD_PATH": str(Path(temp) / "missing")},
	132	+ timeout=20)
	133	+ self.assertNotEqual(result.returncode, 0)
	134	+ self.assertIn("required old-path asset", result.stderr)
	135	+
	136	+ def test_unusable_weight_map_is_rejected_before_any_weight_load(self):
	137	+ with tempfile.TemporaryDirectory() as temp:
	138	+ hf = Path(temp)
	139	+ for mapping in ({"weight": []}, {"weight": None}, None):
	140	+ with self.subTest(weight_map=mapping):
	141	+ (hf / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping}))
	142	+ with self.assertRaises(ValueError):
	143	+ self.assets.validate_hf_assets(hf)
	144	+
	145	+
	146	+if __name__ == "__main__":
	147	+ unittest.main()

Atests/ut_fsdp/models/mistral3/test_repro_config.py+230-0

@@ -0,0 +1,230 @@
	1	+"""CPU-only regression checks for the optional Magistral reproduction entry."""
	2	+
	3	+import importlib.util
	4	+import os
	5	+from pathlib import Path
	6	+import subprocess
	7	+import sys
	8	+from types import SimpleNamespace as NS
	9	+import unittest
	10	+from unittest.mock import patch
	11	+from contextlib import redirect_stdout
	12	+import io
	13	+import json
	14	+import ast
	15	+import re
	16	+import tempfile
	17	+
	18	+
	19	+ROOT = Path(__file__).resolve().parents[4]
	20	+ENTRY = ROOT / "examples/magistral-2509/task3_repro_train.py"
	21	+
	22	+
	23	+def arguments():
	24	+ return NS(
	25	+ training=NS(micro_batch_size=1, gradient_accumulation_steps=1,
	26	+ global_batch_size=4, train_iters=2000, lr_warmup_ratio=0.01,
	27	+ adam_fused=True, save="checkpoint", save_interval=100),
	28	+ parallel=NS(data_parallel_size=2, fsdp_plan=NS(
	29	+ num_to_forward_prefetch=2, num_to_backward_prefetch=2,
	30	+ reshard_after_forward=True)),
	31	+ features=NS(recompute=True),
	32	+ model=NS(attn_implementation="flash_attention_2"),
	33	+ data=NS(dataset_param=NS(basic_parameters=NS(dataset="normal.json", dataset_dir="data", max_samples=None)),
	34	+ dataloader_param=NS(shuffle=True)),
	35	+ )
	36	+
	37	+
	38	+class ReproConfigTests(unittest.TestCase):
	39	+ @classmethod
	40	+ def setUpClass(cls):
	41	+ if not ENTRY.is_file():
	42	+ raise FileNotFoundError(f"Required checked-in reproduction entry is missing: {ENTRY}")
	43	+ spec = importlib.util.spec_from_file_location("magistral_repro", ENTRY)
	44	+ cls.repro = importlib.util.module_from_spec(spec)
	45	+ spec.loader.exec_module(cls.repro)
	46	+
	47	+ def launch(self, **values):
	48	+ env = {"PATH": os.environ["PATH"], "PYTHON_BIN": sys.executable,
	49	+ "HOME": os.environ.get("HOME", ""), **values}
	50	+ return subprocess.run(
	51	+ ["bash", str(ENTRY.with_name("task3_repro.sh")), "--dry-run", "pure"],
	52	+ env=env, capture_output=True, text=True, timeout=15,
	53	+ )
	54	+
	55	+ def test_four_card_default(self):
	56	+ result = self.launch()
	57	+ self.assertEqual(result.returncode, 0, result.stderr)
	58	+ self.assertIn("GRAD_ACC_STEP=1", result.stdout)
	59	+
	60	+ def test_multinode_batch_uses_total_world(self):
	61	+ result = self.launch(NPUS_PER_NODE="2", NNODES="2")
	62	+ self.assertEqual(result.returncode, 0, result.stderr)
	63	+ self.assertIn("GRAD_ACC_STEP=1", result.stdout)
	64	+
	65	+ def test_two_card_requires_explicit_diagnostic_opt_in(self):
	66	+ self.assertEqual(self.launch(NPUS_PER_NODE="2").returncode, 2)
	67	+ result = self.launch(NPUS_PER_NODE="2", ALLOW_NON_REFERENCE_TOPOLOGY="1")
	68	+ self.assertEqual(result.returncode, 0, result.stderr)
	69	+ self.assertIn("GRAD_ACC_STEP=2", result.stdout)
	70	+
	71	+ def test_invalid_batch_and_rank_rejected_before_launch(self):
	72	+ for values in ({"MBS": "bad"}, {"MBS": "0"}, {"NNODES": "08"},
	73	+ {"GRAD_ACC_STEP": "2"}, {"NODE_RANK": "1"},
	74	+ {"GLOBAL_BATCH_SIZE": "3"}, {"CP": "2"}):
	75	+ with self.subTest(values=values):
	76	+ self.assertEqual(self.launch(**values).returncode, 2)
	77	+
	78	+ def test_default_runtime_preserves_config(self):
	79	+ repro = self.repro
	80	+ args = arguments()
	81	+ before = repr(args)
	82	+ with patch.dict(os.environ, {}, clear=True):
	83	+ repro.apply_pure_runtime_overrides(args)
	84	+ self.assertEqual(repr(args), before)
	85	+
	86	+ def test_requested_microbatch_reaches_effective_batch(self):
	87	+ repro = self.repro
	88	+ args = arguments()
	89	+ with patch.dict(os.environ, {"MBS": "2", "TASK3_REPRO_GRAD_ACC_STEP": "2",
	90	+ "GLOBAL_BATCH_SIZE": "8"}, clear=True):
	91	+ repro.apply_pure_runtime_overrides(args)
	92	+ self.assertEqual(args.training.micro_batch_size, 2)
	93	+ self.assertEqual(args.training.global_batch_size, 8)
	94	+
	95	+ def test_reference_data_controls_do_not_change_normal_defaults(self):
	96	+ args = arguments()
	97	+ with patch.dict(os.environ, {"TASK3_REPRO_DATA": "/tmp/reference.json"}, clear=True):
	98	+ self.repro.apply_pure_runtime_overrides(args)
	99	+ self.assertEqual(args.data.dataset_param.basic_parameters.max_samples, 4)
	100	+ self.assertFalse(args.data.dataloader_param.shuffle)
	101	+ self.assertIsNone(arguments().data.dataset_param.basic_parameters.max_samples)
	102	+ self.assertTrue(arguments().data.dataloader_param.shuffle)
	103	+
	104	+ def test_runtime_batch_mismatch_rejected(self):
	105	+ repro = self.repro
	106	+ with patch.dict(os.environ, {"TASK3_REPRO_GRAD_ACC_STEP": "1",
	107	+ "GLOBAL_BATCH_SIZE": "4"}, clear=True):
	108	+ with self.assertRaises(SystemExit):
	109	+ repro.apply_pure_runtime_overrides(arguments())
	110	+
	111	+ def test_invalid_warmup_rejected(self):
	112	+ repro = self.repro
	113	+ for value in ("nan", "inf", "-0.1", "1.1", "bad"):
	114	+ with self.subTest(value=value), patch.dict(os.environ, {
	115	+ "TASK3_REPRO_LR_WARMUP_FRACTION": value}, clear=True):
	116	+ with self.assertRaises(SystemExit):
	117	+ repro.apply_pure_runtime_overrides(arguments())
	118	+
	119	+ def test_profiling_cannot_silently_enter_a_comparison(self):
	120	+ with patch.dict(os.environ, {"TASK3_REPRO_MODE": "pure", "TASK3_REPRO_PROFILE": "1"}, clear=True):
	121	+ with self.assertRaisesRegex(SystemExit, "outside comparison runs"):
	122	+ self.repro.main()
	123	+
	124	+ def test_native_entry_receives_final_config_and_save_policy(self):
	125	+ repro = self.repro
	126	+ args = arguments()
	127	+
	128	+ class ConfigManager:
	129	+ def load_and_parse(self):
	130	+ return args
	131	+
	132	+ modules = {
	133	+ "mindspeed_mm.fsdp.config.config_manager": NS(ConfigManager=ConfigManager),
	134	+ "mindspeed_mm.fsdp.models.mistral3.repro_hooks": NS(patch_pure_trainer=lambda: None),
	135	+ }
	136	+ output = io.StringIO()
	137	+ with patch.dict(sys.modules, modules), patch.dict(os.environ, {
	138	+ "TASK3_REPRO_MODE": "pure", "TASK3_REPRO_TRAIN_ITERS": "5"
	139	+ }, clear=True), redirect_stdout(output), patch.object(
	140	+ repro.runpy, "run_path", side_effect=lambda a, *k: ConfigManager().load_and_parse()
	141	+ ) as native:
	142	+ repro.main()
	143	+ native.assert_called_once()
	144	+ self.assertEqual(args.training.save, "")
	145	+ self.assertEqual(args.training.save_interval, 100)
	146	+ line = next(line for line in output.getvalue().splitlines()
	147	+ if line.startswith("TASK3_PARSED_CONFIG="))
	148	+ snapshot = json.loads(line.split("=", 1)[1])
	149	+ self.assertEqual(snapshot["training"]["train_iters"], 5)
	150	+ self.assertEqual(snapshot["fsdp_plan"]["num_to_forward_prefetch"], 2)
	151	+
	152	+ def test_epoch_config_checkpoint_policy_does_not_require_train_iters(self):
	153	+ for enabled in ("0", "1"):
	154	+ with self.subTest(save_enabled=enabled):
	155	+ args = arguments()
	156	+ args.training.train_iters = None
	157	+ args.training.train_epochs = 2
	158	+
	159	+ class ConfigManager:
	160	+ def load_and_parse(self):
	161	+ return args
	162	+
	163	+ modules = {
	164	+ "mindspeed_mm.fsdp.config.config_manager": NS(ConfigManager=ConfigManager),
	165	+ "mindspeed_mm.fsdp.models.mistral3.repro_hooks": NS(patch_pure_trainer=lambda: None),
	166	+ }
	167	+ output = io.StringIO()
	168	+ with patch.dict(sys.modules, modules), patch.dict(os.environ, {
	169	+ "TASK3_REPRO_MODE": "pure", "TASK3_REPRO_SAVE_CHECKPOINT": enabled,
	170	+ }, clear=True), redirect_stdout(output), patch.object(
	171	+ self.repro.runpy, "run_path",
	172	+ side_effect=lambda a, *k: ConfigManager().load_and_parse(),
	173	+ ):
	174	+ self.repro.main()
	175	+ self.assertIsNone(args.training.train_iters)
	176	+ self.assertEqual(args.training.save_interval, 100)
	177	+ self.assertEqual(args.training.save, "checkpoint" if enabled == "1" else "")
	178	+ snapshot = json.loads(next(line.split("=", 1)[1] for line in
	179	+ output.getvalue().splitlines() if line.startswith("TASK3_PARSED_CONFIG=")))
	180	+ self.assertEqual(snapshot["training"]["save"], args.training.save)
	181	+
	182	+ def test_asset_preflight_annotations_are_deferred(self):
	183	+ shell = ENTRY.with_name("task3_repro.sh").read_text()
	184	+ blocks = re.findall(r"<<'PY'\n(.*?)\nPY", shell, flags=re.S)
	185	+ block = next(value for value in blocks if "def resolve(" in value)
	186	+ tree = ast.parse(block)
	187	+ self.assertIsInstance(tree.body[0], ast.ImportFrom)
	188	+ self.assertEqual(tree.body[0].module, "__future__")
	189	+ self.assertEqual(tree.body[0].names[0].name, "annotations")
	190	+ resolve = next(node for node in tree.body if isinstance(node, ast.FunctionDef)
	191	+ and node.name == "resolve")
	192	+ namespace = {"Path": Path, "repo_root": Path("/tmp")}
	193	+ exec(compile(ast.Module(body=[tree.body[0], resolve], type_ignores=[]),
	194	+ "asset-preflight", "exec"), namespace)
	195	+ self.assertIsInstance(namespace["resolve"].__annotations__["return"], str)
	196	+ self.assertEqual(namespace["resolve"]("model"), Path("/tmp/model"))
	197	+
	198	+ def test_old_launcher_preserves_arguments_and_training_exit_code(self):
	199	+ with tempfile.TemporaryDirectory() as temp:
	200	+ directory = Path(temp)
	201	+ interpreter = directory / "record_python"
	202	+ interpreter.write_text(
	203	+ "#!/usr/bin/env python3\nimport json, os, sys\n"
	204	+ "print('ARGS_JSON=' + json.dumps(sys.argv[1:]))\n"
	205	+ "sys.exit(int(os.environ.get('SIMULATED_EXIT', '0')))\n"
	206	+ )
	207	+ interpreter.chmod(0o700)
	208	+ env = {"PATH": os.environ["PATH"], "PYTHON_BIN": str(interpreter),
	209	+ "SKIP_CANN_ENV": "1", "CANN_ENV_SCRIPT": str(directory / "no-cann"),
	210	+ "LOAD_PATH": str(directory), "LOG_DIR": str(directory / "logs"),
	211	+ "SAVE_PATH": "", "TRAIN_ITERS": "5", "LR_WARMUP_FRACTION": "0"}
	212	+ for status in (0, 17):
	213	+ with self.subTest(exit_code=status):
	214	+ result = subprocess.run(
	215	+ ["bash", str(ENTRY.with_name("finetune_magistral_2509_lora.sh"))],
	216	+ env={**env, "SIMULATED_EXIT": str(status)}, capture_output=True, text=True, timeout=15,
	217	+ )
	218	+ self.assertEqual(result.returncode, status, result.stderr)
	219	+ command = json.loads(next(line.split("=", 1)[1] for line in result.stdout.splitlines()
	220	+ if line.startswith("ARGS_JSON=")))
	221	+ for key, value in {"--train-iters": "5", "--lr-warmup-fraction": "0",
	222	+ "--seq-length": "512", "--global-batch-size": "4",
	223	+ "--seed": "42", "--lora-r": "8"}.items():
	224	+ self.assertEqual(command[command.index(key) + 1], value)
	225	+ self.assertEqual(command.count("--ckpt-format"), 1)
	226	+ self.assertIn(f"TRAIN_EXIT_CODE={status}", result.stdout)
	227	+
	228	+
	229	+if __name__ == "__main__":
	230	+ unittest.main()

Atests/ut_fsdp/models/mistral3/test_review_regressions.py+197-0