跳到主要内容

千卡级训练配置:Llama3-70B

本文介绍如何在 MTT S5000 集群上配置 Llama3-70B 千卡级预训练。内容覆盖容器准备、集群检查、模型与数据准备、并行配置、任务启动、监控和断点恢复。文中的主机名和路径均为占位符,应按实际环境替换。

本文不提供吞吐量、MFU、训练耗时等性能指标。

本文中的命令分为两类:标注“宿主机执行”的命令在 Docker 宿主机运行,其他命令默认在训练容器内运行;带有 <...> 的内容必须替换为当前环境的实际值后再执行。正式训练前请记录镜像 digest、代码 commit、数据版本和启动参数,保证任务可复现。

完整流程如下:

1. 环境准备​

1.1 集群拓扑​

以下配置以 128 个节点、每节点 8 张 GPU、共 1024 张 GPU 为例:

参数示例值说明
节点数128根据实际可用节点调整
每节点 GPU 数8MTT S5000
总 GPU 数1024节点数 × 每节点 GPU 数
TP8优先在节点内组成张量并行组
PP8将 80 个 Transformer 层分配到流水线阶段
CP18K 上下文示例不额外启用 CP
DP161024 ÷ (TP × PP × CP)

启动前必须满足:

WORLD_SIZE = TP × PP × CP × DP

WORLD_SIZE 必须等于所有节点 slots 的总和。调整节点数、TP、PP 或 CP 后,应重新计算 DP,并先执行小规模 dry-run。

1.2 创建容器​

以下命令在每台 Docker 宿主机执行。所有节点必须使用相同镜像摘要,并挂载相同的共享目录。

docker pull registry.mthreads.com/mcctest/training-suite:v2.1.7

# 可选:记录实际使用的镜像摘要,确保所有节点使用同一镜像
docker image inspect registry.mthreads.com/mcctest/training-suite:v2.1.7 \
--format '{{index .RepoDigests 0}}'

docker run -d \
--name training_llama70b \
--runtime=mthreads \
--privileged \
--network=host \
--shm-size=256g \
-e MTHREADS_VISIBLE_DEVICES=all \
-v <SHARED_ROOT>:/data/shared \
-v <CODE_ROOT>:/data/code \
registry.mthreads.com/mcctest/training-suite:v2.1.7 \
sleep infinity

进入容器:

docker exec -it training_llama70b /bin/bash

1.3 检查软件和硬件​

mthreads-gmi
cat /usr/local/musa/version.json

python3 - <<'PY'
import torch
import torch_musa

print("PyTorch:", torch.__version__)
print("torch_musa:", torch_musa.__version__)
print("MUSA available:", torch.musa.is_available())
print("MUSA device count:", torch.musa.device_count())
PY

所有节点还需要检查:

  • GPU 数量、固件、驱动和容器镜像摘要一致;
  • 节点之间可以免密 SSH;
  • 训练网络双向可达,MCCL 集合通信测试通过;
  • 模型、数据、日志和检查点目录在全部节点上可见;
  • 系统时间同步,主机名和 DNS 解析结果一致。

2. 准备目录和 hostfile​

2.1 目录约定​

/data/shared/llama3-70b/
├── data/ # Megatron 格式的 .bin/.idx 数据
├── tokenizer/ # Llama3 Tokenizer
├── checkpoints/ # 模型和优化器检查点
└── logs/ # 每个 rank 的训练日志

/data/code/
├── Megatron-LM/
└── megatron-lm-musa-patch/

2.2 创建 hostfile​

hostfile 每行写一个可解析的主机名。不要将生产 IP 写入公开文档。

train-node-000 slots=8
train-node-001 slots=8
...
train-node-127 slots=8

检查节点数和总 slot:

HOSTFILE=/data/shared/llama3-70b/hostfile

grep -Ev '^#|^$' "${HOSTFILE}" | wc -l
awk -F'slots=' '/slots=/{sum += $2} END{print sum}' "${HOSTFILE}"

