在 AIStudio 训练任务中使用 Megatron

Megatron 提供张量并行、流水线并行、数据并行、上下文并行和专家并行等大模型训练能力。本页面向已经有 Megatron 训练代码或准备使用 NVIDIA 官方示例的机器学习工程师，说明怎样把上游 Megatron 映射到 AIStudio 的镜像、GPU Worker、共享存储和分布式启动方式。

本页面使用一个两 GPU、张量并行度为 2 的小模型完成训练、分布式 checkpoint 保存和重新加载。这个有界场景用于确认 Megatron 与 AIStudio 的运行链路；真实模型的架构、数据、并行策略和训练周期仍由您的训练目标决定。

判断 Megatron 是否适合当前任务

以下场景适合使用 Megatron：

从头预训练、继续预训练或全参数训练较大的 Transformer 模型；
需要组合张量并行、流水线并行、数据并行、上下文并行或专家并行；
已经有 Megatron-LM 启动脚本，希望迁移到 AIStudio 的单 Worker 或多 Worker 训练任务；
正在开发基于 Megatron Core 的训练框架或自定义训练循环。

如果目标是用 LoRA 或 QLoRA 对中小模型做监督微调，优先使用 LLaMA Factory。如果目标是 GRPO 或 PPO 等强化学习后训练，可以使用 verl。

Megatron-Infinigence 是针对特定大规模 MoE 模型和平台优化组合提供的另一条训练路径。需要使用该工具时，请阅读 Megatron-Infinigence 训练框架概述，并为两条路径分别维护镜像、代码目录和训练参数。

选择 Megatron 组件

NVIDIA 的 Megatron 仓库包含以下相互关联的组件：

Megatron-LM：基于 Megatron Core 的参考训练应用和预配置脚本，适合验证训练方案、研究并行策略或从官方示例开始修改。
Megatron Core：提供模型组件、并行策略和分布式 checkpoint 等底层能力，适合维护自定义训练框架或训练循环。
Megatron Bridge：提供 Hugging Face 与 Megatron checkpoint 的双向转换，以及受维护的模型训练配置。需要从已有 Hugging Face 权重继续训练，或把训练结果导回 Hugging Face 格式时使用。

本文的有界场景直接运行 Megatron Core 官方示例。它从随机权重开始，使用 Mock 数据，不需要下载模型或数据集；完成这个场景可以确认平台运行链路，加载或微调已有模型时再替换模型和数据配置。

在目标可用区准备容器镜像

在能够访问源代码、源 Registry 和目标 Registry 的本地设备或 CI Runner 上准备镜像，再把完成的镜像推送到训练资源所在的可用区。从下面两条路径中选择一条执行；两条路径都能生成满足本场景运行要求的目标镜像，不需要重复执行。

路径一：使用 regctl 复制并创建派生镜像

这条路径适合基础镜像较大、传输容易中断，或不希望在构建设备中保存完整基础镜像的场景。regctl 先把基础镜像复制到目标 Registry，再复用其中的镜像层创建一个新的 Megatron 镜像。

按照使用 regctl 将远程镜像复制到目标可用区 Registry安装 regctl 并登录目标 Registry。如果 NGC 要求认证，按照 NVIDIA NGC 登录说明使用 NGC API Key 登录：
language-bash
```
read -rsp 'NGC API key: ' NGC_API_KEY
printf '%s' "$NGC_API_KEY" |
  regctl registry login nvcr.io \
    --user '$oauthtoken' \
    --pass-stdin
unset NGC_API_KEY
```

定义源镜像、目标镜像和 Megatron 提交，把 linux/amd64 基础镜像复制到目标可用区：

language-bash

