NLP——DeepSpeed框架介绍

整体介绍

  • DeepSpeed 是由微软开发的开源的深度学习优化库，旨在提高大规模模型训练的效率和可扩展性
  • DeepSpeed 框架的核心技术有：
    • ZeRO 冗余优化技术 ：通过分布式内存管理，将模型参数、梯度和优化器状态进行分区，大幅降低显存占用，是首次支持千亿级参数模型训练的框架
    • 3D 并行策略 ：支持数据并行、流水线并行和张量切片模型并行，并可灵活组合
    • 混合精度训练 ：自动混合精度训练（AMP）将单精度和半精度浮点数结合使用，降低内存需求的同时提升计算效率
    • 智能推理优化器 ：支持张量并行与异构内存技术，提供低延迟高吞吐的分布式推理服务，可将成本降低 70%
    • 全链路内存管理 ：集成 CPU 卸载与显存碎片整理技术，单卡即可训练百亿级模型，资源利用率提升 6 倍
  • DeepSpeed 框架的组件构成有：
    • Apis ：提供易用的 API 接口
    • 运行时组件 ：管理、执行和性能优化，基于 Python 语言实现，负责部署训练任务到分布式设备、数据分区、模型分区等
    • 底层内核 ：用 C++ 和 CUDA 实现，优化计算和通信
  • DeepSpeed 生态兼容性极好 ：原生兼容 PyTorch 与 Hugging Face 生态，通过简洁 API 可快速迁移项目，开发效率提升 300%

安装 DeepSpeed

  • 仅需一行安装命令即可：

    1
    pip install deepspeed

  • 暗转完成后，可以使用 ds_report 命令验证安装是否成功

    • 这个命令可以查看环境配置信息

DeepSpeed 的使用

  • 代码修改 ：仅仅需要非常少的代码修改即可将原始训练代码切换到 DeepSpeed 框架上（DeepSpeed 与 PyTorch 无缝集成，只需少量修改即可启用加速）

    • 第一步：通过 deepspeed.initialize 将模型包装为 DeepSpeed 引擎，自动应用优化
    • 第二步：使用 model_engine.backward 和 model_engine.step 替换PyTorch原本的 loss.backward() 和 optimizer.step()

  • 配置 DeepSpeed ：DeepSpeed 的优化行为通过 JSON 配置文件（ds_config.json）指定

    • 一般的项目都会自带一些配置好的 .json 文件示例，可直接修改使用

  • 运行训练 ：使用 DeepSpeed 的命令行工具启动训练

    • 单节点训练命令为：

      1
      deepspeed train.py --deepspeed_config ds_config.json

    • 对于多节点集群，使用下面的命令：

      1
      deepspeed --num_gpus 8 --num_nodes 2 train.py --deepspeed_config ds_config.json
      • 其中--num_gpus 指定每节点使用的 GPU 数量，--num_nodes 指定集群中的节点数

    • 其他常用参数：

      • 使用--log_dir 参数启用日志记录，监控内存使用、训练速度等

  • 特别说明：使用了 DeepSpeed 框架的代码需要使用 deepspeed 命令来启动

    • 补充：直接使用 python 命令启动时会出现 deepspeed.initialize 调用的错误
    • 原因：DeepSpeed 作为一个分布式训练库，需要特殊的启动器来管理多个进程和 GPU 之间的通信和资源分配

