A21a_昇腾MindSpeed迁移Megatron并行模型

昇腾MindSpeed迁移Megatron并行模型

作者：MindSpeed[竹林七贤, 夏, CurryRice], 瑾丞

目 录

1. Megatron迁移背景概述
2. 安装环境准备
3. Megatron迁移步骤
4. 启动Megatron多机多卡训练
5. Megatron并行训练验证
6. 常见问题与解决

(这是昇腾知识体系的配套预览材料，转载随意，如反馈bug请移步原文：链接)

第1章 Megatron迁移背景概述

1.1 迁移目的

当前大语言模型训练需要处理海量参数与数据，NVIDIA的Megatron-LM框架通过并行计算技术解决了这一问题。随着自主算力平台的兴起，华为昇腾的硬件架构（NPU、集合通信模式）和软件特性（如CANN软件包、MindSpeed组件等）与GPU存在差异，因此需要处理分布式训练框架的迁移问题：

1. 适配昇腾硬件环境（解决CUDA接口不兼容问题）
2. 激活昇腾专属性能（利用NPU的分布式训练优势）
3. 降低训练成本（通过更高效的显存管理和算子融合技术）

1.2 迁移整体流程图解

1.3 关键步骤清单

1.4 主要的迁移变化点

⚠️ 注意：迁移过程中需特别关注版本匹配。建议使用Python3.9以上版本，避免因类型注解兼容性问题导致脚本报错。

1.5 新手入门建议

1. 单机验证优先：先在一台昇腾服务器上完成模型训练，确保基础流程正确
2. 渐进式扩展：从2台设备开始进行多机训练，再逐步增加节点数量
3. 定期检查通信配置（HCCL_IF_IP等环境变量设置）

通过上述迁移方案，开发者可将基于Megatron-LM的模型无缝转移到昇腾集群，在保证精度不受影响的前提下，充分发挥国产芯片的性能优势。

第2章 安装环境准备

在将Megatron并行框架迁移到昇腾NPU之前，必须完成软件环境配置。只有当这些基础条件都满足时，后续迁移模型的操作才能成功。在实际环境中，大多数启动失败的问题都源于这一步配置不当——因此建议严格按照版本配套表执行操作，并保留环境配置日志以备故障排查。通过本章操作您将配置好一个能直接运行训练脚本的开发环境：

1. Python虚拟环境（conda）提供干净的依赖管理
2. 昇腾软件栈（CANN等）确保NPU硬件与计算框架能通信
3. Megatron-LM代码仓库和MindSpeed插件建立完整模型迁移基础

步骤1：创建Python虚拟环境

建议为昇腾平台独立创建conda虚拟环境，确保后续操作不受其他Python版本干扰。

conda create --name megatron_demo python=3.10  # 创建名为megatron_demo的Python3.10环境
conda activate megatron_demo                 # 激活该环境

补充说明：

• 建议使用独立虚拟环境（而非全局安装）避免版本冲突
• Python 3.9及以上版本支持新特性，若出现报错需升级

步骤2：安装昇腾软件栈基础组件

这个环节包括安装NPU驱动和基础加速库，需确保各软件包的版本互相匹配，属于关键步骤，详情请参考昇腾官方提供的快速体验指南（链接）进行安装，此处不做赘述。

步骤3：准备训练代码仓库

这里需要两套代码：原始Megatron-LM框架（来自NVIDIA）和昇腾适配的插件MindSpeed。

# 安装MindSpeed
git clone https://gitee.com/ascend/MindSpeed.git -b core_r0.10.0
pip install -e MindSpeed

# 获取原版Megatron-LM代码
git clone https://github.com/NVIDIA/Megatron-LM.git # 克隆源码
cd Megatron-LM
git checkout core_r0.10.0 # 切换至适配版本

mkdir gpt-tokenizer        # 创建词典文件夹
mkdir gpt_pretrain_data    # 创建预训练数据目录
mkdir ckpt                 # 创建模型检查点保存路径

版本配套要求：

步骤4：安装数据集依赖的库

# 安装数据处理所需的库（用于parquet转jsonn）
pip3 install nltk pyarrow pandas

