DeepSpeed 深度学习学习笔记：高效训练大型模型

主要参考官网文档，对于具体内容还需参考官方文档

---

1. 引言：为什么需要 DeepSpeed？

大型模型训练的挑战

随着深度学习模型规模的爆炸式增长（从 BERT 的几亿参数到 GPT-3 的千亿参数，再到现在的万亿参数模型），传统的单 GPU 训练方式变得力不从心，即使是多 GPU 训练也面临巨大挑战：

• 内存限制 (Memory Wall):
  • 模型参数： 模型的参数量巨大，例如一个 1750 亿参数的 GPT-3 模型，即使使用 FP16 精度，也需要 350GB 的显存来存储参数。这远超单个 GPU 的显存容量（如 A100 80GB）。
  • 优化器状态： Adam 等优化器需要存储梯度、动量等额外状态，通常是参数量的 10-12 倍。
  • 梯度： 梯度与参数量相同。
  • 激活值： 前向传播产生的激活值在反向传播时需要重新计算或存储，也占用大量显存。
• 计算效率 (Compute Efficiency):
  • 即使能把模型放进显存，如何高效地利用多 GPU 进行并行计算也是一大难题。
• 通信开销 (Communication Overhead):
  • 多 GPU 训练中，参数同步、梯度聚合等操作会产生大量通信，成为性能瓶颈。
• 工程复杂性 (Engineering Complexity):
  • 实现高效的分布式训练，需要深入理解并行策略、通信原语等，对开发者要求高。

DeepSpeed 核心优势

DeepSpeed 应运而生，旨在解决上述挑战，使大型模型训练触手可及：

• 内存优化：
  • 通过 ZeRO (Zero Redundancy Optimizer) 系列技术，极致地减少显存占用，使千亿甚至万亿参数模型在现有硬件上训练成为可能。
• 计算效率：
  • 优化了计算图，减少了通信开销，并支持高效的混合精度训练。
• 易用性：
  • 提供简单易用的 API 和配置文件，用户只需少量代码修改即可应用复杂的优化技术。
• 可扩展性：
  • 支持从单 GPU 到数千个 GPU 的弹性扩展。

2. DeepSpeed 核心概念

理解 DeepSpeed，首先要掌握其背后的几个关键技术。

2.1 ZeRO (Zero Redundancy Optimizer)

ZeRO 是 DeepSpeed 最核心的内存优化技术，通过消除冗余内存来减少显存占用。它将模型状态（优化器状态、梯度、参数）在不同的 GPU 之间进行分片（partitioning）。

• ZeRO-Offload (CPU/NVMe 卸载):

  • 不属于 ZeRO-Stage 的一个阶段，而是一种补充策略。
  • 将部分模型状态（主要是优化器状态和/或梯度）从 GPU 显存卸载到 CPU 内存甚至 NVMe 硬盘，以释放宝贵的 GPU 显存。
  • 优点： 显存节省显著，对于中大型模型非常有效。
  • 缺点： 卸载会引入 CPU 或 NVMe 的 I/O 延迟，可能导致训练速度变慢。

• ZeRO-Stage 1 (Optimizer States Partitioning):

  • 原理： 将优化器状态（如 Adam 的动量和方差）在所有 GPU 上均匀分片。每个 GPU 只存储其分片的部分。

  • 显存节省： 节省大约 4 倍于 FP16 参数的显存（因为 Adam 优化器状态通常是 FP32，占用 8 倍参数显存，分片后变为 8/N 倍，其中 N 是 GPU 数量）。

  • 通信开销： 在梯度更新前需要 AllGather 完整的优化器状态。

  • 示意图：

    ZeRO-Stage 1

    GPU3

    GPU2

    GPU1

    GPU0

    Full Model Parameters

    Full Model Parameters

    Full Model Parameters

    Full Model Parameters

    Full Gradients

    Full Gradients

    Full Gradients

    Full Gradients

    Total Optimizer State

    Total Optimizer State

    Total Optimizer State

    Total Optimizer State

    OS_part0

    OS_part1

    OS_part2

    OS_part3

    P3

    G3

    OS3

    P2

    G2

    OS2

    P1

    G1

    OS1

    P0

    G0

    OS0

    Large Model

    Optimizer

    Parameters

    Gradients

    Optimizer States

    说明：P代表参数，G代表梯度，OS代表优化器状态。ZeRO-Stage 1 将 OS 在各 GPU 间分片。

