使用 Megatron-LM 训练 Llama3-8B 模型：从环境搭建到FP8高效训练
本文将详细介绍如何使用 NVIDIA Megatron-LM 框架训练 Llama3-8B 模型，支持 BF16 和 FP8 混合精度训练、灵活的分布式配置和 Mock 数据快速验证。我们提供了一个高度可配置的训练脚本，让您可以根据硬件资源灵活调整训练参数。
注意
本教程默认使用 BF16 精度和 NVIDIA A100 GPU 进行训练，具有良好的兼容性和稳定性。如果您使用 NVIDIA H100 GPU，可以启用 FP8 精度获得更高的训练性能。教程适用于各种 NVIDIA GPU 平台，在无问芯穹 AIStudio 平台上验证。
学习目标
- 如何搭建 Megatron-LM 训练环境
- 如何配置动态化的训练脚本
- 如何使用 Mock 数据进行快速验证
- 如何配置分布式训练参数
- 如何启用 BF16 混合精度训练（默认）
- 如何在 H100 上启用 FP8 混合精度训练
- 如何监控训练过程和调试问题
- 如何从 Mock 数据切换到真实数据训练
Megatron-LM 优势
Megatron-LM 是 NVIDIA 开发的大规模语言模型训练框架，具有以下优势：
- 高效的模型并行：支持张量并行、流水线并行和上下文并行
- BF16/FP8 混合精度：BF16 提供广泛兼容性，FP8 在 H100 上提供更高的训练效率
- 内存优化：通过激活重计算、梯度检查点等技术减少显存占用
- 可扩展性：支持从单机到数千GPU的大规模分布式训练
- Mock 数据支持：无需准备数据集即可快速验证环境和配置
前提条件
硬件要求
- 默认推荐：NVIDIA A100 GPU（BF16 性能最优）
- 高性能选择：NVIDIA H100 GPU（完整 FP8 支持）
- 最低要求：8GB+ 显存的现代 GPU
- 内存：每个 GPU 建议 64GB+ 系统内存
软件要求
- CUDA 11.8+ 或 12.0+
- Python 3.8+
- PyTorch 2.0+
环境搭建
推荐容器镜像
本教程基于 NVIDIA 官方 PyTorch 容器镜像 nvcr.io/nvidia/pytorch:25.03-py3，该镜像已预装：
- CUDA 12.6+ - GPU 加速计算框架
- Python 3.11 - 编程环境
- PyTorch 2.6+ - 深度学习框架
- NCCL - 分布式通信库
- OpenMPI - 消息传递接口
- cuDNN - CUDA 深度神经网络库
- Transformer Engine - FP8 混合精度支持
- Flash Attention - 高效注意力机制
- 基础 Python 包 - transformers, datasets, accelerate, tensorboard, wandb 等
重要
版本要求：官方文档建议使用特定的 Megatron-LM 版本。在生产环境中，建议检出到稳定版本：
bash
git checkout "core_r0.12.0"AIStudio 平台环境变量适配
重要说明：AIStudio 平台在分布式训练中使用了与 PyTorch 原生不同的环境变量定义。这种设计与 Kubeflow 的 pytorch-operator 保持一致，方便从 Kubeflow 迁移的用户无缝过渡。
平台环境变量定义
AIStudio 平台自动注入的环境变量：
MASTER_ADDR：分布式通信的主节点地址（通常为 worker0 的 pod 名称）MASTER_PORT：主节点端口号，默认值为 29500WORLD_SIZE：pod（节点）数量（注意：不同于 PyTorch 原生的总进程数）RANK：pod 编号（从 0 到 WORLD_SIZE-1，用于设置节点排名）
与 PyTorch 原生的区别
| 环境变量 | AIStudio 平台定义 | PyTorch 原生定义 | 说明 |
|---|---|---|---|
WORLD_SIZE | pod（节点）数量 | 总进程数（节点数×每节点GPU数） | 平台简化了概念 |
RANK | pod 编号 | 进程全局编号 | 平台对应节点级别 |
环境变量转换逻辑
由于环境变量定义的差异，Megatron-LM 训练脚本需要进行适配转换：
bash
# 1. 获取平台注入的环境变量（节点级别）
PLATFORM_WORLD_SIZE=${WORLD_SIZE:-1} # 平台：节点数量
PLATFORM_RANK=${RANK:-0} # 平台：节点编号
PLATFORM_MASTER_ADDR=${MASTER_ADDR:-localhost}
PLATFORM_MASTER_PORT=${MASTER_PORT:-29500}
# 2. 转换为 Megatron-LM 所需的变量
NUM_NODES=$PLATFORM_WORLD_SIZE # 节点数量
NODE_RANK=$PLATFORM_RANK # 节点排名
MASTER_ADDR=$PLATFORM_MASTER_ADDR # 主节点地址
MASTER_PORT=$PLATFORM_MASTER_PORT # 主节点端口
# 3. 计算总 GPU 数量（Megatron-LM 的 WORLD_SIZE）
GPUS_PER_NODE=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
MEGATRON_WORLD_SIZE=$(($GPUS_PER_NODE * $NUM_NODES))注意
为什么需要转换？
- AIStudio 平台：简化概念，
WORLD_SIZE直接表示有多少个节点（pod） - Megatron-LM：需要知道总的 GPU 数量来正确分配模型和数据
- 转换目的：确保 Megatron-LM 能够正确理解分布式环境的规模
示例场景
假设您在 AIStudio 平台上启动 2 个 Worker，每个 Worker 有 4 个 GPU：
平台注入的环境变量：
bash
# Worker 0（第一个 pod）
WORLD_SIZE=2
RANK=0
MASTER_ADDR=worker-0-pod-name
MASTER_PORT=29500
# Worker 1（第二个 pod）
WORLD_SIZE=2
RANK=1
MASTER_ADDR=worker-0-pod-name
MASTER_PORT=29500脚本转换后的变量：
bash
NUM_NODES=2 # 2 个节点
GPUS_PER_NODE=4 # 每个节点 4 个 GPU
MEGATRON_WORLD_SIZE=8 # 总共 8 个 GPU
NODE_RANK=0 或 1 # 节点排名最终 Megatron-LM 分布式配置：
- 总共 8 个 GPU 进程
- Worker 0：GPU 进程编号 0-3
- Worker 1：GPU 进程编号 4-7
容器启动后设置
重要说明：平台的容器根文件系统是持久化的，这意味着您只需要一次性下载和配置 Megatron-LM 代码库，之后的训练任务可以直接使用。
一次性初始化步骤
首次启动容器后，通过 shell 界面手动执行以下初始化步骤：
bash
# 1. 进入工作目录
cd /workspace
# 2. 克隆 Megatron-LM 仓库（仅需执行一次）
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
# 3. 检出到推荐的稳定版本
git checkout "core_r0.12.0"
# 4. 设置环境变量（可写入 ~/.bashrc 以便持久化）
echo 'export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH"' >> ~/.bashrc
source ~/.bashrc
# 5. 验证环境
python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
nvidia-smi后续使用
完成一次性初始化后，后续的训练任务只需要：
bash
# 进入 Megatron-LM 目录
cd /workspace/Megatron-LM
# 直接开始训练（无需重新下载代码）
bash train_llama3_8b_dynamic.sh注意
无需额外安装：nvcr.io/nvidia/pytorch:25.03-py3 镜像已包含所有必需依赖，包括 transformers、datasets、accelerate、tensorboard 等，无需额外的 pip 安装。
配置说明和重要提示
与官方脚本的差异
本教程与官方 train_llama3_8b_h100_fp8.sh 脚本在某些默认配置上有所不同，这是为了提供更好的通用性和学习体验：
| 配置项 | 官方脚本默认值 | 本教程默认值 | 说明 |
|---|---|---|---|
| 精度类型 | fp8 | bf16 | 本教程默认使用 BF16 以确保广泛兼容性 |
| 默认GPU | H100 | A100 | 本教程以 A100 为基准，适配更多硬件环境 |
| 上下文并行 | CP_SIZE=1 | CP_SIZE=1 | 一致 |
| 张量并行 | TP_SIZE=1 | TP_SIZE=1 | 一致 |
| 全局批大小 | 128 | 128 | 一致 |
| Megatron版本 | core_r0.12.0 | main | 官方建议使用稳定版本 |
选择建议
默认配置用户（A100）
- 使用本教程的默认 BF16 + A100 配置
- 适用于学习和生产，性能稳定兼容性好
H100 高性能用户
- 可直接使用官方 FP8 配置获得最佳性能
- 设置
DTYPE="fp8"启用 FP8 训练
其他GPU用户
- 使用本教程的 BF16 配置确保兼容性
- 根据具体硬件调整并行配置
我们基于官方脚本创建一个高度可配置的版本，支持通过环境变量调整所有重要参数：
bash
# train_llama3_8b_dynamic.sh
#!/bin/bash
# =============================================================================
# Megatron-LM Llama3-8B 动态训练脚本
# 支持通过环境变量灵活配置所有训练参数
# =============================================================================
# 性能调优环境变量
export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
export NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-0}
export NCCL_IB_TIMEOUT=${NCCL_IB_TIMEOUT:-23}
export NCCL_P2P_DISABLE=${NCCL_P2P_DISABLE:-0}
export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4}
# AIStudio 平台环境变量适配
# 注意：平台注入的环境变量与 PyTorch 原生定义不同
PLATFORM_WORLD_SIZE=${WORLD_SIZE:-1} # 平台：节点数量
PLATFORM_RANK=${RANK:-0} # 平台：节点编号
PLATFORM_MASTER_ADDR=${MASTER_ADDR:-localhost}
PLATFORM_MASTER_PORT=${MASTER_PORT:-29500}
# 转换为 Megatron-LM 标准变量
NUM_NODES=$PLATFORM_WORLD_SIZE
NODE_RANK=$PLATFORM_RANK
MASTER_ADDR=$PLATFORM_MASTER_ADDR
MASTER_PORT=$PLATFORM_MASTER_PORT
# 自动检测每个节点的 GPU 数量
GPUS_PER_NODE=${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}
# 计算 Megatron-LM 所需的总 GPU 数量
WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
# 基础路径配置
CHECKPOINT_PATH=${1:-${CHECKPOINT_PATH:-"checkpoints/llama3_8b_fp8"}}
TENSORBOARD_LOGS_PATH=${2:-${TENSORBOARD_LOGS_PATH:-"tensorboard_logs/llama3_8b_fp8"}}
TOKENIZER_ARG=${3:-${TOKENIZER_ARG:-"MOCK"}}
DATA_ARG=${4:-${DATA_ARG:-"MOCK"}}
# 实验名称配置
EXP_NAME=${EXP_NAME:-"llama3-8b-fp8-training"}
EXP_DIR=${EXP_DIR:-"./experiments"}
FULL_EXP_PATH="${EXP_DIR}/${EXP_NAME}"
# 如果使用实验目录结构，重新设置路径
if [[ "$USE_EXP_STRUCTURE" == "1" ]]; then
CHECKPOINT_PATH="${FULL_EXP_PATH}/checkpoints"
TENSORBOARD_LOGS_PATH="${FULL_EXP_PATH}/tensorboard"
fi
# 创建必要目录
mkdir -p "$(dirname "$CHECKPOINT_PATH")"
mkdir -p "$(dirname "$TENSORBOARD_LOGS_PATH")"
# 检查 GPU 数量和环境变量适配
AVAILABLE_GPUS=$(nvidia-smi --list-gpus | wc -l)
if [[ $GPUS_PER_NODE -gt $AVAILABLE_GPUS ]]; then
echo "警告: 配置的 GPU 数量 ($GPUS_PER_NODE) 超过可用 GPU 数量 ($AVAILABLE_GPUS)"
echo "自动调整为可用 GPU 数量: $AVAILABLE_GPUS"
GPUS_PER_NODE=$AVAILABLE_GPUS
WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
fi
# 显示环境变量转换信息
echo "==================== AIStudio 平台环境变量适配 ===================="
echo "平台注入变量: WORLD_SIZE=$PLATFORM_WORLD_SIZE (节点数), RANK=$PLATFORM_RANK (节点编号)"
echo "转换后变量: NUM_NODES=$NUM_NODES, NODE_RANK=$NODE_RANK"
echo "GPU 配置: ${GPUS_PER_NODE} GPUs/节点 × ${NUM_NODES} 节点 = ${WORLD_SIZE} 总GPU"
echo "主节点: $MASTER_ADDR:$MASTER_PORT"
echo "=================================================================="
# 默认并行配置（基于官方脚本）
TP_SIZE=${TP_SIZE:-1}
CP_SIZE=${CP_SIZE:-1} # 注意：官方脚本实际默认 CP_SIZE=1，不是2
PP_SIZE=${PP_SIZE:-1}
# 验证并行配置
TOTAL_PARALLEL=$((TP_SIZE * CP_SIZE * PP_SIZE))
if [[ $TOTAL_PARALLEL -gt $WORLD_SIZE ]]; then
echo "错误: 并行配置 (TP=$TP_SIZE × CP=$CP_SIZE × PP=$PP_SIZE = $TOTAL_PARALLEL) 超过总 GPU 数量 ($WORLD_SIZE)"
exit 1
fi
# 训练参数配置
MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-128}
NUM_LAYERS=${NUM_LAYERS:-32}
HIDDEN_SIZE=${HIDDEN_SIZE:-4096}
FFN_HIDDEN_SIZE=${FFN_HIDDEN_SIZE:-14336}
NUM_ATTENTION_HEADS=${NUM_ATTENTION_HEADS:-32}
NUM_QUERY_GROUPS=${NUM_QUERY_GROUPS:-8}
# 序列和位置配置
SEQ_LENGTH=${SEQ_LENGTH:-8192}
MAX_POSITION_EMBEDDINGS=${MAX_POSITION_EMBEDDINGS:-$SEQ_LENGTH}
# 精度配置（官方脚本默认使用 FP8 + BF16）
DTYPE=${DTYPE:-"bf16"} # 教程默认：bf16，官方脚本：fp8
USE_BF16=${USE_BF16:-1} # BF16 始终启用
# 学习率配置
LEARNING_RATE=${LEARNING_RATE:-0.00015}
MIN_LEARNING_RATE=${MIN_LEARNING_RATE:-0.00001}
LR_DECAY_STYLE=${LR_DECAY_STYLE:-"cosine"}
WEIGHT_DECAY=${WEIGHT_DECAY:-0.1}
# 训练轮次配置
TRAIN_SAMPLES=${TRAIN_SAMPLES:-1953125000}
LR_DECAY_SAMPLES=${LR_DECAY_SAMPLES:-1949218748}
LR_WARMUP_SAMPLES=${LR_WARMUP_SAMPLES:-3906252}
# 如果启用快速测试模式
if [[ "$FAST_TEST_MODE" == "1" ]]; then
echo "启用快速测试模式，调整训练参数..."
TRAIN_SAMPLES=1000
LR_DECAY_SAMPLES=800
LR_WARMUP_SAMPLES=100
GLOBAL_BATCH_SIZE=32
fi
# 数据缓存路径
DATA_CACHE_PATH="${PWD}/benchmark_cache_${EXP_NAME}"
mkdir -p "$DATA_CACHE_PATH"
# 脚本路径
PRETRAIN_SCRIPT_PATH=${PRETRAIN_SCRIPT_PATH:-"pretrain_gpt.py"}
# 检查脚本是否存在
if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then
echo "错误: 找不到训练脚本 $PRETRAIN_SCRIPT_PATH"
echo "请确保在 Megatron-LM 根目录下运行此脚本"
exit 1
fi
# 打印配置信息
echo "==================== 训练配置 ===================="
echo "实验名称: $EXP_NAME"
echo "检查点路径: $CHECKPOINT_PATH"
echo "TensorBoard 路径: $TENSORBOARD_LOGS_PATH"
echo "分布式配置: ${NUM_NODES} 节点, 节点排名 ${NODE_RANK}"
echo "GPU 配置: ${GPUS_PER_NODE} GPUs/节点 × ${NUM_NODES} 节点 = ${WORLD_SIZE} 总GPU"
echo "并行配置: TP=${TP_SIZE}, CP=${CP_SIZE}, PP=${PP_SIZE}"
echo "批处理配置: Micro=${MICRO_BATCH_SIZE}, Global=${GLOBAL_BATCH_SIZE}"
echo "模型配置: ${NUM_LAYERS}层, ${HIDDEN_SIZE}隐藏维度, ${SEQ_LENGTH}序列长度"
echo "精度配置: ${DTYPE}${USE_BF16:+, BF16}"
echo "数据模式: ${TOKENIZER_ARG}"
echo "=================================================="
# 分布式参数（适配 AIStudio 平台）
DISTRIBUTED_ARGS=(
--nproc_per_node $GPUS_PER_NODE
--nnodes $NUM_NODES
--node_rank $NODE_RANK
--master_addr $MASTER_ADDR
--master_port $MASTER_PORT
)
# 模型参数
MODEL_ARGS=(
--use-mcore-models
--num-layers $NUM_LAYERS
--hidden-size $HIDDEN_SIZE
--ffn-hidden-size $FFN_HIDDEN_SIZE
--num-attention-heads $NUM_ATTENTION_HEADS
--group-query-attention
--num-query-groups $NUM_QUERY_GROUPS
--kv-channels 128
--seq-length $SEQ_LENGTH
--max-position-embeddings $MAX_POSITION_EMBEDDINGS
--position-embedding-type rope
--rotary-base 1000000
--rotary-percent 1.0
--attention-dropout 0.0
--hidden-dropout 0.0
--swiglu
--init-method-std 0.0134
--attention-backend fused
--apply-layernorm-1p
--untie-embeddings-and-output-weights
--disable-bias-linear
)
# 训练参数
TRAINING_ARGS=(
--micro-batch-size $MICRO_BATCH_SIZE
--global-batch-size $GLOBAL_BATCH_SIZE
--train-samples $TRAIN_SAMPLES
--lr-decay-samples $LR_DECAY_SAMPLES
--lr-warmup-samples $LR_WARMUP_SAMPLES
--lr $LEARNING_RATE
--min-lr $MIN_LEARNING_RATE
--lr-decay-style $LR_DECAY_STYLE
--clip-grad 1.0
--weight-decay $WEIGHT_DECAY
--adam-beta1 0.9
--adam-beta2 0.95
)
# 精度配置
if [[ "$USE_BF16" == "1" ]]; then
TRAINING_ARGS+=(--bf16)
TRAINING_ARGS+=(--grad-reduce-in-bf16)
fi
# 优化参数
TRAINING_ARGS+=(
--cross-entropy-loss-fusion
--calculate-per-token-loss
--manual-gc
--empty-unused-memory-level 1
--exit-duration-in-mins 235
)
# FP8 配置
DTYPE_ARGS=()
if [[ "$DTYPE" == "fp8" ]]; then
DTYPE_ARGS+=(
"--fp8-format hybrid"
"--fp8-amax-history-len 1024"
"--fp8-amax-compute-algo max"
"--fp8-param-gather"
)
fi
# 模型并行参数
MODEL_PARALLEL_ARGS=(
--tensor-model-parallel-size $TP_SIZE
--context-parallel-size $CP_SIZE
--sequence-parallel
)
if [[ $PP_SIZE -gt 1 ]]; then
MODEL_PARALLEL_ARGS+=(--pipeline-model-parallel-size $PP_SIZE)
fi
# 分布式优化器配置
DDP_ARGS=(
--use-distributed-optimizer
--overlap-grad-reduce
--overlap-param-gather
)
TRAINING_ARGS+=("${DDP_ARGS[@]}")
# 数据配置
DATA_ARGS_LIST=()
if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then
echo "使用 Mock 数据模式进行训练..."
DATA_ARGS_LIST+=(
"--mock-data"
"--tokenizer-type NullTokenizer"
"--vocab-size 128256"
"--data-cache-path ${DATA_CACHE_PATH}"
"--tiktoken-pattern v2"
"--split '99,1,0'"
"--no-create-attention-mask-in-dataloader"
"--no-mmap-bin-files"
"--num-workers 1"
)
else
echo "使用真实数据进行训练..."
DATA_ARGS_LIST+=(
"--data-path $DATA_ARG"
"--tokenizer-type HuggingFaceTokenizer"
"--tokenizer-model $TOKENIZER_ARG"
"--data-cache-path ${DATA_CACHE_PATH}"
"--split '99,1,0'"
"--no-create-attention-mask-in-dataloader"
"--no-mmap-bin-files"
"--num-workers 1"
"--vocab-size 128256"
)
fi
# 评估和日志配置
LOG_INTERVAL=${LOG_INTERVAL:-1}
SAVE_INTERVAL=${SAVE_INTERVAL:-1000}
EVAL_INTERVAL=${EVAL_INTERVAL:-100}
EVAL_AND_LOGGING_ARGS=(
--log-interval $LOG_INTERVAL
--eval-iters 32
--eval-interval $EVAL_INTERVAL
--save-interval $SAVE_INTERVAL
--log-throughput
--distributed-timeout-minutes 60
--save "$CHECKPOINT_PATH"
--load "$CHECKPOINT_PATH"
--tensorboard-dir "$TENSORBOARD_LOGS_PATH"
--ckpt-format torch_dist
)
# 性能分析配置
if [[ "$ENABLE_PROFILING" == "1" ]]; then
EVAL_AND_LOGGING_ARGS+=(
--profile
--profile-step-start 4
--profile-step-end 6
)
fi
# 运行训练
echo "开始训练..."
echo "训练命令: torchrun ${DISTRIBUTED_ARGS[@]} $PRETRAIN_SCRIPT_PATH ..."
torchrun ${DISTRIBUTED_ARGS[@]} \
"$PRETRAIN_SCRIPT_PATH" \
${MODEL_ARGS[@]} \
${TRAINING_ARGS[@]} \
${DTYPE_ARGS[@]} \
${MODEL_PARALLEL_ARGS[@]} \
${DATA_ARGS_LIST[@]} \
${EVAL_AND_LOGGING_ARGS[@]}
TRAIN_EXIT_CODE=$?
if [[ $TRAIN_EXIT_CODE -eq 0 ]]; then
echo "训练成功完成!"
echo "检查点保存在: $CHECKPOINT_PATH"
echo "TensorBoard 日志: $TENSORBOARD_LOGS_PATH"
else
echo "训练失败，退出码: $TRAIN_EXIT_CODE"
exit $TRAIN_EXIT_CODE
fi环境配置模板
为了方便管理不同的训练配置，我们提供了针对 AIStudio 平台的配置模板：
AIStudio 平台配置
bash
# config_aistudio_single_node.env - AIStudio 单节点配置
# 说明：在 AIStudio 平台上，WORLD_SIZE 和 RANK 由平台自动注入
# 这些变量表示节点级别的信息，不需要手动设置
export GPUS_PER_NODE=8 # 每个节点的GPU数量（根据实际情况调整）
export TP_SIZE=1 # 张量并行大小
export CP_SIZE=1 # 上下文并行大小
export PP_SIZE=1 # 流水线并行大小
export MICRO_BATCH_SIZE=1 # 微批大小
export GLOBAL_BATCH_SIZE=128 # 全局批大小
export DTYPE="bf16" # 精度类型（A100 推荐）
export EXP_NAME="llama3-8b-aistudio"
export USE_EXP_STRUCTURE=1
# 注意：WORLD_SIZE、RANK、MASTER_ADDR、MASTER_PORT 由 AIStudio 平台自动注入
# 无需手动设置这些变量AIStudio 多节点配置
bash
# config_aistudio_multi_node.env - AIStudio 多节点配置
# 说明：多节点训练时，平台会自动设置 WORLD_SIZE=节点数量，RANK=节点编号
export GPUS_PER_NODE=8 # 每个节点的GPU数量
export TP_SIZE=2 # 张量并行大小
export CP_SIZE=2 # 上下文并行大小
export PP_SIZE=2 # 流水线并行大小
export MICRO_BATCH_SIZE=1 # 微批大小
export GLOBAL_BATCH_SIZE=512 # 全局批大小
export DTYPE="bf16" # 精度类型
export EXP_NAME="llama3-8b-multinode"
export USE_EXP_STRUCTURE=1
# 平台自动注入的环境变量示例：
# 节点0: WORLD_SIZE=2, RANK=0, MASTER_ADDR=worker-0, MASTER_PORT=29500
# 节点1: WORLD_SIZE=2, RANK=1, MASTER_ADDR=worker-0, MASTER_PORT=29500容器启动命令
在 AIStudio 平台上，您可以通过图形界面启动容器。以下是推荐的启动命令配置。
重要
AIStudio 平台特殊说明：
平台会自动注入以下环境变量，无需手动设置：
WORLD_SIZE：节点数量（由平台根据 Worker 数量自动设置）RANK：节点编号（每个 Worker 自动分配 0, 1, 2...）MASTER_ADDR：主节点地址（通常为 worker-0 的 pod 名称）MASTER_PORT：主节点端口（默认 29500）
您只需要设置训练相关的参数即可。
基础启动命令
由于使用的是 nvcr.io/nvidia/pytorch:25.03-py3 镜像且代码库已在初始化时下载，只需少量配置：
bash
# AIStudio 单节点训练（推荐用于首次测试）
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=1 && \
export CP_SIZE=1 && \
export PP_SIZE=1 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=128 && \
export EXP_NAME="llama3-8b-aistudio" && \
export USE_EXP_STRUCTURE=1 && \
export FAST_TEST_MODE=1 && \
bash train_llama3_8b_dynamic.sh可配置环境变量
您可以在启动命令中设置以下环境变量来自定义训练配置：
训练配置参数：
bash
export TP_SIZE=1 # 张量并行大小 (1,2,4,8)
export CP_SIZE=1 # 上下文并行大小 (1,2,4,8)
export PP_SIZE=1 # 流水线并行大小 (1,2,4,8)
export MICRO_BATCH_SIZE=1 # 微批大小 (1,2,4,8)
export GLOBAL_BATCH_SIZE=128 # 全局批大小 (32,64,128,256,512)模型和精度配置：
bash
export DTYPE="bf16" # 精度类型 (bf16,fp8,fp16) - A100默认BF16
export USE_BF16=1 # 启用BF16 (0,1)
export FAST_TEST_MODE=1 # 快速测试模式 (0,1)
export NUM_LAYERS=32 # 模型层数 (32,40,48,80)
export HIDDEN_SIZE=4096 # 隐藏层维度 (4096,5120,8192)
export SEQ_LENGTH=8192 # 序列长度 (2048,4096,8192,16384)实验管理配置：
bash
export EXP_NAME="llama3-8b-training" # 实验名称
export USE_EXP_STRUCTURE=1 # 使用实验目录结构 (0,1)
export LOG_INTERVAL=1 # 日志输出间隔 (1,10,100)
export SAVE_INTERVAL=1000 # 检查点保存间隔 (100,500,1000)常用启动命令模板
单节点快速测试：
bash
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=1 && \
export CP_SIZE=1 && \
export PP_SIZE=1 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=32 && \
export FAST_TEST_MODE=1 && \
export EXP_NAME="test-aistudio" && \
bash train_llama3_8b_dynamic.sh单节点标准训练（A100）：
bash
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=1 && \
export CP_SIZE=1 && \
export PP_SIZE=1 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=128 && \
export DTYPE="bf16" && \
export EXP_NAME="llama3-8b-a100" && \
export USE_EXP_STRUCTURE=1 && \
bash train_llama3_8b_dynamic.sh单节点高性能训练（H100 + FP8）：
bash
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=1 && \
export CP_SIZE=1 && \
export PP_SIZE=1 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=128 && \
export DTYPE="fp8" && \
export USE_BF16=1 && \
export EXP_NAME="llama3-8b-h100-fp8" && \
export USE_EXP_STRUCTURE=1 && \
bash train_llama3_8b_dynamic.sh多节点训练示例：
bash
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=2 && \
export CP_SIZE=2 && \
export PP_SIZE=2 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=512 && \
export DTYPE="bf16" && \
export EXP_NAME="llama3-8b-multinode" && \
export USE_EXP_STRUCTURE=1 && \
bash train_llama3_8b_dynamic.sh内存优化配置：
bash
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=4 && \
export CP_SIZE=2 && \
export PP_SIZE=1 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=64 && \
export SEQ_LENGTH=4096 && \
export DTYPE="bf16" && \
export EXP_NAME="llama3-8b-memory-opt" && \
export USE_EXP_STRUCTURE=1 && \
bash train_llama3_8b_dynamic.sh提示
AIStudio 平台使用建议：
- 首次启动建议使用快速测试模式验证环境
- 平台会自动处理节点间通信，无需手动配置网络参数
- 默认使用 BF16 精度，具有最佳兼容性
- H100 用户可以启用 FP8 精度获得更高性能
- 根据GPU内存大小调整批处理大小和并行策略
- 使用
USE_EXP_STRUCTURE=1来组织实验文件
注意
关于容器镜像的说明：
使用 nvcr.io/nvidia/pytorch:25.03-py3 镜像的优势：
- 开箱即用：预装 CUDA、PyTorch、Transformer Engine 等核心组件
- 依赖完整：包含 transformers、datasets、accelerate、tensorboard、wandb 等所有必需包
- 版本兼容：所有组件版本经过 NVIDIA 优化和测试
- 性能优化：针对 NVIDIA GPU 进行了深度优化
- 稳定可靠：定期更新和安全补丁
无需任何额外的 pip 安装，由于容器的根文件系统是持久化的，您只需要一次性下载 Megatron-LM 代码库，后续训练可直接使用。
启动后管理
容器启动后，您可以通过 shell 界面进行以下操作：
bash
# 检查环境变量是否正确注入
echo "平台环境变量:"
echo "WORLD_SIZE=$WORLD_SIZE (节点数量)"
echo "RANK=$RANK (节点编号)"
echo "MASTER_ADDR=$MASTER_ADDR"
echo "MASTER_PORT=$MASTER_PORT"
# 检查环境是否正确配置
nvidia-smi
python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
# 查看训练状态
tail -f experiments/你的实验名称/tensorboard/events.out.tfevents.*
# 在 AIStudio 平台上，您也可以直接在平台界面中查看 TensorBoard
# 详细使用方法请参考：ai-studio/dev-instance/tensorboard.md
# 监控GPU使用情况
watch -n 1 nvidia-smi
# 查看训练日志和输出
ls -la experiments/你的实验名称/
# 检查 Transformer Engine 是否可用（FP8 支持）
python -c "
try:
import transformer_engine
print('Transformer Engine 可用 - 支持 FP8 训练')
except ImportError:
print('Transformer Engine 不可用')
"提示
AIStudio 平台使用建议：
- 使用推荐的
nvcr.io/nvidia/pytorch:25.03-py3镜像，环境配置最简单 - 平台会自动注入分布式环境变量，无需手动配置网络参数
- 首次启动建议使用单节点测试模式验证环境
- 默认使用 BF16 精度，具有最佳兼容性
- H100 用户可以启用 FP8 精度获得更高性能
- 使用
USE_EXP_STRUCTURE=1来组织实验文件
快速开始
步骤1：环境检查
首先验证您的环境是否正确配置：
bash
# 检查 GPU
nvidia-smi
# 检查 CUDA
nvcc --version
# 检查 PyTorch 和 CUDA 兼容性
python -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}, GPU Available: {torch.cuda.is_available()}')"
# 检查 Megatron-LM
cd /path/to/Megatron-LM
python -c "import megatron; print('Megatron-LM imported successfully')"步骤2：单GPU测试
首先使用单GPU和Mock数据进行快速测试：
bash
# 进入 Megatron-LM 目录
cd /path/to/Megatron-LM
# 加载单GPU配置
source config_single_gpu.env
# 运行训练脚本
bash train_llama3_8b_dynamic.sh预期输出：
==================== 训练配置 ====================
实验名称: test-single-gpu
GPU 配置: 1 GPUs × 1 节点 = 1 总GPU
并行配置: TP=1, CP=1, PP=1
批处理配置: Micro=1, Global=8
模型配置: 32层, 4096隐藏维度, 8192序列长度
精度配置: bf16, BF16
数据模式: MOCK
==================================================
使用 Mock 数据模式进行训练...
开始训练...步骤3：多GPU训练
确认单GPU测试成功后，尝试多GPU训练：
bash
# 加载8GPU配置
source config_8gpu_standard.env
# 运行训练
bash train_llama3_8b_dynamic.sh高级配置
自定义实验管理
启用实验目录结构来更好地组织训练文件：
bash
export USE_EXP_STRUCTURE=1
export EXP_NAME="my-llama3-experiment"
export EXP_DIR="./my_experiments"
# 这将创建以下目录结构：
# ./my_experiments/my-llama3-experiment/
# ├── checkpoints/
# └── tensorboard/性能调优
针对不同硬件配置的优化建议：
A100 GPU 优化（默认推荐）
bash
export DTYPE="bf16" # A100 推荐使用 BF16，默认配置
export USE_BF16=1
export NCCL_IB_TIMEOUT=19H100 GPU 优化
bash
export DTYPE="fp8" # H100 支持 FP8，性能最佳
export USE_BF16=1
export NCCL_IB_TIMEOUT=23
export CUDA_DEVICE_MAX_CONNECTIONS=1其他 GPU 优化
bash
export DTYPE="bf16" # 通用兼容性最好
export USE_BF16=1
# 或者使用 FP16（如果 BF16 不支持）
# export DTYPE="fp16"
# export USE_BF16=0内存优化
bash
# 减少批大小
export MICRO_BATCH_SIZE=1
export GLOBAL_BATCH_SIZE=64
# 启用内存优化
export CUDA_LAUNCH_BLOCKING=1数据配置
使用真实数据
当您准备好真实数据时，可以这样配置：
bash
# 假设您已经预处理了数据
export TOKENIZER_ARG="/path/to/your/tokenizer"
export DATA_ARG="/path/to/your/preprocessed_data"
# 运行训练
bash train_llama3_8b_dynamic.sh推荐测试数据集
根据官方文档，推荐使用以下数据集进行测试：
- WikiText-103：https://huggingface.co/datasets/Salesforce/wikitext
- 适合初学者和测试环境
- 数据质量高，处理简单
数据预处理
如果您需要预处理自己的数据：
bash
# 使用 Megatron-LM 的数据预处理工具（官方推荐方式）
python tools/preprocess_data.py \
--input your_data.json \
--output-prefix my_data \
--tokenizer-type HuggingFaceTokenizer \
--tokenizer-model /path/to/tokenizer \
--append-eod \
--workers 8监控和调试
TensorBoard 监控
您可以使用 TensorBoard 来监控训练进度。
注意
AIStudio 平台 TensorBoard 支持：
AIStudio 平台提供了内置的 TensorBoard 支持，可以直接在平台界面中查看训练指标。详细使用方法请参考：TensorBoard 使用指南
常见问题和解决方案
1. GPU 内存不足
bash
错误: CUDA out of memory
解决方案:
- 减少 MICRO_BATCH_SIZE 或 GLOBAL_BATCH_SIZE
- 增加 TP_SIZE 来分布模型参数
- 启用梯度检查点: --recompute-activations2. 通信超时
bash
错误: NCCL timeout
解决方案:
export NCCL_IB_TIMEOUT=60
export NCCL_ASYNC_ERROR_HANDLING=13. FP8 不支持
bash
错误: FP8 not supported
解决方案:
export DTYPE="bf16" # 降级到 BF16（推荐）
# 或者
export DTYPE="fp16" # 降级到 FP16故障排除
启动前检查清单
bash
# 创建检查脚本 check_environment.sh
#!/bin/bash
echo "=== Megatron-LM 环境检查 ==="
# 检查 GPU
echo "1. GPU 检查:"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader,nounits
# 检查 CUDA
echo -e "\n2. CUDA 检查:"
nvcc --version | grep "release"
# 检查 PyTorch
echo -e "\n3. PyTorch 检查:"
python -c "
import torch
print(f'PyTorch版本: {torch.__version__}')
print(f'CUDA可用: {torch.cuda.is_available()}')
print(f'GPU数量: {torch.cuda.device_count()}')
"
# 检查 Megatron-LM
echo -e "\n4. Megatron-LM 检查:"
if [ -f "pretrain_gpt.py" ]; then
echo "pretrain_gpt.py 存在"
else
echo "pretrain_gpt.py 不存在，请检查路径"
fi
# 检查 Transformer Engine
echo -e "\n5. Transformer Engine 检查:"
python -c "
try:
import transformer_engine
print('Transformer Engine 可用')
except ImportError:
print('Transformer Engine 不可用')
"
echo -e "\n=== 检查完成 ==="调试模式
启用详细的调试输出：
bash
export CUDA_LAUNCH_BLOCKING=1
export NCCL_DEBUG=INFO
export TORCH_DISTRIBUTED_DEBUG=DETAIL
# 运行训练
bash train_llama3_8b_dynamic.sh生产环境部署
多机训练配置
在 AIStudio 平台上进行多机训练时，平台会自动处理节点间通信配置。您只需要：
- 在任务配置界面：设置所需的 Worker 数量
- 在启动命令中：配置训练参数，平台会自动注入分布式环境变量
bash
# 多节点训练启动命令示例
cd /workspace/Megatron-LM && \
export PYTHONPATH="/workspace/Megatron-LM:$PYTHONPATH" && \
export TP_SIZE=2 && \
export CP_SIZE=2 && \
export PP_SIZE=2 && \
export MICRO_BATCH_SIZE=1 && \
export GLOBAL_BATCH_SIZE=512 && \
export DTYPE="bf16" && \
export EXP_NAME="llama3-8b-multinode" && \
export USE_EXP_STRUCTURE=1 && \
bash train_llama3_8b_dynamic.sh注意
平台会自动设置：
WORLD_SIZE：根据配置的 Worker 数量RANK：每个 Worker 的编号（0, 1, 2...）MASTER_ADDR：主节点地址MASTER_PORT：通信端口
最佳实践
1. 渐进式扩展
- 从单GPU Mock数据开始
- 验证成功后扩展到多GPU
- 最后切换到真实数据
2. 参数调优策略
- 先固定模型架构，调优训练参数
- 根据硬件资源调整并行策略
- 监控GPU利用率和内存使用
3. 检查点管理
bash
# 定期保存检查点
export SAVE_INTERVAL=500
# 启用检查点压缩
export CKPT_FORMAT="torch_dist"
# 自动清理旧检查点
export MAX_CHECKPOINTS=54. 日志管理
bash
# 结构化日志输出
export LOG_LEVEL=INFO
export LOG_INTERVAL=10
# 启用 WandB 集成（可选）
export WANDB_PROJECT="megatron-llama3"
export WANDB_RUN_NAME=$EXP_NAME总结
本教程提供了一个完整的 Megatron-LM Llama3-8B 训练解决方案，包括：
- 灵活的环境配置
- 动态化的训练脚本
- Mock 数据快速验证
- 多种硬件配置支持
- 完整的监控和调试方案
- 生产环境部署指南
通过使用提供的动态化脚本和配置模板，您可以快速适配不同的硬件环境和训练需求。建议按照教程步骤逐步验证，确保每个阶段都能正常工作后再进行下一步。
提示
如果您在使用过程中遇到问题，可以查看 Megatron-LM 官方文档 或在社区论坛寻求帮助。记得在提问时提供完整的错误信息和配置参数。