Skip to content
←

필사 모드: DeepSpeed 完全指南：ZeRO 优化与大规模模型训练

中文
0%
정확도 0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

引言

要训练 GPT-4、LLaMA、Falcon 这类大规模语言模型(LLM)，单张 GPU 是不可能做到的。1750 亿参数的 GPT-3 若以 fp16 存储，需要约 350GB 的 GPU 内存，若再算上优化器状态(Adam)，则会达到 1.4TB。为解决这一问题，Microsoft 开发出了 DeepSpeed 这个库。

DeepSpeed 以 ZeRO(Zero Redundancy Optimizer)优化技术为核心，让数百亿乃至数万亿参数的模型也能在普通 GPU 集群上完成训练。本指南将从 DeepSpeed 的核心概念一路讲到实战配置，全面覆盖。

1. DeepSpeed 简介

1.1 背景与动机

深度学习模型越来越大，由此产生三个瓶颈。

  • 内存不足：模型参数、梯度、优化器状态超出 GPU 显存
  • 算力瓶颈：单张 GPU 的 FLOPS 上限
  • 通信瓶颈：多 GPU 环境下梯度同步的开销

传统的数据并行(Data Parallelism)让每张 GPU 都保留一份完整的模型副本，因此无法解决内存问题。模型并行(Model Parallelism)则实现复杂，且往往效率不高。

DeepSpeed 结合了这两种方案的优点，并通过全新的 ZeRO 优化实现了革命性的内存效率。

1.2 与 PyTorch 的集成

DeepSpeed 只需对现有 PyTorch 代码做最小限度的修改即可使用。核心改动就是一次 deepspeed.initialize() 调用，它会一并处理引擎创建、分布式初始化、优化器封装。

import deepspeed

# 现有的 PyTorch 代码
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 切换到 DeepSpeed(最小改动)
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

1.3 安装

pip install deepspeed

# 包含 CUDA 内核编译的安装
DS_BUILD_OPS=1 pip install deepspeed

# 验证安装
ds_report

2. ZeRO 优化分阶段解析

ZeRO(Zero Redundancy Optimizer)是消除数据并行训练中内存冗余的核心技术。在数据并行中，N 张 GPU 都保留着完全相同的模型副本，这是极大的内存浪费。ZeRO 逐阶段地消除这种冗余。

训练中的内存占用分析

设模型参数量为 Ψ，在混合精度(fp16)训练下，每张 GPU 的内存占用如下。

组成部分大小说明
参数 (fp16)2Ψ bytes前向/反向传播
梯度 (fp16)2Ψ bytes反向传播的结果
主参数 (fp32)4Ψ bytes供优化器使用
Adam 动量 (fp32)4Ψ bytes一阶矩
Adam 方差 (fp32)4Ψ bytes二阶矩
合计16Ψ bytes

若是 7B 参数的模型，大约需要 112GB；若是 70B 参数的模型，则需要约 1.1TB。

2.1 ZeRO-1：优化器状态分片

ZeRO-1 将占内存比重最大的优化器状态分散到 N 张 GPU 上。

  • Adam 的情形：fp32 参数 + 动量 + 方差 = 12Ψ bytes
  • 分散到 N 张 GPU 后，每张 GPU 的优化器状态 = 12Ψ/N bytes

内存节省效果：

原始：16Ψ bytes per GPU
ZeRO-1：4Ψ + 12Ψ/N bytes per GPU(N=64 时 ≈ 4.2Ψ bytes，约节省 4 倍)

在更新过程中，每张 GPU 只更新自己负责的参数分片，随后通过 All-gather 恢复出完整参数。

{
  "zero_optimization": {
    "stage": 1
  }
}

2.2 ZeRO-2：梯度分片

ZeRO-2 在优化器状态之外，还进一步分散了梯度。

反向传播过程中，每张 GPU 只在自己负责的参数分片上累积对应的梯度。通过 Reduce-scatter 运算，每张 GPU 只收集自己负责区间的梯度。

原始：16Ψ bytes per GPU
ZeRO-2：2Ψ + (2Ψ + 12Ψ)/N bytes per GPU(N=64 时 ≈ 2.2Ψ bytes，约节省 8 倍)

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