• ZeRO-Stage 2 (Optimizer States + Gradients Partitioning):

  • 原理： 在 Stage 1 的基础上，进一步将梯度也在所有 GPU 上均匀分片。每个 GPU 只存储和处理其分片的部分梯度。

  • 显存节省： 在 Stage 1 的基础上，再节省大约 3 倍于 FP16 参数的显存（梯度通常是 FP16，占用 2 倍参数显存，分片后变为 2/N 倍）。总共节省约 7 倍显存。

  • 通信开销： 在反向传播过程中需要进行 AllReduce 来聚合完整的梯度，但在更新参数前，每个 GPU 只处理自己分片的那部分梯度。

  • 示意图：

    ZeRO-Stage 2 Gradients & OS Partitioning

    GPU3

    GPU2

    GPU1

    GPU0

    Full Model Parameters

    Full Model Parameters

    Full Model Parameters

    Full Model Parameters

    Total Gradients

    Total Optimizer State

    G_part0

    G_part1

    G_part2

    G_part3

    OS_part3

    OS_part2

    OS_part1

    OS_part0

    P3

    G3_part

    OS3_part

    P2

    G2_part

    OS2_part

    P1

    G1_part

    OS1_part

    P0

    G0_part

    OS0_part

    Large Model

    G0_part+G1_part+G2_part+G3_part

    OS0_part+OS1_part+OS2_part+OS3_part

    Optimizer

    说明：ZeRO-Stage 2 将 G 和 OS 在各 GPU 间分片。

• ZeRO-Stage 3 (Optimizer States + Gradients + Parameters Partitioning):

  • 原理： 在 Stage 2 的基础上，进一步将模型参数也在所有 GPU 上均匀分片。每个 GPU 只存储其分片的部分参数。在需要完整参数进行前向/反向传播时，通过 AllGather 操作从其他 GPU 收集。

  • 显存节省： 理论上可以节省高达 N 倍显存（N 为 GPU 数量），因为每个 GPU 只存储 1/N 的模型参数、梯度和优化器状态。

  • 通信开销： 前向传播和反向传播都需要进行 AllGather 操作来获取完整的参数/梯度，通信量最大，但能支持最大模型。

  • 示意图：

    ZeRO-Stage 3 Parameters, Gradients & OS Partitioning

    GPU3

    GPU2

    GPU1

    GPU0

    AllGather for Full Parameters

    AllReduce for Full Gradients

    Total Optimizer State

    P_part0

    P_part1

    P_part2

    P_part3

    G_part3

    G_part2

    G_part1

    G_part0

    OS_part3

    OS_part2

    OS_part1

    OS_part0

    P3_part

    G3_part

    OS3_part

    P2_part

    G2_part

    OS2_part

    P1_part

    G1_part

    OS1_part

    P0_part

    G0_part

    OS0_part

    P0_part+P1_part+P2_part+P3_part

    (Large Model)

    G0_part+G1_part+G2_part+G3_part

    OS0_part+OS1_part+OS2_part+OS3_part

    (Optimizer)

    说明：ZeRO-Stage 3 将 P, G, OS 全部在各 GPU 间分片。

• ZeRO-Infinity:

  • 原理： 在 ZeRO-Stage 3 的基础上，结合了 NVMe 卸载、动态卸载、异构内存管理等技术。
  • 目标： 实现几乎无限的内存拓展，将模型训练扩展到集群上最大的模型，甚至超过单节点 NVMe 容量。
  • 显存节省： 能够训练 TB 级别的模型参数，超越单个节点的所有显存和内存限制。