NODE_COUNT=$(grep -Ev '^#|^$' "${HOSTFILE}" | wc -l)
SLOT_COUNT=$(awk -F'slots=' '/slots=/{sum += $2} END{print sum}' "${HOSTFILE}")
test "${NODE_COUNT}" -eq 128
test "${SLOT_COUNT}" -eq 1024

hostfile 的顺序通常会影响 rank 分配。第一个节点是否作为 master、节点名是否可被所有节点解析,以及 launcher 使用的 SSH 用户和端口,必须以实际启动脚本为准。

3. 准备模型和数据​

3.1 模型结构​

Llama3-70B 的核心结构参数如下:

MODEL_ARGS=(
--num-layers 80
--hidden-size 8192
--ffn-hidden-size 28672
--num-attention-heads 64
--group-query-attention
--num-query-groups 8
--seq-length 8192
--max-position-embeddings 8192
--position-embedding-type rope
--normalization RMSNorm
--swiglu
--untie-embeddings-and-output-weights
--bf16
)

模型结构应与所使用的 Llama3-70B 配置文件保持一致。不要仅复制参数后直接加载不匹配的检查点。

3.2 数据准备​

3.2.1 下载训练资源​

本文使用的数据已经转换为 Megatron IndexedDataset 格式,由同名的 .bin 数据文件和 .idx 索引文件组成。训练脚本可以直接读取这组文件,不需要再执行 Parquet 转换或分词预处理。

DATASET_DIR=/data/shared/llama3-70b
DOWNLOAD_DIR=$DATASET_DIR/downloads
DATA_DIR=$DATASET_DIR/dataset

mkdir -p "$DOWNLOAD_DIR" "$DATA_DIR" "$DATASET_DIR/tokenizer"
cd "$DOWNLOAD_DIR"

wget \
https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/llama3-datasets/wudao_llama3bpe_content_document.bin \
-O "$DATA_DIR/dedup-md5-pile-pile-cc_text_document.bin"

wget \
https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/llama3-datasets/wudao_llama3bpe_content_document.idx \
-O "$DATA_DIR/dedup-md5-pile-pile-cc_text_document.idx"

curl -fL \
https://mt-ai-infra.tos-cn-beijing.volces.com/train_llama3_70b_multi_nodes.sh \
-o train_llama3_70b_multi_nodes.sh

chmod +x train_llama3_70b_multi_nodes.sh

下载时将文件重命名为训练脚本固定使用的 dedup-md5-pile-pile-cc_text_document 前缀。该前缀不能包含 .bin 或 .idx 扩展名。

检查两个文件是否同时存在且不为空:

DATA_PREFIX=/data/shared/llama3-70b/dataset/dedup-md5-pile-pile-cc_text_document

test -s "${DATA_PREFIX}.bin"
test -s "${DATA_PREFIX}.idx"
ls -lh "${DATA_PREFIX}.bin" "${DATA_PREFIX}.idx"

数据链接只提供 .bin/.idx 文件,不包含训练脚本要求的 Hugging Face Tokenizer。应将与该数据分词方式和 Llama3 模型词表一致的 Tokenizer 文件放入 /data/shared/llama3-70b/tokenizer,并确认其中包含 tokenizer.json、tokenizer_config.json 等必要文件。

检查训练脚本内容和参数说明:

sed -n '1,220p' "$DOWNLOAD_DIR/train_llama3_70b_multi_nodes.sh"
bash "$DOWNLOAD_DIR/train_llama3_70b_multi_nodes.sh" --help || true

脚本可能包含与当前集群不同的节点数、路径或并行参数。正式运行前应确认工作目录、hostfile、Tokenizer、检查点路径和并行参数均与当前环境一致,并先执行小规模 dry-run。

训练脚本会根据 --dataset_dir 自动生成以下路径:

DATA_PATH=/data/shared/llama3-70b/dataset/dedup-md5-pile-pile-cc_text_document
TOKENIZED_MODEL=/data/shared/llama3-70b/tokenizer

3.2.2 验证数据路径​

DATASET_DIR=/data/shared/llama3-70b
DATA_PREFIX=$DATASET_DIR/dataset/dedup-md5-pile-pile-cc_text_document