DeepSpeed 使用示例（基于 PyTorch）

  • DeepSpeed 使用示例： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader, Dataset
    import deepspeed

    class DiyModel(nn.Module):
        def __init__(self, input_dim, output_dim):
            super(DiyModel, self).__init__()
            self.fc = nn.Sequential(
                nn.Linear(input_dim, 128),
                nn.ReLU(),
                nn.Linear(128, output_dim)
            )
        
        def forward(self, x):
            return self.fc(x)

    class RandomDataset(Dataset):
        def __len__(self):
            return 1000
        
        def __getitem__(self, idx):
            x = torch.randn(32)
            y = torch.randint(0, 10, (1,)).item()
            return x, y

    # 初始化模型和数据（无需为 DeepSpeed 特别处理）
    model = DiyModel(32, 10)
    train_dataset = RandomDataset()
    train_loader = DataLoader(train_dataset, batch_size=32)

    # 初始化 DeepSpeed 引擎（仅需使用 deepspeed.initialize 初始化模型得到 model_engine 即可）
    # 注：这一步会自动将模型参数移动到 GPU 上，下面使用的数据也需要将数据移动到对应的 GPU 上才能运行，否则会出现 设备不一致的错误（CPU vs GPU）
    # zero_optimization 字段的 stage 键值对应如下效果
    ## stage: 0：不使用 ZeRO 优化（默认值是 0）
    ## stage: 1：优化器状态分片
    ## stage: 2：优化器状态和梯度分片
    ## stage: 3：优化器状态、梯度和参数分片（最高内存效率）
    model_engine, optimizer, _, _ = deepspeed.initialize(
        args=None,
        model=model,
        model_parameters=model.parameters(),
        config={
            "train_batch_size": 32,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.001,
                    "betas": [0.9, 0.999]
                }
            },
            "fp16": {
                "enabled": True,
                "loss_scale": 0,
                "loss_scale_window": 1000,
                "initial_scale_power": 16
            },
            "zero_optimization": {
                "stage": 2,  # 这里指定Zero层级（0、1、2、3）
                "offload_optimizer": {
                "device": "cpu"  # 可选：指定优化器卸载设备
                }
            }
        }
    )

    # 训练过程（训练时不再使用原来的模型，使用 model_engine）
    for epoch in range(10):
        for batch_idx, (data, target) in enumerate(train_loader):
            # 将数据挪到和模型相同的 GPU 上
            device = model_engine.device
            model_dtype = next(model_engine.parameters()).dtype # 通过模型的第一个参数获取dtype
            data = data.to(device, dtype=model_dtype) # 将数据移至模型所在设备并转换为与模型相同的dtype
            target = target.to(device)

            outputs = model_engine(data) # 使用（deepspeed.initialize 初始化得到的）model_engine 来进行前向过程
            loss = nn.CrossEntropyLoss()(outputs, target) 
            model_engine.backward(loss) # 使用 model_engine 来进行后向过程
            model_engine.step() # 使用 model_engine 来更新模型参数（注：此时不再需要显示调用 optimizer)
            if batch_idx % 100 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}")

附录：如何指定目标 GPU

启动时指定 CUDA_VISIBLE_DEVICES 环境变量

  • 在运行命令前设置环境变量来限制 DeepSpeed 可见的 GPU：

    1
    CUDA_VISIBLE_DEVICES=0,1,2,3 deepspeed your_script.py --args ...

  • 该方案是最常用的方法 ，且适用于常见的很多框架

在 Python 脚本中设置环境变量

  • 也可以在Python脚本中通过os.environ设置这个环境变量：

    1
    2
    3
    4
    5
    6
    7
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  # 指定使用GPU 0和1

    import deepspeed
    import torch

    # 训练逻辑

  • 注：需要在导入DeepSpeed或PyTorch之前设置

  • 该方案同样适用于常见的很多框架，但因为需要修改代码，不常用

使用 deepspeed 命令的 --include 参数

  • 如果使用的是 DeepSpeed 的 launcher，也可以通过--include参数指定使用的 GPU： 
    1
    deepspeed --include localhost:0,1 your_script.py --args ...

多机多卡如何指定 GPU

  • 如果你使用的是单机多卡，以上方法都能很好地工作
  • 对于多机多卡训练，通常需要结合 deepspeed 命令的其他参数如 --hostfile 等一起使用

附录：数据加载位置管理

  • 由于 deepspeed 会自动将模型参数加载到指定 GPU 上，所以数据也要加载到指定 GPU，否则会出现设备不一致的错误
  • 加载命令如下（亲测解决方案）： 
    1
    2
    device = model_engine.device
    data = data.to(device)

附录：混合精度训练数据格式管理

  • 由于 deepspeed 在启动混合精度训练时，可能会按照指定格式来指定参数形式，此时数据也需要转换为指定类型
  • 解决方式如下（亲测遇到错误时的解决方案）： 
    1
    2
    model_dtype = next(model_engine.parameters()).dtype # 通过模型的第一个参数获取类型，注：写这么复杂的原因是，当前还不支持直接调用 model_engine.dtype()
    data = data.to(device, dtype=model_dtype)