使用 AWS Trainium 加速芯片进行 Llama 2 继续预训练

最近一年多来，大语言模型（LLM）引起了跨多个行业的开发者、科学家、技术专家、企业家和高管的极大兴趣和关注。这些模型可用于问答、总结、翻译等应用中，以及客户支持的会话式代理，用于营销的内容创作和编码助手，极大地提高了相关领域的工作效率。

作为垂直领域数据提供商，通联数据也在积极探索大语言模型训练的各种可能路径，通过和亚马逊云科技的深入合作，尝试了基于亚马逊云科技的自研芯片，使用自有数据对 Llama2 进行了继续预训练的工作，并取得了良好的结果。这次给大家分享一下使用亚马逊云科技的 Trainium 芯片对 Llama2 进行继续预训练的整个流程，供大家参考。

名词解释

AWS Trainium：AWS Trainium 是亚马逊云科技专门为超过 1000 亿个参数模型的深度学习训练打造的第二代机器学习（ML）加速芯片。每个 Amazon EC2 Trn1 实例有 16 个 AWS Trainium 加速芯片，每个 Trainium 加速芯片具有 32GB 的高带宽内存，提供高达 190 TFLOPS 的 FP16/BF16 计算能力，并采用了实例内、超高速非阻塞互连技术 NeuronLink。https://aws.amazon.com/cn/machine-learning/trainium/
Neuronx-nemo-megatron：Neuronx-nemo-megatron （也称为 AWS Neuron Reference for NeMo Megatron）是包含了 NeMo 和 Apex 开源软件包的修改版本，经过适配以在 AWS Neuron 和 AWS EC2 Trn1 实例上使用。目前 Neuronx-nemo-megatron 已经实现了使用成千上万个 Trainium 加速芯片进行千亿参数规模的模型预训练，并支持高级训练功能，如 3D 并行、序列并行和检查点激活。https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nemo-megatron/index.html
Llama2：Llama2 是 Meta AI 推出的一系列预训练和微调的大型语言模型（LLMs），参数规模从 70 亿到 700 亿不等。Llama2 可以用于多种自然语言处理任务，比如文本生成、问答、对话生成、情感分析等，它可以通过预训练和微调来适应不同的任务和应用场景。
继续预训练：基于基础模型（比如 Llama2，Baichuan），继续从大量的无标签文本数据使用自监督学习的方式继续学习文本中的潜在知识，比较常见的场景是中文能力增强、垂直领域知识增强。

Trainium

每个 Neuron 核心均由如下所示的四个专用引擎和大容量片上 SRAM 存储器组成，利用每个引擎高效地执行计算。

张量引擎：基于脉动阵列的张量引擎，针对卷积等矩阵运算进行了优化。使用 FP16/BF16 的计算性能超过 100 TFLOPS，达到 Inferentia Neuron 核心性能的 6 倍。
矢量引擎：针对批量归一化和池化处理进行优化的引擎。
标量引擎：针对 ReLU 等激活函数进行优化的引擎。
通用 SIMD 引擎：Trainium 中新的内置 SIMD 引擎，支持自定义运算符。现在可以用 C++ 编写自定义运算符。

Neuron SDK

AWS Neuron 是一种 SDK，可优化在 AWS Inferentia 和 Trainium 上执行的复杂神经网络模型的性能。AWS Neuron 包括深度学习编译器、运行时和工具，这些工具与 TensorFlow 和 PyTorch 等流行框架原生集成，它预装在 AWS Deep Learning AMI 和 Deep Learning Containers 中，供客户快速开始运行高性能且经济高效的推理。

Neuron 编译器接受多种格式（TensorFlow、PyTorch、XLA HLO）的机器学习模型，并优化它们以在 Neuron 设备上的运行。Neuron 编译器在机器学习框架内调用，其中模型由 Neuron Framework 插件发送到编译器。生成的编译器工件称为 NEFF 文件（Neuron 可执行文件格式），该文件又由 Neuron 运行时加载到 Neuron 设备。

Neuron 运行时由内核驱动程序和 C/C++ 库组成，后者提供 API 来访问 Inferentia 和 Trainium Neuron 设备。TensorFlow 和 PyTorch 的 Neuron ML 框架插件使用 Neuron 运行时在 NeuronCores 上加载和运行模型。Neuron 运行时将编译的深度学习模型（也称为 Neuron 可执行文件格式，NEFF）加载到 Neuron 设备，并针对高吞吐量和低延迟进行了优化。

为何要选择 AWS Trainium 芯片进行训练

