计算加速套件 TACO Kit
有奖捉虫：行业应用 & 管理与支持文档专题 HOT
文档中心 > 计算加速套件 TACO Kit > TACO Train AI 训练加速引擎 > 使用 TACO Train 部署分布式集群 > 大模型预训练最佳实践

大模型预训练最佳实践

最近更新时间：2023-12-21 21:03:42

我的收藏

本页目录：

背景信息

最近 ChatGPT 凭借其强大的语言理解能力、文本生成能力、对话能力等在多个领域均取得了巨大成功，掀起了新一轮的人工智能浪潮。ChatGPT，GPT3，以及 GPT3.5 都是基于 Transformer 架构堆叠而成，研究发现随着训练数据量和模型容量的增加可以持续提高模型的泛化能力和表达能力，研究大模型成为了近两年的趋势。国内外头部科技公司均有布局，发布了若干千亿规模以上的大模型，如下图所示。
?
?
?
近几年 NLP 预训练模型规模的发展，模型已经从亿级发展到了万亿级参数规模。具体来说，2018 年 BERT 模型最大参数量为3.4亿，2019 年 GPT-2 为十亿级参数的模型。2020 年发布的百亿级规模有 T5 和 T-NLG，以及千亿参数规模的 GPT-3。在 2021 年末，Google 发布了 SwitchTransformer，首次将模型规模提升至万亿。然而GPU硬件发展的速度难以满足 Transformer 模型规模发展的需求。近四年中，模型参数量增长了十万倍，但 GPU 的显存仅增长了4倍。
?
?
?
举例来说，万亿模型的模型训练仅参数和优化器状态便需要 1.7TB 以上的存储空间，至少需要 425 张 A100-40G，这还不包括训练过程中产生的激活值所需的存储。在这样的背景下，大模型训练不仅受限于海量的算力,?更受限于巨大的存储需求。
业内大模型预训练的解决方案主要包括微软的 DeepSpeed 和英伟达的 Megatron-LM。DeepSpeed 引入 ZeRO（Zero Redundancy Optimizer?）优化器，将模型参数、梯度、优化器状态按需分配到不同的训练卡上，满足大模型对存储的极致要求；Megatron-LM 基于3D并行（张量并行、流水并行、数据并行）将模型参数进行切分，满足大模型在有限显存资源下的训练诉求。
?
腾讯内部也有非常多的大模型预训练业务，为了以最小的成本和最快的性能训练大模型，太极机器学习平台对 DeepSpeed 和 Megatron-LM 进行了深度定制优化，推出了AngelPTM 训练框架。2022 年 4 月腾讯发布的混元（HunYuan-1T）AI大模型便是基于 AngelPTM 框架训练而来。鉴于最近大模型的火热趋势，我们决定将内部成熟落地的 AngelPTM 框架推广给广大公有云用户，帮助业务降本增效。

AngelPTM简介

?AngelPTM 基于 ZERO策略，将模型的参数、梯度、优化器状态以模型并行的方式切分到所有GPU，并自研ZeRO-Cache框架把内存作为二级存储offload参数、梯度、优化器状态到CPU内存，同时也支持把SSD作为第三级存储。ZeRO-Cache为了最大化最优化的利用内存和显存进行模型状态的缓存，引入了显存内存统一存储视角，将存储容量的上界由内存扩容到内存+显存总和。同时将多流异步化做到了极致，在GPU计算的同时进行数据IO和NCCL通信，使用异构流水线均衡设备间的负载，最大化提升整个系统的吞吐。ZeRO-Cache将GPU显存、CPU内存统一视角管理，减少了冗余存储和内存碎片，增加了内存的利用率，将机器的存储空间压榨到了极致。

部署步骤

本文指导用户基于腾讯云高性能计算集群和TACO AngelPTM训练框架，搭建大模型训练环境，展示AngelPTM的强大能力。

购买高性能计算实例

购买实例，其中实例、存储及镜像请参考以下信息选择，其余配置请参考 通过购买页创建实例 按需选择。
实例：选择 GPU 型 HCCPNV4h。
镜像：建议选择公共镜像，公共镜像当中已安装 RDMA 网卡驱动，且支持自动安装 GPU 驱动。
操作系统请使用 CentOS 7.6、Ubuntu 18.04?或 TencentOS 2.4（TK4）。
若您选择公共镜像，则请勾选“后台自动安装GPU驱动”，实例将在系统启动后预装对应版本驱动。如下图所示：
?
若您选择自定义镜像，则需要自行安装 RDMA 网卡驱动和 GPU 驱动，请通过 联系我们 获取腾讯云售后支持。

安装docker和NVIDIA docker

1. 参考 使用标准登录方式登录 Linux 实例，登录实例。
2. 执行以下命令，安装 docker。
curl -s -L http://mirrors.tencent.com/install/GPU/taco/get-docker.sh | sudo bash
若您无法通过该命令安装，请尝试多次执行命令，或参考 Docker 官方文档 Install Docker Engine 进行安装。
3. 执行以下命令，安装 nvidia-docker2。
curl -s -L http://mirrors.tencent.com/install/GPU/taco/get-nvidia-docker2.sh | sudo bash
若您无法通过该命令安装，请尝试多次执行命令，或参考 NVIDIA 官方文档 Installation Guide & mdash 进行安装。

