【AIGC】大模型训练框架-DeepSpeed
一、训练框架概述
一个好的分布式训练框架可以大大提高训练效率，节省时间和成本。但是如何选择一个适合自己的框架呢？如何验证分布式训练的效率？应该如何应用分布式训练？此外，如何避免那些常见的坑。
人工智能领域的ChatGPT模型成为了各家公司争相研究的焦点，然而，训练这样的模型需要大量的时间和资源成本，随着模型和训练数据规模的增大，训练成本也在逐渐增加。而开源系统对于机器的利用率不高，通常也不能让机器达到最大效率，更有甚者吞吐量无法提高，不能让训练效率大幅提高。
针对上述问题，进行了广泛的业界调研，现有的分布式训练框架对比：
训练框架地址：
Megatron-LM: https://github.com/NVIDIA/Megatron-LM
DeepSpeed: https://github.com/microsoft/DeepSpeed
mesh-tensorflow: https://github.com/tensorflow/mesh
fairscale: https://github.com/facebookresearch/fairscale
- 具备高效、有效、易使用等三个方面的特点。
- 作为一款开源系统，Deepspeed框架能够在提高训练效率的同时，也能保证开发生产力的高效。
- 使用该框架可以大大节省时间和成本，让我们能够更加高效地完成建模和训练任务，从而提升项目的整体效果。
二、Deepspeed框架介绍
Deepspeed是由微软开发的一款开源深度学习优化库，旨在提高大规模模型训练的效率和可扩展性。该框架采用多种技术手段来加速训练，其中包括模型并行化、梯度累积、动态精度缩放、本地模式混合精度等。
此外，Deepspeed还提供了一些辅助工具，如分布式训练管理、内存优化和模型压缩等，以帮助开发者更好地管理和优化大规模深度学习训练任务。该框架是基于PyTorch构建的，因此可以简单修改以便进行迁移使用。Deepspeed已经在许多大规模深度学习项目中得到了应用，包括语言模型、图像分类、目标检测等领域。在深度学习模型软件体系架构中，Deepspeed扮演着重要的角色，如下图所示：
Deepspeed软件架构主要分为三个部分，如下图所示:
-
APIs
提供简单易用的api接口，使用者只需要调用几个接口就能够完成模型训练和推理任务。其中最重要的接口是initialize，用于初始化引擎并配置训练参数和优化技术。配置参数通常保存在config.json文件中。
-
Runtime
运行时组件，是Deepspeed的核心组件，主要用于管理、执行和性能优化。该组件包括将训练任务部署到分布式设备、数据分区和模型分区、系统优化、微调、故障检测、checkpoints保存和加载等功能，并使用Python语言实现。
-
Ops
底层内核实现，使用C++和CUDA优化计算和通信。此组件涵盖了许多关键技术。
Deepspeed是一款高效、可扩展的深度学习优化库，其中一些核心技术如下：
-
ZeRO（Zero Redundancy Optimizer）
ZeRO是一种内存优化技术，用于大规模分布式深度学习。该技术可以消除数据并行进程中的内存冗余，通过在数据并行进程之间划分模型状态参数、梯度和优化器状态，而不是复制它们。此外，ZeRO还使用动态通信调度在分布式设备之间共享必要的状态，以保持数据并行的计算粒度和通信量。基于ZeRO，DeepSpeed实现了数据并行、流水线并行和张量切片模型并行等方式的训练，以提高显存和计算效率，并能够训练具有万亿个参数的模型。
-
3D并行
数据并行和模型并行可能会导致内存冗余，因此DeepSpeed采用基于ZeRO的3D并行来优化显存利用和计算效率。该技术将模型状态参数、梯度和优化器状态按照3D方式划分，并使用动态物理内存分配来减少内存占用。
-
梯度累积
DeepSpeed还使用梯度累积来提高批量大小，从而提高训练效率。梯度累积可以通过多次前向传递和反向传递来累积梯度，从而实现较大的批量大小。
通过以上技术和方法，DeepSpeed实现了高效、可扩展的深度学习训练，提高了训练速度和模型性能。
三、DeepSpeed框架实践
下面以DeepSpeed框架与ChatGLM微调模型相结合的实践案例为例，演示分布式训练的过程（注意这个案例仅用于说明，实际应用时可能会有所不同）
环境准备
conda install pytorch==1.12.0 \
torchvision==0.13.0 \
torchaudio==0.12.0 \
cudatoolkit=11.3 -c pytorch
pip install deepspeed==0.8.1
sudo yum install openmpi-bin libopenmpi-dev
pip install mpi4py
下载预训练模型
https://huggingface.co/THUDM/chatglm-6b/tree/main
准备训练样本
训练数据通常由三个部分组成：prompt、response和history。引入用户特征数据可以提升模型训练效果。
样本数据示例：
cat /data/dev.json | head -1
{
"prompt": "什么事。\n[姓名]：张三\n[年龄]：45.0\n[性别]：男士\n[省份]：湖北",
"response": "打扰您了，给您介绍下近期的优惠活动？",
"history": [
[
"",
"您好，请问是张先生吗？"
],
[
"是。",
"喂，您好，我这是您专属客服顾问。"
]
]
}
准备训练代码
git clone https://github.com/THUDM/ChatGLM-6B.git
微调训练核心流程伪代码
import torch
import deepspeed
from torch.utils.data import RandomSampler, DataLoader
from torch.utils.data.distributed import DistributedSampler
...
# 一、初始化DeepSpeed引擎
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,model=model,model_parameters=params)
# 二、分布式环境设置
deepspeed.init_distributed()
....
# 三、预加载模型和训练数据
model = ChatGLMForConditionalGeneration.from_pretrained(args.model_dir)
tokenizer = ChatGLMTokenizer.from_pretrained(args.model_dir)
# DataLoaders creation:
train_dataloader = DataLoader(train_dataset,
batch_size=conf["train_micro_batch_size_per_gpu"],
sampler=RandomSampler(train_dataset),
collate_fn=coll_fn,
drop_last=True,
num_workers=0)
....
# 四、训练循环
for step, batch in enumerate(train_dataloader):
#用于向前传播和损失计算
loss = model_engine(batch)
#向后传播
model_engine.backward(loss)
#优化器更新
model_engine.step()
启动训练前需要配置机器之间免登、.deepspeed_env环境变量、调整训练参数
cd ./ChatGLM-6B/ptuning/
nohup deepspeed --hostfile=myhostfile --master_port 9000 main.py \
--deepspeed deepspeed.json \
--do_train \
--train_file /data/train.json \
--test_file /data/dev.json \
--prompt_column prompt \
--response_column response \
--history_column history \
--overwrite_cache \
--model_name_or_path /data/pre_model/chatglm/chatglm-6b \
--output_dir ./output/out-chatglm-6b-ft-le-4 \
--overwrite_output_dir \
--max_source_length 64 \
--max_target_length 64 \
--per_device_train_batch_size 4 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 1 \
--predict_with_generate \
--max_steps 50000 \
--logging_steps 10 \
--save_steps 10000 \
--learning_rate $LR \
--fp16 &
训练loss迭代示意图：
下面是介绍使用ChatGLM模型进行推理服务的示例。请注意，此处的示例仅用于演示，实际应用可能会有所不同，生产环境通常使用的是TorchServing方式部署的分布式推理服务。
首先需要安装额外的依赖 pip install fastapi uvicorn，然后运行仓库中的 api.py：
pip install fastapi uvicorn
python api.py
默认部署在本地的 8000 端口，通过 POST 方法进行调用
curl -X POST "http://127.0.0.1:8000" \
-H 'Content-Type: application/json' \
-d '{"prompt": "你好", "history": []}'
得到的返回值为
{
"response":"你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。",
"history":[["你好","你好👋！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。"]],
"status":200,
"time":"2023-06-03 21:38:40"
}
四、注意事项
实践过程可能会遇到如下问题，下面也给了相应解决方法。
1、多机多卡跑不起来，多机间不能通信，报错如下，
Call to connect returned Connection refused, retrying
解决方案：
deepspeed环境依赖问题，重新装pytorch依赖；
conda install pytorch==1.12.0 \
torchvision==0.13.0 \
torchaudio==0.12.0 \
cudatoolkit=11.3 -c pytorch
2、使用官方ChatGLM-6B 代码训练的时候 loss不降不明显
解决方案：
通过换为ChatGLM-Finetuning的训练框架，loss下降平稳
git clone https://github.com/liucongg/ChatGLM-Finetuning.git
3、因为ChatGLM-Finetuning官方代码采样器的时候用的是随机采样，导致8卡单机只能跑6卡;
解决方案：
改源码，调整成支持对训练数据分布式采样
# 改成分布式采样
# DataLoaders creation:
if args.local_rank == -1:
train_sampler = RandomSampler(train_dataset)
else:
train_sampler = DistributedSampler(train_dataset)
train_dataloader = DataLoader(train_dataset,
batch_size=conf["train_micro_batch_size_per_gpu"],
sampler=train_sampler,
collate_fn=coll_fn,
drop_last=True,
num_workers=0)
4、将代码改成支持分布式采样之后，分布式跑10万级训练数据没问题，扩大到100万级训练数据之后依旧只能跑6卡的问题
问题原因：
在训练100万级训练数据时，向量化数据需要4个小时左右，向量化后的数据占用内存在140G左右，默认情况下分布式训练每个进程都需要加载完整数据集，在当前单机8卡机器内存总共只有1T的情况下，最多只能使用6卡，有1/4的资源空闲。
解决方案：
改源码，重写dataset逻辑为：
1.均分完整数据集到所有进程;
2.训练时每个进程同一时间只加载单个分段大小(batch_size)数据集；
3.token转向量后的内存数据dump到文件，重新训练的时候加载向量化后的数据。
4.模型训练过程中，训练数据质量比较关键，扩充训练样本的特征提示，可以提升模型训练效果。
5.多机训练时，每台机器的训练代码和环境务必一致。
五、总结
本小节介绍如何选择适合自己的分布式训练框架，并介绍DeepSpeed框架的优势。介绍了如何使用DeepSpeed框架来构建分布式训练实践，并解决实际训练过程中的各种问题和注意事项。
本小节也详细介绍DeepSpeed框架的能力，包括其核心技术ZeRO、3D并行和梯度累积的使用方法和技术细节。
六、拓展
（一）ZeRO（零冗余优化器）
推荐阅读：图解大模型训练之：数据并行( DeepSpeed ZeRO，零冗余优化)
ZeRO（Zero Redundancy Optimizer） 是一种用于优化大规模深度学习模型训练的技术。它的主要目标是降低训练期间的内存占用、通信开销和计算负载，从而使用户能够训练更大的模型并更高效地利用硬件资源。
ZERO论文首先分析了模型训练中内存主要消耗在两个方面：
-
model states：模型状态，包括包括优化器参数（例如Adam的动量和方差）、梯度、模型参数
-
residual states：剩余状态，包括包括激活函数、临时缓冲区、内存碎片
ZERO分别使用ZeRO-DP和ZeRO-R来优化model states和residual states。
- ZeRO-DP包括三个阶段（以64GPU的混合精度训练举例，采用Adam优化器）：
- ZeRO-R通过激活值分割来优化激活值内存；固定缓冲区大小以防止随着模型规模的增加而导致缓冲区过大；根据张量的不同生命周期来管理内存，以防止内存碎片化。
通过上述一系列措施，ZeRO优化了模型训练过程中的内存消耗，并使用动态通信来保持计算和通信效率，使得我们可以用有限的计算资源来训练更大规模的模型。
-
基本背景
在整个训练中，有很多states并不会每时每刻都用到，举例来说：
1.Adam优化下的optimizer states只在最终做update时才用到 2.数据并行中，gradients只在最后做AllReduce和updates时才用到 3.参数W只在做forward和backward的那一刻才用到 4......... -
主要思想
如果数据算完即废，等需要的时候，我再想办法从之前暂存的地方拿回来，那不就省了一笔存储空间吗？
-
具体方法
（1）ZeRO-Offload
核心思想
- 显存不够，内存来凑。
- 如果我把要存储的大头卸载(offload)到CPU上，而把计算部分放到GPU上，这样比起跨机，是不是能既降显存，也能减少一些通讯压力呢？
基本原理
- forward和backward计算量高，因此和它们相关的部分，例如参数W（fp16），activation，就全放入GPU。
- update的部分计算量低，因此和它相关的部分，全部放入CPU中。例如W(fp32)，optimizer states（fp32）和gradients(fp16)等。
（2）ZeRO-Infinity
与ZeRO-Offload原理类似，找个除GPU之外的地方存数据，需要使用的时候再获取。
（二）DP、MP、PP
推荐阅读：深度学习模型训练显存占用分析及DP、MP、PP分布式训练策略
（1）Data Parallelism （DP）数据并行
主要思想：
将模型复制到多个GPU设备，每张GPU当中都存放了一个复制的GPU版本。每个设备都可以并行接受输入的 data batches 。具体过程如下所示：
主要特点：
- 数据并行策略并不能节省总的训练内存，这是好理解的，因为在训练过程中，模型被分配到了每个 GPU 一份，这就意味着产生了很多额外的内存消耗。另外，如果模型本身的权重大小难以保存在一张 GPU 时。数据并行的算法将不能使用，无论你扩大自己的计算资源到 8张 还是 100张，都无法展开训练。
- 使用数据并行带来的好处在于，数据并行将极大提高训练速度，因为每个设备只需在总训练数据的一小部分上进行训练。不过，由于设备间梯度交换会带来通信开销，速度提升并不与设备数量成线性关系。
（2）Model Parallelism （MP）模型并行
当模型自身参数量很大时，数据并行过程将无法使用，因此， Model Parallel (MP) 模型并行在这样的背景下应运而生，MP 算法不同于 DP 将模型复制到每个设备上，MP 选择将模型拆开，每个 GPU 设备只保存模型权重的一部分，在 forward 前向传递过程中，数据将依次通过各个设备，一个设备的输出作为另外一个设备的输入。
主要特点：
- 模型并行的优势在于面对参数量很大的模型，能够有效减少模型对显存的占用。
- 缺点很明显，假设我们有三个 GPU 设备，数据并行会同时使用这三个 GPU 进行训练，而模型并行用这三个设备执行一次训练，带来的时间损耗很大。
（3）Pipeline Parallelism （PP）流水线并行
流水线并行（PP）是模型并行的一种变体，通过将每批输入数据拆分成若干较小的 “micro-batches”，来减少设备的空闲时间。只有在整个模型处理完每个micro-batches后，才会更新模型参数，这意味着当其他设备仍在处理上一个 micro-batches 时，每个设备就可以开始处理下一个 micro-batches 。
(1) 它只在模型作为 torch.nn.Sequential 模块实现时起作用；
(2) 它要求每个模块的输入和输出要么是单个张量，要么是张量的元组
扫一扫
461
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
