526互联

DeepSpeed分布式训练

发布时间 2023-12-20 12:56:49作者: Abyss_J

一、DeepSpeed 总纲

官方文档：DeepSpeed 官方英文文档

cpoy下来的原版

DeepSpeed 模型的训练是使用DeeSpeed引擎完成的

DeepSpeed引擎可以包装任何 torch.nn.module 类型的模型

二、训练

1、通过 deepspeed.initialize 初始化

# deepspeed.initialize确保分布式数据并行或混合精度训练所需的所有必要设置都在后台适当完成。
# 除了包装模型外，DeepSpeed还可以根据传递给DeepSpeed.initialize的参数
# 和DeepSpeed配置文件来构建和管理训练优化器、数据加载器和学习率调度器

# NOTE：通过 deepspeed.initialize 初始化，默认使用 nccl 后端，
#       若需要替换其他如 gloo， mpi 后端， 需要使用 deepspeed.init_distributed()进行初始化
deepspeed.initialize(args=None, 
                     model: Optional[Module] = None, 
                     optimizer: Optional[Union[Optimizer, Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer]]] = None, 
                     model_parameters: Optional[Module] = None, 
                     training_data: Optional[Dataset] = None, 
                     lr_scheduler: Optional[Union[_LRScheduler, Callable[[Optimizer], _LRScheduler]]] = None, 
                     mpu=None, 
                     dist_init_required: Optional[bool] = None, 
                     collate_fn=None, 
                     config=None, 
                     config_params=None)
# return: engine, optimizer, training_daaloader, lr_scheduler 

model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
                                                     model=net,
                                                     model_parameters=net.parameters())

2、分布式初始化

deespeed 默认使用nccl 后端， 如果已经使用了分布式环境，需将 torch.distributed.init_process_group(...)

替换为：

deepspeed.init_distributed()

若在调用 deepspeed.intialize() 函数之后，则无需调用 deepspeed.init_distributed()

因为 deepspeed.intialize() 会自动初始化分布式环境。

3、训练

DeepSpeed 的训练非常简单，具体使用三个 API 即可。

在引擎盖下，DeepSpeed通过预定义的学习率调度器，

以混合精度自动执行分布式数据并行训练所需的必要操作

for step, batch in enumerate(data_loader):
    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

NOTE

需要注意一下 Learning Rate Scheduler

1）当使用 DeepSpeed 的 learning rate scheduler （在 ds_config.json 文件中） 时，当 model_engine.step() 执行时， 即在每个 training step 进行更新。

2）当使用用户自定义 learning rate scheduler 时，若需要在每个 training step 进行更新，应

在进行 deepspeed.initialize 初始化时，将 scheduler 传递给 deepspeed.initialize , 通过 Deeppeed 进行 update 或 save/restore

3）当 learning rate scheduler 需要在其他时间步进行更新， 如 training epoch, 不能传递给deepspeed.initialize ，需要手动管理

4、保存

模型保存主要通过 DeepSpeed 的 save_checkpoint 和 load_checkpoint 这两个API进行管理，需要两个额外参数：ckpt_dir 和 ckpt_id

Args：

ckpt_dir：保存目录

ckpt_id：唯一标识目录中检查点的标识符，下面例子里使用 loss 值

client_sd ：用户希望保存的额外数据，字典形式保存。

#load checkpoint
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)
step = client_sd['step']

#advance data loader to ckpt step
dataloader_to_step(data_loader, step + 1)

for step, batch in enumerate(data_loader):

    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

    #save checkpoint
    if step % args.save_interval:
        client_sd['step'] = step
        ckpt_id = loss.item()
        model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sd = client_sd)

NOTE：

所有进程都必须调用此方法，而不仅仅是 rank = 0 的进程。这是因为每个进程都需要保存其weights 和 scheduler+optimizer 的状态。如果只为 rank = 0 的进程调用此方法，则此方法将挂起，等待与其他进程同步。

5、DeepSpeed Configuration

DeepSpeed 的 Config 配置，json 格式，传递给 deepspeed.initialize 的 args

  1. 训练批次大小 (train_batch_size):在配置文件中，可以通过指定一个整数值来设置训练批次的大小。这个值代表每个训练步骤中用于训练的样本数。
  2. 梯度累积步数 (gradient_accumulation_steps):通过设置这个参数，可以定义梯度累积的步数。这意味着在执行优化器步骤之前，模型将进行多少次前向传播和反向传播。这对于处理大批量训练数据而内存有限的情况很有用。
  3. 优化器类型和参数 (optimizer):在配置文件中，可以指定优化器的类型（如"Adam"）并定义相应的参数，比如学习率 (lr)。这样可以配置模型训练中所使用的优化器及其超参数。
  4. FP16（半精度浮点数） (fp16):可以通过设置这个参数来启用或禁用FP16混合精度训练。将其设置为true表示启用FP16，以减少模型训练时的内存占用。
  5. 零优化 (zero_optimization):这个参数用于启用或禁用零优化技术，即在模型训练中将零梯度忽略以减少计算。将其设置为true表示启用零优化。
# 一个示例
{
    "train_batch_size": 8,
    "gradient_accumulation_steps": 4,
    "optimizer": {
    "type": "Adam",
    "params": {
        "lr": 0.0001,
        "weight_decay": 0.01
    }
    },
    "fp16": {
    "enabled": true
    },
    "zero_optimization": true
}

6、资源配置

单节点

在仅在单节点运行时，DeepSpeed需要注意哪些不同的配置？CUDA_VISIBLE_DEVICES与DeepSpeed的使用有何异同？

在仅在单节点运行 DeepSpeed 时，需要注意以下几点不同的配置和使用情况：

  • 不需要指定hostfile：在单节点运行时，不需要使用 hostfile 来指定计算资源。DeepSpeed会自动检测本地机器上的GPU数量，并使用这些GPU进行训练。
  • -include和--exclude参数的使用：虽然在单节点情况下不需要 hostfile，但仍可以使用 --include和 --exclude参数来指定使用的GPU资源。此时应将主机名指定为'localhost'，例如：deepspeed --include localhost:1 ...
  • CUDA_VISIBLE_DEVICES与DeepSpeed的异同：

异同点：

  • CUDA_VISIBLE_DEVICES 是环境变量，用于指定哪些GPU设备可以被程序看到和使用。它可以限制程序在单节点内使用的GPU设备。
  • DeepSpeed 提供了 --include 和 --exclude 参数，用于在启动时指定要使用的GPU资源，类似于CUDA_VISIBLE_DEVICES的功能，但更加灵活且可以在命令行中指定不同的节点和GPU。

相同点：

  • 两者都可以用于限制DeepSpeed或其他深度学习框架在单节点上使用的GPU设备。

在单节点情况下，DeepSpeed的配置更多地集中在指定的节点内，因此不需要显式指定hostfile，而可以通过命令行参数更灵活地控制使用的GPU资源。相比之下，CUDA_VISIBLE_DEVICES是一个环境变量，需要在运行程序之前设置，控制整个进程可见的GPU设备。

三、推理

TODO