企业宣传,产品推广,广告招商,广告投放联系seowdb

上 优雅谈大模型 白话ZeRO

并行策略

正如和分布式一样,如何利用多设备和多硬件也是很重要的一个环节。大模型训练也是如此,如今训练大模型离不开各种分布式并行策略,常用的并行策略包括:

下图为其中一种算法,将所有的梯度分为五份,然后按照圆圈的方式传播(而不是广播数据)累加,直到每个GPU都完成一个部分的累加,然后再批量同步数据。

经过五次循环之后,每个GPU再将各自负责的那部分数据直接同步给其他显卡,避免多次广播运算。

由于数据并行简单易实现,应用最为广泛。然而每张显卡都存储完整的模型,此时显存大小成为了瓶颈。例如若存在2张显卡,那么系统中就存在2份模型参数,如果有4张卡,那么系统中就存在4份模型参数,如果有N张卡,系统中就存在N份模型参数,而其中N-1份其实都是冗余。那么有新的办法来解决这个问题没有,ZeRO就是其中的一种方法。

初识ZeRO

零冗余优化器(Zero Redundancy Optimizer,简称ZeRO)是一种用于大规模分布式深度学习的新型内存优化技术。ZeRO可以在当前一代GPU集群上训练具有100B参数的深度学习模型,吞吐量是当前最佳系统的吞吐量的三到五倍。它还为训练具有数万亿个参数的模型提供了一条清晰的路径,展示了深度学习系统技术的前所未有的飞跃。同时ZeRO也被集成到了DeepSpeed,它是用于加速分布式深度学习训练的高性能类库。

在搞清楚ZeRO之前需要先明白在模型训练时候,显卡的占用情况。一般而言,模型在训练的时候会不断地更新参数,目前采用的都是梯度下降法。而混合精度训练(mixed precision training)和Adam优化器基本上已经是目前训练模型的标配。

Adam在SGD基础上,为每个参数梯度增加了一阶动量(momentum)和二阶动量(variance)1。混合精度训练,字如其名,同时存在fp16和fp32两种格式的数值,其中模型参数、模型梯度都是fp16,此外还有fp32的模型参数,如果优化器是Adam,则还有fp32的momentum和variance。

回过头来看看显卡中的数据分类,一般而言可以分为两类:

假设模型参数量为Φ ,则共需要2Φ+2Φ+(4Φ+4Φ+4Φ)=4Φ+12Φ=16Φ 字节存储空间。从这个统计中可以发现,其实Adam优化器占据了75%的空间。<若读者要微调3B的模型,那么则需要48B的空间!!!>

ZeRO是为了克服数据并行性和模型并行性的局限性,同时实现两者的优点。它通过在数据并行进程中对模型状态(参数、梯度和优化器状态)进行分区,而不是复制它们,从而消除了数据并行进程之间的内存冗余。同时采用动态通信计划在分布式训练设备之间共享必要的状态,以保持数据并行性的计算粒度和通信量。

如何来解读上图,蓝色就代表模型参数数Φ(若每个参数类型是FP16,则为2Φ字节),橙色代表模型梯度数Φ,K为优化器Adam的参数量系数,即KΦ。

在DeepSpeed的训练框架中, Pos对应ZeRO-1, Pos+g对应ZeRO-2, Pos+g+p对应ZeRO-3,而一般在生产中ZeRO-1就足够使用。

本文转载自​​,作者:​​

© 版权声明
评论 抢沙发
加载中~
每日一言
不怕万人阻挡,只怕自己投降
Not afraid of people blocking, I'm afraid their surrender