目录

一、单卡场景

1. 混合精度训练

2. 梯度检查点-以时间换空间

二、多卡场景

1. 模型并行

1. 1 张量并行（Tensor Parallelism）

1.2. 流水线并行（Pipeline Parallelism）

1.3 ZeRO(Zero Redundancy Optimizer)— DeepSpeed 的核心优化

三、前沿优化方案

1. 内存高效Attention

2. 参数高效微调（PEFT）

⚠️ 避坑指南

📊 业界案例

---

一、单卡场景

1. 混合精度训练

核心思想：降低模型权重的数值精度，从 FP32 → FP16/BF16 → INT8 → INT4，大幅减少显存占用

FP16 / BF16 半精度推理

• 显存减半（FP32 7B 模型约 28GB → FP16 约 14GB）

• 使用 torch_dtype 设置：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "meta-llama/Llama-2-7b-hf"

# 加载模型时指定精度
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # 或者 torch.bfloat16
    device_map="auto",          # 自动分配设备
    trust_remote_code=True      # 部分模型需要
)

2. 梯度检查点-以时间换空间

核心思想：不保存所有中间激活值，反向传播时重新计算，大幅降低训练时的显存占用

• 适用于微调（SFT、DPO）场景

• 显存节省 50%~80%，但训练速度降低 20%~30%

model.gradient_checkpointing_enable()

或在from_pretrained 中设置：

model = AutoModelForCausalLM.from_pretrained(
    ...,
    gradient_checkpointing=True
)

二、多卡场景

1. 模型并行

核心思想：将一个大模型的层或参数切分到多个 GPU 上，每张卡只存储一部分模型。

1. 1 张量并行（Tensor Parallelism）

• 将单个层的权重矩阵（如注意力头、FFN）按维度拆分

• 例如：将 7B 模型的 q_proj 矩阵拆到 2 张卡上

• 需要频繁的 GPU 间通信（All-Reduce）

• 适用场景：模型参数量 < 单卡显存容量

✅ 工具：Megatron-LM

1.2. 流水线并行（Pipeline Parallelism）

• 将模型按层拆分，每张卡负责一部分层

• 数据像“流水线”一样依次通过各卡

• 显著降低单卡显存占用

• 适用场景：模型参数量 > 单卡显存容量

✅ 工具：DeepSpeed、Accelerate

1.3 ZeRO(Zero Redundancy Optimizer)— DeepSpeed 的核心优化

核心思想：消除数据并行中冗余的优化器状态、梯度和参数副本

ZeRO 技术分为三个阶段：

阶段	优化对象	显存节省
ZeRo-1	将优化器状态分片到各GPU	3~5x
ZeRo-2	梯度	8~10x
ZeRo-3	模型参数	15x+

示例：

// deepspeed_config.json
{
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu"
    }
  },
  "train_micro_batch_size_per_gpu": 1
}

# 启动训练
deepspeed --num_gpus=4 train.py \
  --deepspeed deepspeed_config.json

三、前沿优化方案

1. 内存高效Attention

FlashAttention-2

• 通过 IO 感知算法减少显存读写

• 速度提升 2-4 倍

2. 参数高效微调（PEFT）

LoRA

• 冻结原参数，训练低秩适配矩阵

• 显存占用降低70%

⚠️ 避坑指南

• 梯度检查点会导致约30%训练速度下降

• FP16训练可能出现梯度下溢，需配合Loss Scaling

• 模型并行需要改写网络结构，调试复杂

📊 业界案例

• LLaMA-2 70B：采用8路张量并行+16路流水线并行

• GPT-3：梯度检查点+FP16节省显存78%