MindSpeed MM模型迁移指南

核心接口适配

MindSpeed MM训练逻辑整体沿用Megatron风格，各类模型使用统一的pretrain_xxx.py文件启动，通过配置文件对训练过程进行调控。examples目录下提供各类模型启动的训练脚本。

pretrain_xxx.py中提供了模型训练过程中涉及的主要（回调）函数，用户需要根据具体任务需求，在这些函数中实现相应的业务逻辑。这些函数将在训练过程中被自动调用：

​​函数	描述
model_provider	构建模型
get_batch	构建模型的前向输入数据
loss_func	计算模型前向损失
forward_step	模型前向，并计算损失
train_valid_test_datasets_provider	构造数据加载器

核心接口调用流程

sequenceDiagram
    participant M as pretrain_xxx
    participant E as Training
    participant D as Data
    participant MP as ModelProvider
    participant LC as LossCalculator
    participant Model

    M->>E: 初始化训练
    activate E

    E->>MP: 调用model_provider构建模型
    activate MP
    MP->>Model: 创建模型实例
    MP-->>E: 返回模型
    deactivate MP

    E->>D: 调用train_valid_test_datasets_provider
    activate D
    D-->>E: 返回DataLoader
    deactivate D



    loop 每个训练步骤
        E->>D: 调用get_batch获取数据
        activate D
        D-->>E: 返回batch数据
        deactivate D

        E->>Model: 调用forward_step前向计算
        activate Model
        Model->>LC: 调用loss_func计算损失
        activate LC
        LC-->>E: 返回损失值
        deactivate LC
        deactivate Model

        E->>E: 反向传播与参数更新
        Note over E: 梯度计算、优化器步进<br/>学习率调整
    end

    E-->>M: 训练完成

模型迁移

整体迁移的要点围绕上述核心接口展开。

环境搭建

【模型开发时推荐使用配套的环境版本】

请参考安装指南，完成昇腾软件安装。

数据模块迁移

迁移基本原则

MindSpeed MM 框架对数据模块提供了高度兼容性，用户现有的 DataSet 和 DataLoader 实现可以无缝迁移。 快速迁移步骤：

1. ​保持现有实现​：用户自定义的 DataSet、DataLoader 等数据模块无需修改
2. ​注册数据提供函数​：通过入口脚本的 train_valid_test_datasets_provider 函数返回DataLoader数据对象
3. ​获取训练批次​：在 get_batch 函数中处理当前训练批次数据

迁移示例：

def train_valid_test_datasets_provider():
    """Build train, valid, and test datasets."""
    # 返回您现有的数据集实例
    train_dataloader= build_dataloader(CustomDataset(...))
    valid_dataloader= build_dataloader(CustomDataset(...))
    test_dataloader= build_dataloader(CustomDataset(...))

    return train_dataloader, valid_dataloader, test_dataloader

def get_batch(data_iterator, args):
    """Generate a batch."""
    if data_iterator is not None:
        batch = next(data_iterator)
    else:
        raise ValueError("Data iterator is None. Unable to retrieve batch.")
    move_to_device(batch, get_args().params_dtype)

    return batch

原生数据组件使用

MindSpeed MM 同时提供了一套优化的多模态数据集处理模块，包括build_mm_dataset、build_mm_dataloader等，用户可通过 data.json 配置文件灵活定义数据源、预处理流程及加载策略。

模型结构迁移

在 MindSpeed MM 框架中，所有训练模型都通过标准化的入口函数进行构建和执行：model_provider构造模型，forward_step执行前向结果，loss_func计算训练损失。

def model_provider(*args, **kwargs):
    model = CustomModel(config)

    return model

def loss_func(output_tensor):
   # 基于模型输出计算损失
   loss = compute_loss(output_tensor)
   return loss

def forward_step(data_iterator, model):
    """Forward step."""
    batch_data = get_batch(data_iterator)
    output_tensor = model(**batch_data)
    return output_tensor, loss_func

FSDP2 训练必备条件

用户自定义模型

对于用户自行开发的模型，使用 FSDP2 进行分布式训练只需简单继承指定的 Mixin 类，​模型本体的实现无需任何修改​：

from mindspeed_mm.models.common.module import MultiModalModule
from mindspeed_mm.models.transformers.base_model import FSDP2Mixin, WeightInitMixin

class CustomModel(MultiModalModule, FSDP2Mixin, WeightInitMixin):
 """您的自定义模型类"""
    ...

三方库模型适配

对于从第三方库（如 Transformers）导入的模型，MindSpeed MM 提供了统一的适配方案：

1. ​创建适配类​：将第三方模型继承 FSDP2 Mixin 类
2. ​注册到模型库​：将适配后的类添加到 ModelHub
3. ​通过标准接口使用：使用 TransformersModel 加载适配后的模型

下面以Internvl-3.5为例

步骤1：创建适配类

# mindspeed_mm/models/transformers/internvl3.5.py
from internvl.modeling_internvl_chat import InternVLChatModel
from mindspeed_mm.models.transformers.base_model import FSDP2Mixin, WeightInitMixin