ls -lh "${DATA_PREFIX}.bin" "${DATA_PREFIX}.idx"
test -f "$DATASET_DIR/tokenizer/tokenizer.json"

正式启动前检查数据来源、授权、分片清单和 checksum,并在少量样本上验证 Tokenizer 编解码结果。

4. 配置训练任务​

4.1 并行参数​

TP_SIZE=8
PP_SIZE=8
CP_SIZE=1
WORLD_SIZE=1024
DP_SIZE=$((WORLD_SIZE / (TP_SIZE * PP_SIZE * CP_SIZE)))

test $((TP_SIZE * PP_SIZE * CP_SIZE * DP_SIZE)) -eq "${WORLD_SIZE}"

对应的 Megatron 参数:

PARALLEL_ARGS=(
--tensor-model-parallel-size "${TP_SIZE}"
--pipeline-model-parallel-size "${PP_SIZE}"
--context-parallel-size "${CP_SIZE}"
--sequence-parallel
--use-distributed-optimizer
)

4.2 Batch 参数​

不要从性能记录中复制 batch。根据显存和数据规模先完成小规模验证,再逐步放大:

MICRO_BATCH_SIZE=<MICRO_BATCH_SIZE>
NUM_MICROBATCHES=<NUM_MICROBATCHES>
GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * DP_SIZE))

4.3 检查启动脚本​

不同 TrainingSuite 镜像中的脚本名称可能不同,先查找 Llama3-70B 启动入口:

find /data/code -type f \
\( -iname '*llama*70b*.sh' -o -iname '*llama3*70b*.py' \)

确认脚本使用的代码版本、模型参数、数据路径、Tokenizer、hostfile、并行参数和检查点路径与本文一致。

如果启动脚本使用 torchrun、DeepSpeed 或其他 launcher,还必须确认 MASTER_ADDR、MASTER_PORT、NNODES、NODE_RANK、NPROC_PER_NODE 等变量的来源。不要在多个节点上手工重复执行同一个启动命令,除非脚本明确要求这样做。

5. 启动训练​

5.1 小规模 dry-run​

先使用单节点或一个模型并行组运行少量 step,检查:

  • Forward、Backward 和 optimizer step 正常;
  • loss 为有限值,没有 NaN/Inf;
  • TP、PP 和 DP 通信正常;
  • 数据读取没有长时间阻塞;
  • 检查点可以保存并重新加载。

5.2 启动千卡任务​

建议通过统一入口脚本提交,不要逐台手工运行不同命令:

DATASET_DIR=/data/shared/llama3-70b
LAUNCH_SCRIPT=$DATASET_DIR/downloads/train_llama3_70b_multi_nodes.sh

bash "$LAUNCH_SCRIPT" \
--dataset_dir "$DATASET_DIR" \
--data_format bf16

训练脚本从当前工作目录读取 ./hostfile 和 ./pretrain_gpt.py,因此必须在包含这两个文件的训练工作目录中执行。脚本中的 CHECKPOINT_PATH 当前仍是固定路径,正式运行前应修改为实际检查点目录。

建议先验证启动入口能够打印参数并退出,再提交正式任务:

bash "$LAUNCH_SCRIPT"

不传 --dataset_dir 时脚本会直接退出,可用该行为检查参数解析;不要把无参数命令当作正式启动命令。

6. 监控训练​

tail -F /data/shared/llama3-70b/logs/<RANK0_LOG>

监控重点包括:

  • 所有 rank 是否完成初始化;
  • loss、梯度范数和学习率是否正常;
  • 是否出现 NaN、Inf、OOM、通信超时或数据读取错误;
  • 节点、GPU、网络和共享存储是否持续健康;
  • 检查点是否按计划写入且可以校验。

本文不记录或展示性能指标。

7. 断点恢复和停止任务​

恢复前确认最新检查点包含模型、优化器、scheduler、RNG 和数据位置状态。恢复后检查 step 和 loss 是否连续。

恢复命令中的参数名以当前训练脚本为准,常见形式如下:

