预训练大语言模型面对来自于模型规模和数据规模的多重挑战，为了解决这些问题，PAI提供了在DLC上使用Megatron-LM进行大语言模型预训练的最佳实践。该实践通过准备镜像环境与训练数据、调整训练策略、提交与管理任务等流程，帮助您在DLC上轻松高效地预训练大语言模型。本文为您介绍该实践的使用流程。
背景信息
基本概念
语言模型（LM）
是一种利用自然文本来预测词（Token）顺序的机器学习方法。
大语言模型（LLM）
是指参数数量达到亿级别的神经网络语言模型，例如：GPT-3、GPT-4、PaLM、PaLM2等。
Megatron-LM
是由NVIDIA的应用深度学习研究团队开发的Transformer模型训练框架，它可以高效利用算力、显存和通信带宽，大幅提升了大语言模型大规模预训练的效率。Megatron-LM已经成为许多大语言模型预训练任务的首选框架。
使用流程
首先准备训练任务所需的镜像环境和训练数据，并将训练数据创建为PAI数据集。
配置一系列优化策略来优化大语言模型的训练过程，以提高模型的性能和准确度。
提交大语言模型预训练任务，并对任务进行管理。
前提条件
已创建工作空间，具体操作，请参见创建工作空间。
已准备专有资源组，本案例使用灵骏智算资源gu7xf机型（例如：资源规格选择ml.gu7xf.c96m1600.8-gu108），关于如何创建灵骏智算资源，请参见新建与管理灵骏智算资源。
已了解GPT的基本概念，详情请参见Language Models are Few-Shot Learners。
使用限制
仅支持在华北6（乌兰察布）使用灵骏智算资源。如果您在其他地域提交训练任务，请选择使用通用训练资源专有资源组，关于如何创建通用训练资源，请参见新建及管理通用训练资源。
准备镜像环境与训练数据
准备镜像环境
目前支持使用以下几种类型的镜像环境：
PyTorch镜像和自定义Megatron
在DLC提供的公共镜像中选择一PyTorch类型的镜像；下载Megatron并创建为代码集，具体操作，请参见代码配置。
集成了Megatron的镜像
PAI提供了集成社区PyTorch和Megatron最新版本的镜像：
registry.${region}.aliyuncs.com/pai-dlc/llmacc:23.05-gpu-py310-cu121-ubuntu22.04其中：${region}需要替换为当前地域ID，目前支持的地域为cn-beijing、cn-shanghai、cn-hangzhou、cn-shenzhen和cn-wulanchabu，您可以根据需要选择其中一个。
本案例使用位于华北6（乌兰察布）的灵骏智算资源，因此使用 registry-vpc.cn-wulanchabu.aliyuncs.com/pai-dlc/llmacc:23.05-gpu-py310-cu121-ubuntu22.04作为演示镜像。
准备训练数据
首先将词表文件和合并规则表文件下载到本地。
本案例使用BPE 方式进行分词，使用GPT-2的词表和合并规则表：
对维基百科语料或自定义语料文件进行处理生成Megatron-LM支持的索引数据集文件。
任意选择一种语料即可，本案例使用维基百科语料生成索引数据集文件。
维基百科语料
维基百科数据集是自然语言处理的常见语料，在词表文件所在目录使用WikiExtractor产生Megatron-LM支持的索引数据集文件：
# 1. Install wikiextractor pip install wikiextractor # 2. Download & extract to loose json wget https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2 python -m wikiextractor.WikiExtractor --json ./enwiki-latest-pages-articles.xml.bz2 # 3. Use nltk/spacy/ftfy to clean the corpus # 4. Convert the corpus into indexed datasets find text/ -type d | sed 's|text|mkdir -p datasets/text|' | sh find text/ -type f -name "wiki_*" | xargs -P $(nproc) -I{} \ bash -c "python env/megatron_lm/src/tools/preprocess_data.py --input {} \ --output-prefix datasets/{} \ --vocab-file gpt2-vocab.json \ --dataset-impl mmap \ --tokenizer-type GPT2BPETokenizer \ --merge-file gpt2-merges.txt \ --append-eod \ --workers 1 --chunk-size 1024" # 5. Merge indexed datasets mkdir -p flatten_datasets find datasets/text/ -type f | sed 's|\(.*\)/wiki\(.*\)|cp \1/wiki\2 ../../flatten_datasets/\1_wiki\2|' | sh python env/megatron_lm/src/tools/merge_datasets.py --output-prefix ./enwiki-latest-pages-articles --input flatten_datasets/自定义语料
使用loose JSON格式的语料，例如：
{"text": "A platform that provides enterprise-level data modeling services based on machine learning algorithms to quickly meet your needs for data-driven operations.", "type": "Eng", "id": "0", "title": "Machine Learning Platform for AI", "src": "www.alibabacloud.com"} {"text": "This topic describes how to use the Deep Learning Containers (DLC) of Machine Learning Platform for AI (PAI) to train transfer learning models based on the PyTorch framework.", "type": "Eng", "id": "42", "title": "Submit a standalone training job that uses PyTorch", "src": "www.alibabacloud.com"}使用Megatron-LM中提供的数据预处理工具对语料文件进行处理，可以产生Megatron-LM支持的mmap格式的索引数据集文件（my-gpt2_text_document.bin和my-gpt2_text_document.idx）：
python /path/to/megatron-lm/tools/preprocess_data.py \ --input my-corpus.json \ --output-prefix my-gpt2 \ --vocab-file gpt2-vocab.json \ --dataset-impl mmap \ --tokenizer-type GPT2BPETokenizer \ --merge-file gpt2-merges.txt \ --append-eod \ --workers 2 \ --chunk-size 1执行后会打印出类似如下的日志：
Opening my-corpus.json > building GPT2BPETokenizer tokenizer ... > padded vocab (size: 50257) with 47 dummy tokens (new size: 50304) > building GPT2BPETokenizer tokenizer ... > building GPT2BPETokenizer tokenizer ... Vocab size: 50257 Output prefix: my-gpt2 Time to startup: 0.07754158973693848 > padded vocab (size: 50257) with 47 dummy tokens (new size: 50304) Done! Now finalizing. > padded vocab (size: 50257) with 47 dummy tokens (new size: 50304)创建数据集。
生成索引数据集文件后，您可以将词表文件、合并规则文件和索引数据集文件上传到OSS Bucket，具体操作，请参见控制台上传文件。
将上传到OSS Bucket的数据创建为PAI数据集，具体操作，请参见创建及管理数据集。
说明如果OSS Bucket目录中已产生了npy文件，您可以开启数据集加速，提升数据访问效率，具体操作，请参见在PAI平台使用数据集加速器。
调整大规模训练策略
参考附录：PTD-P并行技术原理介绍，来配置训练任务的优化策略，后续提交DLC容器训练任务时，配置到执行命令中。目前支持使用以下两种方式调整训练策略。
手动调整训练策略
在训练大语言模型时，DLC支持在执行命令中配置以下参数手动调整训练策略，来优化整体训练性能。
参数 | 描述 |
--fp16 | 在Megatron-LM中支持使用fp16方式训练模型。 |
--micro-batch-size |
|
--tensor-model-parallel-size | 用来配置tensor parallel size参数值。
|
--pipeline-model-parallel-size | 用来配置pipeline parallel size参数值。 您可以尝试从1开始逐渐增加pipeline parallel size值，直到模型可以正常开始训练，不出现OOM错误。 说明 pipeline parallel size参数值需要小于模型层数，且可以被模型层数整除。 |
--data-parallel-size | 您可以根据尝试获得的tensor parallel size和pipeline parallel size的值，以及准备训练的总卡数，算出可行的data parallel size值： |
--use-distributed-optimizer | 将参数配置为true，分别用来打开优化器状态并行技术、序列并行技术、激活重算技术，来进一步优化显存占用。 |
--sequence-parallel | |
--recompute-activations | |
--use-flash-attn | 将参数配置为true，用来打开flash attention，进一步提升计算效率。 |
--global batch size | 在模型效果允许的前提下增加global batch size，从而降低流水并行时的层间等待开销。 |
自动调整训练策略
如果训练GPT-like模型，您可以使用准备镜像环境与训练数据章节中的集成了Megatron的镜像，并只需提供 global batch size或者最大允许的global batch size，就可以使用如下命令自动拉起Megatron预训练作业：
torchrun \
--nnodes $WORLD_SIZE \
--node_rank $RANK \
--nproc_per_node ${LOCAL_WORLD_SIZE:-4} \
--master_addr $MASTER_ADDR \
--master_port $MASTER_PORT \
-m llmacc.pretrain \
--task-type llmacc.megatron.GPT \
--max-global-batch-size 16 \
# megatron arguments此外，也可以直接将llmacc作为工具算出可能的并行配置：
python \
-m llmacc.pretrain \
--task-type llmacc.megatron.GPT \
--world-size 64 # 64 GPUs in total \
--local-world-size 8 # 8 GPUs for each node \
--max-global-batch-size 16 \
# megatron arguments对于非GU50或GU100机型，您可以通过配置DEVICE_MEMORY_GB环境变量来配置显存大小。
提交与管理任务
进入创建任务页面，具体操作，请参见创建训练任务。
在创建任务页面，配置以下参数，其他参数配置详情，请参见创建训练任务。
本案例以训练GPT-22B模型为例：
参数
描述
资源组
选择已创建的灵骏智算资源。
优先级
配置为1。
节点镜像
选择镜像地址，并配置：
registry-vpc.cn-wulanchabu.aliyuncs.com/pai-dlc/llmacc:23.05-gpu-py310-cu121-ubuntu22.04。任务类型
选择PyTorch。
数据集配置
选择已创建的数据集。
执行命令
配置为以下命令，执行命令配置详情，请参见调整大规模训练策略。
torchrun \ --nnodes $WORLD_SIZE \ --node_rank $RANK \ --nproc_per_node gpu \ --master_addr $MASTER_ADDR \ --master_port $MASTER_PORT \ -m llmacc.pretrain \ --task-type llmacc.megatron.GPT \ --num-layers 48 \ --num-attention-heads 64 \ --hidden-size 6144 \ --seq-length 2048 \ --max-position-embeddings 2048 \ --train-iters 100 \ --lr 2.0e-4 \ --min-lr 2.0e-4 \ --vocab-file /mnt/data/gpt2-vocab.json \ --merge-file /mnt/data/gpt2-merges.txt \ --data-path /mnt/data/enwiki-20230526-pages-articles \ --log-interval 10 \ --max-global-batch-size 32任务资源配置
节点数量：配置为2。
CPU（核数）：80。
内存（GB）：800。
共享内存（GB）：800。
GPU（卡数）：8。
单击提交。
页面自动跳转到容器训练任务页面。在该页面，您可以单击任务名称进入任务详情页面，来观察任务执行状态。
在任务详情页面下方单击资源视图，查看任务中使用的资源。
更多管理训练任务的操作，请参见创建及管理容器训练任务。
附录：PTD-P并行技术原理介绍
大规模并行训练的一个关键因素是并行训练策略的选择。在Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM和Reducing Activation Recomputation in Large Transformer Models这两篇论文中介绍了Megatron-LM常见的大规模训练策略。
常见的大模型训练技术包括：数据并行技术、模型并行技术（包括张量并行技术和流水并行技术）、优化器状态并行技术、序列并行技术、激活重算技术等：
数据并行技术（Data Parallel）：在多个GPU组上有相同的模型参数副本，但读取不同的样本。在参数更新前使用allreduce来全局平均梯度，从而加速模型训练。
优化器状态并行技术（Distributed Optimizer / ZERO-DP）：在使用数据并行技术的同时，将模型参数对应的优化器状态切分到不同的GPU上，从而支持训练更大的模型。
激活重算技术（Activation Recomputation）：在反向传播时重新计算部分激活，避免占用显存存储这部分激活，从而支持训练更大的模型。
序列并行技术（Sequence Parallel）：一个Transformer层内的dropout和layer-norm等部分参数切分到不同GPU上，从而支持训练更大的模型。
模型并行技术（Model Parallel）：在多个GPU上存放一套模型参数的不同分片，从而支持训练更大的模型。
流水并行技术（Pipeline Parallel）：模型内不同Transformer层切分到不同的GPU上；
张量并行技术（Tensor Parallel）：一个Transformer层内的Attention和Linear部分参数切分到不同的GPU上；
模型并行的示意图如下，来源：Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM和Reducing Activation Recomputation in Large Transformer Models。
上述每种技术都有不同的权衡，如果互相结合需要仔细推敲才能获得良好的训练性能。Megatron-LM结合上述这些技术，提出了一种称为PTD-P并行的技术，从而支持在千卡规模上以较好的计算性能来训练大语言模型。