01
—
为什么需要模型并行？
在深度学习训练中,我们面临着一个"大胃王"问题 - 模型越来越大,单个GPU的内存总是不够用:
GPT-3有1750亿参数 PaLM有5400亿参数 MT-NLG有5300亿参数
当模型太大，无法装入单个GPU显存 需要更快的训练和推理速度 想要更有效地利用硬件资源
所谓的“分层次烤蛋糕”就是指将模型的参数分布到多个设备上，每个设备只存储一部分模型。
02
—
什么是ZeRO 优化？
一个拥有 1.5B 参数的 GPT-2 在使用 FP16 训练时只需要 3G 来存储模型的参数，但是却无法在单张 32GB 的 GPU 上进行1.5B 参数的训练。所以内存去哪了呢？
通过实验发现，训练期间大部分的内存被模型状态（Model States），也就是：优化器参数（Optimizer states），梯度（Gradients）和模型（Parameters）所消耗。除此之外，残余状态（Residual States）消耗了剩余的内存，残余状态包括：前向传播时得到的 Activations，进行计算通信的临时缓冲区和还有没有被妥善管理的内存碎片。
ZeRO 是 DeepSpeed 引入的一种优化方法，专门用于分片大型模型的参数、梯度和优化器状态。
核心思想
将模型参数、梯度和优化器状态在所有设备之间分片存储，而不是每个设备存储完整副本。 ZeRO 有三个阶段：
参数分片（Params: Stage 3）：除了拆分优化器和梯度之外，每个设备只存储部分模型参数。
梯度分片（Gradient:Stage 2）：除了拆分优化器之外之外，继续拆分梯度每个设备只存储部分梯度，减少梯度同步的内存需求。
优化器状态分片（Optimizer:Stage 1）：拆分优化器，让每个设备只存储部分优化器状态。
import torchimport torch.distributed as distfrom torch.nn.parallel import DistributedDataParallel as DDPfrom deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizerfrom deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer as ZeROStage3# 假设我们有一个简单的模型class SimpleModel(torch.nn.Module):def __init__(self):super().__init__()self.linear1 = torch.nn.Linear(1024, 1024)self.linear2 = torch.nn.Linear(1024, 1024)def forward(self, x):x = self.linear1(x)x = self.linear2(x)return x# =============== 传统数据并行 ===============def traditional_ddp():# 初始化模型model = SimpleModel()# 将模型移到GPUmodel = model.cuda()# 包装成DDP模型model = DDP(model)# 创建优化器optimizer = torch.optim.Adam(model.parameters())# 在每个GPU上都保存完整的模型副本、梯度和优化器状态# 这就是为什么内存消耗大的原因# =============== ZeRO Stage 1 ===============def zero_stage_1():model = SimpleModel().cuda()model = DDP(model)# 创建ZeRO优化器，只对优化器状态进行分片# 通常通过json 配置文件来完成配置。optimizer_config = {"zero_optimization": {"stage": 1,"allgather_partitions": True,"allgather_bucket_size": 5e8,"overlap_comm": True}}optimizer = DeepSpeedZeroOptimizer(model.parameters(),optimizer_config,timers=None)# 优化器状态（如Adam的动量）被分片到不同GPU上# 但模型参数和梯度仍完整保存在每个GPU上# =============== ZeRO Stage 2 ===============def zero_stage_2():model = SimpleModel().cuda()model = DDP(model)optimizer_config = {"zero_optimization": {"stage": 2,"allgather_partitions": True,"allgather_bucket_size": 5e8,"reduce_scatter": True,"overlap_comm": True}}optimizer = DeepSpeedZeroOptimizer(model.parameters(),optimizer_config,timers=None)# 优化器状态和梯度都被分片# 在反向传播时动态收集梯度# 模型参数仍完整保存在每个GPU上# =============== ZeRO Stage 3 ===============def zero_stage_3():model = SimpleModel().cuda()optimizer_config = {"zero_optimization": {"stage": 3,"allgather_partitions": True,"allgather_bucket_size": 5e8,"reduce_scatter": True,"overlap_comm": True,"contiguous_gradients": True,"stage3_max_live_parameters": 1e9}}optimizer = ZeROStage3(model.parameters(),optimizer_config,timers=None)# 模型参数、梯度和优化器状态都被分片# 在前向和反向传播时动态收集需要的参数# 实现了真正的模型并行def training_step():# 在需要时收集参数with optimizer.gather_parameters():outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()# 参数自动释放，节省内存optimizer.step()
—
对比ZeRO的三种方式
以一个大小为1G的模型并且有4个GPU为例，来说明三个不同stage所占有的内存。
1）如果采用数据并行的方式：
训练需要至少16G的内存。
# =============== 内存使用比较 ===============def memory_comparison():# 假设模型大小为1GB，有4个GPU# 传统DDP：每个GPU需要# - 模型参数: 1GB# - 梯度: 1GB# - 优化器状态: 2GB (Adam优化器需要两个状态)# 总共: 每个GPU 4GB，总共16GB
2）如果采用ZeRO Stage 1 的方式，也就是采用拆分优化器的方式。
训练需要至少10G的内存。
# ZeRO-1：每个GPU需要# - 模型参数: 1GB# - 梯度: 1GB# - 优化器状态: 0.5GB (分片后)# 总共: 每个GPU 2.5GB，总共10GB
3）如果采用ZeRO Stage 2的方式，也就是采用拆分梯度和优化器的方式。
训练需要至少7G的内存。
# ZeRO-2：每个GPU需要# - 模型参数: 1GB# - 梯度: 0.25GB (分片后)# - 优化器状态: 0.5GB (分片后)# 总共: 每个GPU 1.75GB，总共7GB
4）如果采用ZeRO Stage 3的方式，也就是采用拆分梯度和优化器以及参数的方式。
训练至少需要4G的内存。
# ZeRO-3：每个GPU需要# - 模型参数: 0.25GB (分片后)# - 梯度: 0.25GB (分片后)# - 优化器状态: 0.5GB (分片后)# 总共: 每个GPU 1GB，总共4GB
进技术交流群请添加AINLP小助手微信（id: ainlp2)
请备注具体方向+所用到的相关技术点 关于AINLP
AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括LLM、预训练模型、自动生成、文本摘要、智能问答、聊天机器人、机器翻译、知识图谱、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP小助手微信(id：ainlp2)，备注工作/研究方向+加群目的。