export PLATFORM="linux/amd64"
export DST_REGISTRY="<目标可用区 Registry 公网访问地址>"
export DST_NAMESPACE="<租户 ID>"
export SRC_BASE_IMAGE="nvcr.io/nvidia/pytorch:25.03-py3"
export BASE_IMAGE="$DST_REGISTRY/$DST_NAMESPACE/nvidia-pytorch:25.03-py3"
export MEGATRON_IMAGE="$DST_REGISTRY/$DST_NAMESPACE/megatron-core:0.18.2-cu128"
export MEGATRON_COMMIT="571370c829ca768fe37244f4e2e7f28d8accc4ab"

SRC_BASE_DIGEST="$(
  regctl image digest \
    --platform "$PLATFORM" \
    "$SRC_BASE_IMAGE"
)"

regctl image copy \
  --platform "$PLATFORM" \
  "$SRC_BASE_IMAGE" \
  "$BASE_IMAGE"

test "$(regctl image digest "$BASE_IMAGE")" = "$SRC_BASE_DIGEST"

获取固定提交的 Megatron 源码，并导出不包含 Git 历史的镜像层目录。如果团队维护了代码镜像，请把 GitHub 地址替换为该镜像地址。

language-bash

git clone --filter=blob:none https://github.com/NVIDIA/Megatron-LM.git
git -C Megatron-LM checkout --detach "$MEGATRON_COMMIT"
test "$(git -C Megatron-LM rev-parse HEAD)" = "$MEGATRON_COMMIT"

mkdir -p megatron-rootfs/opt/Megatron-LM
git -C Megatron-LM archive "$MEGATRON_COMMIT" |
  tar -x -C megatron-rootfs/opt/Megatron-LM

从基础镜像创建新的 Megatron 镜像。--create 只写入 MEGATRON_IMAGE，不会修改 BASE_IMAGE：

language-bash

regctl image mod "$BASE_IMAGE" \
  --create "$MEGATRON_IMAGE" \
  --layer-add "dir=megatron-rootfs,platform=$PLATFORM" \
  --env "MEGATRON_COMMIT=$MEGATRON_COMMIT" \
  --env "PYTHONPATH=/opt/Megatron-LM" \
  --label "org.opencontainers.image.source=https://github.com/NVIDIA/Megatron-LM" \
  --label "org.opencontainers.image.revision=$MEGATRON_COMMIT"

test "$(regctl image digest "$BASE_IMAGE")" = "$SRC_BASE_DIGEST"
regctl image get-file \
  "$MEGATRON_IMAGE" \
  /opt/Megatron-LM/megatron/core/package_info.py \
  >/dev/null
regctl image digest "$MEGATRON_IMAGE"

信息

regctl image mod 的适用范围

regctl 将 image mod 标记为实验性命令。固定并验证所用 regctl 版本，在工具升级后重新执行摘要和文件检查。这条路径适合添加已经准备好的文件和镜像配置；需要在镜像中安装依赖、编译源码或执行其它构建命令时，请使用下面的 Docker Buildx 路径。

路径二：使用 Docker Buildx 完整构建并推送

这是标准的 Dockerfile 构建路径，适合能够稳定访问 NGC 和 GitHub，并已经安装 Docker Engine 与 Buildx 的本地设备或 CI Runner。构建设备会下载并解压基础镜像。

定义目标镜像和 Megatron 提交，登录 NGC 与目标可用区 Registry：

language-bash

export PLATFORM="linux/amd64"
export DST_REGISTRY="<目标可用区 Registry 公网访问地址>"
export DST_USERNAME="<镜像中心提供的用户名>"
export DST_NAMESPACE="<租户 ID>"
export MEGATRON_IMAGE="$DST_REGISTRY/$DST_NAMESPACE/megatron-core:0.18.2-cu128"
export MEGATRON_COMMIT="571370c829ca768fe37244f4e2e7f28d8accc4ab"

docker login nvcr.io --username '$oauthtoken'
docker login "$DST_REGISTRY" --username "$DST_USERNAME"

在第一个登录命令的密码提示中输入 NGC API Key；在第二个登录命令中输入镜像中心提供的目标 Registry 密码。