超级计算机级性能：Trainium 支持的实例可按需访问超级计算机级性能，提供高达 6 exaflops 的计算能力。
成本效益：通过 Trn1 实例，AWS 显著降低成本，最高可达 50%。客户能够利用生成式 AI 的优势的同时，无需支付过高的费用。
生态系统支持：Trainium 由 AWS Neuron SDK 支持，该 SDK 提供与 TensorFlow 和 PyTorch 等流行 ML 框架的本机集成，从而简化了开发人员的过渡。

在 AWS ParallelCluster 中进行 Llama2 的继续预训练

使用亚马逊云科技的自研 Trainium 芯片进行 Llama2 的继续预训练目前可以使用两种训练集群，一种是使用 AWS ParallelCluster，另外一种是使用 Amazon SageMaker HyperPod。本次分享使用的训练集群是 AWS ParallelCluster 的训练流程。整体架构如下：

训练集群准备

一．首先我们需要在 AWS 上准备部署 AWS ParallelCluster 的网络环境，包含一个公有子网，一个私有子网，公有子网里部署有 NAT Gateway 用于私有子网访问互联网。

二．在云端跳板机上安装 AWS CLI 以及 AWS ParallelCluster CLI 命令行工具

参照以下文档安装 AWS CLI ，然后运行 aws configure 配置

https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html

参照以下文档安装 AWS ParallelCluster CLI

https://docs.aws.amazon.com/parallelcluster/latest/ug/install-v3-virtual-environment.html

三. 新建 AWS ParallelCluster 集群

准备 pcluster.yaml 文件，子网（<Public Subnet ID>和<Private Subnet ID>）需要改成步骤一里准备的子网的 ID，<Key Name>改成 AWS 对应区域新建的密钥对的名字：

Region: us-east-1
Image:
  Os: ubuntu2004
HeadNode:
  InstanceType: c5.4xlarge
  Networking:
    SubnetId: <Public Subnet ID>
  Ssh:
    KeyName: <Key Name>
  LocalStorage:
    RootVolume:
      Size: 1024
  CustomActions:
    OnNodeConfigured:
      Script: s3://neuron-s3/pcluster/post-install-scripts/neuron-installation/v15/u20/pt/install_neuron.sh
  Iam:
    S3Access:
       - BucketName: neuron-s3
         EnableWriteAccess: false
Scheduling:
  Scheduler: slurm
  SlurmQueues:
    - Name: compute1
      CapacityType: ONDEMAND
      ComputeSettings:
        LocalStorage:
          RootVolume:
            Size: 1024
          EphemeralVolume:
            MountDir: /local_storage
      ComputeResources:
        - Efa:
            Enabled: true
          InstanceType: trn1.32xlarge
          MaxCount: 4
          MinCount: 0
          Name: queue1-i1
      Networking:
        SubnetIds:
          - <Private Subnet ID>
        PlacementGroup:
          Enabled: true
      CustomActions:
        OnNodeConfigured:
          Script: s3://neuron-s3/pcluster/post-install-scripts/neuron-installation/v15/u20/pt/install_neuron.sh
      Iam:
        S3Access:
          - BucketName: neuron-s3
            EnableWriteAccess: false
SharedStorage:
- FsxLustreSettings:
    DeploymentType: SCRATCH_2
    StorageCapacity: 1200
  MountDir: /fsx
  Name: pclusterfsx
  StorageType: FsxLustre

运行以下命令新建 ParallelCluster 集群

pcluster create-cluster --cluster-configuration pcluster.yaml -n llama2-pretrain-apc

等集群部署完成以后，查看一下 AWS ParallelCluster 的 Head Node 的公网 IP 地址用于后续访问。

Llama2 继续预训练

用 SSH 客户端登录到 AWS ParallelCluster 集群的 Head Node
```
ssh -i YOUR_KEY.pem ubuntu@HEAD_NODE_IP_ADDRESS
```

激活 PyTorch Neuron 的 virtual environment

cd ~
source ./aws_neuron_venv_pytorch/bin/activate

克隆 neuronx-nemo-megatron 仓库到 Head Node（使用 2.15 版本）

cd ~
git clone -b 2.15_release  https://github.com/aws-neuron/neuronx-nemo-megatron.git
cd neuronx-nemo-megatron

安装 wheel Python 包
```
pip3 install wheel
./build.sh
```

安装 neuronx-nemo-megatron packages 包和依赖