2.2 混合精度训练 (Mixed Precision Training)

• 原理：

  在训练过程中同时使用 FP32（单精度浮点数）和 FP16（半精度浮点数）。

  • FP16 优势： 显存占用减半，部分 GPU 硬件（如 NVIDIA Tensor Cores）对 FP16 有原生加速，提高计算速度。
  • FP32 优势： 精度更高，避免小梯度下溢。

• DeepSpeed 实现： 自动管理 FP16 和 FP32 的转换、梯度缩放 (GradScaler) 以防止梯度下溢，以及参数副本的维护。

2.3 模型并行 (Model Parallelism)

当模型过大，即使 ZeRO-Stage 3 也无法完全容纳在一个 GPU 上时，就需要模型并行。模型并行将模型的不同部分放在不同的 GPU 上。

• 张量并行 (Tensor Parallelism):
  • 原理： 将单个层的张量（如权重矩阵）沿某个维度拆分到不同 GPU 上，每个 GPU 只负责计算一部分。
  • 优点： 减少单层计算的显存和计算压力。
  • 缺点： 引入大量 GPU 间通信，实现复杂。
  • DeepSpeed 配合： DeepSpeed 自身不直接实现张量并行，但可以与 Megatron-LM 等提供张量并行的框架结合使用。

2.4 管道并行 (Pipeline Parallelism)

• 原理： 将模型的不同层或层组划分到不同的 GPU 上，形成一个“流水线”。每个 GPU 负责模型的一个阶段。数据批次被进一步拆分成更小的“微批次”，在流水线中流动。
• 优点：
  • 减少了每个 GPU 上的显存占用，因为每个 GPU 只存储模型的一部分。
  • 通过流水线效应，可以提高多 GPU 利用率。
• 缺点：
  • 引入“气泡”（bubble）效应，即流水线中的空闲时间，可能降低 GPU 效率。
  • 实现复杂，需要精心划分模型。
• DeepSpeed 实现： DeepSpeed 提供了对管道并行的原生支持。

2.5 优化器 (Optimizers)

DeepSpeed 提供了自己的优化器实现，这些优化器通常是高度优化的，并且能够与 ZeRO 技术无缝集成。例如，它有 DeepSpeedCPUAdam (将 Adam 的计算卸载到 CPU) 和 FusedAdam (GPU 上更快的 Adam 实现)。

3. 安装与环境配置

3.1 Python 环境准备

使用 Conda 或 Virtualenv 创建隔离的 Python 环境。

conda create -n deepspeed_env python=3.10
conda activate deepspeed_env

3.2 GPU 驱动与 CUDA/cuDNN

查询 PyTorch 官方网站以获取兼容性信息

3.3 DeepSpeed 安装

# 确保安装了正确版本的 PyTorch (例如，CUDA 12.1)
# 详情参考 PyTorch 官网，这里以 CUDA 12.1 为例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# 安装 DeepSpeed
pip install deepspeed
# 如果需要从源码安装以获取最新功能或进行开发：
# git clone https://github.com/microsoft/DeepSpeed.git
# cd DeepSpeed
# pip install . --global-config-path

验证安装：

python -c "import deepspeed; print(deepspeed.__version__)"

4. DeepSpeed 基本使用

DeepSpeed 的核心是其配置文件 (Configuration File)。将大部分优化选项写入一个 JSON 文件中，然后在训练脚本中通过 deepspeed.init_distributed() 和 deepspeed.initialize() 来加载和应用。

4.1 配置文件 (Configuration File) 详解

创建一个名为 ds_config.json 的 DeepSpeed 配置文件：

