前情提要

在此之前，我对lavis(github.com/salesforce/...%25E8%25BF%259B%25E8%25A1%258C%25E4%25BF%25AE%25E6%2594%25B9%25EF%25BC%258C%25E5%259C%25A8%25E6%25AD%25A4%25E5%259F%25BA%25E7%25A1%2580%25E4%25B8%258A%25E5%25B0%2586BLIP2%25E6%258E%25A5%25E5%2585%25A5%25E9%2598%25BF%25E9%2587%258C%25E9%2580%259A%25E4%25B9%2589%25E5%258D%2583%25E9%2597%25AE%25E7%259A%2584Qwen-7B-Chat%25E4%25B8%25AD%25E6%259E%2584%25E6%2588%2590%25E4%25BA%2586MiniGPT4Qwen%25E9%25A1%25B9%25E7%259B%25AE "https://github.com/salesforce/LAVIS)%E8%BF%9B%E8%A1%8C%E4%BF%AE%E6%94%B9%EF%BC%8C%E5%9C%A8%E6%AD%A4%E5%9F%BA%E7%A1%80%E4%B8%8A%E5%B0%86BLIP2%E6%8E%A5%E5%85%A5%E9%98%BF%E9%87%8C%E9%80%9A%E4%B9%89%E5%8D%83%E9%97%AE%E7%9A%84Qwen-7B-Chat%E4%B8%AD%E6%9E%84%E6%88%90%E4%BA%86MiniGPT4Qwen%E9%A1%B9%E7%9B%AE")

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需要使用以下方法进行分布式多进程的初始化

ini 复制代码

deepspeed.init_distributed(
        dist_backend='nccl',
        init_method='env://',
        distributed_port=8080,
    )

其他操作：都可以使用torch.distributed的操作，如：torch.distributed.barrier(),等价于使用deepspeed.utils.dist.barrier()

初始化模型

ini 复制代码

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的时候才可用。

ini 复制代码

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如下：

python 复制代码

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)

训练循环

ini 复制代码

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

ini 复制代码

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

ini 复制代码

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数，如下：

scss 复制代码

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联动：

ini 复制代码

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

请在代码中加入

python 复制代码

parser.add_argument('--local_rank',default=-1,type=int)
parser.add_argument('--deepspeed_config', default=None, type=str,required=True)

批量大小相关参数 (Batch size)

json 复制代码

"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的

优化器

json 复制代码

"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

json 复制代码

"scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 1e-4,
            "warmup_num_steps": 5
        }
    },

与optimizer类似
详情可以看deepspeed文档

梯度裁剪

json 复制代码

"gradient_clipping": 1.0,

gradient_clipping: 启用梯度剪裁，剪裁阈值为指定值(默认值：1.0)

logging相关

json 复制代码

"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）

yaml 复制代码

"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）

json 复制代码

"zero_optimization": {
    "stage": 0
}

Zero-1（optimizer-state shard）

json 复制代码

"zero_optimization": {
    "stage": 1
}

一些通信方面的操作比较共通，详见zero-2（zero-3有所不同）

Zero-2（optimizer-state + gradient shard）

json 复制代码

"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）

json 复制代码

"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，会有一些额外的时间开销

设置

json 复制代码

"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开启后才能使用！

json 复制代码

"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

json 复制代码

"tensorboard": {
    "enabled": true, 
    "output_path": "log/", 
    "job_name": "2023-12-15"
}

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

前情提要

ZERO理论部分参考

简单的DeepSpeed Tutorials

分布式方面

初始化模型

zero3直接分片初始化（如果模型太大，无法在一张卡上加载）

DataLoader

优化器

学习率调度器（Scheduler）

训练循环

保存checkpoint

加载checkpoint

其他DeepSpeed踩的一些坑的记录（可能会不断更新）

数据类型问题

在EVA-ViT-G（1B模型）上的ZERO实验

赋予MiniGPT4Qwen以DeepSpeed的翅膀

DeepSpeed Config的参数介绍（按需～）

参考

批量大小相关参数 (Batch size)

关于训练步数（有关打印log和optimizer/scheduer的行为）

优化器

Scheduler

梯度裁剪

logging相关

混合精度训练（支持bf16和fp16）

Zero相关

Zero-0（等价于DDP）

Zero-1（optimizer-state shard）

Zero-2（optimizer-state + gradient shard）

Zero-3（optimizer-state + gradient + model-params shard）

offload相关

offload解释

设置

NVMe

启用tensorboard

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