获取固定提交的 Megatron 源码，并创建不包含 Git 历史的构建上下文：

language-bash

git clone --filter=blob:none https://github.com/NVIDIA/Megatron-LM.git
git -C Megatron-LM checkout --detach "$MEGATRON_COMMIT"
test "$(git -C Megatron-LM rev-parse HEAD)" = "$MEGATRON_COMMIT"

mkdir -p megatron-image/Megatron-LM
git -C Megatron-LM archive "$MEGATRON_COMMIT" |
  tar -x -C megatron-image/Megatron-LM

在 megatron-image 目录创建以下 Dockerfile：

language-dockerfile

FROM nvcr.io/nvidia/pytorch:25.03-py3

ARG MEGATRON_COMMIT=571370c829ca768fe37244f4e2e7f28d8accc4ab

COPY Megatron-LM /opt/Megatron-LM

ENV MEGATRON_COMMIT=${MEGATRON_COMMIT} \
    PYTHONPATH=/opt/Megatron-LM

构建 linux/amd64 镜像并推送到目标可用区 Registry：

language-bash

docker buildx build \
  --platform "$PLATFORM" \
  --build-arg MEGATRON_COMMIT="$MEGATRON_COMMIT" \
  --tag "$MEGATRON_IMAGE" \
  --push \
  megatron-image

docker buildx imagetools inspect "$MEGATRON_IMAGE"

确认镜像可用于任务

完成任一条路径后，返回镜像中心，在目标可用区确认 megatron-core:0.18.2-cu128 可用于任务。

这个镜像固定 Megatron Core 0.18.2 的发布提交，并使用包含 Python 3.12、PyTorch 2.7 和 CUDA 12.8 的 NVIDIA PyTorch 25.03 基础镜像。固定发布版可以复现同一套 Megatron API；CUDA 12.x 对节点驱动的要求低于 CUDA 13.x，更容易适配不同的 AIStudio 资源池。

创建页找不到镜像时，先检查镜像的可用服务和可用区；镜像未覆盖训练任务或目标可用区时，再调整镜像配置。

准备共享存储输出目录

为训练任务挂载一份可写共享存储，并把容器内路径设置为 /mnt/megatron。每次运行使用独立目录，例如：

language-text

/mnt/megatron/
└── runs/
    └── mcore-tp2-01/
        ├── ckpt/
        ├── runtime.txt
        └── train.log

把 checkpoint 和日志写入可写共享存储，以便任务退出后继续读取结果。/infini-data 是公共只读目录，任务容器根文件系统则随 Worker 生命周期结束而释放。

运行两 GPU 张量并行训练

本场景直接运行镜像中的 examples/run_simple_mcore_train_loop.py。该示例构建两层小型 GPT 模型，使用 Mock 数据完成 5 个 optimizer step，然后通过 Megatron 分布式 checkpoint API 保存并重新加载模型。

Step 1 创建训练任务

创建训练任务时使用以下配置：

Worker 数量：1；
Worker 规格：选择包含 2 张 NVIDIA GPU 的规格；
分布式框架：单机；
镜像：选择上一节构建的 megatron-core:0.18.2-cu128；
训练变慢检测：本例只运行 5 个 step，可以保持关闭。扩展为正式的长时间 Megatron-LM 训练时开启该选项；平台发现 step 耗时持续增加后，会在容错日志中输出告警。告警的读取和排查方法见查看训练变慢检测结果；
共享存储容器内路径：/mnt/megatron。

Step 2 填写启动命令

将本次运行目录作为显式变量写在启动命令开头。再次运行时更换 RUN_ROOT，为每次运行保留独立输出。

language-bash

set -euo pipefail

RUN_ROOT=/mnt/megatron/runs/mcore-tp2-01
MEGATRON_ROOT=/opt/Megatron-LM

test "$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l | tr -d ' ')" -eq 2
mkdir -p "$RUN_ROOT"

