Megatron-LM 支持

Megatron-LM 支持#

本指南将清晰地引导你如何使用 Megatron-LM 来训练模型。

备注

本指南假设你已经按照 安装指南 中的源码安装方式配置好了环境。如果还没有，请先参考该指南。

步骤 1：安装 Megatron-LM 支持#

安装 Megatron-LM 相关依赖：

pip install -e ".[megatron]"

# for uv
# uv sync -extra megatron

另外还需要从源码安装 NVIDIA 的 Apex 库以支持混合精度训练：

pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
    --config-settings "--build-option=--cpp_ext" \
    --config-settings "--build-option=--cuda_ext" \
    --resume-retries 10 git+https://github.com/NVIDIA/apex.git

# for uv
# uv pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
#    --config-settings "--build-option=--cpp_ext" \
#    --config-settings "--build-option=--cuda_ext" \
#    --resume-retries 10 git+https://github.com/NVIDIA/apex.git

替代方案：使用 Docker#

我们提供了 Docker 配置以简化环境管理。

构建 Docker 镜像#

Trinity-RFT 提供了专门用于 Megatron-LM 的 Dockerfile，位于 scripts/docker_for_megatron/Dockerfile。 可以使用以下命令构建镜像：

docker build -f scripts/docker_for_megatron/Dockerfile -t trinity-rft-megatron:latest .

💡 你可以在构建前自定义 Dockerfile —— 例如添加 pip 镜像源或设置 API 密钥。

运行容器#

docker run -it \
    --gpus all \
    --shm-size="64g" \
    --rm \
    -v $PWD:/workspace \
    -v <your_data_and_checkpoints_path>:/data \
    trinity-rft-megatron:latest

请将 <your_data_and_checkpoints_path> 替换为你机器上存储数据集和模型检查点的实际路径。

步骤 2：配置并运行训练#

大多数配置设置已在 快速入门指南 中涵盖。此处我们仅关注 Megatron-LM 特有 的配置。

Megatron 配置示例#

以下是将 actor、reference model 和 critic 配置为使用 Megatron-LM 的示例：

actor_rollout_ref:
  ...
  actor:
    strategy: megatron  # 为保持向后兼容性保留
    megatron:
      # 模型并行设置
      tensor_model_parallel_size: 2
      pipeline_model_parallel_size: 1
      expert_model_parallel_size: 1

      # offload 设置（除非内存受限，否则设为 false）
      param_offload: false
      grad_offload: false
      optimizer_offload: false

      # 使用 mBridge 进行参数导入/导出（可选）
      use_mbridge: false

      # 使用 Megatron checkpointing
      use_dist_checkpointing: false
      dist_checkpointing_path: null

      # 重计算设置（有助于训练期间节省内存）
      override_transformer_config:
        recompute_granularity: full
        recompute_method: uniform
        recompute_num_layers: 1
  ...
  ref:
    megatron:
      tensor_model_parallel_size: 2
      pipeline_model_parallel_size: 1
      expert_model_parallel_size: 1
      param_offload: false
      grad_offload: false
      optimizer_offload: false
      use_mbridge: false
      use_dist_checkpointing: false
      dist_checkpointing_path: null
      override_transformer_config:
        recompute_granularity: full
        recompute_method: uniform
        recompute_num_layers: 1
  ...

critic:
  strategy: megatron
  megatron:
    tensor_model_parallel_size: 2
    pipeline_model_parallel_size: 1
    expert_model_parallel_size: 1
    param_offload: false
    grad_offload: false
    optimizer_offload: false
    use_mbridge: false
    use_dist_checkpointing: false
    dist_checkpointing_path: null
    override_transformer_config:
      recompute_granularity: full
      recompute_method: uniform
      recompute_num_layers: 1
  ...

训练 Mixture-of-Experts (MoE) 模型#

如果你正在训练像 Qwen/Qwen3-30B-A3B 这样的 MoE 模型，则需要采用以下两种方法之一，以确保其正常工作：

  1. 使用 MBridge（推荐）： 只需在配置文件中设置 use_mbridge: true。这将直接启用对 MoE 模型所需的支持。

  2. 手动转换模型： 如果你不希望使用 MBridge，请设置 use_mbridge: false。在训练前，你必须先使用 verl 仓库中的 Hugging Face 到 MCore 转换器 将 Hugging Face 模型转换为 MCore 格式。转换完成后，在配置中更新：

    • use_dist_checkpointing: true

    • dist_checkpointing_path: /PATH/TO/CONVERTED/MODEL/

⚠️ 重要提示：如果跳过上述任一方法，MoE 模型可能无法正确加载或训练。请务必选择以上两种方式之一。

目录