2.3 ZeRO-3：参数分片

ZeRO-3 连模型参数本身也一并分散。这是最强力的方案，但实现也最复杂。

运行原理：

  1. 每张 GPU 平时只保留全部参数的 1/N
  2. 前向传播时：在每层执行前，通过 All-gather 收集出完整参数
  3. 计算完成后：立即释放该参数
  4. 反向传播同样按需(on-demand)收集参数
ZeRO-3：(2Ψ + 2Ψ + 12Ψ)/N bytes per GPU
N=64 时：16Ψ/64 ≈ 0.25Ψ bytes(约节省 64 倍！)

缺点： All-gather 通信在每一层都会发生，通信开销随之增加。不过由于与流水线执行相重叠，实际延迟并不大。

{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

ZeRO-3 的 Python 初始化：

import deepspeed
from transformers import AutoModelForCausalLM

# 在 ZeRO-3 下，模型创建方式同样重要
with deepspeed.zero.Init(config_dict_or_path="ds_config.json"):
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# 这样做能让参数从一开始就处于分片状态，避免内存峰值

3. ZeRO-Offload 与 ZeRO-Infinity

3.1 ZeRO-Offload

ZeRO-Offload 把优化器状态和梯度转移到 CPU 内存，大幅减轻 GPU 内存压力。

  • GPU：fp16 参数 + fp16 梯度(前向/反向传播)
  • CPU：fp32 主参数 + Adam 状态(优化器更新)

这样一来，即便只有单张 GPU，也能训练数十亿参数的模型。

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

参数卸载(ZeRO-3 + Offload)：

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    }
  }
}

3.2 ZeRO-Infinity(NVMe 卸载)

ZeRO-Infinity 进一步利用 NVMe SSD，支持事实上不受限制的模型规模。

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 4,
      "fast_init": false
    },
    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 5,
      "buffer_size": 1e8
    },
    "aio": {
      "block_size": 1048576,
      "queue_depth": 8,
      "thread_count": 1,
      "single_submit": false,
      "overlap_events": true
    }
  }
}

带宽考量：

  • GPU-CPU 带宽：以 PCIe 4.0 x16 为准，约 32 GB/s
  • CPU-NVMe 带宽：以 NVMe SSD 为准，约 7 GB/s
  • NVMe 比 CPU 更慢，因此会拖慢训练速度，但能让原本不可能的模型规模成为可能

4. DeepSpeed 配置文件完全指南

4.1 ds_config.json 基本结构

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": 1.0,

  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "bf16": {
    "enabled": "auto"
  },

  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  },

  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },

  "activation_checkpointing": {
    "partition_activations": false,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": null,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  },

  "wall_clock_breakdown": false,
  "steps_per_print": 100
}

4.2 混合精度设置详解

fp16 设置：

{
  "fp16": {
    "enabled": true,
    "auto_cast": false,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "consecutive_hysteresis": false,
    "min_loss_scale": 1
  }
}
  • loss_scale: 0：启用动态损失缩放
  • initial_scale_power: 16：初始损失缩放 = 2^16 = 65536
  • loss_scale_window：缩放系数上调前所需的成功步数
  • hysteresis：缩放系数下调前所需的溢出计数

bf16 设置(Ampere 及以上架构 GPU)：

{
  "bf16": {
    "enabled": true
  }
}

bf16 的动态范围比 fp16 更宽，因此不需要损失缩放。推荐在 A100/H100 上使用。

4.3 梯度累积与批大小

批大小的关系式：

train_batch_size = train_micro_batch_size_per_gpu × gradient_accumulation_steps × world_size

{
  "train_batch_size": 2048,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 64
}

上述设置在 8 张 GPU 环境下：4 × 64 × 8 = 2048 的全局批大小。

使用 "auto" 时，由 Transformers Trainer 自动填充相应的值。

4.4 学习率调度器选项

{
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "last_batch_iteration": -1,
      "total_num_steps": 100000,
      "warmup_min_lr": 0,
      "warmup_max_lr": 3e-4,
      "warmup_num_steps": 2000,
      "warmup_type": "linear"
    }
  }
}

