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 完整的优化器状态。

  • 示意图：

    说明：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 将 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 将 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.pyimport 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 argsdef 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.pymain()

训练多个模型

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() methodloss = 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, LayerSpecclass MyPipelineModel(PipelineModule):def __init__(self, **kwargs):# 定义模型层，这里用LayerSpec来表示每个阶段包含哪些层# 实际模型会更复杂，例如 Transformer Blockssuper().__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()