本文为稀土掘金技术社区首发签约文章，30天内禁止转载，30天后未获授权禁止转载，侵权必究！
随着 ChatGPT 的现象级走红，引领了 AI 大模型时代的变革，从而导致 AI 算力日益紧缺。与此同时，中美贸易战以及美国对华进行AI芯片相关的制裁导致 AI 算力的国产化适配势在必行。之前讲述了基于昇腾910使用ChatGLM-6B进行模型训练和推理，主要是针对 MindSpore 框架进行大模型训练，这也是华为自家研发的 AI 框架。 在昇腾上面除了使用 MindSpore 进行大模型训练，我们也可以使用 PyTorch 进行大模型训练。这两个框架也是在昇腾 NPU 上华为花主要精力维护的两款 AI 框架。
本文主要针对 MindSpore 和 Pytorch 分布式 AI 框架在进行多机多卡训练（双机16卡），为了文章具有更好的阅读体验，具体代码放置在GitHub：llm-action。
数据准备及数据格式转换
本文多机多卡训练采用斯坦福羊驼提供的数据集：下载地址。该数据集是利用OpenAI提供的GPT模型API生成质量较高的指令数据。然而，Alpaca 格式不适合用于多轮对话。参考 Vicuna 进行多轮对话训练的数据集格式，下面使用 alpaca_converter.py 脚本添加 prompts 模板，同时，转换为多轮对话格式。
kotlin
> python alpaca_converter.py \
--data_path /home/workspace/data/alpaca_data.json \
--output_path /home/workspace/data/alpaca-data-conversation.json
> tree -h /home/workspace/data
/home/workspace/data
├── [ 35M] alpaca-data-conversation.json
└── [ 22M] alpaca_data.json转换前格式：
json
[
{
"instruction": "Give three tips for staying healthy.",
"input": "",
"output": "1.Eat a balanced diet and make sure to include plenty of fruits and vegetables. \n2. Exercise regularly to keep your body active and strong. \n3. Get enough sleep and maintain a consistent sleep schedule."
},
{
"instruction": "What are the three primary colors?",
"input": "",
"output": "The three primary colors are red, blue, and yellow."
},
]转换后格式：
json
[
{
"id": "1",
"conversations": [
{
"from": "human",
"value": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\nGive three tips for staying healthy.\n\n### Response:"
},
{
"from": "gpt",
"value": "1.Eat a balanced diet and make sure to include plenty of fruits and vegetables. \n2. Exercise regularly to keep your body active and strong. \n3. Get enough sleep and maintain a consistent sleep schedule."
}
]
},
{
"id": "2",
"conversations": [
{
"from": "human",
"value": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\nWhat are the three primary colors?\n\n### Response:"
},
{
"from": "gpt",
"value": "The three primary colors are red, blue, and yellow."
}
]
},
]模型准备及模型格式转换
本文使用 LLaMA-13B 作为多机多卡训练的基座模型，关于 LLaMA 的介绍已经太多，这里不再赘述，下面下载 LLaMA-13B 模型权重文件并换成 HuggingFace Transformers 格式，有以下两种方案。
- 方案一：从 Meta 官方下载原始的LLaMA模型权重文件使用 HuggingFace Transformers 中提供的模型转换工具进行模型格式转换。
- 方案二：直接从 HuggingFace 模型仓库下载转换好的模型：decapoda-research / llama-13b-hf。
模型和数据集准备好之后，我们先来看一下基于昇腾 910的 AI 集群的网络通信。
网络通信
昇腾 910 通过自研的集合通信库（HCCL）整合了三种高速互联接口：PCIe、RoCE v2 以及 HCCS。 通常情况下，NPU 与 CPU 之间采用 PCIe 进行通信，服务器内的 NPU 之间则采用 HCCS（组间） 或 PCIe（组间） 进行通信，而服务器之间使用 RoCE 进行通信。
下面我们使用 MindSpore 进行多机多卡训练。
使用 MindSpore 进行多机多卡训练
基础环境安装
- 操作系统版本/架构：EulerOS release 2.0 (SP8)/aarch64
- NPU：8卡 910B1 64G
- Python：3.7.5
- NPU 驱动 ：23.0.rc1，下载
- NPU 固件 ：6.3.0.1.241，下载
- CANN 工具包 ：6.3.RC1.alpha003，下载
- MindSpore：2.0.0
- MindFormers：dev
对于驱动、固件以及 CANN 工具集的安装和之前在昇腾910上面安装方式基本一致，找到对应版本的驱动、固件以及CANN工具集进行安装即可。
创建新的虚拟环境：
ini
conda create -n llama_cluster python=3.7.5
conda activate llama_cluster安装 MindSpore 和 MindFormers（MindSpore Transformers 的简称）。
bash
pip install ./mindspore-2.0.0-cp37-cp37m-linux_aarch64.whl
cd /home/workspace/code
git clone https://gitee.com/mindspore/mindformers.git
cd mindformers
bash build.sh如果模型权重为 Huggingface Transformers 模型权重格式，还需要下载Pytorch 和 Huggingface Transformers 库用于将权重转换为 MindSpore 格式。
pip install torch transformers
模型权重格式转换
将 Huggingface Transformers 模型权重格式转换成 MindFormers 权重格式。
bash
cd /home/workspace/code/mindformers
python mindformers/models/llama/convert_weight.py \
--torch_ckpt_dir /home/workspace/model/vicuna-13B \
--mindspore_ckpt_path /home/worspace/model/vicuna-13B-ms数据集格式转换
使用 MindFormers 进行模型训练时，需先将带有 Prompt 模板的数据集转换为 MindRecord 格式。
css
cd /home/workspace/code/mindformers/mindformers/tools/dataset_preprocess/llama
python llama_preprocess.py \
--input_glob '/home/workspace/data/alpaca-data-conversation.json' \
--dataset_type qa \
--model_file /home/workspace/model/tokenizer.model \
--seq_length 1024 \
--repeat 3 \
--shuffle 1 \
--output_file /home/workspace/data/alpaca-fastchat1024.mindrecord参数说明：
- input_glob：输入数据集路径。
- dataset_type：数据集类型，qa（对话）/wiki（生成）。
- model_file：tokenizer模型权重。
- seq_length：最大序列长度。
- repeat: 数据重复次数。
- shuffle: 是否对数据集进行shuffle，值为0或1。
- output_file：输出数据集路径。
模型训练
首先，在每台服务器上面生成用于Ascend芯片分布式通信的芯片资源信息配置文件（RANK_TABLE_FILE）。Ascend HCCL 的 RANK_TABLE_FILE 文件提供 Ascend 分布式训练作业的集群信息。
bash
python ./mindformers/tools/hccl_tools.py --device_num "[0,8)" --server_ip=90.90.3.xx
python ./mindformers/tools/hccl_tools.py --device_num "[0,8)" --server_ip=90.90.3.yy然后，合并每个机器上的RANK_TABLE_FILE文件。可以使用官网提供的脚本merge_hccl.py进行合并，也可以手动进行合并。
bash
python ./mindformers/tools/merge_hccl.py hccl*.json合并后的文件如下所示。
json
{
"version": "1.0",
"server_count": "2",
"server_list": [
{
"server_id": "90.90.3.xx",
"device": [
{
"device_id": "0",
"device_ip": "10.0.0.36",
"rank_id": "0"
},
...
{
"device_id": "7",
"device_ip": "10.0.0.48",
"rank_id": "7"
}
],
"host_nic_ip": "reserve"
},
{
"server_id": "90.90.3.yy",
"device": [
{
"device_id": "0",
"device_ip": "10.0.0.37",
"rank_id": "8"
},
...
{
"device_id": "7",
"device_ip": "10.0.0.49",
"rank_id": "15"
}
],
"host_nic_ip": "reserve"
}
],
"status": "completed"
}之后，将合并后的RANK_TABLE_FILE文件分别复制到所有的机器上。
下面修改配置文件configs/llama/run_llama_13B.yaml 将 only_save_strategy 改为 True；同时，修改并行策略，目前支持纯数据并行（类似DeepSpeed的ZeRO3，会对模型进行切分）进行训练。
yaml
only_save_strategy: True
parallel_config:
data_parallel: 16
model_parallel: 1
pipeline_stage: 1然后，启动训练任务生成模型切分策略文件。
bash
# server1：
cd scripts
bash run_distribute.sh ../hccl_2s_16p.json ../configs/llama/run_llama_13b.yaml [0,8] train 16
# server2：
cd scripts
bash run_distribute.sh ../hccl_2s_16p.json ../configs/llama/run_llama_13b.yaml [8,16] train 16根据策略文件，切分模型。
bash
python ./mindformers/tools/transform_ckpt.py \
--src_ckpt_strategy '' \
--dst_ckpt_strategy output/strategy/ckpt_strategy_rank_0.ckpt \
--src_ckpt_dir /home/zk/workspace/model/ \
--dst_ckpt_dir /home/zk/workspace/model/ckpt-multi-node \
--prefix "checkpoint_"参数说明：
- src_ckpt_strategy：原权重切分策略，空代表不切分
- dst_ckpt_strategy：目标切分策略路径
- src_ckpt_dir：原权重路径，需按照
checkpoint_root_dir/rank_{}/*.ckpt格式 - dst_ckpt_dir：目标权重路径 prefix：生成权重的文件名前缀
模型切分完成之后，接下来，修改配置文件 configs/llama/run_llama_13B.yaml 中，模型权重文件路径、数据集路径、模型超参数等配置信息。
yaml
load_checkpoint: '/home/llama-sft/ckpt/13B-vicuna-dp'
train_dataset: &train_dataset
data_loader:
type: MindDataset
dataset_dir: "/home/ms_llama/datasets/alpaca_1024"
shuffle: True
...配置修改完成之后，在每一台服务器上面运行训练脚本。
bash
# 配置环境变量，使能GE，MS_GE_ATOMIC_CLEAN_POLICY算子内存复用，并将HCCL通信超时增大到600s
export MS_ENABLE_GE=1
export MS_GE_TRAIN=1
export MS_GE_ATOMIC_CLEAN_POLICY=1
export HCCL_CONNECT_TIMEOUT=600
# server1：
cd scripts
bash run_distribute.sh ../hccl_2s_16p.json ../configs/llama/run_llama_13b.yaml [0,8] train 16
# server2：
cd scripts
bash run_distribute.sh ../hccl_2s_16p.json ../configs/llama/run_llama_13b.yaml [8,16] train 16运行日志：
yaml
...
2023-08-07 17:27:21,717 - mindformers - INFO - .........Build context config..........
2023-08-07 17:27:21,718 - mindformers - INFO - initial recompute_config from dict: {'recompute': True, 'parallel_optimizer_comm_recompute': False, 'mp_comm_recompute': True, 'recompute_slice_activation': True}
2023-08-07 17:27:21,718 - mindformers - INFO - initial parallel_config from dict: {'data_parallel': 16, 'model_parallel': 1, 'pipeline_stage': 1, 'optimizer_shard': True, 'micro_batch_num': 16, 'vocab_emb_dp': True, 'gradient_aggregation_group': 4}
2023-08-07 17:27:21,719 - mindformers - INFO - context config is: [ParallelConfig]
...
2023-08-07 19:12:16,935 - mindformers - INFO - Epoch:[ 1/ 1], step:[ 804/ 812], loss:[0.762/0.762], time:27421.790 ms, lr:[1.2624264e-09], overflow cond: False, loss_scale: 65536.0
2023-08-07 19:12:16,935 - mindformers - INFO - Per sink_size step time: 27443.713 ms, per step time: 6860.928 ms, avg loss: 0.762
2023-08-07 19:12:44,368 - mindformers - INFO - Epoch:[ 1/ 1], step:[ 808/ 812], loss:[0.702/0.702], time:27411.658 ms, lr:[3.08156e-10], overflow cond: False, loss_scale: 65536.0
2023-08-07 19:12:44,369 - mindformers - INFO - Per sink_size step time: 27433.661 ms, per step time: 6858.415 ms, avg loss: 0.702
2023-08-07 19:13:11,812 - mindformers - INFO - Epoch:[ 1/ 1], step:[ 812/ 812], loss:[0.739/0.739], time:27419.873 ms, lr:[0.], overflow cond: False, loss_scale: 65536.0
2023-08-07 19:13:11,813 - mindformers - INFO - Per sink_size step time: 27443.428 ms, per step time: 6860.857 ms, avg loss: 0.739
2023-08-07 19:14:21,468 - mindformers - INFO - .........Training Over!.............至此，基于 MindSpore 和 MindFormers 进行多机多卡训练结束。 下面我们继续来看下使用 Pytorch 进行多机多卡训练。
使用 PyTorch 进行多机多卡训练
基础环境安装
- 操作系统版本/架构：EulerOS release 2.0 (SP8)/aarch64
- NPU：8卡 910B1 64G
- Python：3.7.5
- NPU 驱动 ：23.0.rc2.b090，下载
- NPU 固件 ：6.3.0.1.241，下载
- CANN 工具包 ：6.3.RC1.alpha003，下载
- Pytorch：1.11.0
- Huggingface Transformers：4.28.1
- DeepSpeed：0.6.0
目前，在昇腾上面针对 MindSpore 和 PyTorch 框架的支持是不同的团队支持。因此，驱动、固件以及 CANN 工具集目前还没有办法使用同一套。因此，需要根据官方的要求选择对应的版本进行重新安装，安装方式和之前在昇腾910上面安装方式基本一致。
为了避免之前的环境干扰，创建新的虚拟环境：
bash
conda create -n llama_pytorch_cluster python=3.7.5
conda activate llama_pytorch_cluster安装依赖。
css
pip3 install --upgrade pip
pip3 install einops sympy regex decorator scipy setuptools scm prompt toolkit安装 PyTorch 和 PyTorch 的适配插件 torch_npu。
bash
# 安装 1.11.0 版本
wget https://repo.huaweicloud.com/kunpeng/archive/Ascend/PyTorch/torch-1.11.0-cp37-cp37m-linux_aarch64.whl
pip3 install torch-1.11.0-cp37-cp37m-linux_aarch64.whl
# 安装 1.11.0 版本
wget https://gitee.com/ascend/pytorch/releases/download/v3.0.0-pytorch1.11.0/torch_npu-1.11.0-cp37-cp37m-linux_aarch64.whl
pip3 install torch_npu-1.11.0-cp37-cp37m-linux_aarch64.whl安装Apex。
bash
yum install -y patch libjpeg-turbo-devel dos2unix openblas git
# gcc 7.3.0版本及以上，cmake 3.12.0版本及以上。若用户要安装1.11.0版本PyTorch，则gcc需为7.5.0版本以上。
yum install -y gcc==7.3.0 cmake==3.12.0
# 获取昇腾适配的APEX源码。
git clone -b master https://gitee.com/ascend/apex.git
# 进入昇腾适配的APEX源码目录，获取原生APEX代码。
cd apex
git clone https://github.com/NVIDIA/apex.git
# 进入原生APEX代码目录，切换对应分支。
cd apex
git checkout 4ef930c1c884fdca5f472ab2ce7cb9b505d26c1a
cd ..
# 在昇腾适配的APEX源码目录下的scripts中执行命令生成NPU适配的全量代码。
cd scripts
bash gen.sh
进入原生APEX执行命令编译生成二进制安装包。
cd ../apex
python3 setup.py --cpp_ext --npu_float_status bdist_wheel
# 执行如下命令安装。如果使用非root用户安装，需要在命令后加--user。
cd dist
pip3 install apex-0.1_ascend-cp37-cp37m-linux_aarch64.whl在昇腾上面使用 Pytorch 训练 LLaMA-13B 是基于 FastChat 项目进行的适配。下面安装相关的依赖包及Huggingface Transformers库。
bash
git clone https://gitee.com/ascend/ModelZoo-PyTorch.git
cd ModelZoo-PyTorch/PyTorch/built-in/foundation/LLaMA-13B
pip3 install -e .安装 DeepSpeed 以及 Ascend NPU 适配 DeepSpeed 的插件。
bash
pip3 install deepspeed==0.6.0
git clone https://gitee.com/ascend/DeepSpeed.git
cd DeepSpeed
git checkout ad2347e
python setup.py develop安装完成之后，修改将deepspeed_npu适配包导入到deepspeed中。首先，使用 whereis 命令查看 deepspeed 安装路径，找到路径后，在.../deepspeed/bin/deepspeed文件中增加以下内容。
arduino
import deepspeed_npu下面，替换transformers库中相关文件。将源码包根目录下transformers_modify文件夹中的各个文件分别替换到transformers 安装目录下的对应位置（基于transformers 4.28.1版本）：
bash
training_args.py -> transformers/training_args.pu
trainer.py -> transformers/trainer.py
versions.py -> utils/versions.py
modeling_llama.py -> transformers/models/llama/modeling_llama.py模型训练
下面进行模型训练，本文使用 Pytorch 原生方式启动训练任务（除此之外，还可以采用DeepSpeed命令进行启动，需先安装pdsh，同时配置hostfile。）
bash
# server1：
export HCCL_CONNECT_TIMEOUT=1200
cd /home/workspace/code/ModelZoo-PyTorch/PyTorch/built-in/foundation/LLaMA-13B
nohup python -m torch.distributed.launch --nproc_per_node 8 --nnodes 2 --node_rank 0 --master_addr='90.90.3.xx' --master_port='29500' \
fastchat/train/train_mem.py \
--model_name_or_path /home/workspace/model/vicuna-13B \
--data_path /home/workspace/data/alpaca-data-conversation.json \
--fp16 True \
--output_dir /home/workspace/model/output \
--num_train_epochs 30 \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 16 \
--evaluation_strategy "no" \
--save_strategy "steps" \
--save_steps 500 \
--save_total_limit 2 \
--learning_rate 2e-5 \
--weight_decay 0. \
--warmup_ratio 0.03 \
--lr_scheduler_type "cosine" \
--logging_steps 1 \
--tf32 False \
--model_max_length 1024 \
--gradient_checkpointing True \
--lazy_preprocess True \
--deepspeed ./deepspeed_config_13B.json > train_13B.log &
# server2：
export HCCL_CONNECT_TIMEOUT=1200
cd /home/workspace/code/ModelZoo-PyTorch/PyTorch/built-in/foundation/LLaMA-13B
nohup python -m torch.distributed.launch --nproc_per_node 8 --nnodes 2 --node_rank 1 --master_addr='90.90.3.xx' --master_port='29500' \
fastchat/train/train_mem.py \
--model_name_or_path /home/workspace/model/vicuna-13B \
--data_path /home/workspace/data/alpaca-data-conversation.json \
--fp16 True \
--output_dir /home/workspace/model/output \
--num_train_epochs 30 \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 16 \
--evaluation_strategy "no" \
--save_strategy "steps" \
--save_steps 500 \
--save_total_limit 2 \
--learning_rate 2e-5 \
--weight_decay 0. \
--warmup_ratio 0.03 \
--lr_scheduler_type "cosine" \
--logging_steps 1 \
--tf32 False \
--model_max_length 1024 \
--gradient_checkpointing True \
--lazy_preprocess True \
--deepspeed ./deepspeed_config_13B.json > train_13B.log &参数说明：
- --model_name_or_path：预训练参数路径
- --data_path：数据集路径
- --fp16：参数使用fp16保存
- --num_train_epochs：训练epoch数
- --per_device_train_batch_size：每张卡上的训练batch size
- --per_device_eval_batch_size：每张卡上的评估batch size
- --gradient_accumulation_steps：梯度累积的步数
- --evaluation_strategy：评估策略
- --save_strategy：ckpt保存策略
- --save_steps：ckpt保存间隔步数
- --save_total_limit：ckpt最大保存数量
- --learning_rate：学习率
- --weight_decay：weight decay策略
- --warmup_ratio：warmup步数的比例
- --lr_scheduler_type：学习率衰减方式
- --logging_steps：训练日志打印间隔步数
- --tf32 False：使用tf32训练，npu暂不支持
- --model_max_length：模型训练的sequence length
- --gradient_checkpointing：是否开启重计算
- --deepspeed：deepspeed配置脚本路径
显存占用：
diff
+------------------------------------------------------------------------------------------------+
| npu-smi 23.0.rc2.b090 Version: 23.0.rc2.b090 |
+---------------------------+---------------+----------------------------------------------------+
| NPU Name | Health | Power(W) Temp(C) Hugepages-Usage(page)|
| Chip | Bus-Id | AICore(%) Memory-Usage(MB) HBM-Usage(MB) |
+===========================+===============+====================================================+
| 0 910B1 | OK | 327.1 54 0 / 0 |
| 0 | 0000:C1:00.0 | 57 0 / 0 64293/ 65536 |
+===========================+===============+====================================================+
...
+===========================+===============+====================================================+
| 7 910B1 | OK | 451.6 76 0 / 0 |
| 0 | 0000:42:00.0 | 68 0 / 0 64114/ 65536 |
+===========================+===============+====================================================+
+---------------------------+---------------+----------------------------------------------------+
| NPU Chip | Process id | Process name | Process memory(MB) |
+===========================+===============+====================================================+
| 0 0 | 2441447 | python | 59618 |
+===========================+===============+====================================================+
...
+===========================+===============+====================================================+
| 7 0 | 2441454 | python | 59802 |
+===========================+===============+====================================================+至此，基于昇腾910采用 Pytorch 进行多机多卡训练结束。
结语
本文讲述了在昇腾910上面使用 LLaMA-13B 进行多机多卡模型训练；但目前在昇腾910上面，无论是基于 MindSpore 还是基于 Pytorch 进行大模型训练，总体都还不是很完善，就当前时间节点而言，经常会有一些奇奇怪怪的坑；虽然，从硬件层面上来说昇腾910具备训练大模型的能力，但是软件层面（无论是生态还是AI框架的易用性，对于业界最新大模型的支持和适配等）上还有很长的路要走。
当然，目前来看，AI框架趋于收敛，大模型算法架构也趋于收敛，对于像华为昇腾这样的AI软硬件平台发展其生态来说是好事，不需要疲于拼命做太多AI框架和模型的适配工作。
参考文档：