支持的调度器类型：

  • LRRangeTest
  • OneCycle
  • WarmupLR
  • WarmupDecayLR
  • WarmupCosineLR(自定义)

5. DeepSpeed + PyTorch 集成

5.1 基本训练循环

import torch
import deepspeed
from torch.utils.data import DataLoader

def main():
    # 初始化分布式环境
    deepspeed.init_distributed()

    # 创建模型
    model = GPT2LMHeadModel.from_pretrained("gpt2")

    # 初始化 DeepSpeed 引擎
    model_engine, optimizer, train_loader, lr_scheduler = deepspeed.initialize(
        model=model,
        training_data=train_dataset,
        config="ds_config.json"
    )

    # 训练循环
    for epoch in range(num_epochs):
        for batch in train_loader:
            input_ids = batch["input_ids"].to(model_engine.device)
            labels = batch["labels"].to(model_engine.device)

            # 前向传播
            outputs = model_engine(input_ids=input_ids, labels=labels)
            loss = outputs.loss

            # 反向传播(缩放/累积由 DeepSpeed 处理)
            model_engine.backward(loss)

            # 参数更新(包含梯度裁剪)
            model_engine.step()

    print(f"Step: {model_engine.global_steps}, Loss: {loss.item():.4f}")

5.2 检查点的保存与加载

# 保存检查点
def save_checkpoint(model_engine, save_dir, tag=None):
    # 以 DeepSpeed 格式保存(包含分布式状态)
    model_engine.save_checkpoint(save_dir, tag=tag)

# 例：每 1000 步保存一次
if model_engine.global_steps % 1000 == 0:
    save_checkpoint(model_engine, "./checkpoints", tag=f"step_{model_engine.global_steps}")

# 加载检查点
_, client_sd = model_engine.load_checkpoint("./checkpoints", tag="step_1000")

# 在 ZeRO-3 环境下提取为普通 PyTorch 格式的权重
if args.zero_stage == 3:
    state_dict = model_engine._zero3_consolidated_16bit_state_dict()
    if model_engine.local_rank == 0:
        torch.save(state_dict, "model_weights.pt")

5.3 Transformers Trainer 集成

将 DeepSpeed 与 HuggingFace Transformers 集成的最简便方式：

from transformers import TrainingArguments, Trainer, AutoModelForCausalLM

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    fp16=True,
    deepspeed="ds_config.json",  # DeepSpeed 配置文件路径
    logging_steps=10,
    save_steps=1000,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

命令行执行：

deepspeed --num_gpus=8 train.py \
    --deepspeed ds_config.json \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --dataset_name openwebtext \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 3 \
    --fp16

5.4 Accelerate 集成

from accelerate import Accelerator
from accelerate.utils import DeepSpeedPlugin

deepspeed_plugin = DeepSpeedPlugin(
    zero_stage=2,
    gradient_accumulation_steps=8,
    gradient_clipping=1.0,
)

accelerator = Accelerator(
    mixed_precision="fp16",
    deepspeed_plugin=deepspeed_plugin,
)

model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)

for batch in train_dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

6. DeepSpeed 流水线并行

6.1 流水线并行的概念

流水线并行把模型的各层依序分布到多张 GPU 上。GPU 0 负责靠前的层，GPU 1 负责中间的层，GPU N 负责最后的层。

1F1B 调度(One Forward, One Backward)： 将流水线气泡降到最低的调度方式。每张 GPU 在完成一个微批次的前向传播后，紧接着立即执行反向传播。

6.2 使用 PipelineModule

from deepspeed.pipe import PipelineModule, LayerSpec

class GPT2PipelineModel(PipelineModule):
    def __init__(self, num_layers, hidden_size, num_heads, vocab_size):
        layers = [
            LayerSpec(EmbeddingLayer, vocab_size, hidden_size),
        ]
        for i in range(num_layers):
            layers.append(LayerSpec(TransformerBlock, hidden_size, num_heads))
        layers.append(LayerSpec(LMHead, hidden_size, vocab_size))

        super().__init__(
            layers=layers,
            loss_fn=cross_entropy_loss,
            num_stages=4,  # 流水线阶段数(= GPU 数量)
            topology=PipeDataParallelTopology(num_pp=4, num_dp=2)
        )

