已开启
docs: add Megatron backend model migration guide #2773
Ruiyu_Qiu创建于 7月4日
docs: add Megatron backend model migration guide #2773
已开启
共 1 个文件变更+562-0
| @@ -0,0 +1,562 @@ | |||
| 1 | +# MindSpeed MM Megatron 后端迁移指南 | ||
| 2 | + | ||
Y | |||
| 3 | +本文面向需要将新模型、新数据集或第三方训练流程接入 MindSpeed MM Megatron 后端的开发者，说明训练入口、模型接入、数据接入、配置与启动脚本、权重转换和运行验证的开发约定。 | ||
| 4 | + | ||
| 5 | +Megatron 后端基于 [MindSpeed](https://gitcode.com/Ascend/MindSpeed) 与 [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) 构建。模型可复用仓内并行组件，也可新增编码器、layer spec 或组装类。文中以 Qwen3-VL 说明组件选择和配置。 | ||
| 6 | + | ||
| 7 | +## 1. 先识别 Megatron 路线 | ||
| 8 | + | ||
| 9 | +仓内提供 Megatron、Megatron-FSDP2 和 FSDP2 三套训练后端。选择入口时，需同时核对模型构建方式和并行配置。 | ||
| 10 | + | ||
| 11 | +| 项目 | Megatron 后端（本文主线） | Megatron-FSDP2 后端 | FSDP2 后端 | | ||
| 12 | +| --- | --- | --- | --- | | ||
| 13 | +| 常见入口 | `pretrain_vlm.py`、`pretrain_sora.py` 等 | `pretrain_transformers.py`，部分其他 `pretrain_*` 入口 | `mindspeed_mm/fsdp/train/trainer.py` 或任务专用入口 | | ||
| 14 | +| 模型构建 | `model_provider` 构建由 Megatron 组件组成的模型 | Transformers 风格模型配合 FSDP2 分片 | 模型插件注册，由 `ModelHub.build` 构建 | | ||
| 15 | +| 并行配置 | TP、PP、CP、VPP 等模型并行参数 | `use_torch_fsdp2` 与额外 `fsdp2_config.yaml` | 顶层 YAML 的 `parallel` 配置 | | ||
| 16 | +| 数据接入 | 数据 provider、`get_batch` 与仓内数据构建函数 | 复用 Megatron 风格数据回调 | 数据插件注册与 FSDP2 数据构建函数 | | ||
| 17 | + | ||
| 18 | +本文以 `pretrain_vlm.py` 的模型接入为主线。视频生成、音频等任务需选择相应入口，核对模型调用和 loss 返回约定。FSDP2 的接入方式见 [FSDP2 迁移指南](./fsdp2_model_migration_guide.md)。 | ||
| 19 | + | ||
| 20 | +Megatron VLM 的调用链如下： | ||
| 21 | + | ||
| 22 | +```text | ||
| 23 | +torchrun → pretrain_vlm.py → mindspeed_mm.training.pretrain | ||
| 24 | + → merge_mm_args 读取 model / data / tools 配置 | ||
| 25 | + → model_provider 构建模型 | ||
| 26 | + → train_valid_test_datasets_provider 构建数据迭代对象 | ||
| 27 | + → forward_step → get_batch → model.forward → loss_func | ||
| 28 | +``` | ||
| 29 | + | ||
| 30 | +## 2. 迁移前准备 | ||
| 31 | + | ||
| 32 | +环境安装参考 [安装指导](../installation/install_guide.md)，模型和数据的准备方法参考对应案例文档。开始开发前，先确认以下信息： | ||
| 33 | + | ||
| 34 | +| 准备项 | 需要确认什么 | | ||
| 35 | +| --- | --- | | ||
| 36 | +| 源仓入口 | 模型构建、权重加载、数据集构建和训练入口的位置 | | ||
| 37 | +| 模型 I/O | `forward` 接收的 batch 字段、输出格式、loss 计算位置 | | ||
| 38 | +| 运行资产 | 本地模型权重、tokenizer/processor、训练样本与模态数据根目录 | | ||
| 39 | +| 并行需求 | TP/PP/CP 等需求，以及视觉、文本等子模块是否需要分别切分 | | ||
| 40 | +| 权重格式 | 能否复用转换器，原始权重与训练模型的名称、布局和精度是否一致 | | ||
| 41 | +| 参考实现 | 可复用的仓内组件、相近案例，以及用于精度对比的已验证实现 | | ||
| 42 | + | ||
| 43 | +优先检查源模型与已有实现的差异。结构一致时通常只需补配置；有新的注意力、位置编码或模态连接方式时，再增加组件或调整组装类。现有案例的后端可能不同，复用前应核对其入口和模型构建路径。 | ||
| 44 | + | ||
| 45 | +例如 Qwen3-VL 已有视觉编码器、layer spec、DeepStack 处理和 Megatron 权重转换器，但 `examples/qwen3vl` 的现成训练配置使用 Megatron-FSDP2 或 FSDP2，迁入 Megatron 后端时需另配模型参数和启动脚本。 | ||
| 46 | + | ||
| 47 | +## 3. 训练入口与核心接口 | ||
| 48 | + | ||
| 49 | +### 3.1 五个核心接口 | ||
| 50 | + | ||
| 51 | +以下五个接口均定义在 [pretrain_vlm.py](../../../../pretrain_vlm.py) 中。`model_provider`、数据 provider 和 `forward_step` 传给共享训练循环 `pretrain`；`get_batch` 由 `forward_step` 调用，`loss_func` 由 `forward_step` 返回给训练框架。 | ||
| 52 | + | ||
| 53 | +| 接口 | 作用 | | ||
| 54 | +| --- | --- | | ||
| 55 | +| `model_provider` | 构建模型，处理流水线首尾 stage 标志与参数冻结 | | ||
| 56 | +| `train_valid_test_datasets_provider` | 构建并返回训练、验证、测试所需的数据迭代对象 | | ||
| 57 | +| `get_batch` | 取 batch、搬运设备并适配字段 | | ||
| 58 | +| `forward_step` | 调用模型，返回 `(output_tensor, loss_func)` | | ||
| 59 | +| `loss_func` | 按约定处理训练 loss、有效 token 统计和日志聚合 | | ||
| 60 | + | ||
| 61 | +入口末尾的 `pretrain` 调用见 [pretrain_vlm.py](../../../../pretrain_vlm.py) 的 `__main__` 部分。新入口可参考下面的骨架，按任务实现五个接口，补充 batch 搬运、模型调用和 PP 分支： | ||
| 62 | + | ||
| 63 | +```python | ||
| 64 | +def forward_step(data_iterator, model): | ||
| 65 | + batch = get_batch(data_iterator) | ||
| 66 | + output_tensor = model(**batch) | ||
| 67 | + return output_tensor, loss_func | ||
| 68 | + | ||
| 69 | + | ||
| 70 | +if __name__ == "__main__": | ||
| 71 | + train_valid_test_datasets_provider.is_distributed = True | ||
| 72 | + pretrain( | ||
| 73 | + train_valid_test_datasets_provider, | ||
| 74 | + model_provider, | ||
| 75 | + ModelType.encoder_or_decoder, | ||
| 76 | + forward_step, | ||
| 77 | + extra_args_provider=mm_extra_args_provider, | ||
| 78 | + args_defaults={"dataloader_type": "external"}, | ||
| 79 | + ) | ||
| 80 | +``` | ||
| 81 | + | ||
| 82 | +### 3.2 模型构造 | ||
| 83 | + | ||
| 84 | +模型构造由 [pretrain_vlm.py](../../../../pretrain_vlm.py) 中的 `model_provider` 完成： | ||
| 85 | + | ||
| 86 | +```text | ||
| 87 | +args.mm.model → deepcopy → 设置 pre_process / post_process | ||
| 88 | + → _configure_modules：通过 get_model_config 转换各子模块配置 | ||
| 89 | + → VLMModel（或新组装类）→ _apply_freezing | ||
| 90 | +``` | ||
| 91 | + | ||
| 92 | +`_configure_modules` 分别处理视觉、音频和文本模块的配置，再调用 [transformer_model_config.py](../../../../mindspeed_mm/utils/transformer_model_config.py) 中的 `get_model_config`。该函数用子模块字段和全局并行参数构造 `TransformerConfig`，将校验后的字段写回 `ConfigReader`。 | ||
| 93 | + | ||
| 94 | +[vlm_model.py](../../../../mindspeed_mm/models/vlm_model.py) 中的 `VLMModel.__init__` 接收整体模型配置，按各子模块配置是否非空创建编码器和解码器。模块实现由子模块的 `model_id` 选择，登记位置见 4.2。 | ||
| 95 | + | ||
| 96 | +构造完成后，`_apply_freezing` 调用模型的 `freeze`。训练时，`forward_step` 还会读取 `add_image_encoder` 和 `image_encoder.post_process`。更换组装类时，需保留这些属性和接口，或同步修改入口。 | ||
| 97 | + | ||
| 98 | +### 3.3 模型输出与 loss | ||
| 99 | + | ||
| 100 | +复用 `pretrain_vlm.py` 时，模型返回包含 `loss_dict` 的字典。具体要求随 loss 配置变化： | ||
修正文档路径。 PR 描述中写的是 docs/zh/features/megatron_developer_migration_guide.md，但本次实际新增文件为 docs/zh/guides/development/megatron_model_migration_guide.md，请统一描述中的路径。