多模态大模型实战-DeepSpeed x MiniGPT4Qwen系列4:给Lavis带上DeepSpeed翅膀

前情提要

在此之前,我对lavis(https://github.com/salesforce/LAVIS)进行修改,在此基础上将BLIP2接入阿里通义千问的Qwen-7B-Chat中构成了MiniGPT4Qwen项目

https://zhuanlan.zhihu.com/p/664612306

接下来两期,我对大模型训练的各个组件进行了详细的分析,如:Trainer+Registry机制构成一套灵活、易配置的代码框架( https://zhuanlan.zhihu.com/p/670572461 )以及 大模型训练时的混合精度(amp)和gradient-checkpointing技术的分析和实验( https://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的项目哈:

完整项目地址: https://github.com/Coobiw/MiniGPT4Qwen

我的DeepSpeed Tutorials: https://github.com/Coobiw/MiniGPT4Qwen/tree/master/deepspeed_tutorials

ZERO理论部分参考

建议看这篇博客,写的不错~

https://zhuanlan.zhihu.com/p/663517415

简单的DeepSpeed Tutorials

根据DeepSpeed的官方Docs( https://deepspeed.readthedocs.io/en/latest/ ),DeepSpeed使用上和DDP对模型的封装区别并不大,下面会分别介绍每个组件的定义,完整的代码见https://github.com/Coobiw/MiniGPT4Qwen/blob/master/deepspeed_tutorials/test_zero_no_optim.py

分布式方面

初始化:不同于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踩的一些坑的记录(可能会不断更新)

建议参考博客:https://mp.weixin.qq.com/s/mn47BK9IFYwq70ofbyv6FQ

数据类型问题

使用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地址:https://github.com/Coobiw/MiniGPT4Qwen/tree/master/deepspeed_tutorials/ds_configs

首先,借用我个人上篇博客的分析采用混合精度时的显存占用:

在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(见https://github.com/Coobiw/MiniGPT4Qwen/blob/master/lavis/runners/deepspeed_runner.py),整体代码比较复杂,这里就不介绍代码层面的事情(上面已经讨论了那么多了哈哈哈),大家具体去github仓库看吧

这里想说几个点,用MiniGPT4Qwen,由于只训练中间的一个linear projection层(几M的参数量),导致实际上ZERO系列对显存的优化并不明显(尤其是ZERO-1和ZERO-2,几乎不会有提升,甚至由于通信的缘故导致实际显存占用有较少的增加,ZERO-3我还在调),但如果想多训练一些部分(如:把BLIP2的Q-former也打开,给Qwen上LoRA,甚至想训练ViT),在3090上,你不开ZERO优化就是不可能的哈。

这里放一个使用示例吧~额滴麦麦~

DeepSpeed Config的参数介绍(按需~)

参考

https://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”
}

关于auto(一般需要通过命令行参数传递)

​ 

Read More 

發佈留言

發佈留言必須填寫的電子郵件地址不會公開。 必填欄位標示為 *