# 在 ds_config 中添加流水线设置
pipeline_config = {
    "train_batch_size": 64,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 8,
    "pipeline": {
        "seed_layers": true,
        "activation_checkpoint_interval": 1
    }
}

6.3 流水线 + ZeRO 组合

model_engine, _, _, _ = deepspeed.initialize(
    model=pipeline_model,
    config={
        "train_micro_batch_size_per_gpu": 2,
        "gradient_accumulation_steps": 8,
        "zero_optimization": {"stage": 1},  # 与 ZeRO-1 组合
        "fp16": {"enabled": True}
    }
)

# 流水线训练
for step in range(num_steps):
    loss = model_engine.train_batch()
    if model_engine.global_steps % 100 == 0:
        print(f"Step {model_engine.global_steps}, Loss: {loss.item():.4f}")

7. 张量并行

7.1 Megatron-DeepSpeed 集成

张量并行把单个层分布到多张 GPU 上，把矩阵乘法沿列或行方向进行切分。

# Megatron-DeepSpeed 设置
megatron_config = {
  "tensor_model_parallel_size": 4,
  "pipeline_model_parallel_size": 2,
  "data_parallel_size": 4,  # world_size / (tp_size × pp_size)
}

在 ds_config.json 中设置张量并行：

{
  "tensor_parallel": {
    "tp_size": 4,
    "mpu": null,
    "tp_grain_size": 64
  }
}

执行命令：

deepspeed --num_nodes=4 --num_gpus=8 \
    pretrain_gpt.py \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 4 \
    --num-layers 96 \
    --hidden-size 12288 \
    --num-attention-heads 96 \
    --seq-length 2048 \
    --global-batch-size 1024 \
    --train-iters 500000 \
    --deepspeed \
    --deepspeed_config ds_config.json

8. 激活值检查点(Gradient Checkpointing)

8.1 概念与权衡

激活值检查点(activation checkpointing)在前向传播中不保存中间激活值，而是在反向传播时重新计算。

  • 内存节省：把与层数成正比的激活值内存，降为 O(sqrt(层数))
  • 速度损失：约多出 33% 的额外计算成本
{
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true,
    "contiguous_memory_optimization": true,
    "number_checkpoints": 4,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  }
}

# 在 PyTorch 层面同样可以启用
from deepspeed.runtime.activation_checkpointing import checkpointing

class TransformerLayer(nn.Module):
    def forward(self, x):
        return checkpointing.checkpoint(self._forward, x)

    def _forward(self, x):
        # 实际运算
        return self.attention(self.norm1(x)) + x

9. 混合专家(MoE)与 DeepSpeed

9.1 MoE 架构概览

Mixture of Experts 让每个输入 token 只经由一部分"专家"子网络处理，从而在保持计算量不变的同时增加参数量。

from deepspeed.moe.layer import MoE

class MoETransformerBlock(nn.Module):
    def __init__(self, hidden_size, num_experts, top_k=2):
        super().__init__()
        self.attention = MultiHeadAttention(hidden_size)
        self.moe_layer = MoE(
            hidden_size=hidden_size,
            expert=FeedForward(hidden_size, hidden_size * 4),
            num_experts=num_experts,
            ep_size=1,          # 专家并行大小
            k=top_k,            # 激活的专家数量
            capacity_factor=1.25,
            eval_capacity_factor=2.0,
            min_capacity=4,
            use_residual=False
        )

    def forward(self, x, attention_mask=None):
        x = x + self.attention(x, attention_mask)
        x, _, _ = self.moe_layer(x)
        return x

9.2 专家并行设置

{
  "zero_optimization": {
    "stage": 2
  },
  "moe_expert_parallel_size": 4,
  "moe": {
    "enabled": true,
    "ep_size": 4,
    "moe_param_group": true
  }
}

# 分离 MoE 参数组(重要！)
def create_moe_param_groups(model):
    parameters = {
        "params": [],
        "name": "parameters"
    }
    moe_parameters = {
        "params": [],
        "moe": True,
        "name": "moe_parameters"
    }
    for module_name, module in model.named_modules():
        if isinstance(module, MoE):
            moe_parameters["params"].extend(
                [p for n, p in module.named_parameters() if "expert" in n]
            )
        else:
            parameters["params"].extend(module.parameters(recurse=False))
    return [parameters, moe_parameters]