# 官方 Mock 数据示例会使用这个小型 C++ helper。
make -C "$MEGATRON_ROOT/megatron/core/datasets"

{
  python -c "import torch, megatron.core; print('torch:', torch.__version__); print('cuda:', torch.version.cuda); print('gpu_count:', torch.cuda.device_count())"
  printf 'megatron_commit: %s\n' "${MEGATRON_COMMIT:?MEGATRON_COMMIT is required}"
} > "$RUN_ROOT/runtime.txt"

cd "$RUN_ROOT"
torchrun \
  --nproc_per_node=2 \
  "$MEGATRON_ROOT/examples/run_simple_mcore_train_loop.py" \
  > >(tee -i "$RUN_ROOT/train.log") 2>&1

任务获得 GPU 后会在一个 Worker 内启动两个 PyTorch 进程。官方示例把张量并行度固定为 2，因此两个进程共同训练一个张量并行模型，而不是分别训练两份独立模型。

Step 3 核对训练和 checkpoint

任务结束后核对以下结果：

任务状态为“运行成功”，启动命令没有用 sleep 或 || true 隐藏退出状态。
train.log 包含 Iteration 0 到 Iteration 4，并在末尾显示 Successfully loaded the model。
runtime.txt 记录的 GPU 数量为 2，Megatron 提交与镜像中固定的提交一致。
ckpt 目录包含 Megatron 分布式 checkpoint 元数据和两个 rank 写入的分片。
在任务外通过同一共享存储读取 runtime.txt、train.log 和 ckpt，确认结果不依赖已经退出的 Worker。

这个官方示例保存并重新加载模型权重，但不保存 optimizer 和学习率调度器。生产训练请使用 Megatron-LM 或 Megatron Bridge 的正式 checkpoint 入口，把模型、optimizer、调度器、随机状态和已完成迭代一起保存，从而生成可继续训练的恢复点。

扩展到多 Worker 训练

多 Worker 训练仍然运行同一份 Megatron 代码，但需要把 AIStudio 注入的 Worker 数量和编号传给 torchrun。先完成上面的单 Worker 场景，再使用一个新的运行目录验证多 Worker 启动。

创建 2 个 Worker、每个 Worker 2 张 GPU，并在分布式框架中选择 PyTorch DDP。保持 RDMA 关闭即可验证基本的多节点 NCCL 通信；只有训练方案确实需要并且所选资源支持时，再单独配置 RDMA。

使用以下启动命令：

language-bash

set -euo pipefail

RUN_ROOT=/mnt/megatron/runs/mcore-tp2-dp2-01
MEGATRON_ROOT=/opt/Megatron-LM

NNODES=${WORLD_SIZE:?AIStudio WORLD_SIZE is required}
NODE_RANK=${RANK:?AIStudio RANK is required}
MASTER_ADDR=${MASTER_ADDR:?AIStudio MASTER_ADDR is required}
MASTER_PORT=${MASTER_PORT:-29500}
GPUS_PER_NODE=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l | tr -d ' ')

test "$NNODES" -eq 2
test "$GPUS_PER_NODE" -eq 2
mkdir -p "$RUN_ROOT/nodes/node-$NODE_RANK"

# 镜像文件系统按 Worker 隔离，因此每个 Worker 都准备本地 helper。
make -C "$MEGATRON_ROOT/megatron/core/datasets"

cd "$RUN_ROOT"
torchrun \
  --nproc_per_node="$GPUS_PER_NODE" \
  --nnodes="$NNODES" \
  --node_rank="$NODE_RANK" \
  --master_addr="$MASTER_ADDR" \
  --master_port="$MASTER_PORT" \
  "$MEGATRON_ROOT/examples/run_simple_mcore_train_loop.py" \
  > >(tee -i "$RUN_ROOT/nodes/node-$NODE_RANK/train.log") 2>&1