// ds_config.json
{
  "train_batch_size": 16,            // 训练批次大小 (全局批次大小 = train_batch_size)
  "train_micro_batch_size_per_gpu": 2, // 每个 GPU 上的微批次大小 (Gradient Accumulation Steps = train_batch_size / (num_gpus * train_micro_batch_size_per_gpu))
  "gradient_accumulation_steps": 8, // 梯度累积步数 (如果未设置，DeepSpeed 会根据前两个自动计算)

  "optimizer": {                     // 优化器配置
    "type": "AdamW",                 // 优化器类型，例如 AdamW
    "params": {
      "lr": 1e-5,                    // 学习率
      "betas": [0.9, 0.999],         // AdamW 参数
      "eps": 1e-8,
      "weight_decay": 0.01           // 权重衰减
    }
  },

  "scheduler": {                     // 学习率调度器配置
    "type": "WarmupLR",              // 调度器类型，例如 WarmupLR
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 1e-5,
      "warmup_num_steps": 1000
    }
  },

  "fp16": {                          // 混合精度训练配置
    "enabled": true,                 // 启用混合精度
    "loss_scale": 0,                 // 0表示DeepSpeed自动管理，否则为固定值
    "initial_scale_power": 12,       // 初始梯度缩放因子 (2^12 = 4096)
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "zero_optimization": {             // ZeRO 优化器配置
    "stage": 2,                      // ZeRO 阶段 (0, 1, 2, 3)
    "offload_optimizer": {           // 优化器状态卸载到 CPU
      "device": "cpu",               // 卸载到 CPU
      "pin_memory": true             // 锁定 CPU 内存，提高传输速度
    },
    "offload_param": {               // 参数卸载到 CPU (仅 Stage 3 配合使用)
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,            // 开启通信与计算重叠
    "contiguous_gradients": true,    // 尝试使梯度内存连续
    "sub_group_size": 1e9,           // 较大的值表示禁用子组优化
    "stage3_prefetch_bucket_size": 0, // stage3预取大小，0为禁用
    "stage3_param_persistence_threshold": 1e5, // 仅对Stage3，参数持久化阈值
    "stage3_max_live_parameters": 1e9, // Stage3 最大驻留参数
    "stage3_max_act_size": 1e9       // Stage3 最大激活值大小
  },

  "gradient_clipping": 1.0,          // 梯度裁剪，防止梯度爆炸
  "prescale_gradients": false,       // 启用 prescale_gradients 可能会提高 Stage3 的性能
  "wall_clock_breakdown": false,     // 是否详细记录训练时间分解
  "sparse_attention": false,         // 稀疏注意力 (针对稀疏模型)

  "flops_profiler": {                // FLOPS 分析器
    "enabled": false,
    "profile_step": 1,
    "module_depth": -1,
    "top_modules": 1,
    "detailed": true,
    "output_file": null
  },

  "activation_checkpointing": {       // 激活值检查点 (类似于梯度检查点，节省显存)
    "enabled": false,
    "cpu_checkpointing": false,      // 激活值检查点是否卸载到 CPU
    "partition_activations": false   // 激活值是否分片 (只适用于 ZeRO Stage 3)
  },

  "pipeline": {                      // 管道并行配置 (如果启用)
    "seed_layers": false,
    "micro_batch_size": 0
  },

  "steps_per_print": 10,             // 每隔多少步打印一次日志
  "seed": 1234,                      // 随机种子

  "tensorboard": {                   // TensorBoard 日志
    "enabled": true,
    "output_path": "tensorboard_logs/",
    "job_name": "my_deepspeed_experiment"
  }
}

• train_batch_size vs train_micro_batch_size_per_gpu vs gradient_accumulation_steps:
  • train_batch_size 是全局等效的训练批次大小，它定义了模型在一个优化步骤中处理的数据量。
  • train_micro_batch_size_per_gpu 是每个 GPU 上的微批次大小。
  • gradient_accumulation_steps 是梯度累积的步数。
  • 它们之间的关系是：train_batch_size = num_gpus * train_micro_batch_size_per_gpu * gradient_accumulation_steps。如果你设置了前两个，DeepSpeed 会自动计算 gradient_accumulation_steps。这是控制实际更新模型参数频率的关键。

4.2 启动 DeepSpeed 训练

使用 deepspeed 命令行工具来启动训练脚本：

deepspeed --num_gpus=4 your_training_script.py --deepspeed_config ds_config.json

• --num_gpus: 指定使用的 GPU 数量。
• your_training_script.py: 你的 PyTorch 训练脚本。
• --deepspeed_config: 指定 DeepSpeed 配置文件的路径。

4.3 简单的 PyTorch 训练脚本改造

只需要少量修改即可将 PyTorch 脚本转换为 DeepSpeed 脚本。

# your_training_script.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import deepspeed
import argparse

# 1. 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(20, 1)

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

def get_args():
    parser = argparse.ArgumentParser(description="DeepSpeed Training Example")
    parser.add_argument('--deepspeed_config', default='ds_config.json', type=str,
                        help='DeepSpeed config file path')
    parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from DeepSpeed')
    args = parser.parse_args()
    return args

def main():
    args = get_args()

    # 2. DeepSpeed 初始化
    # model, optimizer, _, lr_scheduler = deepspeed.initialize(
    #     model=model,
    #     model_parameters=model.parameters(),
    #     config_params=ds_config_dict # 可以直接传递字典
    # )
    # 或者从配置文件加载
    model = SimpleModel()
    model_engine, optimizer, _, lr_scheduler = deepspeed.initialize(
        args=args,
        model=model,
        model_parameters=model.parameters()
    )

    # 3. 准备数据 (示例数据)
    # 确保数据在正确的设备上 (DeepSpeed 会自动处理数据到对应GPU)
    # 对于DeepSpeed，通常数据加载无需手动.to(device)
    dummy_data = torch.randn(100, 10).cuda()
    dummy_labels = torch.randn(100, 1).cuda()
    dataset = TensorDataset(dummy_data, dummy_labels)
    # DeepSpeed 会处理分布式采样器
    train_dataloader = DataLoader(dataset, batch_size=model_engine.train_micro_batch_size_per_gpu)

    # 4. 定义损失函数
    criterion = nn.MSELoss()

    # 5. 训练循环
    print(f"Rank {model_engine.local_rank}: Starting training...")
    for epoch in range(10):
        for i, (inputs, labels) in enumerate(train_dataloader):
            # inputs, labels = inputs.cuda(), labels.cuda() # DeepSpeed会自动移动数据

            outputs = model_engine(inputs)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            model_engine.backward(loss)
            model_engine.step()

            # 打印日志 (可选，DeepSpeed 会有自己的打印逻辑)
            if (i + 1) % model_engine.steps_per_print() == 0:
                print(f"Rank {model_engine.local_rank}, Epoch [{epoch+1}/10], Step [{i+1}/{len(train_dataloader)}], Loss: {loss.item():.4f}")

        if model_engine.local_rank == 0: # 只有主进程保存
            print(f"Epoch {epoch+1} finished.")
            # 保存检查点
            # model_engine.save_checkpoint(f"checkpoints/epoch_{epoch+1}")

    # 保存最终模型 (所有进程都会保存自己的分片)
    if model_engine.local_rank == 0:
        model_engine.save_checkpoint(save_dir="final_model_checkpoint", client_state={"epoch": epoch})
        print("Training finished and model saved.")

if __name__ == "__main__":
    # DeepSpeed 使用 torch.distributed.launch 类似的机制
    # 需要在命令行通过 deepspeed 命令启动，而不是直接 python your_script.py
    main()

训练多个模型

model_engines = [engine for engine, _, _, _ in [deepspeed.initialize(m, ...,) for m in models]]
for batch in data_loader:
    losses = [engine(batch[0], batch[1]) for engine in model_engines]
    loss = sum(l / (i + 1) for i, l in enumerate(losses))
    loss.backward()

    for engine in model_engines:
        engine._backward_epilogue()

    for engine in model_engines:
        engine.step()

    for engine in model_engines:
        engine.optimizer.zero_grad()

除了使用多个 DeepSpeedEngine 之外，上述用法与典型用法在两个关键方面有所不同：

1. 反向传播调用是使用共同损失值而非各自模型引擎进行的。
2. 在 loss.backward() 之后，会在模型引擎上调用 _backward_epilogue。

4.4 推理

deepspeed.init_inference() 返回一个 推理引擎，其类型为 InferenceEngine

deepspeed.init_inference()传入config参数，对于config 支持以下几种

• class deepspeed.inference.config.InferenceCheckpointConfig
• deepspeed.inference.config.DeepSpeedInferenceConfig
• deepspeed.inference.config.DeepSpeedTPConfig
• deepspeed.inference.config.DeepSpeedMoEConfig
• deepspeed.inference.config.QuantizationConfig

对每个config具体参数参考

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

5. 深入理解 ZeRO 优化器

5.1 ZeRO-Offload：CPU 和 NVMe 卸载

• 配置： 在 ds_config.json 中配置 zero_optimization.offload_optimizer 和 zero_optimization.offload_param。
• 适用场景：
  • 优化器卸载到 CPU： 适用于显存稍有不足，但 CPU 内存充足且 I/O 速度较快的场景。
  • 参数卸载到 CPU/NVMe： 主要用于 ZeRO-Stage 3，当即使参数分片后 GPU 显存仍然不够时。NVMe 卸载（需要 nvme_path 配置）比 CPU 卸载能支持更大的模型，但速度更慢。
• 权衡： 显存节省与训练速度之间的权衡。

5.2 ZeRO-Stage 1, 2, 3：优化器状态、梯度、参数分片

• 配置： 关键在于 zero_optimization.stage 参数。
• 选择原则：
  • Stage 1 (最小开销，适度显存节省)： 适用于优化器状态是主要显存瓶颈的场景。
  • Stage 2 (中等开销，显著显存节省)： 适用于梯度和优化器状态共同造成瓶颈的场景，是常用且平衡的选项。
  • Stage 3 (最大开销，极致显存节省)： 用于训练现有硬件能容纳的最大模型，通信开销最高。需要 activation_checkpointing 等技术配合以进一步节省显存。
• 通信优化： overlap_comm (通信与计算重叠) 和 contiguous_gradients (梯度内存连续性) 对性能至关重要，建议在 Stage 2/3 中启用。

5.3 ZeRO-Infinity：无限内存拓展

• 配置： 通常是在 Stage 3 的基础上，通过更激进的卸载策略和动态管理实现。
• 目标： 挑战单节点和甚至多节点集群的显存极限。
• 使用： 相比标准 ZeRO Stage 3 更复杂，需要根据具体硬件和模型进行调整。通常在需要训练万亿参数模型时考虑。

6. 混合精度训练 (AMP)

6.1 原理：FP16 与 FP32

• FP16 (Half-precision)：占用 2 字节，范围小，精度低，但部分硬件支持加速。
• FP32 (Single-precision)：占用 4 字节，范围大，精度高。
• 混合精度策略：大部分计算使用 FP16，但保持参数的 FP32 副本以确保精度；梯度在累积时转换为 FP32；损失计算使用 FP32。

6.2 DeepSpeed 中的 AMP 配置

在 ds_config.json 中设置 fp16.enabled: true。DeepSpeed 会自动管理：

• 模型参数的类型转换（FP16 用于计算，FP32 用于参数更新）。
• 梯度缩放 (GradScaler)。

6.3 GradScaler (梯度缩放)

• 目的： 防止 FP16 精度下浮点数表示范围不足导致的小梯度下溢为 0。
• 原理： 在反向传播前，将损失函数乘以一个大的缩放因子，使梯度也相应地变大，避免下溢。梯度计算完成后，再除以相同的缩放因子恢复正常值，用于参数更新。
• DeepSpeed 自动管理： 用户无需手动实现 GradScaler，DeepSpeed 会根据配置文件中的 loss_scale、initial_scale_power 等参数自动处理。

7. 模型并行与管道并行

DeepSpeed 主要关注数据并行和内存优化（通过 ZeRO）。虽然它本身不直接提供像 Megatron-LM 那样的复杂张量并行模块，但可以通过与这些库的集成来支持模型并行。DeepSpeed 提供了原生管道并行。

7.1 张量并行 (Tensor Parallelism) 概述

• 将一个层的权重和输入张量在 GPU 之间拆分。例如，一个大的矩阵乘法被拆分成多个小的矩阵乘法在不同 GPU 上并行计算。
• 显存： 每个 GPU 只存储部分权重，节省显存。
• 通信： 每个操作后都需要进行 AllReduce 或 AllGather 等通信。

7.2 管道并行 (Pipeline Parallelism) 概述

• 将模型垂直切分，不同层在不同 GPU 上。

• DeepSpeed 中的配置： 在 ds_config.json 中配置 pipeline.stages（如果使用）。

• 代码示例（概念性）：

  # DeepSpeed 管道并行示例（概念性，实际实现更复杂）
  # 通常需要定义好各个子模块，并指定哪些模块属于哪个pipeline stage
  # DeepSpeed会帮你处理forward/backward pass在不同GPU上的调度
  import deepspeed
  from deepspeed.pipe import PipelineModule, LayerSpec
  
  class MyPipelineModel(PipelineModule):
      def __init__(self, **kwargs):
          # 定义模型层，这里用LayerSpec来表示每个阶段包含哪些层
          # 实际模型会更复杂，例如 Transformer Blocks
          super().__init__(layers=[
              LayerSpec(nn.Linear, 10, 20),
              LayerSpec(nn.ReLU),
              LayerSpec(nn.Linear, 20, 30),
              LayerSpec(nn.ReLU),
              LayerSpec(nn.Linear, 30, 1)
          ], **kwargs)
  
  # 实际使用时，需要在配置文件中开启 pipeline 模式，并设置 micro_batch_size
  # "pipeline": {
  #   "stages": N_STAGES, // N_STAGES 个阶段，对应 N_STAGES 个 GPU
  #   "micro_batch_size": YOUR_MICRO_BATCH_SIZE
  # }
  # 然后在 initialize 时传入 PipeLineModule 实例
  # model_engine, optimizer, _, _ = deepspeed.initialize(
  #     args=args,
  #     model=MyPipelineModel(loss_fn=nn.MSELoss()), # 需要指定损失函数
  #     model_parameters=[p for p in model.parameters() if p.requires_grad]
  # )
  

8. 优化器与调度器

8.1 DeepSpeed 支持的优化器

DeepSpeed 封装了许多标准优化器，并提供了一些定制优化器，例如：

• AdamW (默认推荐): 通常在 DeepSpeed 中使用 FusedAdam 或 OneBitAdam（如果启用）。
• Adam: 标准 Adam。
• DeepSpeedCPUAdam: 将 Adam 的计算卸载到 CPU，进一步节省 GPU 显存。当 zero_optimization.offload_optimizer.device 设置为 “cpu” 时，DeepSpeed 会自动选择使用它。
• OneBitAdam/OneBitSGD: 高效通信的优化器，用于减少梯度通信量（但可能牺牲一些精度）。

8.2 学习率调度器

DeepSpeed 支持多种学习率调度器（如 WarmupLR, ReduceLROnPlateau 等），可以在配置文件中直接配置。

9. Checkpointing (检查点)

训练大型模型通常需要很长时间，因此保存和加载检查点以防止训练中断以及进行实验是非常重要的。

• DeepSpeed 的检查点： DeepSpeed 能够以分布式方式保存和加载模型参数、优化器状态和学习率调度器状态。对于 ZeRO-Stage 3，它会智能地处理分片后的参数。

• 保存：

  model_engine.save_checkpoint(save_dir, client_state=None)
  

  • save_dir: 保存检查点的目录。
  • client_state: 可以保存任何自定义状态信息（如当前 epoch 数）。

• 加载：

  load_path, client_state = model_engine.load_checkpoint(load_dir, tag=None, load_module_only=False)
  

  • load_dir: 检查点所在的目录。
  • tag: 如果有多个检查点，可以指定标签。
  • load_module_only: 如果只加载模型参数而不加载优化器和调度器状态，设为 True。

10. 故障排查与最佳实践

10.1 常见错误及解决方案

• CUDA OOM (Out of Memory):
  • 原因： 显存不足。
  • 解决方案：
    • 降低 train_micro_batch_size_per_gpu。
    • 增加 gradient_accumulation_steps。
    • 将 zero_optimization.stage 提高到 2 或 3。
    • 启用 zero_optimization.offload_optimizer 和/或 offload_param 到 CPU/NVMe。
    • 启用 activation_checkpointing。
• 通信死锁/hang:
  • 原因： 分布式训练中进程间通信问题。
  • 解决方案：
    • 确保所有进程都正确启动并连接。
    • 检查防火墙设置。
    • 确认 torch.distributed.init_process_group 等初始化正确。
    • 有时是由于某个 GPU 上 OOM 导致进程崩溃，其他进程等待。
• 性能下降：
  • 原因： CPU/NVMe 卸载开销过大，通信成为瓶颈，微批次过小。
  • 解决方案：
    • 优化 zero_optimization.offload_optimizer 配置，权衡卸载量。
    • 调整 overlap_comm 和 contiguous_gradients。
    • 适当增加 train_micro_batch_size_per_gpu。
    • 考虑使用更快的互联（如 InfiniBand）。
    • 使用 DeepSpeed 的 wall_clock_breakdown 和 flops_profiler 进行性能分析。

10.2 性能调优技巧

• 从小到大： 先尝试 ZeRO-Stage 1/2，如果显存仍然不足再考虑 Stage 3。
• 梯度累积： 充分利用 gradient_accumulation_steps 来增加等效批次大小，同时保持较小的微批次大小以节省显存。
• 通信与计算重叠： 开启 overlap_comm: true 可以有效隐藏通信延迟。
• 激活值检查点： 对于深度模型，启用 activation_checkpointing 是一个非常有效的显存优化手段，它会重新计算激活值而不是存储它们。
• FloPS Profiler： 使用 flops_profiler 来分析模型的计算瓶颈。
• 数据加载： 确保数据加载器不会成为瓶颈，使用 num_workers 适当并行加载。

10.3 日志与监控

DeepSpeed 提供了丰富的日志输出。结合 TensorBoard 可以可视化训练过程中的各种指标（损失、学习率、内存使用等）。

11. 实际应用案例与高级特性

11.1 结合 Hugging Face Transformers

DeepSpeed 与 Hugging Face 的 transformers 库高度集成。Hugging Face 提供了可以直接使用 DeepSpeed 的 Trainer 类。

• 使用方式：

  # 使用 transformers 命令行工具
  transformers-cli run --model_name_or_path <model_path> \
  --deepspeed <deepspeed_config_file> \
  --num_gpus <num_gpus> \
  --do_train \
  # ... 其他训练参数
  

  或者在 Python 脚本中：

  from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
  # ... 定义模型、tokenizer、数据集
  
  training_args = TrainingArguments(
      output_dir="./output",
      deepspeed="ds_config.json", # 指定 DeepSpeed 配置文件
      # ... 其他参数
  )
  
  trainer = Trainer(
      model=model,
      args=training_args,
      # ... 其他 Trainer 参数
  )
  
  trainer.train()