对比:手动 DDP (Trainer + 4bit 量化)
import os
import torch
from transformers import (
AutoTokenizer, AutoModelForCausalLM,
TrainingArguments, Trainer, DataCollatorForLanguageModeling
)
from transformers import BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
# ==========================================
# 第一部分:手动环境初始化 (手动 DDP 的标志)
# ==========================================
# 1. 显式获取当前进程的排名
# 当你使用 torchrun --nproc_per_node=2 启动时,
# 进程 0 的 LOCAL_RANK 为 0,进程 1 的 LOCAL_RANK 为 1
local_rank = int(os.environ.get("LOCAL_RANK", 0))
# 2. 显式设置当前进程使用的 GPU
# 这一步至关重要,防止所有进程都默认去操作 cuda:0
torch.cuda.set_device(local_rank)
# ==========================================
# 第二部分:模型加载 (关键:手动 device_map)
# ==========================================
model_name = "/path/to/your/model"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True
)
# 3. 加载模型时手动指定 device_map
# {"": local_rank} 的意思是:把整个模型完整地放到 local_rank 这张卡上
# 如果不写这一行,accelerate 可能会尝试把模型切分到多张卡 (模型并行),破坏 DDP
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map={"": local_rank}, # <--- 核心手动配置
torch_dtype=torch.float16
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# ... (数据集处理代码略) ...
# 配置 LoRA
lora_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
model.enable_input_require_grads() # 量化训练必须开启
# ==========================================
# 第三部分:Trainer 参数配置 (显式开启 DDP 优化)
# ==========================================
training_args = TrainingArguments(
output_dir="./output",
num_train_epochs=3,
per_device_train_batch_size=1,
gradient_accumulation_steps=16,
# --- 以下是 DDP 相关的关键手动配置 ---
# 1. 指定通信后端 (通常为 nccl)
ddp_backend="nccl",
# 2. LoRA 训练建议设为 False,提升效率
# 如果为 True,DDP 会反向遍历图找没用到的参数,非常慢
ddp_find_unused_parameters=False,
# 3. 梯度检查点配置 (非重入式,避免与 DDP 冲突)
gradient_checkpointing_kwargs={"use_reentrant": False},
gradient_checkpointing=True,
# 4. 混合精度
fp16=True,
# 其他常规参数
logging_steps=10,
save_steps=100,
save_total_limit=2,
report_to="none",
remove_unused_columns=False
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset, # 假设已定义
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
# 开始训练
trainer.train()
为什么这种写法叫"手动 DDP"?
通常使用
Trainer时,你不需要写前两行代码(local_rank和set_device),只需要在TrainingArguments里什么都不写,Trainer会自动检测是否是多卡环境并自动分配。但在这个特定场景下,手动接管了分配权,原因如下:
- 量化导致的复杂度 :
BitsAndBytesConfig(4bit 加载) 依赖accelerate库的device_map机制。- 避免模型并行 :如果不手动指定
device_map={"": local_rank},accelerate可能会认为你的模型很大,尝试把它拆分到 GPU 0 和 GPU 1 上(模型并行)。- 强制数据并行 :你想做的是 DDP(数据并行),即 GPU 0 有一份完整模型,GPU 1 也有一份完整模型。所以你必须手动告诉代码:"不管模型多大,请把它完整 地塞到
local_rank这张卡里,不要拆分。"总结
这种写法是
Trainer(高级封装) 和PyTorch DDP(底层逻辑) 的一种混合体。它既保留了Trainer的易用性,又解决了量化模型在多卡环境下的设备分配冲突问题。
Accelerate 是 Hugging Face 提供的一个轻量级库,用于简化 单机多卡 (或分布式)训练的代码编写。它底层封装了 PyTorch 的 DistributedDataParallel(DDP),让你无需手动处理进程启动、设备分配、梯度同步等细节。
1. 核心特点
-
无需手动管理设备 :不需要写
torch.cuda.set_device(local_rank)。 -
自动封装模型 :不需要手动
DistributedDataParallel(model)。 -
自动处理梯度同步 :直接用
accelerator.backward(loss)替代loss.backward()。from accelerate import Accelerator
Step 1: 初始化
accelerator = Accelerator()
Step 2: 定义你的 model, optimizer, dataloader, loss
model = ...
optimizer = ...
dataloader = ...
loss_fn = ...Step 3: 用 accelerator.prepare() 包装所有组件
model, optimizer, dataloader, loss_fn = accelerator.prepare(
model, optimizer, dataloader, loss_fn
)Step 4: 训练循环(注意 backward 要用 accelerator)
for batch in dataloader:
optimizer.zero_grad()
outputs = model(batch)
loss = loss_fn(outputs, labels)
accelerator.backward(loss) # ← 关键!不是 loss.backward()
optimizer.step()启动训练
单机多卡(自动检测 GPU 数量)
accelerate launch train.py
指定 4 卡
accelerate launch --num_processes=4 train.py
CPU / MPS / TPU 也支持
accelerate launch --cpu train.py
下面是一个完整的 使用 accelerate 实现单机多卡 DDP 训练 的示例,包含:
- 数据加载
- 模型定义
- 优化器与训练循环
- 自动设备放置与梯度同步
前提条件
- 安装:
pip install accelerate torch torchvision - 环境:单台机器,配备 ≥2 张 GPU(如 2x A100 / RTX 3090)
- 启动方式:直接运行 Python 脚本 (
accelerate会自动处理torch.distributed初始化)
示例代码:图像分类训练(ResNet + CIFAR-10)
# train_ddp_accelerate.py
from accelerate import Accelerator
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
import torch
import torch.nn as nn
import torch.optim as optim
def main():
# 1. 初始化 Accelerator(自动处理 DDP、设备、混合精度等)
accelerator = Accelerator()
# 2. 准备数据
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
# 注意:Accelerator 会自动对 sampler 进行分布式处理
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
# 3. 模型、优化器、损失函数
model = models.resnet18(num_classes=10)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
# 4. 使用 accelerator.prepare() 包装所有组件
model, optimizer, train_loader, criterion = accelerator.prepare(
model, optimizer, train_loader, criterion
)
# 5. 训练循环
model.train()
for epoch in range(10):
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
# 关键:用 accelerator.backward() 替代 loss.backward()
accelerator.backward(loss)
optimizer.step()
if batch_idx % 100 == 0 and accelerator.is_main_process:
print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")
if __name__ == "__main__":
main()
如何运行?
方法一:直接运行(推荐)
# 单机多卡(例如 2 卡)
accelerate launch --multi_gpu train_ddp_accelerate.py
accelerate launch会自动设置RANK,WORLD_SIZE,MASTER_ADDR等环境变量,并启动多个进程。
方法二:指定 GPU 数量(可选)
accelerate launch --num_processes=2 --multi_gpu train_ddp_accelerate.py
accelerator.prepare() 做了什么?
| 组件 | 处理效果 |
|---|---|
model |
自动包装为 DistributedDataParallel(DDP) |
optimizer |
适配 DDP 下的参数引用 |
DataLoader |
自动插入 DistributedSampler,确保每个 GPU 加载不同数据子集 |
loss |
通常不需要 prepare,但可参与自动设备迁移 |
高级功能(按需启用)
1. 混合精度训练(FP16)
accelerator = Accelerator(mixed_precision="fp16") # 或 "bf16"
2. 梯度累积(模拟大 batch)
accelerator = Accelerator(gradient_accumulation_steps=4)
# 在 backward 前加:
if (step + 1) % 4 == 0:
optimizer.zero_grad()
accelerator = Accelerator(gradient_accumulation_steps=8)
for step, batch in enumerate(dataloader):
with accelerator.accumulate(model): # 自动处理 zero_grad / step / sync
loss = model(batch)
accelerator.backward(loss)
optimizer.step()
# scheduler.step() 也可放这里
3. 保存/加载模型(仅主进程保存)
if accelerator.is_main_process:
unwrapped_model = accelerator.unwrap_model(model)
torch.save(unwrapped_model.state_dict(), "model.pth")
4. 日志只在主进程打印
if accelerator.is_main_process:
print("Only printed once!")
注意事项
- 不要手动调用
model.to(device)------accelerator.prepare()会自动处理。 - 不要用
loss.backward()------ 必须用accelerator.backward(loss)。 - 数据集无需手动划分 ------
DistributedSampler已自动插入。 - 验证/测试时也要用
accelerator.prepare()包装 DataLoader。
优势总结
| 传统 DDP | 使用 Accelerate |
|---|---|
需写 init_process_group、DistributedSampler、model = DDP(model) |
一行 Accelerator() 全搞定 |
多进程启动需 torchrun 或 mp.spawn |
直接 accelerate launch |
| 混合精度、梯度累积代码复杂 | 参数化配置即可 |
保存模型需处理 module. 前缀 |
unwrap_model() 自动处理 |
总结
accelerate= PyTorch 分布式训练的"胶水层"它不改变你的训练逻辑,只帮你自动处理设备、并行、精度、同步等底层细节,让你专注模型和算法本身。
=========================================================================
完整代码示例
假设要微调一个 BERT 模型做文本分类:
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup
from datasets import load_dataset
from accelerate import Accelerator
from tqdm.auto import tqdm
# 1. 初始化 Accelerator (这是核心!)
# fp16/bf16: 开启混合精度
# gradient_accumulation_steps: 梯度累积
accelerator = Accelerator(
mixed_precision="fp16", # 或者 "bf16"
gradient_accumulation_steps=4
)
# 2. 加载模型和数据 (注意:这里不需要手动 .to(device))
model_name = "bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
# 准备数据
dataset = load_dataset("csv", data_files="your_data.csv", split="train")
def preprocess_function(examples):
return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
tokenized_dataset = dataset.map(preprocess_function, batched=True)
tokenized_dataset = tokenized_dataset.remove_columns(["text"])
tokenized_dataset.set_format("torch")
dataloader = DataLoader(tokenized_dataset, shuffle=True, batch_size=16)
# 优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
# 3. 关键步骤:使用 prepare 将模型、优化器、数据加载器包装起来
# 这一步会自动:
# - 将模型放到对应的 GPU
# - 用 DDP 包装模型
# - 将数据分发到对应的 GPU
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
# 训练循环
num_epochs = 3
num_training_steps = num_epochs * len(dataloader)
lr_scheduler = get_linear_schedule_with_warmup(
optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)
print(f"Start training on {accelerator.num_processes} GPUs")
for epoch in range(num_epochs):
model.train()
progress_bar = tqdm(dataloader, disable=not accelerator.is_local_main_process) # 只在主进程显示进度条
for batch in progress_bar:
# 4. 前向传播 (不需要手动 .to(device))
outputs = model(**batch)
loss = outputs.loss
# 5. 反向传播 (关键变化!)
# Accelerate 会自动处理梯度同步和缩放
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.set_description(f"loss: {loss.item():.4f}")
# 6. 保存模型 (只让主进程保存)
# accelerator.wait_for_everyone():确保所有 GPU 都跑完这一轮,再让主进程去保存
accelerator.wait_for_everyone()
if accelerator.is_main_process:
# unwrap_model:剥去 Accelerate 的包装,获取原始模型进行保存
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(f"./output/bert_ft_epoch_{epoch}")
tokenizer.save_pretrained(f"./output/bert_ft_epoch_{epoch}")
accelerator.end_training()
如何启动训练
使用 Accelerate 不需要写复杂的 torchrun 命令,而是使用 accelerate launch:
# 使用所有可见 GPU 进行训练
accelerate launch train.py
# 指定使用 2 张卡
accelerate launch --num_processes=2 train.py
# 指定使用特定的卡 (例如卡0和卡1)
CUDA_VISIBLE_DEVICES=0,1 accelerate launch train.py
关键点解析(对比之前的代码)
| 特性 | 手动 DDP (之前的代码) | Accelerate DDP (现在的代码) |
|---|---|---|
| 设备管理 | 需要写 local_rank = ... 和 torch.cuda.set_device(...) |
不需要 。accelerator.prepare 自动处理。 |
| 模型封装 | 需要手动处理 device_map={"": local_rank} |
不需要 。prepare 自动封装 DDP 并分配设备。 |
| 反向传播 | loss.backward() |
accelerator.backward(loss) (支持混合精度自动缩放)。 |
| 多卡同步 | 需要手动 dist.all_reduce 或 gather_for_metrics |
提供了 accelerator.gather, accelerator.wait_for_everyone 等便捷 API。 |
| 保存模型 | 需要判断 if local_rank == 0 |
使用 if accelerator.is_main_process 和 unwrap_model。 |
为什么推荐用 Accelerate?
对于场景(大模型 QLoRA 微调):
- 代码迁移成本低 :你只需要把普通的 PyTorch 训练脚本拿来,加上
Accelerator()初始化,把.backward()换成accelerator.backward(),再把对象prepare一下,就能跑多卡了。 - 兼容性极好 :
Accelerate是 Hugging Face 生态的基石,Trainer底层就是用它写的。它能无缝配合DeepSpeed、FSDP等高级并行策略,以后你想改用 DeepSpeed ZeRO-3,只需要改一行启动参数,不需要改训练代码。
=========================================================================
改用 DeepSpeed ZeRO-3
改用 DeepSpeed ZeRO-3,只需要改一行启动参数,不需要改训练代码
第一步:创建一个 DeepSpeed 配置文件
创建一个名为 ds_zero3_config.json 的文件(文件名随意),写入 ZeRO-3 的配置:
{
"train_batch_size": 16,
"train_micro_batch_size_per_gpu": 1,
"gradient_accumulation_steps": 16,
"optimizer": {
"type": "AdamW",
"params": {
"lr": "2e-5",
"betas": [0.9, 0.95],
"eps": "1e-8",
"weight_decay": "0.01"
}
},
"scheduler": {
"type": "WarmupDecayLR",
"params": {
"total_num_steps": 10000,
"warmup_min_lr": "0",
"warmup_max_lr": "2e-5",
"warmup_num_steps": 500
}
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"bf16": {
"enabled": false
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto"
},
"gradient_clipping": 1.0,
"steps_per_print": 10
}
第二步:修改启动命令(只改这一行!)
原本的启动命令(DDP 模式):
accelerate launch train.py
现在的启动命令(DeepSpeed ZeRO-3 模式):
只需要加上 --config_file 参数指向你的 JSON 文件:
accelerate launch --config_file ds_zero3_config.json train.py
第三步:Python 代码完全不用变
train.py 内容依然是这样的,完全不用改动:
# train.py 内容(与 DDP 时完全一致)
from accelerate import Accelerator
# 初始化 Accelerator
# 即使你加了 ds_config,这里也不需要传任何参数,accelerate 会自动读取启动命令里的配置
accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
# ... 后面的训练循环代码完全一样 ...
为什么这么神奇?
- 自动识别 :当你运行
accelerate launch --config_file ds_zero3_config.json时,Accelerate库会读取这个 JSON 文件。 - 底层替换 :看到
zero_optimization.stage = 3后,Accelerate会自动将你的模型包装成DeepSpeedEngine,而不是普通的DistributedDataParallel。 - 参数分片:DeepSpeed 引擎会接管模型参数,将它们切分并散落在各个 GPU(甚至 CPU)上,而你的 Python 代码依然觉得自己是在操作一个完整的模型。
总结
- DDP 模式:每张卡存一份完整模型副本(显存占用大)。
- ZeRO-3 模式:模型参数被切分到多张卡(显存占用极小,能训更大的模型)。
- 代码工作量 :对于开发者来说,工作量为 0。你只需要准备好配置文件,然后改一行启动命令。
=========================================================================
以上内容总结:
核心结论:
对于 train.py(训练逻辑代码) ,Accelerate DDP 和 Accelerate DeepSpeed 的代码是 100% 完全一样 的。
区别仅在于:
- 外部配置文件 :DeepSpeed 需要一个
ds_config.json,DDP 不需要。 - 启动命令:DeepSpeed 需要指定配置文件,DDP 不需要。
- (可选)初始化方式 :如果你不想用 JSON 文件,想把配置写在 Python 里,那么
Accelerator()的初始化代码会有一点点不同。
1. 训练代码对比 (完全相同)
无论用 DDP 还是 DeepSpeed,你的 train.py 内容通常长这样,完全不需要修改:
# train.py (DDP 和 DeepSpeed 共用)
import torch
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_scheduler
from accelerate import Accelerator
# --- 初始化部分 ---
# 注意:这里代码是一样的!
# Accelerate 会自动检测你是用 DDP 还是 DeepSpeed
accelerator = Accelerator()
# --- 数据和模型准备 ---
model = AutoModelForSequenceClassification.from_pretrained("bert-base-chinese")
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
dataloader = DataLoader(...)
# prepare() 会自动根据后端包装模型
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
# --- 训练循环 ---
for batch in dataloader:
outputs = model(**batch)
loss = outputs.loss
accelerator.backward(loss) # 这一行对两者都通用
optimizer.step()
optimizer.zero_grad()
2. 真正的区别在于:配置与启动
虽然代码没变,但它们跑起来的方式截然不同。
模式 A:Accelerate DDP (标准模式)
特点:不需要额外文件,配置简单。
-
启动命令:
accelerate launch train.py -
或者配置混合精度:
accelerate launch --mixed_precision=fp16 train.py -
代码(可选) :如果你想显式指定参数,可以直接写在
Accelerator()里:accelerator = Accelerator(mixed_precision="fp16", gradient_accumulation_steps=2)
模式 B:Accelerate DeepSpeed (ZeRO 模式)
特点:需要一个 JSON 配置文件来告诉 DeepSpeed 怎么切分参数。
-
编写配置文件
ds_config.json(这是 DDP 没有的):{ "train_batch_size": 16, "train_micro_batch_size_per_gpu": 2, "gradient_accumulation_steps": 8, "optimizer": { "type": "AdamW", "params": { "lr": "2e-5" } }, "fp16": { "enabled": true }, "zero_optimization": { "stage": 3 // 开启 ZeRO-3 } } -
启动命令 (必须加上
--config_file):accelerate launch --config_file ds_config.json train.py -
代码(此时
train.py内部):# 当 DeepSpeed 启动时,这行代码会读取 config 文件,而不是括号里的参数 accelerator = Accelerator()
3. 进阶区别:代码中的 DeepSpeedPlugin
如果不想用 JSON 文件 ,坚持要把所有配置都写在 Python 代码里,那么 Accelerator 的初始化会有区别:
DDP 的 Python 写法:
from accelerate import Accelerator
accelerator = Accelerator(
mixed_precision="fp16",
gradient_accumulation_steps=4
)
DeepSpeed 的 Python 写法 (代码有区别):
需要引入 DeepSpeedPlugin 并将其传给 Accelerator。
from accelerate import Accelerator
from accelerate.utils import DeepSpeedPlugin
# 1. 定义 DeepSpeed 配置字典 (对应 JSON 内容)
ds_config = {
"train_batch_size": 16,
"train_micro_batch_size_per_gpu": 2,
"gradient_accumulation_steps": 8,
"optimizer": {
"type": "AdamW",
"params": { "lr": "2e-5" }
},
"fp16": {"enabled": True},
"zero_optimization": {
"stage": 3
}
}
# 2. 创建 Plugin 对象
ds_plugin = DeepSpeedPlugin(ds_config)
# 3. 传给 Accelerator (这里和 DDP 不同)
accelerator = Accelerator(deepspeed_plugin=ds_plugin)
# 后面的 prepare 和训练循环完全一样
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
总结
| 特性 | Accelerate DDP | Accelerate DeepSpeed |
|---|---|---|
| Python 训练逻辑 | 完全相同 (prepare, backward) |
完全相同 (prepare, backward) |
| 配置方式 | CLI 参数 或 Accelerator() 参数 |
必须用 JSON 文件 或 DeepSpeedPlugin 对象 |
| 显存优化 | 有限 (主要靠混合精度/梯度累积) | 极强 (ZeRO 切分优化器/梯度/参数) |
| 推荐写法 | accelerator = Accelerator() |
accelerator launch --config_file ds.json |
一句话概括 :
Accelerate 的魔法在于解耦 。train.py 只是描述"怎么训练",而 DDP 还是 DeepSpeed 只是"怎么分配硬件资源",两者通过 accelerator.prepare 这个适配器连接,所以核心代码可以保持不变。
=========================================================================
补充:DDP 和 DeepSpeed区别
DDP (Distributed Data Parallel) 和 DeepSpeed 都是为了解决多卡并行训练的问题,但它们的核心思路 和适用场景完全不同。
一句话总结:
- DDP :数据并行 。每张卡存一份完整模型,显存占用大,适合中小模型。
- DeepSpeed (ZeRO) :状态/参数切分 。模型参数被拆散存放在多张卡上,显存占用极小,适合超大模型。
1. 核心区别对比表
| 特性 | DDP (PyTorch 标准) | DeepSpeed (主要是 ZeRO 策略) |
|---|---|---|
| 核心思想 | 复制 (Replication) | 切分 (Partitioning / Sharding) |
| 模型存储 | 每张卡都存一份完整的模型参数、梯度和优化器状态。 | 将优化器状态、梯度、参数切分存储在不同的卡上。 |
| 显存占用 | 高。O(N)×GPU数。 | 低。单卡显存占用随 GPU 数量增加而减少。 |
| 通信量 | 较低。只同步梯度。 | 较高。除了同步梯度,还需要在计算时动态获取/同步参数。 |
| 计算速度 | 快。通信开销小,计算效率高。 | 稍慢。因为增加了参数收集的通信开销,但换来了能训练大模型的能力。 |
| 适用模型 | 7B、13B 等参数量较小、能塞进单卡显存的模型。 | 30B、70B、175B+ 等参数量巨大、单卡塞不下的模型。 |
2. 生动的比喻:看书训练
假设你要"阅读"(训练)一本 1000 页 的书(模型),你有 4 个学生(4 张 GPU)。
DDP 方式(复印机模式)
- 做法 :把这本书复印 4 份,每个学生手里都有 完整的 1000 页。
- 分工:学生 1 读第 1-250 页,学生 2 读第 251-500 页...
- 缺点:如果书特别厚(比如 10000 页),每个学生手里的书都拿不动(显存溢出,OOM)。
- 优点:大家各自读各自的,不需要互相借书,速度很快。
DeepSpeed ZeRO-3 方式(拼图模式)
- 做法:不复印,把书拆开。撕成 4 份,学生 A 拿第 1-250 页,学生 B 拿第 251-500 页...
- 分工 :
- 学生 A 需要看第 260 页时,必须大喊一声:"谁有第 260 页?"
- 学生 B 说:"我有,发给你!"
- 学生 A 看完计算完后,把结果传回去,或者把这一页擦除掉。
- 优点:无论书有多厚(10000 页甚至 100000 页),只要人数(卡数)够多,每个人手里只拿几十页,完全拿得动。
- 缺点:学生 A 频繁向学生 B 借书(通信),会花费一些时间,导致阅读速度比 DDP 慢一点。
3. 技术细节上的区别
DDP 的问题:显存被"三巨头"吃光了
训练一个模型,显存主要被这三样东西占用(按比例大致为 2:1:1):
- 模型参数:比如 FP16 下,7B 模型约占 14GB。
- 优化器状态:Adam 优化器需要存储动量等,通常占用 2 倍于参数的显存(约 28GB)。
- 梯度:占用 1 倍于参数的显存(约 14GB)。
DDP 的痛点 :
在 DDP 模式下,每张卡都要存这三样东西。
如果模型是 7B:
- 单卡需要:14GB (参数) + 28GB (优化器) + 14GB (梯度) = 56GB。
- 这就是为什么单张 A100 (40GB/80GB) 用 DDP 有时也会显存不足,或者你无法在单卡上训练太大的模型。
DeepSpeed 的解决方案:ZeRO (Zero Redundancy Optimizer)
DeepSpeed 引入了 ZeRO 技术,把上面的"三巨头"进行切分:
- ZeRO Stage 1 :切分优化器状态 。
- 节省显存:4倍。
- 例子:上面的 56GB -> 28GB。
- ZeRO Stage 2 :切分优化器状态 + 梯度 。
- 节省显存:8倍。
- 例子:上面的 56GB -> 14GB。
- ZeRO Stage 3 :切分优化器状态 + 梯度 + 模型参数 。
- 节省显存:与 GPU 数量成正比。
- 例子:如果有 4 张卡,每张卡只需存 1/4 的模型。这就是**之前说的"只改一行代码就能跑 70B 模型"**的秘密。
4. 应该怎么选?
-
如果你训练的是 7B 以下的模型:
- 首选 DDP。
- 代码简单,速度快,不容易出 Bug。如果显存不够,加
gradient_accumulation或用4bit量化(QLoRA)。
-
如果你想训练 30B、70B 甚至更大的模型:
- 必须用 DeepSpeed (ZeRO-3)。
- DDP 根本塞不进去。只有 DeepSpeed 能把模型切碎了存。
-
如果你想用 QLoRA (4bit 量化) 跑大模型:
- 实际上你不需要 ZeRO-3。
- 因为 QLoRA 已经把模型压缩到了 4bit,显存占用很小了,普通的 DDP 就能跑 70B 模型。只有在不量化的全量微调场景下,ZeRO 才是大杀器。
=========================================================================
DDP 和 DeepSpeed(尤其是 ZeRO-3)确实都需要进行通信(GPU 之间互相传数据)。
既然都要通信,为什么 DeepSpeed 能塞进去大模型,而 DDP 不行?
核心区别在于:通信的"内容" 、通信的"时机" 以及 通信的"量"。
可以用一个形象的比喻来对比:学生做作业(计算)与对答案(通信)。
1. DDP:只对"答案" (通信量小,频率低)
假设有 4 个学生(4 张 GPU),大家手里都有一本完整的书(模型)。
- 通信内容 :只传递梯度(即"答案的修正值")。
- 通信时机 :在大家做完题(反向传播结束)之后。
- 过程 :
- 大家各自算各自的题(前向+反向)。
- 算完后,大家聚在一起(All-Reduce),互相核对一下:"第 1 题的梯度是 0.5 吗?""是的,平均一下。"
- 核对完,大家各自修改自己书上的知识点(更新参数)。
- 下一轮继续。
总结:
- 优点 :通信量较小(只有梯度),且只在最后通信一次,所以计算效率最高。
- 缺点:每个人都必须拿一本完整的书(存完整模型)。如果书太厚(模型太大),书包(显存)就装不下了。
2. DeepSpeed (ZeRO-3):连"书页"都要借 (通信量大,频率高)
还是 4 个学生,但书太厚了,书包根本装不下。于是大家决定把书撕了,每人只拿几页。
- 通信内容 :模型参数 + 梯度 + 优化器状态(即"书页"、"答案修正值"、"草稿纸")。
- 通信时机 :在做题的整个过程中持续发生。
- 过程 :
- 学生 A 做题做到第 10 页,但他手里只有第 1-5 页。
- 通信发生:学生 A 必须大喊:"谁有第 10 页?"
- 学生 B 说:"我有,传给你。"(获取参数)
- 学生 A 拿到第 10 页,开始做题。
- 做完后,学生 A 算出第 10 页的修正值。
- 通信再次发生 :学生 A 把修正值传给学生 B:"这是第 10 页的修改意见,你记一下。"(更新梯度/状态)
总结:
- 优点:每个人手里只需要拿几张纸(极省显存),大家可以一起读一本超级厚的书(超大模型)。
- 缺点 :做题过程中,不停地要"借书页"、"还书页",通信量变得非常巨大,且非常频繁。
3. 深度对比:DDP vs DeepSpeed (ZeRO-3)
| 特性 | DDP (DataParallel) | DeepSpeed ZeRO-3 |
|---|---|---|
| 手里有什么 | 完整模型 (Parameters) | 1/N 的模型碎片 (Sharded Parameters) |
| 传什么数据 | 只传梯度 (Gradients) | 传参数 (Forward时) + 传梯度 (Backward时) |
| 什么时候传 | 一个 Batch 结束后传一次 | 每一层计算前/后都要传 |
| 通信压力 | 小 (通信带宽占用低) | 极大 (极度依赖 NVLink/高速网络) |
| 为什么快/慢 | 计算快 (主要时间在计算,通信时间短) | 相对慢 (大量时间花在等待借书页上) |
| 为什么能省显存 | 不省 (每卡存完整模型) | 省 (每卡只存一部分模型) |
4. 结论:时间换空间
- DDP 是为了速度。它的通信是"低成本"的,但它对显存容量要求高。
- DeepSpeed 是为了容量。它牺牲了通信效率(变慢了),通过疯狂地传递数据碎片,换取了"显存可以无限叠加"的能力。
所以,如果单卡显存够用,大家都会选 DDP (因为它通信少,跑得快);只有当模型大到 DDP 塞不进去时,才会被迫使用 DeepSpeed(虽然通信多,但至少能跑起来)。
=========================================================================
补充:DDP与DP区别
DP (DataParallel) 和 DDP (DistributedDataParallel) 都是 PyTorch 提供的数据并行训练方式(即:每张卡复制一份完整模型,处理不同的数据)。
虽然目的相同,但DP 是老一代的简易方案,DDP 是新一代的工业级标准方案。
一句话总结: 除非你只是想在单机上用 2 张卡快速跑通代码验证想法,否则永远不要用 DP,请直接使用 DDP。
1. 核心区别对比表
| 特性 | DP (DataParallel) | DDP (DistributedDataParallel) |
|---|---|---|
| 并行方式 | 单进程多线程 (Single Process, Multi-Threaded) | 多进程 (Multi-Process) |
| 性能瓶颈 | 严重受限于 Python GIL (全局解释器锁) | 无 GIL 限制,真正的并行计算 |
| 通信效率 | 低(梯度传输必须经过 CPU) | 高(GPU 之间直接通过 NCCL 通信) |
| 代码复杂度 | 极低(一行代码搞定) | 较高(需要处理进程组、初始化等) |
| 适用场景 | 单机多卡,快速调试 | 单机多卡 & 多机多卡,生产环境 |
| 显存均衡 | 不均衡(主卡显存占用通常比其他卡高) | 均衡(所有卡显存占用一致) |
2. 生动的比喻:搬砖
假设你要搬 10000 块砖(训练数据),你有 4 个工人(4 张 GPU)。
DP 的模式(包工头模式)
- 结构 :只有一个包工头(主进程) ,他雇佣了 3 个临时工(子线程)。
- 工作流 :
- 包工头把砖分给临时工。
- 临时工干完活,把结果(梯度)拿回来交给包工头。
- 包工头自己负责汇总结果,计算更新方法,然后再把新方法告诉临时工。
- 缺点 :包工头(CPU/主线程)累得半死,而因为 Python 的 GIL 锁,同一时间只能有一个人在说话,其他人得排队。导致临时工经常在等包工头,效率极低。
DDP 的模式(合作小组模式)
- 结构 :4 个独立的合伙人(独立的进程),大家地位平等。
- 工作流 :
- 每个人各领一部分砖(通过 DistributedSampler)。
- 大家自己干自己的活。
- 干完后,大家聚在一起开个会(All-Reduce),核对一下笔记,统一更新一下知识。
- 散会,继续干活。
- 优点:没有包工头卡脖子,大家同步交流,效率极高。
3. 为什么 DP 性能差?(技术细节)
-
Python GIL 锁:
- DP 使用多线程。Python 的多线程由于 GIL 的存在,同一时刻只能有一个线程在执行 Python 字节码。这意味着在梯度同步和参数更新时,CPU 是串行的,无法利用多核优势。
- DDP 使用多进程。每个进程有自己独立的 Python 解释器和 GIL,互不干扰,真正实现了并行。
-
梯度传输路径:
- DP:GPU 0, 1, 2 的梯度必须先传输到 CPU 内存,由 CPU 汇总计算平均,然后再发回 GPU。这增加了 GPU -> CPU -> GPU 的数据拷贝开销。
- DDP:使用 NVIDIA 的 NCCL 库,允许 GPU 0 直接通过 NVLink 或 PCIe 网络与其他 GPU 通信,不需要 CPU 中转。
-
显存溢出风险:
- 在 DP 中,主卡(GPU 0)除了存模型,还要负责损失计算、梯度汇总和更新,因此主卡的显存占用通常比其他卡高。容易导致"其他卡没事,主卡先爆显存(OOM)"的情况。
4. 代码对比
DP 的写法(极其简单,但慢)
import torch.nn as nn
# 定义模型
model = MyModel()
# 只需要这一行!
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model) # 自动包裹模型
model.to("cuda")
# 后面正常训练,PyTorch 会自动切分数据输入到不同卡
DDP 的写法(稍微复杂,但快)
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 1. 初始化进程组
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
# 2. 配置当前进程使用的 GPU
torch.cuda.set_device(local_rank)
device = torch.device(f"cuda:{local_rank}")
# 3. 创建模型并移动到 GPU
model = MyModel().to(device)
# 4. 用 DDP 包裹模型
model = DDP(model, device_ids=[local_rank])
# 5. 数据加载器需要使用 DistributedSampler
# (需要手动切分数据集,确保每个进程拿到的数据不同)
# ... (后续训练代码)
# 6. 记得最后要销毁进程组
dist.destroy_process_group()
*(注:实际上我们通常会使用 Trainer 或 Accelerate 来自动处理 DDP 的繁琐步骤,而不需要手写上面那么多代码)*
5. 总结与建议
-
DP (DataParallel):
- 优点:API 简单,改动代码少。
- 缺点:慢,受限于 GIL,显存不均,不支持多机训练。
- 适用:只是想在双卡机器上跑个小 demo,不关心速度,或者显卡非常少(1-2张)且模型很小。
-
DDP (DistributedDataParallel):
- 优点:速度快,支持大模型,支持多机多卡,显存均衡。
- 缺点:上手概念稍多。
- 适用:几乎所有正式的科研训练和工业级部署。
结论: 只要涉及到多卡训练,默认使用 DDP (或者封装了 DDP 的工具,如 Trainer / Accelerate)。请忘掉 DP 的存在。