pip3 install ./build/*.whl
pip3 install -r requirements.txt torch==1.13.1 protobuf==3.20.3

构建 Megatron helper 模块

cd ~
python3 -c "from nemo.collections.nlp.data.language_modeling.megatron.dataset_utils import compile_helper; \
compile_helper()"

安装 git-lfs

curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get install git-lfs

下载 Llama2 的模型及权重

cd ~/data
git clone https://huggingface.co/meta-llama/Llama-2-7b-hf

将准备用于增量预训练的数据进行分词预处理

cd ~/neuronx-nemo-megatron/
python3 nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \
    --input=/home/ubuntu/data/datayes/data/2023_07-12.jsonl \
    --json-keys=text \
    --tokenizer-library=huggingface \
    --tokenizer-type=/home/ubuntu/data/Llama-2-7b-hf \
    --dataset-impl=mmap \
    --output-prefix=/home/ubuntu/data/datayes/tokenized/2023_07-12_tokenized \
    --append-eod \
    --need-pad-id \
--workers=32

将 Llama2 HF 的模型转为 Nemo 格式

cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling/checkpoint_conversion
python3 convert_hf_checkpoint_to_nemo_llama.py \
    --path_to_checkpoint /home/ubuntu/data/Llama-2-7b-hf \
    --config_file /home/ubuntu/data/Llama-2-7b-hf/config.json \
    --tp_degree 8 \
    --pp_degree 1 \
--output_path /home/ubuntu/data/Llama-2-7b-nemo

修改 test_llama.sh 文件

cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling
vim test_llama.sh

修改 test_llama.sh 中的 TOKENIZER_PATH 和 DATASET_PATH 指向对应的路径：

: ${TOKENIZER_PATH=$HOME/data/Llama-2-7b-hf}
: ${DATASET_PATH=$HOME/data/datayes/tokenized/2023_07-12_tokenized_text_document}

增加以下部分到 test_llama.sh 最后，model.resume_from_checkpoint 指向到前面转换 Llama2 Nemo 格式的 checkpoint 文件：

model.use_cpu_initialization=False \
+model.load_xser=True \
model.resume_from_checkpoint='/home/ubuntu/data/Llama-2-7b-nemo/mp_rank_07/model_optim_rng.ckpt' 2>&1  | tee  $LOG_PATH/log

修改 megatron_llama_config.yaml 文件
```
cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling/conf
vim megatron_llama_config.yaml
```
修改 megatron_llama_config.yaml 文件中的 save_top_k 为 -1, save_last 为 True，这样会在继续预训练完成以后保存训练完成的 checkpoint 文件；修改文件中的 convert_to_hf 为 True，output_dir 和 config_path 到指定位置，这样可以在继续预训练完成以后将 Nemo 格式的 checkpoint 转成 hugging face 格式的模型文件：
```
checkpoint_callback_params:
    monitor: step
    save_top_k: -1
    mode: max
    save_last: True

  convert_to_hf: True 
  output_dir: /home/ubuntu/data/Llama-2-7b-trained/ 
  config_path: /home/ubuntu/data/Llama-2-7b-trained/config.json 
```

运行以下命令启动预编译任务

cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling
sbatch --nodes 1 compile.slurm ./llama_7b.sh

运行以下命令检查队列任务情况
```
squeue  
```
等任务状态显示 R 以后，使用 tail 命令检查日志（预编译任务的日志命名是 slurm-compile.slurm-<任务id>.out）
```
tail -f slurm-compile.slurm-1.out
```

等日志中显示类似下面的信息的时候说明预编译已经完成

2023-11-06 21:26:11.000884: 14376 INFO ||NEURON_PARALLEL_COMPILE||: Total graphs: 32 2023-11-06 21:26:11.000884: 14376 INFO ||NEURON_PARALLEL_COMPILE||: Total successful compilations: 32 2023-11-06 21:26:11.000884: 14376 INFO ||NEURON_PARALLEL_COMPILE||: Total failed compilations: 0

运行以下命令启动继续预训练任务

cd ~/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling
sbatch --nodes 1 run.slurm ./llama_7b.sh

运行以下命令检查队列任务情况
```
squeue
```
等任务状态显示 R 以后，使用 tail 命令检查日志确定训练进度（预训练任务的日志命名是 slurm-run.slurm–<任务id>.out)）
```
tail -f slurm-run.slurm-2.out
```

在日志中看到类似下面的消息说明模型已经训练完成

[NeMo I 2023-11-05 22:57:43 megatron_gpt_pretraining:112] Finished converting Llama checkpoints

训练好的模型会放在第 12 步定义的 output_dir 定义的路径下，后续可以使用该模型进行推理评估

训练过程中的问题定位

一般来说，如果在继续预训练过程中遇到问题，可以从以下两个方向进行 Troubleshooting：

如果训练任务已经启动并生成了相应的 slurm*.log 日志，则可以分析 slurm 日志检查问题原因
如果训练任务没有启动，则可以参考下面两个文档检查节点的初始化情况：

https://docs.aws.amazon.com/parallelcluster/latest/ug/troubleshooting-fc-v3-compute-node-initialization-v3.html

https://docs.aws.amazon.com/parallelcluster/latest/ug/troubleshooting-v3-scaling-issues.html#troubleshooting-v3-node-init

亚马逊AWS官方博客