安装过程可能遇到的报错及应对方案：

完整流程总结

第3章 Megatron迁移步骤详解

本章节介绍模型迁移到昇腾平台的核心操作流程，建议在阅读前先通读相关背景知识，了解NPU与GPU架构差异后，再按步骤执行。

3.1 数据处理全流程

大模型需要结构化数据格式才能运行，此步骤完成原始数据的预处理适配。

获取NVIDIA GPT-3.5 Turbo项目中包含的两个关键文件： vocab.json（词汇表）和 merges.txt（合并规则）。下载 vocab.json 和 merges.txt（链接）并放到 Megatron-LM/gpt-tokenizer 目录内，然后分别重命名为 gpt2-vocab.json 和 gpt2-merges.txt。

下载Alpaca数据集链接，放到服务器上任意目录内，示例目录为 /home/datasets/Alpaca。

语料格式转换：读取 parquet 格式的原始语料，并将其转换为 JSON 格式，以便后续处理

# 依赖安装
pip3 install nltk pyarrow pandas

cd /home/datasets/Alpaca/
python convert_parquet.py

convert_parquet.py 脚本代码如下：

import json
import pandas as pd
data_df = pd.read_parquet("train-00000-of-00001-a09b74b3ef9c3b56.parquet")
data_df['text'] = data_df['text'].apply(lambda v: json.dumps({"text": v}))
with open("alpaca_json.json", encoding='utf-8', mode='w') as f:
    for i, row in data_df.iterrows():
        f.write(row['text'])
        f.write('\n')

数据预处理：

1. 修改 Megatron-LM/tools/preprocess_data.py 脚本，在 import torch 下新增一行 import mindspeed.megatron_adaptor，如下所示：

import torch
import mindspeed.megatron_adaptor # 新增代码
import numpy as np

2. 在 Megatron-LM 目录下执行预处理脚本

python tools/preprocess_data.py \
   --input /home/datasets/Alpaca/alpaca_json.json \
   --output-prefix ./gpt_pretrain_data/alpaca  \
   --tokenizer-type GPT2BPETokenizer \
   --vocab-file ./gpt-tokenizer/gpt2-vocab.json \
   --merge-file ./gpt-tokenizer/gpt2-merges.txt \
   --append-eod \
   --log-interval 1000 \
   --workers 8

执行成功后，将在 gpt_pretrain_data 目录下生成两个文件：alpaca_text_document.bin（二进制训练数据）和alpaca_文本前缀.idx（索引文件），这两个文件共同构成模型输入源。

3.2 启动昇腾平台预训练

这是整个迁移的关键步骤：通过单行代码修改激活NPU优化特性

第一步：激活NPU加速功能

在Megatron-LM/pretrain_gpt.py文件中增加：

import torch
# 下一行插入MindSpeed适配器
import mindspeed.megatron_adaptor # 启用NPU专属计算优化模块
from functools import partial

第二步：调整分布式训练参数

示例脚本参考原GPU版本pretrain_gpt_distributed.sh（脚本链接），并修改其中的配置：

# 修改前
...
CHECKPOINT_PATH=<Specify path>
VOCAB_FILE=<Specify path to file>/gpt2-vocab.json
MERGE_FILE=<Specify path to file>/gpt2-merges.txt
DATA_PATH=<Specify path and file prefix>_text_document
...
OUTPUT_ARGS="
    --log-interval 100 \
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters 10
"

torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $CHECKPOINT_PATH \
    --load $CHECKPOINT_PATH

# 修改后
...
CHECKPOINT_PATH=./ckpt # 修改为实际路径，下同
VOCAB_FILE=./gpt-tokenizer/gpt2-vocab.json
MERGE_FILE=./gpt-tokenizer/gpt2-merges.txt
DATA_PATH=./gpt_pretrain_data/alpaca_text_document
...

GPT_ARGS="
    ...
    --transformer-impl local \ # 新增配置
    --no-masked-softmax-fusion # 新增配置
"
...

OUTPUT_ARGS="
    --log-interval 1 \ # 从 100 改为 1
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters 10
"

torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $CHECKPOINT_PATH \
    --ckpt-format torch # 删去 --load，新增 --ckpt-format torch

3.3 执行训练与效果验证

启动训练

在 Megatron-LM 目录下直接执行：

bash pretrain_gpt_distributed.sh

预期输出：

• 若出现OOM报错，可尝试减少--micro-batch-size的数值
• 日志显示迭代次数递增（如下图所示），说明迁移成功，即可停止训练。

完整的 pretrain_gpt_distributed.sh 示例代码如下：

#!/bin/bash

# Runs the "345M" parameter model

export CUDA_DEVICE_MAX_CONNECTIONS=1

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=localhost
MASTER_PORT=6000
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

CHECKPOINT_PATH=./ckpt
VOCAB_FILE=./gpt-tokenizer/gpt2-vocab.json
MERGE_FILE=./gpt-tokenizer/gpt2-merges.txt
DATA_PATH=./gpt_pretrain_data/alpaca_text_document

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

GPT_ARGS="
    --num-layers 24 \
    --hidden-size 1024 \
    --num-attention-heads 16 \
    --seq-length 1024 \
    --max-position-embeddings 1024 \
    --micro-batch-size 8 \
    --global-batch-size 64 \
    --lr 0.00015 \
    --train-iters 500000 \
    --lr-decay-iters 320000 \
    --lr-decay-style cosine \
    --min-lr 1.0e-5 \
    --weight-decay 1e-2 \
    --lr-warmup-fraction .01 \
    --clip-grad 1.0 \
    --fp16 \
    --transformer-impl local \
    --no-masked-softmax-fusion
"

DATA_ARGS="
    --data-path $DATA_PATH \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
    --split 949,50,1
"

OUTPUT_ARGS="
    --log-interval 1 \
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters 10
"

torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $CHECKPOINT_PATH \
    --ckpt-format torch

3.4 迁移过程关键步骤

至此，您已完成将Megatron分布式训练框架适配到昇腾平台的关键步骤。下一步可进行多机训练扩展或使用MindSpeed提供的优化特性进一步提升训练效率。

若在迁移过程中遇到疑难问题，建议参考MindSpeed开源仓的介绍提交ISSUE（开源仓链接）。

第4章 启动Megatron多机多卡训练

在单台服务器上进行大模型训练时，可能因算力不足或内存瓶颈导致效率低下。此时需要多台设备协同工作——即多机多卡模式（Multi-Node Multi-GPU/NPU），通过统一调度多个硬件的计算资源来突破单点限制。本章节将以双节点为例，完整展示如何将Megatron-LM训练任务扩展到跨服务器协作场景。

在训练前需要确保两台服务器之间通信正常，无其它进程干扰。确保环境配置一致（包括 conda 环境、cann 环境等），代码一致，且均能正常进行单机训练。

4.1 核心配置逻辑解析

在昇腾平台运行多机训练时需完成三项基础条件：

4.2 配置操作步骤

多机脚本适配

修改 pretrain_gpt_distributed.sh 脚本中的 MASTER_ADDR（使用真实的主节点 IP）、NNODES（从 1 改成 2）和 NODE_RANK（主节点为 0，从节点为 1）。各参数说明及修改后的值见下表：

此外，还需将 --global-batch-size 64 改为 --global-batch-size 128。

如果未使用共享存储，需要修改 megatron/core/datasets/gpt_dataset.py，目的是为了让从节点也能生成数据集缓存，避免从节点获取不到数据集缓存：

# 原始
if not path_to_cache or (
    not cache_hit
    and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0) # 删除这行
):

# 修改后
if not path_to_cache or (
    not cache_hit
):

修改后的多机多卡训练脚本为（以主节点为例）：

#!/bin/bash

# Runs the "345M" parameter model

export CUDA_DEVICE_MAX_CONNECTIONS=1

GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=xxx.xxx.xxx.xxx
MASTER_PORT=6000
NNODES=2
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

CHECKPOINT_PATH=./ckpt
VOCAB_FILE=./gpt-tokenizer/gpt2-vocab.json
MERGE_FILE=./gpt-tokenizer/gpt2-merges.txt
DATA_PATH=./gpt_pretrain_data/alpaca_text_document

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