export CHECKPOINT_DIR=/data/shared/llama3-70b/checkpoints/<STEP_DIR>
bash <LLAMA3_70B_LAUNCH_SCRIPT> --load "${CHECKPOINT_DIR}"

如果改变 TP、PP、CP、数据顺序或代码 commit,不能直接假设原 checkpoint 可以恢复;应先在副本目录执行加载验证。

停止任务时先终止调度入口,再清理残留进程:

pkill -f '<LLAMA3_70B_LAUNCH_SCRIPT>'

pkill 可能误杀同一节点上的其他任务。生产环境应优先使用 launcher 返回的任务 ID 或 PID 文件停止任务,并在所有节点确认没有残留训练进程。

不要直接删除正在写入的检查点目录。

FAQ​

Q1: 如何确认所有节点的软件版本完全一致?

千卡训练开始前,应统一检查操作系统、驱动、MUSA SDK、PyTorch、torch_musa、训练镜像和代码提交。可在管理节点按 hostfile 逐台采集:

HOSTFILE=/data/shared/llama3-70b/hostfile

while read -r host _; do
echo "===== ${host} ====="
ssh "${host}" '
uname -r
cat /usr/local/musa/version.json
python3 -c "import torch, torch_musa; print(torch.__version__, torch_musa.__version__)"
git -C /data/code/Megatron-LM rev-parse HEAD
git -C /data/code/megatron-lm-musa-patch rev-parse HEAD
'
done < <(grep -Ev "^#|^$" "${HOSTFILE}")

发现任一节点版本、镜像摘要或代码提交不同,应先修复或隔离该节点,再启动正式任务。

Q2: 如何确认全部 GPU 都可用?

先在每个节点运行:

mthreads-gmi

然后在训练容器内检查 PyTorch 能看到的设备数量:

python3 - <<'PY'
import torch
import torch_musa

assert torch.musa.is_available()
count = torch.musa.device_count()
print("MUSA device count:", count)
assert count == 8, f"expected 8 devices, got {count}"
PY

同时确认没有不可恢复硬件错误、异常温度、掉卡或残留训练进程。存在问题的节点不要加入 hostfile。

Q3: 如何确认 MCCL 通信正常?

通信测试应按以下顺序逐级扩大:

  1. 单节点 8 卡;
  2. 两节点;
  3. 一个模型并行组;
  4. 全部训练节点。

至少覆盖 AllReduce、AllGather、ReduceScatter 和点对点通信。使用当前环境提供的 MCCL 测试工具:

<MCCL_TEST_COMMAND>

如果小规模正常、全规模失败,重点检查训练网卡、路由、交换网络、端口、防火墙、节点拓扑和各节点 MCCL 环境变量是否一致。

Q4: 启动后部分 rank 一直等待怎么办?

优先查看最早出现异常的 rank,不要只看最终的超时信息。依次检查:

  1. hostfile 中的节点和 slot 数是否正确;
  2. 节点之间是否可以免密 SSH;
  3. WORLD_SIZE = TP × PP × CP × DP 是否成立;
  4. rendezvous 地址和端口是否一致;
  5. 每个容器中的代码、数据和工作目录是否存在;
  6. MCCL 通信测试是否覆盖了相同节点。
Q5: 训练出现 OOM、NaN 或 Inf 怎么办?

先回到小规模 dry-run,并按顺序检查:

  1. 降低 micro batch;
  2. 启用或增加 activation checkpointing;
  3. 检查 BF16、FP8 和 Transformer Engine 配置是否匹配;
  4. 检查数据中是否存在异常样本;
  5. 检查梯度裁剪、学习率和 loss scale;
  6. 确认所有节点使用同一份模型配置。

一次只修改一类参数,并保留修改前后的配置和日志。

Q6: 恢复训练后 loss 不连续怎么办?

确认模型、优化器、scheduler、RNG 和数据位置来自同一次完整保存。不要混用不同 step、不同代码提交或不同并行拓扑生成的检查点。恢复前可先在较小规模执行一次加载验证。

相关资料​