class InternVLChatModelGeneration(InternVLChatModel, FSDP2Mixin, WeightInitMixin):
 """InternVL 模型的 FSDP2 适配版本"""
    def __init__(self, config, vision_model=None, language_model=None, use_flash_attn=True):
        super().__init__(config)

步骤 2：注册到模型库

# mindspeed_mm/models/common/modelzoo.py

from mindspeed_mm.models.transformers.internvl3_5 import InternVLChatModelGeneration

class ModelHub:
    MODEL_MAPPINGS = {
        'internvl': InternVLChatModelGeneration,
    }

步骤 3：通过标准接口使用

# pretrain_transformers.py
from mindspeed_mm.models.transformers_model import TransformersModel

def model_provider(*args, **kwargs):
 """Builds the model."""
 args = get_args()
 vlm_config = deepcopy(args.mm.model)
 model = TransformersModel(vlm_config)

 return model

FSDP2配置文件驱动训练（推荐）

MindSpeed MM 支持通过 YAML 配置文件灵活管理 FSDP2 训练策略，实现训练配置与模型代码的完全解耦。

配置文件参数说明

参数分类	参数名称	描述	取值	默认值	注意事项
基本配置	sharding_size	模型并行分片大小	auto或整数值	1	auto表示world_size大小
	param_dtype	参数存储和计算数据类型	bf16, fp16, fp32	模型dtype	训练精度设置
	reduce_dtype	梯度通信数据类型	-	-	通信精度设置
	output_dtype	前向输出数据类型	-	-	输出精度控制
	cast_forward_inputs	前向输入自动类型转换	true/false	-	确保输入类型匹配
模块包装	sub_modules_to_wrap	FSDP分片子模块路径	模块路径字符串列表	-	模式语法: • model.layers.{*}: 匹配所有子模块 • model.layers.{0-23}: 匹配层数范围 • model.layers.{1,3,5}: 匹配指定层数
	ignored_modules	排除FSDP管理的模块	模块路径字符串列表	-	格式同sub_modules_to_wrap
内存优化	recompute_modules	激活值重计算模块	模块路径字符串列表	-	格式同sub_modules_to_wrap 冲突避免: 需关闭Megatron重计算功能
	use_reentrant	检查点实现类型	true/false	true	是否可重入
	reshard_after_forward	参数重新聚合时机	true/false	-	true: ZeRO3(省内存) false: ZeRO2(高性能)
	offload_to_cpu	参数卸载到CPU	true/false	false	启用时需要设置--distributed-backend npu:hccl,cpu:gloo
	pin_memory	锁定CPU内存	true/false	false	仅offload_to_cpu=true时生效
性能调优	num_to_forward_prefetch	前向预取层数	整数值	-	通信与计算重叠优化
	num_to_backward_prefetch	反向预取层数	整数值	1	通信与计算重叠优化

配置示例如下，fsdp2_config.yaml

sharding_size: auto
sub_modules_to_wrap:
  - "text_decoder.output_layer"
  - "text_decoder.embedding"
  - "text_decoder.rotary_pos_emb"
  - "text_decoder.decoder.layers.{*}"
param_dtype: "bf16"
reduce_dtype: "fp32"
cast_forward_inputs: True
ignored_modules:
  - "image_encoder"
recompute_modules:
  - "text_decoder.decoder.layers.{*}"
num_to_forward_prefetch: 2
num_to_backward_prefetch: 2
offload_to_cpu: False

自定义切分策略（可选）

针对模型结构复杂或无法通过 YAML 配置满足需求的场景，用户可通过 FSDP2Mixin 提供的接口实现自定义切分策略。此时 YAML 配置文件中仅需提供基础配置即可。

自定义fully_shard实现示例

from mindspeed_mm.models.common.module import MultiModalModule
from mindspeed_mm.models.transformers.base_model import FSDP2Mixin, WeightInitMixin

class YourModel(MultiModalModule, FSDP2Mixin, WeightInitMixin):
    def _fully_shard(self, fsdp2_kwargs=None, fsdp2_config=None):
        """
        自定义fully_shard实现
        """
        # 自定义重计算模块（可选）
        set_recompute_modules_to_wrap()

        # 自定义fully_shard包装模块
        set_fullyshard_modules_to_wrap()

        # 自定义预取策略（可选）
        num_to_forward_prefetch = getattr(self.fsdp2_config, "num_to_forward_prefetch", 0)
        num_to_backward_prefetch = getattr(self.fsdp2_config, "num_to_backward_prefetch", 0)
        set_modules_to_prefetch(num_to_forward_prefetch, num_to_backward_prefetch)

如需更深入的定制，请参考 FSDP2Mixin 类的完整实现，了解各个方法的详细用法和扩展点。

启动命令配置

启动 FSDP2 训练时，需要在标准 Megatron 训练命令的基础上添加以下参数

export CUDA_DEVICE_MAX_CONNECTIONS=2

--use-torch-fsdp2 \
--fsdp2-config-path ./fsdp2_config.yaml \
--ckpt-format torch_dcp \
--untie-embeddings-and-output-weights \

关键参数说明：

• --use-torch-fsdp2：启用FSDP2训练模式
• --fsdp2-config-path：指定FSDP2配置文件路径