启动AngelPTM运行环境

#!/bin/bash
?
docker run \\
    -itd \\
    --gpus all \\
    --privileged --cap-add=IPC_LOCK \\
    --ulimit memlock=-1 --ulimit stack=67108864 \\
    --net=host \\
    --ipc=host \\
    --name=ptm \\
    ccr.ccs.tencentyun.com/qcloud/taco-train:torch112-cu116-bm-0.6.1
    
docker exec -it ptm bash
该镜像包含的版本信息如下：
OS：Ubuntu 20.04.5 LTS
python：3.8.10
CUDA toolkit: V11.6.124
cuDNN: 8.4.0
tencent-lightcc: 3.1.1
pytorch: 1.12.1+cu116
AngelPTM DeepSpeed: 0.6.1+93ef832c
AngelPTM Megatron-DeepSpeed:1.1.5
AngelPTM ptm: 0.1

开始测试

28B GPT模型

社区版本 Zero Stage3，单机 A100 40G 所能容纳的最大模型规模
cd /workspace/examples/ptm
bash start.sh 28

55B GPT模型

AngelPTM Zero Cache，单机 A100 40G 所能容纳的最大模型规模
cd /workspace/examples/ptm
bash start.sh 55

评测其他模型规模

可以通过修改start.sh添加其他的模型规模配置：
if [ $MODEL_SIZE == "28" ]; then
        NUM_LAYERS=26
        HIDDEN_SIZE=8192
        NUM_ATTN_HEADS=128
        FFN_HIDDEN_SIZE=32768
        ATTN_HEAD_SIZE=128
        BATCH_SIZE=38
        CONFIG_PREFETCH_CACHE_SUB_GROUP_RATE=0.0
elif [ $MODEL_SIZE == "55" ]; then
        NUM_LAYERS=68
        HIDDEN_SIZE=8192
        NUM_ATTN_HEADS=64
        FFN_HIDDEN_SIZE=32768
        ATTN_HEAD_SIZE=128
        BATCH_SIZE=10
        CONFIG_PREFETCH_CACHE_SUB_GROUP_RATE=0.1
else
        echo "ERROR: Please supplement new model configuration to test!"
        exit -1
fi

配置参数说明

相比原始的 ds config，AngelPTM提供了一些额外的控制参数（参考/workspace/examples/ptm/ds_config_gpt_zero3_maxscale.json），其功能说明说下：
"pipeline_optimizer": true,
打开模型更新过程中流水操作，默认打开，无需配置
"prefetch_cache_sub_group_rate": 0.0
值范围为0.0~1.0，表示放置模型状态到 GPU 的比例：
0.1表示10%的模型状态存储到显存，这样可以支持更大的模型尺寸；
0.0表示模型状态不占用显存资源，这样可以支持更大的 batch，获取更好的性能
"max_contigous_params_size": 2e9,
ZeRO Cache 管理的显存大小，和 stage3_prefetch_bucket_size 以及 stage3_max_live_parameters 配合使用，如果预取参数变大，其值也需要调大，在显存比较充足的情况下可以考虑关闭，设置为-1
"max_param_reduce_events": 0,
梯度 reduce_scatter 同时进行的 event 个数，建议1或者0，配置太大会导致显存 OOM
"is_communication_time_profiling": false,
调试相关，用于打印通信时间，默认关闭
"save_large_model_multi_slice": true,
分片模型保存，避免内存 OOM，默认打开，无需配置
说明：
模型测试脚本来自Megatron-DeepSpeed 训练脚本?。
第一次测试启动较慢主要是由于数据集下载和预处理。
如果想要终止测试，启动新一轮测试之前建议先`ps`确认上次测试程序已经完全退出（由于AngelPTM会使用大量的内存资源cache模型参数和优化器状态，cache的释放需要一段时间才能完成，通常几分钟就好）。

性能结果

?
DeepSpeed
Megatron-DeepSpeed
社区方案
0.8.1+258d2831
7212b58
AngelPTM方案
0.6.1+474caa20
c5808e0
注意：其他环境，例如OS/python/CUDA/cuDNN/pytorch等版本二者一致。

最大容纳模型规模

?
?
?

同等模型规模训练性能

?
?
?

多机扩展比

?
?
?

千亿模型性能

?
?
?
注意：4机32卡训练性能提升比例（26.8%）相比单机有所下降主要是由于网络带宽限制。这里使用的是100Gbps RDMA 网络，未来腾讯云会推出更高带宽高性能计算集群，预期性能提升会与单机接近。

监控信息

下图是单机 A100 40G 训练550亿参数的 GPU 和内存监控信息
?

结论

本文简单介绍了 AngelPTM 产生的背景和技术能力，然后指导用户如何在腾讯云高性能计算集群下搭建 AngelPTM 训练环境。可以看出，相比社区方案：
AngelPTM将单机 A100 40G 容纳的模型规模提升了94.71%
基于社区方案能容纳的最大模型规模，AngelPTM 性能提升了44.42%
千亿模型规模下，AngelPTM 多机扩展比接近线性
?
?
中国站


http://www.vxiaotou.com