10. DeepSpeed Inference

10.1 推理优化

DeepSpeed Inference 能大幅提升已训练模型的推理速度。

import deepspeed
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 转换为 DeepSpeed 推理引擎
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=1,              # 张量并行大小
    dtype=torch.float16,    # 推理 dtype
    checkpoint=None,
    replace_with_kernel_inject=True,  # 注入优化内核
    max_tokens=2048,
    replace_method="auto"
)

model = ds_engine.module

# 执行推理
inputs = tokenizer("DeepSpeed 是 ", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7
    )
print(tokenizer.decode(outputs[0]))

10.2 张量并行推理

# 多 GPU 推理(张量并行)
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=4,              # 分散到 4 张 GPU
    dtype=torch.float16,
    replace_with_kernel_inject=True,
    injection_policy={
        LlamaDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")
    }
)

10.3 ZeRO-Inference

面向单张或少数 GPU 的大规模模型推理：

# 基于 ZeRO-3 的推理(参数从 CPU/NVMe 流式加载)
ds_engine = deepspeed.init_inference(
    model=model,
    config={
        "tensor_parallel": {"tp_size": 1},
        "dtype": "fp16",
        "enable_cuda_graph": False,
        "zero": {
            "stage": 3,
            "offload_param": {
                "device": "cpu"
            }
        }
    }
)

11. DeepSpeed Chat(RLHF)

11.1 RLHF 流水线概览

DeepSpeed-Chat 提供了完整的 RLHF(Reinforcement Learning from Human Feedback)流水线。

三阶段训练：

  1. SFT(Supervised Fine-Tuning)：有监督微调
  2. Reward Model Training：训练奖励模型
  3. RLHF/PPO：以 PPO 算法进行强化学习

11.2 SFT 阶段

# Step 1: SFT
deepspeed main.py \
    --data_path Dahoas/rm-static \
    --data_split 2,4,4 \
    --model_name_or_path facebook/opt-1.3b \
    --per_device_train_batch_size 8 \
    --max_seq_len 512 \
    --learning_rate 9.65e-6 \
    --weight_decay 0.1 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --lr_scheduler_type cosine \
    --num_warmup_steps 0 \
    --seed 1234 \
    --zero_stage 2 \
    --deepspeed \
    --output_dir /output/sft

11.3 奖励模型训练

from deepspeed_chat.reward_model import RewardModel

# 奖励模型 = SFT 模型 + 线性头
reward_model = RewardModel(
    base_model=sft_model,
    tokenizer=tokenizer,
    num_padding_at_beginning=0
)

# 用偏好对(chosen, rejected)数据训练
def compute_reward_loss(reward_model, chosen_ids, rejected_ids):
    chosen_reward = reward_model(chosen_ids)
    rejected_reward = reward_model(rejected_ids)
    # 更受偏好的回答应得到更高的奖励
    loss = -torch.mean(torch.log(torch.sigmoid(chosen_reward - rejected_reward)))
    return loss

11.4 PPO 训练

from deepspeed_chat.ppo_trainer import DeepSpeedPPOTrainer

trainer = DeepSpeedPPOTrainer(
    rlhf_engine=rlhf_engine,
    args=args
)

for epoch in range(args.num_train_epochs):
    for step, batch in enumerate(prompt_dataloader):
        out = trainer.generate_experience(batch["input_ids"])
        # 更新 Actor/Critic 模型
        actor_loss, critic_loss = trainer.train_rlhf(out)

12. 实战示例：LLM 训练

12.1 完整的 Llama-2 微调示例

import torch
import deepspeed
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq
)
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, TaskType

