使用ROCm Megatron-LM训练大语言模型的技术指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00266/article/details/148505791

使用ROCm Megatron-LM训练大语言模型的技术指南

ROCm AMD ROCm™ Software - GitHub Home 项目地址: https://gitcode.com/gh_mirrors/ro/ROCm

概述

ROCm Megatron-LM是AMD针对其Instinct™ MI300X计算卡优化的开源大语言模型训练框架，基于NVIDIA的Megatron-LM项目进行深度定制。该框架专为AMD GPU硬件设计，支持高效训练Llama系列等主流大语言模型，提供出色的扩展性和性能表现。

核心特性

ROCm Megatron-LM具备以下关键技术特性：

高效并行训练：支持3D并行策略（TP+SP+CP）
优化计算内核：集成Transformer Engine和Fused Kernels
内存优化：支持Flash Attention 2技术
混合精度训练：提供FP8/BP16支持
自动调优：包含GEMM内核自动调优功能

支持模型

当前版本已针对以下模型进行优化：

Llama 2系列：7B/70B参数版本
Llama 3系列：8B/70B参数版本
Llama 3.1系列：8B/70B参数版本

环境准备

系统配置

NUMA设置优化：

sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'

GPU性能设置：
```
rocm-smi --setperfdeterminism 1900
```

容器环境

AMD提供预配置的Docker镜像，包含：

ROCm 6.1
PyTorch 2.4.0
Megatron Core 0.9.0
Flash Attention v2.6

获取镜像：

docker pull rocm/megatron-lm:24.12-dev

启动容器：

docker run -it --device /dev/dri --device /dev/kfd --network host \
  --ipc host --group-add video --cap-add SYS_PTRACE \
  --security-opt seccomp=unconfined --privileged \
  -v $CACHE_DIR:/root/.cache \
  --name megatron-dev-env rocm/megatron-lm:24.12-dev /bin/bash

网络性能验证

RCCL带宽测试

分布式训练前应验证节点间通信性能：

# 单节点8GPU测试
./build/all_reduce_perf -b 8 -e 10G -f 2 -g 8

# 多节点测试示例（4节点×8GPU）
mpirun -np 32 -H node1:8,node2:8,node3:8,node4:8 \
  --mca pml ucx --mca btl ^openib \
  -x NCCL_SOCKET_IFNAME=ens50f0np0 \
  -x NCCL_IB_HCA=rdma0:1,rdma1:1,rdma2:1,rdma3:1,rdma4:1,rdma5:1,rdma6:1,rdma7:1 \
  -x NCCL_IB_GID_INDEX=3 \
  $HOME/rccl-tests/build/all_reduce_perf -b 8 -e 8g -f 2 -g 1

关键环境变量说明：

NCCL_IB_HCA：指定RDMA接口
NCCL_IB_GID_INDEX：RoCE模式下的全局ID索引
TORCH_NCCL_HIGH_PRIORITY：提升通信流优先级

数据准备

数据集预处理

示例使用GPT格式数据预处理：

python tools/preprocess_data.py \
    --input my-corpus.json \
    --output-prefix my-gpt2 \
    --vocab-file gpt2-vocab.json \
    --tokenizer-type GPT2BPETokenizer \
    --merge-file gpt2-merges.txt \
    --append-eod

生成文件：

my-gpt2_text_document.bin：二进制数据文件
my-gpt2_text_document.idx：索引文件

模型训练

环境配置

网络接口设置：

export NCCL_SOCKET_IFNAME=ens50f0np0
export GLOO_SOCKET_IFNAME=ens50f0np0

Tokenizer选择：
- Llama 2：Llama2Tokenizer
- Llama 3/3.1：HuggingFaceTokenizer（需指定模型路径）

启动训练

Llama 2训练示例：

bash examples/llama/train_llama2.sh

Llama 3训练示例：

TOKENIZER_MODEL=meta-llama/Llama-3.1-8B \
bash examples/llama/train_llama3.sh

关键参数说明

| 参数 | 选项 | 说明 | |------|------|------| | MBS | 数值 | 微批次大小 | | TP | 1/2/4/8 | 张量并行度 | | TE_FP8 | 0/1 | FP8混合精度开关 | | GEMM_TUNING | 0/1 | GEMM内核调优开关 | | USE_FLASH_ATTN | 0/1 | Flash Attention开关 |