文章
文章用户网址导航

深度学习框架DeepSpeed使用教程及配置参数分享

释放双眼，带上耳机，听听看~！
本文介绍了深度学习框架DeepSpeed的使用教程及配置参数分享，包括分布式训练的初始化、模型初始化等内容，适合对深度学习框架DeepSpeed感兴趣的读者。

前情提要

在此之前，我对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…

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

深度学习框架DeepSpeed使用教程及配置参数分享

在ViT模型使用时，一般情况激活值还是会占很大比例的（经验性的结论可能会占50%，当然也和batch_size有关），且这里并不开启Gradient-Checkpointing

深度学习框架DeepSpeed使用教程及配置参数分享

  • 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使用教程及配置参数分享

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

深度学习框架DeepSpeed使用教程及配置参数分享

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（一般需要通过命令行参数传递）

深度学习框架DeepSpeed使用教程及配置参数分享

本网站的内容主要来自互联网上的各种资源，仅供参考和信息分享之用，不代表本网站拥有相关版权或知识产权。如您认为内容侵犯您的权益，请联系我们，我们将尽快采取行动，包括删除或更正。
AI教程Distributed Training深度学习深度学习框架配置参数
AI教程

PyTorch 2.0.0新手教程：注册算子的主要方式

2023-12-22 14:01:00

AI教程

AIGC技术前瞻与应用

2023-12-22 14:05:14

猜你喜欢
❯

解锁会员权限

开通会员

解锁海量优质VIP资源

立刻开通
个人中心
购物车
优惠劵
今日签到
有新私信 私信列表
搜索

  • 扫码打开当前页

返回顶部