GPT_ARGS="
    --num-layers 24 \
    --hidden-size 1024 \
    --num-attention-heads 16 \
    --seq-length 1024 \
    --max-position-embeddings 1024 \
    --micro-batch-size 8 \
    --global-batch-size 128 \
    --lr 0.00015 \
    --train-iters 500000 \
    --lr-decay-iters 320000 \
    --lr-decay-style cosine \
    --min-lr 1.0e-5 \
    --weight-decay 1e-2 \
    --lr-warmup-fraction .01 \
    --clip-grad 1.0 \
    --fp16 \
    --transformer-impl local \
    --no-masked-softmax-fusion
"

DATA_ARGS="
    --data-path $DATA_PATH \
    --vocab-file $VOCAB_FILE \
    --merge-file $MERGE_FILE \
    --split 949,50,1
"

OUTPUT_ARGS="
    --log-interval 1 \
    --save-interval 10000 \
    --eval-interval 1000 \
    --eval-iters 10
"

torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $CHECKPOINT_PATH \
    --ckpt-format torch

4.3 常见优化建议

• 全局批大小：将单机的--global-batch-size 64调整为128
• 并行策略：根据硬件规模选择适配的并行组合（如TP=2 PP=2 CP=1）
• 日志监控：设置--log-interval 1以实时跟踪训练进程

4.4 小结

本章节通过三个关键动作实现多机协同训练：

1. 网络参数配置统一（IP设置+通信库优化）
2. 训练脚本的并行拓扑定义
3. 数据处理逻辑适配

在完成以上步骤后，用户即可在昇腾平台体验与GPU环境相媲美的分布式训练能力。如遇网络异常，请优先检查硬件连接状态，并联系华为技术支持进行进一步诊断。

第5章 Megatron并行训练验证

在完成模型迁移后，本章节将指导您如何验证Megatron-LM是否成功适配到昇腾平台，并确保多机/单机训练任务能稳定运行。重点包括以下步骤：

1. 配置分布式通信参数（主节点和从节点）
2. 启动训练并确认日志
3. 排查可能的网络或环境问题

5.1 网络环境准备

为了确保多机训练过程中的通信稳定，需要在所有参与节点执行以下操作：

export HCCL_IF_IP=$(ifconfig查询到的本机IP)
export GLOO_SOCKET_IFNAME=$(ifconfig查到的网卡名称)

示例代码：

# 主从节点均需配置以下环境变量
import os

os.environ["HCCL_IF_IP" = "192.168.0.1"]
os.environ["GLOO_SOCKET_IFNAME" = "enp189s0f0")

5.2 启动并行训练

Step1、 先在主节点上执行多机多卡训练脚本：

bash pretrain_gpt_distributed.sh

主节点处于下图状态，说明在等待从节点响应：

Step2、 然后在从节点上执行多机多卡训练脚本。在从节点上观察到下图所示训练日志，说明多机多卡训练适配完成，可以停止训练：

5.3 验证流程图解

5.4 验证结果确认标准

5.5 验证流程总结

本章通过环境变量设置、脚本参数调整以及网络检查三个维度，确保了Megatron-LM模型在昇腾平台上的能成功启动并行训练。实际操作中需注意以下关键点：

• 每个节点的网卡名和ip地址必须严格一致
• 环境配置建议优先使用source set_env.sh
• 多机训练中所有节点应同时启动训练脚本

通过以上验证，您将能够确认模型迁移已成功，并具备进行更大规模集群训练的能力。

第6章 常见问题与解决指南

提示：分布式训练涉及多个节点通信协同，若出现网络或软件栈异常，请按本章指引逐步排查。我们整理了三大典型故障场景及其修复方法，建议先阅读"故障分类说明"再尝试操作。

6.1 主要故障类型说明

6.2 分场景解决方案

问题1：节点间无法建立通信

错误特征 执行训练时主节点卡空等待，从节点报错"RuntimeError: Gloo connectFullMesh failed with … no error"

操作说明

1. 获取网卡信息（主从节点均需执行）：

