deepspeed&zero及大模型显存占用

我们先来看模型训练过程中的显存占用,大模型的显存占用主要分布在两个地方:

1. 模型状态维护上

2. 激活值

这是ZeRO的论文:https://arxiv.org/pdf/1910.02054

先来看下混合精度训练, 根据ZeRO的论文,混合精度训练(Mixed-Precision Training)是当前训练大模型的SOTA的方式。这个方式上:

* 模型参数,及参与运算后计算出的激活值,以及反向传播时候的梯度值等,全部都是以fp16来存储。

* 模型梯度算出来之后,要用梯度来更新模型参数的时候,会用fp32来进行更新,也就是说,原始的参数有一份一直存着的fp32的副本。  然后一次迭代的流程应该就是(猜测):

      *  fp32的参数副本,转一份为fp16的参数

      *  fp16的参数参与运算,算出fp16的激活值,以及fp16的梯度

      *  fp16的梯度,转为fp32的梯度

      *  fp32的梯度被用到更新fp32参数的过程中

基于上述内容,来看

先来看第一块儿:模型状态维护

1. 参与训练forward及backward过程的部分:

      模型参数本身,例如模型大小是7B,用fp16/bf16一个参数2个字节

### 使用 Deepspeed Zero3 进行自定义模型训练 为了在自定义深度学习模型中应用 Deepspeed 的 Zero3 来优化训练过程,从而提高效率和减少内存占用,需遵循特定配置方法[^2]。 #### 安装依赖库 首先安装必要的 Python 库来支持 Deepspeed 和 PyTorch: ```bash pip install deepspeed transformers torch ``` #### 配置 Deepspeed 参数文件 创建一个 JSON 文件 `ds_config.json` 用于指定 Deepspeed 设置。对于 Zero3 而言,此文件应包含如下内容: ```json { "fp16": { "enabled": true, "loss_scale": 0, "initial_scale_power": 16, "loss_scale_window": 1000, "hysteresis": 2, "min_loss_scale": 1 }, "optimizer": { "type": "AdamW", "params": { "lr": 3e-5, "betas": [ 0.9, 0.999 ], "eps": 1e-8, "weight_decay": 0.0 } }, "zero_optimization": { "stage": 3, "offload_optimizer": { "device": "cpu" }, "overlap_comm": true, "contiguous_gradients": true, "reduce_bucket_size": 5e8, "allgather_partitions": true, "allgather_bucket_size": 5e8, "max_live_parameters": 1e7, "sub_group_size": 1e9 }, "gradient_accumulation_steps": 1, "steps_per_print": 2000, "wall_clock_breakdown": false } ``` 上述配置启用了 FP16 训练模式以及第三阶段的 ZeRO 技术,并允许将优化器状态卸载到 CPU 内存以节省 GPU 显存资源。 #### 修改训练脚本引入 Deepspeed 支持 接下来调整现有的训练代码以便利用 Deepspeed 提供的功能。这里展示了一个简单的例子说明怎样集成这些改动: ```python import torch from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments import deepspeed model_name_or_path = 'bert-base-cased' train_dataset = ... # 用户自己的数据集加载逻辑 eval_dataset = ... # 同上 training_args = TrainingArguments( output_dir='./results', num_train_epochs=3, per_device_train_batch_size=16, per_device_eval_batch_size=16, warmup_steps=500, weight_decay=0.01, logging_dir='./logs', ) # 初始化带有 Deepspeed 的 trainer 对象 trainer = Trainer( model_init=lambda: AutoModelForSequenceClassification.from_pretrained(model_name_or_path), args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, # 添加 Deepspeed 参数路径 deepspeed='ds_config.json' ) if __name__ == "__main__": # 开始训练前初始化 Deepspeed 环境变量 deepspeed.init_distributed() # 执行实际训练流程 trainer.train() ``` 这段代码展示了如何通过设置 `Trainer` 类中的参数来启用 Deepspeed 功能,并指定了之前准备好的配置文件位置。此外,在主程序入口处调用 `deepspeed.init_distributed()` 方法完成分布式环境初始化操作[^4]。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值