def main():
    # 加载模型与分词器
    model_name = "meta-llama/Llama-2-7b-hf"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token

    # 在 ZeRO-3 语境下加载模型(提升内存效率)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map=None  # 由 DeepSpeed 直接管理
    )

    # LoRA 设置(可选：参数高效微调)
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.1,
        bias="none"
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # 加载数据集
    dataset = load_dataset("tatsu-lab/alpaca", split="train")

    def preprocess(example):
        prompt = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
        tokens = tokenizer(prompt, truncation=True, max_length=512, padding="max_length")
        tokens["labels"] = tokens["input_ids"].copy()
        return tokens

    tokenized_dataset = dataset.map(preprocess, remove_columns=dataset.column_names)

    # 带 DeepSpeed 的 TrainingArguments
    training_args = TrainingArguments(
        output_dir="./llama2-finetuned",
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        learning_rate=2e-4,
        warmup_steps=100,
        logging_steps=10,
        save_steps=500,
        bf16=True,
        deepspeed="ds_zero2_config.json",
        remove_unused_columns=False,
        report_to="wandb",
        run_name="llama2-alpaca-deepspeed"
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True)
    )

    trainer.train()
    trainer.save_model("./llama2-finetuned-final")

if __name__ == "__main__":
    main()

ZeRO-2 配置文件(ds_zero2_config.json)：

{
  "bf16": {
    "enabled": "auto"
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto",
      "total_num_steps": "auto"
    }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 100,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}

执行：

deepspeed --num_gpus=8 train_llama.py \
    --deepspeed ds_zero2_config.json

12.2 性能基准测试

在 8 张 A100 80GB 环境下的实测对比：

配置最大模型规模吞吐量 (tokens/sec)
基础 DDP~7B12,000
ZeRO-1~14B11,500
ZeRO-2~35B10,800
ZeRO-3~175B9,200
ZeRO-3 + CPU Offload~350B4,100
ZeRO-Infinity~1T+1,800

13. 调试与问题排查

13.1 常见问题与解决方案

OOM(内存不足)错误：

# 内存占用剖析
ds_config = {
    "memory_breakdown": True,
    "wall_clock_breakdown": True,
}

# 分步排查：
# 1. 提高 ZeRO stage(1 → 2 → 3)
# 2. 启用 CPU Offload
# 3. 启用激活值检查点
# 4. 减小批大小
# 5. 缩短序列长度

梯度溢出(fp16)：

{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 500,
    "initial_scale_power": 12,
    "hysteresis": 2,
    "min_loss_scale": 1e-8
  }
}

ZeRO-3 参数收集错误：

# 在 ZeRO-3 下直接访问参数时
from deepspeed import zero

with zero.GatheredParameters(model.parameters()):
    # 仅在此代码块内可使用完整参数
    weight_norm = model.weight.norm()

13.2 性能调优技巧

# 1. 通信与计算重叠
zero_config = {
    "overlap_comm": True,
    "contiguous_gradients": True,
}

# 2. 优化桶大小(依据 GPU 内存与通信带宽调整)
zero_config["allgather_bucket_size"] = 5e8  # 500MB
zero_config["reduce_bucket_size"] = 5e8

# 3. ZeRO-3 预取
zero_config["stage3_prefetch_bucket_size"] = 5e7
zero_config["stage3_param_persistence_threshold"] = 1e6

# 4. 编译优化(PyTorch 2.0+)
model = torch.compile(model)

结语

DeepSpeed 是现代 LLM 训练的核心工具。借助 ZeRO 优化，原本单张 GPU 无法企及的大规模模型，如今也能以合理的成本完成训练。

核心要点：

  • ZeRO-1/2：分散优化器/梯度，提升数据并行的效率
  • ZeRO-3：连参数也一并分散，内存节省随 GPU 数量成比例增长
  • ZeRO-Offload：卸载到 CPU/NVMe，让小规模集群也能训练大型模型
  • 流水线 + 张量并行：水平/垂直方向的模型分布，将通信效率最大化
  • DeepSpeed Inference：通过内核优化，将推理速度提升 2-5 倍

在实际项目中，ZeRO-2 + CPU Offload 的组合是不错的起点。若需要更大的模型，再转向 ZeRO-3；推理场景则可利用 DeepSpeed Inference 的 kernel injection。

参考资料

현재 단락 (1/710)

要训练 GPT-4、LLaMA、Falcon 这类大规模语言模型(LLM)，单张 GPU 是不可能做到的。1750 亿参数的 GPT-3 若以 fp16 存储，需要约 350GB 的 GPU 内存，若...

작성 글자: 0자원문 글자: 19,267자작성 단락: 0/710