ifconfig  # 查看本机IP地址及对应网卡

1. 设置通信绑定策略（主节点示例）:

export HCCL_IF_IP=10.136.0.199
# 若使用Gloo通信库时，需额外绑定网卡名：
export GLOO_SOCKET_IFNAME=enp189s0f0

问题2：HCCL通信库初始化失败

错误特征 日志包含"hcclCommInitRootInfoConfig…error code is 4"，主节点持续等待从节点响应

RuntimeError: create_config:build/CMakeFiles/torch_npu.dir/compiler_depend.ts:201 HCCL function error: hcclCommInitRootInfoConfig(numRanks, &rootInfo, rank, config, &(comm->hcclComm_)), error code is 4
[ERROR] 2025-04-21-19:16:14 (PID:325336, Device:0, RankID:8) ERR02200 DIST call hccl api failed.

解决步骤表

注意：若使用GPU版本Megatron-LM时的NCCL配置，需替换为昇腾平台专用HCCL通信库（Ascend Communication Library for NPU, HCCL）。请通过安装CANN软件栈后自动完成兼容。

问题3：硬件链路异常

错误特征

报错"error code is 19 … The network port is down"并附带hccn_tool诊断建议，如：

RuntimeError: create_config:build/CMakeFiles/torch_npu.dir/compiler_depend.ts:201 HCCL function error: hcclCommInitRootInfoConfig(numRanks, &rootInfo, rank, config, &(comm->hcclComm_)), error code is 19
[ERROR] 2025-04-21-19:11:46 (PID:202188, Device:6, RankID:14) ERR02200 DIST call hccl api failed.
EJ0002: [PID: 202188] 2025-04-21-19:11:44.470.610 The network port is down.Suggest: Run the following hccn_tool commands: 1. for i in {0..7}; do hccn_tool -i $i -optical -g ; done | grep present: Check whether the optical module is in position.2. for i in {0..7}; do hccn_tool -i $i -ip -g ; done. Check whether the IP address is configured. 3. for i in {0..7}; do hccn_tool -i $i -lldp -g: Check whether the switch is connected

修复步骤说明

1. 执行网络诊断（主从节点分别操作）:

# 检查光模块状态（示例）
for i in {0..7}; do hccn_tool -i $i -optical -g ; done | grep present

# 获取IP地址配置
for i in {0..7}; do hccn_tool -i $i -ip -g; done

# 检查交换机连接状态
for i in {0..7}; do hccn_tool -i $i -lldp -g ; done

1. 结果解读与处理
   • 若光模块检测失败（present字段异常）：拔插HDMI线缆，重启服务器
   • 若IP未配置成功：检查set_env.sh中MASTER_ADDR是否指向正确主节点
   • 若LLDP协议无响应：确认交换机已安装华为认证固件
2. 联系技术支持（若上述操作无法解决）：
   • 提供错误日志截图（含HCCL_IF_IP值和hccn_tool输出）
   • 附带服务器硬件序列号和CANN版本信息

6.3 高频故障修复总结表

6.4 预防性检查清单

执行以下命令确保训练环境无误：

# 检查节点间连通性（主节点）
ping 主节点IP && ping 从节点IP

# 确认CANN软件栈版本一致性
cat $CANN_INSTALL_PATH/version.txt

# 核心依赖检查
pip show mindspeed torch-npu

# 通信库状态诊断
nccl-tests/1byteAllReduceTest -c hccl -n 8

初学者建议：首次迁移时，可使用 from torch_npu.contrib import transfer_to_npu 替换原生 CUDA，该接口会自动处理大部分兼容性问题。

6.5 联系技术支持指引

若遇到以下情况，请进行故障排查，如有必要可联系华为技术提交工单：

1. 执行hccn_tool后仍提示端口down
2. 多节点训练启动30秒内未出现"waiting for workers"日志
3. 出现未知算子报错（非NCCL/GLOO范畴）

提交工单时，请提供以下故障信息清单：

• 完整错误截图（含RankID和Device编号）
• 环境配置详情（Python版本、MindSpeed分支、CANN版本）
• 通信拓扑结构图（建议使用hccn_tool -topo获取）