前情提要
在此之前,我对lavis(github.com/salesforce/…
zhuanlan.zhihu.com/p/664612306
接下来两期,我对大模型训练的各个组件进行了详细的分析,如:Trainer+Registry机制构成一套灵活、易配置的代码框架( zhuanlan.zhihu.com/p/670572461 )以及 大模型训练时的混合精度(amp)和gradient-checkpointing技术的分析和实验( zhuanlan.zhihu.com/p/671165275 )。 然而,lavis框架的分布式使用的是最基本的pytorch的DDP,虽然简单易用,但如今还是算是有些out-of-date了。FSDP、DeepSpeed、Megatron等各种分布式并行训练框架更加受到青睐。我尝试过pytorch原生的fsdp(来源自fairscale),个人不是很喜欢去使用,所以这期想介绍一下更加简单易用、Plug-in的DeepSpeed。
本项目将给出一个我自己参考DeepSpeed文档书写的简单tutorials,再介绍一下我踩的一些坑,然后我将DeepSpeed支持进了原本的MiniGPT4Qwen项目中,给出了ZERO-0(等价于DDP)、ZERO-1、ZERO-2的配置。至于一些DeepSpeed的参数配置,我参考和总结了一些放在文章最后,兄弟们按需自取~
项目代码和文档的相关链接如下,大家多提意见呀,有帮助的话麻烦点个star呜呜呜,想有第一个100+stars的项目哈:
完整项目地址: github.com/Coobiw/Mini…
我的DeepSpeed Tutorials: github.com/Coobiw/Mini…
ZERO理论部分参考
建议看这篇博客,写的不错~
zhuanlan.zhihu.com/p/663517415
简单的DeepSpeed Tutorials
根据DeepSpeed的官方Docs( deepspeed.readthedocs.io/en/latest/ ),DeepSpeed使用上和DDP对模型的封装区别并不大,下面会分别介绍每个组件的定义,完整的代码见github.com/Coobiw/Mini…
分布式方面
初始化:不同于torch.distributed使用torch.distributed.init_process_group
,deepspeed需要使用以下方法进行分布式多进程的初始化
deepspeed.init_distributed(
dist_backend='nccl',
init_method='env://',
distributed_port=8080,
)
其他操作:都可以使用torch.distributed
的操作,如:torch.distributed.barrier()
,等价于使用deepspeed.utils.dist.barrier()
初始化模型
model = create_eva_vit_g(
img_size=224,
drop_path_rate=0.,
use_checkpoint=False,
precision="fp32").cuda(args.local_rank)
- 可以直接放到cuda上,但需要是对应的
local_rank
- 也可以不
.cuda()
,在deepspeed.initialize
的时候会自动放到对应的gpu上
zero3直接分片初始化(如果模型太大,无法在一张卡上加载)
如果模型太大,无法在一张显卡上加载,可以加载在多卡上,需要在deepspeed.zero.Init()
下进行初始化,但只有采用zero3的时候才可用。
with deepspeed.zero.Init():
model = xxxx
DataLoader
既可以使用DDP的分布式DataLoader(即带上DistributedSampler
),也可以直接将Dataset输入给deepspeed.initialize()
得到DataLoader
优化器
目前发现如果打开zero_optimization
的offload_optimizer
,是得使用DeepSpeed内部的fused的optimizer的,直接使用torch自带的optimizer会报错:
deepspeed.runtime.zero.utils.ZeRORuntimeException:
You are using ZeRO-Offload with a client provided optimizer (<class ‘torch.optim.adamw.AdamW’>) ….
学习率调度器(Scheduler)
用deepspeed_config
自然是不会出错的方法,但deepspeed支持的scheduler还是太简单了,甚至不好满足简单的linear warmup + 后续cosine decay。如果使用torch的scheduler,由于需要提供optimizer参数,假如optimizer是用的deepspeed实现的optimizer,会报错(因为不是torch里的optimizer类):
TypeError: DeepSpeedZeroOptimizer is not an Optimizer
可以自定义scheduler如下:
import math
class CosineAnnealingLR:
def __init__(self, optimizer, T_max, base_lr,eta_min=0):
self.optimizer = optimizer
self.T_max = T_max
self.base_lr = base_lr
self.eta_min = eta_min
self.current_step = 0
def step(self):
self.current_step += 1
new_lr = self.eta_min + (self.base_lr - self.eta_min) * (1 + math.cos(math.pi * self.current_step / self.T_max)) / 2
for param_group in self.optimizer.param_groups:
param_group['lr'] = new_lr
def state_dict(self):
"""Returns the state of the scheduler as a :class:`dict`.
It contains an entry for every variable in self.__dict__ which
is not the optimizer.
"""
return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}
def load_state_dict(self, state_dict):
"""Loads the schedulers state.
Args:
state_dict (dict): scheduler state. Should be an object returned
from a call to :meth:`state_dict`.
"""
self.__dict__.update(state_dict)
训练循环
model = create_eva_vit_g(img_size=224,drop_path_rate=0.,use_checkpoint=False,precision="fp32").cuda()
trainset = ExampleDataset()
model_engine, optimizer , trainloader , _ = deepspeed.initialize(
args,
model=model,
model_parameters=model.parameters(),
training_data=trainset,
# optimizer=optimizer,
)
scheduler = CosineAnnealingLR(optimizer=optimizer,T_max=10,base_lr=1e-4,eta_min=1e-5)
for epoch in range(start_epoch+1,10):
start_time = time.time()
for global_step, batch in tqdm(enumerate(trainloader,start=start_global_step+1)):
batch = batch.cuda()
batch = batch.to(next(model_engine.parameters()).dtype) # data的dtype可能和模型不对应,比如开了bf16
outputs = model_engine(batch)
loss = outputs.mean()
model_engine.backward(loss)
model_engine.step()
scheduler.step()
-
model_engine.backward(loss)会完成以下操作:
- 如果开了amp且需要对loss进行rescale,会rescale loss(类似torch.cuda.amp.GradScaler)
- 反向传播(loss.backward())
- 梯度积累(gradient accumulation)
-
model_engine.step()完成以下操作
- optimizer.step()
- 如果有定义scheduler,会进行scheduler.step()
- 如果开启了loss rescale,进行scaler.update()
- optimizer.zero_grad()
保存checkpoint
不同于DDP只需要在主进程上save,deepspeed直接在所有进程上save即可(因为zero系列会有进程间的通信,不是所有参数、梯度、优化器状态都在一张卡上)
- 如果只在主进程
save_checkpoint
,其他进程会被hang住,导致无法完成,最后timeout
if epoch % save_interval == 0:
client_sd['global_step'] = global_step
client_sd['epoch'] = epoch
client_sd['scheduler'] = scheduler.state_dict()
ckpt_id = loss.item()
model_engine.save_checkpoint(save_dir=save_dir,tag=f'epoch_{epoch}',client_state = client_sd)
-
client_state:可以多存一下自定义的内容(但不要和一些常规的key的名字相同,如:
module
,optimizer
,lr_scheduler
-
tag:可以指定存储的
.pt
文件的目录名(在save_dir
底下)
加载checkpoint
def load_ckpt(model_engine, scheduler, ckpt_dir, ckpt_tag):
_, client_sd = model_engine.load_checkpoint(load_dir=ckpt_dir,tag=ckpt_tag)
global_step = client_sd['global_step']
start_epoch = client_sd['epoch']
scheduler.load_state_dict(client_sd['scheduler'])
return start_epoch, global_step
后续需要在for循环处修改开始的epoch和dataloader
的step数,如下:
for epoch in range(start_epoch+1,10):
for global_step, batch in tqdm(enumerate(trainloader,start=start_global_step+1)):
...
其他DeepSpeed踩的一些坑的记录(可能会不断更新)
建议参考博客:mp.weixin.qq.com/s/mn47BK9IF…
数据类型问题
使用deepspeed的fp16或bf16,在数据输入、中间算子等位置常常会出现数据类型的问题,在输入处直接改dtype
可能还行,但其他位置一直修改dtype
终究有些不美观,甚至会出现错误,解决方法:可以直接和torch.cuda.amp.autocast
联动:
model_dtype = next(model.parameters()).dtype
with (torch.cuda.amp.autocast(dtype=model_dtype,cache_enabled=False) if model_dtype != torch.float32 else contextlib.nullcontext()):
loss, loss_dict = self.train_step(model=model, samples=samples)
在EVA-ViT-G(1B模型)上的ZERO实验
ZERO的Configs地址:github.com/Coobiw/Mini…
首先,借用我个人上篇博客的分析采用混合精度时的显存占用:
在ViT模型使用时,一般情况激活值还是会占很大比例的(经验性的结论可能会占50%,当然也和batch_size有关),且这里并不开启Gradient-Checkpointing
-
ZERO-1进行优化器状态分片,相当于上面的8x + 4x = 12*x(DeepSpeed的fp32副本算在Optimizer内部)进行分片,可以省大量显存(ZERO-0的 bs=32会OOM,而ZERO-1完全有余量)
-
ZERO-2、ZERO-3的分片虽然也节省显存,但由于通信量,实际上似乎没怎么省,甚至由于通信量导致实际使用显存大于ZERO-1
-
ZERO-2(无offload optimizer)、ZERO-3(无offload)速度很快,比DDP更快,可能是因为通信的缘故,尤其是我是单机多卡的情况,通信速度快
-
offload会降低显存占用,因为会把一些tensor放在CPU,但速度上会变慢许多(因为受限于cpu和gpu间IO速度)
赋予MiniGPT4Qwen以DeepSpeed的翅膀
这里在原有的MiniGPT4Qwen上实现了DeepSpeed的Runner(见github.com/Coobiw/Mini…
这里想说几个点,用MiniGPT4Qwen,由于只训练中间的一个linear projection层(几M的参数量),导致实际上ZERO系列对显存的优化并不明显(尤其是ZERO-1和ZERO-2,几乎不会有提升,甚至由于通信的缘故导致实际显存占用有较少的增加,ZERO-3我还在调),但如果想多训练一些部分(如:把BLIP2的Q-former也打开,给Qwen上LoRA,甚至想训练ViT),在3090上,你不开ZERO优化就是不可能的哈。
这里放一个使用示例吧~额滴麦麦~
DeepSpeed Config的参数介绍(按需~)
参考
zhuanlan.zhihu.com/p/650824387
请在代码中加入
parser.add_argument('--local_rank',default=-1,type=int)
parser.add_argument('--deepspeed_config', default=None, type=str,required=True)
批量大小相关参数 (Batch size)
"train_batch_size": 16,
"gradient_accumulation_steps": 1,
"train_micro_batch_size_per_gpu": 8,
-
计算公式:
train_batch_size = micro_batch_per_gpu * gradient_acc_step * world_size(GPU个数)
-
train_micro_batch_size_per_gpu
:单个GPU在一个步骤中处理的微批量大小(不算梯度累积)- 如果同时提供train_batch_size和gradient_accumulation_steps,可以忽略train_micro_batch_size_per_gpu。
- 默认值:train_batch_size的值
-
gradient_accumulation_steps
:在计算平均并应用梯度之前累积梯度的训练步骤数- 如果同时提供train_batch_size和train_micro_batch_size_per_gpu,可以忽略gradient_accumulation_steps
- 默认值:1
-
train_batch_size
:有效的训练批量大小。这指的是每次模型更新所涉及的数据样本数量- 默认值:32
关于训练步数(有关打印log和optimizer/scheduer的行为)
deepspeed里有micro_step
和global_step
,前者不管梯度积累,后者管
- scheduler运行step方法和optimizer运行step方法都是看global step的
- 打印log的步数是按micro_step的
优化器
"optimizer": {
"type": "Adam",
"params": {
"lr": 1e-4,
"betas": [
0.9,
0.99
],
"eps": 1e-7,
"weight_decay": 0,
"torch_adam": false,
"adam_w_mode": true
}
},
type
: 优化器名称。DeepSpeed原生支持Adam、AdamW等优化器,并可以从torch导入其他优化器params
: 参数字典,用于实例化优化器,参数名称必须与优化器的type
相匹配torch_adam
: 使用torch的Adam实现,而不是fused的Adam实现。 (默认值:false)adam_w_mode
: 应用L2正则化(也称为AdamW)。 (默认值:true)
Scheduler
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 1e-4,
"warmup_num_steps": 5
}
},
- 与optimizer类似
- 详情可以看deepspeed文档
梯度裁剪
"gradient_clipping": 1.0,
gradient_clipping
: 启用梯度剪裁,剪裁阈值为指定值(默认值:1.0)
logging相关
"steps_per_print": 1,
"wall_clock_breakdown": false,
"dump_state":false
-
steps_per_print
: 每过多少个train_step打印进度报告。- 报告内容包括训练的iterations,由于混合精度训练中的溢出而跳过的优化器更新数,当前学习率以及当前动量(包含一阶和二阶)
- 默认值:10
-
wall_clock_breakdown
: 启用前向、反向和更新训练阶段的时序计时,以分析时间延迟(默认值:false) -
dump_state
: 在初始化后打印出DeepSpeed对象的状态信息(默认值:false)
混合精度训练(支持bf16和fp16)
"fp16": {
"enabled": false,
"auto_cast": false,
"loss_scale": 0,
"initial_scale_power": 16,
"loss_scale_window": 1000,
"hysteresis": 2,
"consecutive_hysteresis": false,
"min_loss_scale": 1
}
"bf16": {
"enabled": true
}
-
auto_cast
: 是否将输入强制转换为fp16数据类型 (默认值:false) -
loss_scale
: 表示FP16训练的损失缩放值- 启用动态损失缩放
- 默认值:0.0
-
initial_scale_power: 表示初始动态损失比例值的功率,实际损失规模计算为
2**(initial_scale_power)
(默认值:16) -
loss_scale_window
: 代表动态损失缩放值上升/下降的窗口范围。(默认值:1000) -
hysteresis
: 表示动态损失缩放中的延迟偏移 (默认值:2)(没太去理解) -
consecutive_hysteresis
: 表示是否在达到不会溢出的迭代时重新填充hysteresis值(默认值:false)(没太去理解) -
min_loss_scale
: 表示最小动态损失比例值 (默认值:1)
Zero相关
Zero-0(等价于DDP)
"zero_optimization": {
"stage": 0
}
Zero-1(optimizer-state shard)
"zero_optimization": {
"stage": 1
}
一些通信方面的操作比较共通,详见zero-2(zero-3有所不同)
Zero-2(optimizer-state + gradient shard)
"zero_optimization": {
"stage": 2,
"allgather_partitions": true,
"allgather_bucket_size": 3e8,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 3e8,
"contiguous_gradients": true
}
-
allgather_partitions
: 在每个步骤结束时,从所有GPU中选择使用all-gather的操作或者一系列广播集体操作之间的方式,以收集更新后的参数 (默认值:true) -
allgather_bucket_size
: 用于调节all-gather操作的对张量的分桶大小- 将张量分成较小的桶有助于在通信过程中更高效地传输数据
- 较大的allgather_bucket_size值会导致每个桶的尺寸增大,可能加速通信操作,但也需要更多内存来存储中间结果
- 通信速度和显存(/内存)占用的trade-off
- 默认值:5e8
-
overlap_comm
: 控制通信与计算是否交叠执行- 当设置为True时,DeepSpeed将尝试在梯度计算期间并行进行梯度通信,这有效地缩短通信时间,从而加速整个训练过程
- 默认值:false
-
reduce_scatter
: 使用reduce或reduce-scatter来替代all-reduce以平均梯度。(默认值:true) -
reduce_bucket_size
: 用于控制All-reduce操作的分桶大小- 通信速度和显存(/内存)占用的trade-off
- 默认值:5e8
-
contiguous_gradients
: 在梯度产生时将其复制到一个连续的缓冲区中。在反向传播过程中避免了内存碎片化问题。(默认值:true)
Zero-3(optimizer-state + gradient + model-params shard)
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": 1e6,
"stage3_prefetch_bucket_size": 4e6,
"stage3_param_persistence_threshold": 1e4,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
ZeRO-3 中不使用 allgather_partitions、allgather_bucket_size 和 reduce_scatter 配置参数
-
sub_group_size
: 控制在优化器步骤中参数更新的粒度- 参数被分组到大小为
sub_group_size
的桶中,每个桶依次进行一次更新 - 当与ZeRO-Infinity中的NVMe offload同时使用时,
sub_group_size
决定了在优化器步骤期间从NVMe迁移到CPU内存的模型状态的粒度。这有助于避免超大模型对CPU内存的过度占用 - 在不使用NVMe offload时,请保持其默认值
- 若遇到内存不足(OOM)情况,可以考虑减小
sub_group_size
- 当优化器迭代较缓慢时,也可以考虑增大
sub_group_size
- 默认值:1e9
- 参数被分组到大小为
-
stage3_prefetch_bucket_size
: prefetch参数的固定缓冲区大小- 较小的值使用的内存较少,但可能会因通信而增加停顿
- 默认值:5e8
-
stage3_max_live_parameters
: 保留在GPU上的完整参数数量的上限。(默认值:1e9) -
stage3_max_reuse_distance
: 根据参数在未来何时再次使用的指标来决定是舍弃还是保留参数- 如果一个参数在不久的将来会再次被使用(小于
stage3_max_reuse_distance
),则会保留该参数以减少通信开销 - 在遇到内存不足(OOM)的情况下,可以降低
stage3_max_live_parameters
和stage3_max_reuse_distance
的值 - 默认值:1e9
- 如果一个参数在不久的将来会再次被使用(小于
-
stage3_gather_16bit_weights_on_model_save
: 在保存模型时启用模型FP16权重合并- 对于大型模型和多GPU环境,这是一项在内存和速度方面代价较高的操作
- 默认值:false
offload相关
offload解释
在中间变量产生时,将中间变量移动到 CPU/NVMe 上,在需要使用中间变量时移动到 GPU 上。通过这种方式,可以减小中间变量的显存占用。Zero的Offload优化通常更适用于资源受限,但是又要训练大模型的情况。通过时间换空间。比如把optimizer state、parameters offload到 CPU/NVMe,会有一些额外的时间开销
设置
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "nvme",
"pin_memory": true
}
-
在开启ZeRO第一阶段后,可以使用offload_optimizer
-
在开启ZeRO第三阶段后才可以同时使用offload_optimizer与offload_param
-
offload to NVMe 只在stage 3开启后才能使用!
-
pin_memory
:转移到页面锁定的CPU内存- 这可能会提升吞吐量,但代价是增加了额外的内存开销
- 默认值:false
NVMe
offload to NVMe 只在stage 3开启后才能使用!
"offload_optimizer": {
"device": "nvme",
"nvme_path": "/dev/shm",
"buffer_count": 4,
"fast_init": false
},
"offload_param": {
"device": "nvme",
"nvme_path": "/dev/shm",
"buffer_count": 5,
"buffer_size": 1e8,
"max_in_cpu": 1e9
}
-
nvme_path
: 用于卸载优化器/参数的NVMe设备的文件系统路径。 -
buffer_count
(offload_optimizer): 用于将优化器状态卸载到NVMe的缓冲池中的缓冲区数量。这个数量至少应该是优化器每个参数维护的状态数。例如,Adam优化器有4个状态(参数、梯度、动量和方差)。 (默认值:5) -
fast_init
: 启用在卸载至NVMe时的快速优化器初始化。 (默认值:false) -
buffer_count
(offload_param): 将参数卸载到NVMe的缓冲池中的缓冲区数量。 (默认值:5) -
buffer_size
: 将参数卸载到NVMe的缓冲池中的缓冲区大小。 (默认值:1e8) -
max_in_cpu
: 启用卸载至NVMe时在CPU内存中保留的参数元素数量。 (默认值:1e9)
启用tensorboard
"tensorboard": {
"enabled": true,
"output_path": "log/",
"job_name": "2023-12-15"
}