这个配置创建 4 个 PyTorch 进程。官方示例使用 TP=2，Megatron 会把其余并行维度组成 DP=2。核对两个 Worker 使用不同的 node_rank、4 个 rank 都进入训练、共享 checkpoint 写入完成，并确认两个 Worker 都正常退出。

AIStudio Worker 变量的完整含义和通用启动规则见编写 PyTorch DDP torchrun 启动脚本。这次有界运行只验证 2 Worker × 2 GPU 的基本启动、训练和 checkpoint 链路，不代表任意拓扑的扩展效率。

把有界场景替换为真实训练

完成平台链路验证后，再按训练目标替换模型和数据：

从随机权重开始预训练：使用 Megatron-LM 的正式训练入口，明确模型架构和并行参数，并把文本预处理为 Megatron 使用的 .bin 和 .idx 文件。
从已有权重继续训练：优先使用 Megatron Bridge 从 Hugging Face 目录转换并加载权重。模型可以先从 ModelScope 下载到共享存储，再把该本地目录交给 Bridge。
迁移现有脚本：保留原有模型和并行配置，只把镜像、挂载路径、torchrun 节点参数、checkpoint 路径和退出状态映射到 AIStudio。

真实训练开始前，分别估算模型权重、梯度、optimizer、激活和 checkpoint 所需的 GPU 显存与共享存储容量，再根据目标模型和并行策略选择 GPU 型号、GPU 数量和训练时间。

排查 Megatron 训练任务

导入 Megatron 失败：核对镜像 tag、/opt/Megatron-LM、PYTHONPATH 和 MEGATRON_COMMIT。如需更新版本，请重新构建带有明确提交的镜像。
C++ helper 编译失败：确认镜像包含编译器、Python headers 和 pybind11；也可以在镜像构建阶段提前执行同一条 make 命令。
两个进程在初始化阶段退出：确认 Worker 实际提供 2 张 GPU，并检查 TP=2 是否与进程总数兼容。
多 Worker 初始化超时：先核对 NNODES、NODE_RANK、MASTER_ADDR、MASTER_PORT 和每节点 GPU 数，再检查训练网络和 RDMA 配置。
checkpoint 没有持久化：确认当前工作目录位于可写共享存储，并为每次运行使用独立目录。
任务显示成功但日志不完整：检查启动命令是否保留 torchrun 的真实退出码，并确认所有 Worker 都已经退出。

简介

AICoder 功能

租户管理

凭证管理

费用管理

算力资源池管理

存储管理

镜像中心

通知与告警

存储

开发机

任务

推理服务

资源管理

简介

集群管理

简介

实例管理

开始使用

模型与能力

开发工具与集成

计费、限制与用量

排查与支持

更新与参考

Megatron-Infinigence 训练框架

典型模型部署

判断 Megatron 是否适合当前任务 ​

选择 Megatron 组件 ​

在目标可用区准备容器镜像 ​

路径一：使用 regctl 复制并创建派生镜像 ​

路径二：使用 Docker Buildx 完整构建并推送 ​

确认镜像可用于任务 ​

准备共享存储输出目录 ​

运行两 GPU 张量并行训练 ​

Step 1 创建训练任务 ​

Step 2 填写启动命令 ​

Step 3 核对训练和 checkpoint ​

扩展到多 Worker 训练 ​

把有界场景替换为真实训练 ​

排查 Megatron 训练任务 ​

判断 Megatron 是否适合当前任务

选择 Megatron 组件

在目标可用区准备容器镜像

路径一：使用 regctl 复制并创建派生镜像

路径二：使用 Docker Buildx 完整构建并推送

确认镜像可用于任务

准备共享存储输出目录

运行两 GPU 张量并行训练

Step 1 创建训练任务

Step 2 填写启动命令

Step 3 核对训练和 checkpoint

扩展到多 Worker 训练

把有界场景替换为真实训练

排查 Megatron 训练任务