好的，各位观众老爷们，欢迎来到今天的“DeepSpeed/FairScale：大规模分布式训练与模型并行优化”专场！今天咱们不搞虚的，直接上干货，聊聊如何用DeepSpeed和FairScale这两个神器，把那些动不动就几百亿、几千亿参数的大模型给喂饱，让它们跑得飞起！
一、引言：模型越来越大，显存越来越小？
话说啊，这年头，模型参数量蹭蹭往上涨，恨不得一天一个亿。但咱们手里的显卡，显存就那么点，捉襟见肘啊！单卡训练？那得等到猴年马月！所以，分布式训练是唯一的出路。但是，分布式训练也不是那么容易的，各种问题等着你：
- 显存不够用？ 一个模型几百G，一张卡才几十G，怎么塞得下？
- 通信开销太大？ 几百张卡一起训练，数据传来传去，网络带宽不够啊！
- 训练效率不高？ 卡多了，但效率反而下降了，感觉白花了钱！
别慌！DeepSpeed和FairScale就是来拯救世界的！它们提供了各种模型并行技术，帮你解决这些问题，让你的大模型训练事半功倍。
二、DeepSpeed：微软出品，必属精品？
DeepSpeed是微软开源的一个深度学习优化库，专注于大规模分布式训练。它的目标是让每个人都能轻松训练拥有数十亿甚至数万亿参数的模型。DeepSpeed的核心技术包括：
-
ZeRO (Zero Redundancy Optimizer)
ZeRO是DeepSpeed的核心技术，它通过将优化器的状态（optimizer states）、梯度（gradients）和模型参数（model parameters）分片到不同的GPU上，来减少每个GPU的显存占用。简单来说，就是把一个完整的模型拆成几块，分给不同的GPU来存储。
ZeRO有三个阶段：
- ZeRO-1 (Optimizer State Partitioning): 将优化器的状态分片到不同的GPU上。
- ZeRO-2 (Gradient Partitioning): 将梯度分片到不同的GPU上。
- ZeRO-3 (Parameter Partitioning): 将模型参数分片到不同的GPU上。
ZeRO-3是DeepSpeed中最强大的技术，它可以显著减少显存占用，但也会增加通信开销。
代码示例：
from deepspeed import initialize import torch import torch.nn as nn from torch.optim import Adam # 定义一个简单的模型 class SimpleModel(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(10, 10) def forward(self, x): return self.linear(x) # 初始化模型、优化器 model = SimpleModel() optimizer = Adam(model.parameters(), lr=0.001) data = torch.randn(1, 10) target = torch.randn(1, 10) # 初始化DeepSpeed引擎 model_engine, optimizer, _, _ = initialize( model=model, optimizer=optimizer, config_params={ "train_batch_size": 1, # Global batch size "train_micro_batch_size_per_gpu": 1, # Batch size per GPU "optimizer": { "type": "Adam", "params": { "lr": 0.001 } }, "zero_optimization": { "stage": 3, # 使用 ZeRO-3 "offload_param": { "device": "cpu", # 将模型参数 offload 到 CPU "pin_memory": True }, "offload_optimizer": { "device": "cpu", # 将优化器状态 offload 到 CPU "pin_memory": True }, "overlap_comm": True, "contiguous_gradients": True, "reduce_bucket_size": 1e8, "stage3_prefetch_bucket_size": 1e8, "stage3_param_persistence_threshold": 1e4 } } ) # 训练循环 for i in range(10): output = model_engine(data) loss = torch.nn.functional.mse_loss(output, target) model_engine.backward(loss) model_engine.step() print(f"Step {i}, Loss: {loss.item()}")代码解释：
zero_optimization: 这个配置项是DeepSpeed ZeRO的关键。stage: 设置为3，表示使用ZeRO-3。offload_param: 将模型参数 offload 到 CPU，可以进一步减少 GPU 显存占用。offload_optimizer: 将优化器状态 offload 到 CPU，同上。overlap_comm: 允许计算和通信重叠，提高效率。contiguous_gradients: 确保梯度是连续的，可以提高内存访问效率。reduce_bucket_size和stage3_prefetch_bucket_size: 调整通信的bucket size，可以优化通信性能。stage3_param_persistence_threshold: 控制参数持久化到CPU的阈值。
-
混合精度训练 (Mixed Precision Training)
DeepSpeed支持混合精度训练，可以使用半精度浮点数（FP16）或BrainFloat16（BF16）来训练模型。与单精度浮点数（FP32）相比，FP16和BF16可以减少显存占用，并提高计算速度。
代码示例：
from deepspeed import initialize import torch import torch.nn as nn from torch.optim import Adam class SimpleModel(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(10, 10) def forward(self, x): return self.linear(x) model = SimpleModel() optimizer = Adam(model.parameters(), lr=0.001) data = torch.randn(1, 10) target = torch.randn(1, 10) model_engine, optimizer, _, _ = initialize( model=model, optimizer=optimizer, config_params={ "train_batch_size": 1, "train_micro_batch_size_per_gpu": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.001 } }, "fp16": { "enabled": True, # 启用 FP16 "loss_scale": 0, "loss_scale_window": 1000, "initial_scale_power": 32, "hysteresis": 2, "min_loss_scale": 1 } } ) for i in range(10): data = data.half() # Convert input to half precision target = target.half() # Convert target to half precision output = model_engine(data) loss = torch.nn.functional.mse_loss(output, target) model_engine.backward(loss) model_engine.step() print(f"Step {i}, Loss: {loss.item()}")代码解释：
fp16: 这个配置项是DeepSpeed FP16训练的关键。enabled: 设置为True，表示启用FP16训练。loss_scale: 用于防止梯度下溢，一般设置为0，让DeepSpeed自动调整。- 其他的参数用于更精细地控制loss scaling的行为。
-
Pipeline 并行 (Pipeline Parallelism)
Pipeline 并行将模型分成多个阶段（stage），每个阶段由不同的GPU负责。数据像流水线一样在各个阶段之间传递。这样可以减少每个GPU的显存占用，并提高吞吐量。
代码示例：
import torch import torch.nn as nn from deepspeed import initialize from deepspeed.pipe import PipelineModule, LayerSpec, TiedLayerSpec # 定义一个简单的模型，分成两个阶段 class Stage1(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(10, 20) def forward(self, x): return self.linear(x) class Stage2(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(20, 10) def forward(self, x): return self.linear(x) # 将模型封装成 PipelineModule class PipelineModel(PipelineModule): def __init__(self, num_stages): super().__init__(num_stages=num_stages) self.stage1 = Stage1() self.stage2 = Stage2() def forward(self, x): x = self.stage1(x) x = self.stage2(x) return x # 初始化模型、优化器 model = PipelineModel(num_stages=2) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) data = torch.randn(1, 10) target = torch.randn(1, 10) # 初始化 DeepSpeed 引擎 model_engine, optimizer, _, _ = initialize( model=model, optimizer=optimizer, config_params={ "train_batch_size": 1, "train_micro_batch_size_per_gpu": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.001 } }, "pipeline": { "enabled": True, "num_stages": 2 } } ) # 训练循环 for i in range(10): output = model_engine(data) loss = torch.nn.functional.mse_loss(output, target) model_engine.backward(loss) model_engine.step() print(f"Step {i}, Loss: {loss.item()}")代码解释：
PipelineModule: DeepSpeed 提供的用于封装 pipeline 并行模型的基类。pipeline: 这个配置项是DeepSpeed pipeline 并行的关键。enabled: 设置为True，表示启用 pipeline 并行。num_stages: 指定 pipeline 的阶段数。
-
数据并行 (Data Parallelism)
数据并行是最常见的分布式训练方式。它将数据分片到不同的GPU上，每个GPU训练模型的一个副本。DeepSpeed也支持数据并行，并提供了优化，可以提高数据并行的效率。通常与 ZeRO 结合使用。
三、FairScale：Facebook出品，也不赖！
FairScale是Facebook开源的一个用于大规模训练的库。它提供了各种模型并行技术，与DeepSpeed类似，但也有一些自己的特点。FairScale的核心技术包括：
-
Fully Sharded Data Parallel (FSDP)
FSDP是FairScale的核心技术，它类似于DeepSpeed的ZeRO-3。FSDP将模型参数、梯度和优化器状态分片到不同的GPU上，可以显著减少每个GPU的显存占用。FSDP与ZeRO-3的主要区别在于，FSDP在每个GPU上保留一部分完整的模型参数，这样可以减少通信开销。
代码示例：
import torch import torch.nn as nn from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP from fairscale.optim.oss import OSS from fairscale.optim.grad_scaler import ShardedGradScaler # 定义一个简单的模型 class SimpleModel(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(10, 10) def forward(self, x): return self.linear(x) # 初始化模型、优化器 model = SimpleModel() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) data = torch.randn(1, 10).cuda() target = torch.randn(1, 10).cuda() # 将模型封装成 FSDP model = FSDP(model).cuda() # 使用 OSS 优化器 optimizer = OSS(params=model.parameters(), optim=optimizer) # 使用 ShardedGradScaler 进行混合精度训练 scaler = ShardedGradScaler() # 训练循环 for i in range(10): optimizer.zero_grad() with torch.cuda.amp.autocast(): output = model(data) loss = torch.nn.functional.mse_loss(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() print(f"Step {i}, Loss: {loss.item()}")代码解释：
FullyShardedDataParallel: FairScale 提供的 FSDP 类，用于封装模型。OSS: FairScale 提供的优化器封装类，可以与 FSDP 配合使用。ShardedGradScaler: FairScale 提供的梯度缩放类，用于混合精度训练。torch.cuda.amp.autocast(): PyTorch 提供的自动混合精度上下文管理器。
-
Pipeline Parallelism
FairScale 也支持 Pipeline 并行，与DeepSpeed类似。
-
Data Parallelism
FairScale 也支持数据并行，并提供了优化，可以提高数据并行的效率。
四、DeepSpeed vs FairScale：选哪个？
DeepSpeed和FairScale都是非常强大的分布式训练库，它们各有优缺点。
| 特性 | DeepSpeed | FairScale |
|---|---|---|
| 核心技术 | ZeRO, 混合精度训练, Pipeline 并行, 数据并行 | FSDP, 混合精度训练, Pipeline 并行, 数据并行 |
| 易用性 | 相对容易上手，配置项较多，需要一定的调参经验 | 相对复杂，API 比较底层，需要对分布式训练有深入的理解 |
| 社区支持 | 微软官方支持，文档完善，社区活跃度高 | Facebook 官方支持，文档相对简单，社区活跃度一般 |
| 适用场景 | 适用于各种规模的模型，尤其是超大规模模型，对显存优化有较高要求 | 适用于大规模模型，对分布式训练的性能有较高要求 |
| 混合精度训练 | 集成度高，配置灵活 | 需要配合 PyTorch 的 torch.cuda.amp 使用，相对繁琐 |
| 优化器集成 | 集成度高，支持多种优化器 | 需要使用 FairScale 提供的 OSS 优化器封装类，有一定的限制 |
总结：
- 如果你的模型非常大，显存非常紧张，并且希望快速上手，那么DeepSpeed可能更适合你。
- 如果你对分布式训练有深入的理解，并且希望对性能进行极致优化，那么FairScale可能更适合你。
五、实战案例：训练一个简单的GPT-2模型
为了让大家更直观地了解DeepSpeed和FairScale的使用，我们来用它们训练一个简单的GPT-2模型。由于篇幅限制，这里只给出DeepSpeed的示例代码。
import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from torch.optim import AdamW
from deepspeed import initialize
import deepspeed
import os
# 定义模型
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
model = GPT2LMHeadModel.from_pretrained(model_name)
model.resize_token_embeddings(len(tokenizer))
model.config.pad_token_id = model.config.eos_token_id
# 定义数据集 (这里用一个简单的例子)
train_data = [
"The quick brown fox jumps over the lazy dog.",
"The cat sat on the mat.",
"Hello, world!",
"This is a test sentence."
]
def encode_data(tokenizer, data, max_length=128):
encoded_texts = tokenizer(data, truncation=True, padding="max_length", max_length=max_length, return_tensors="pt")
return encoded_texts["input_ids"], encoded_texts["attention_mask"]
input_ids, attention_mask = encode_data(tokenizer, train_data)
labels = input_ids.clone()
# 初始化 DeepSpeed
config_params = {
"train_batch_size": 1,
"train_micro_batch_size_per_gpu": 1,
"optimizer": {
"type": "AdamW",
"params": {
"lr": 2e-5
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 2e-5,
"warmup_num_steps": 100
}
},
"fp16": {
"enabled": True
},
"zero_optimization": {
"stage": 3,
"offload_param": {
"device": "cpu",
"pin_memory": True
},
"offload_optimizer": {
"device": "cpu",
"pin_memory": True
},
"overlap_comm": True,
"contiguous_gradients": True,
"reduce_bucket_size": 1e8,
"stage3_prefetch_bucket_size": 1e8,
"stage3_param_persistence_threshold": 1e4
},
"gradient_clipping": 1.0
}
model = model.cuda() # Move model to GPU before initializing DeepSpeed
optimizer = AdamW(model.parameters(), lr=2e-5) # Create optimizer
model_engine, optimizer, _, _ = initialize(
model=model,
optimizer=optimizer,
config_params=config_params
)
# 训练循环
num_epochs = 3
for epoch in range(num_epochs):
for i in range(len(input_ids)):
input_ids_batch = input_ids[i].unsqueeze(0).cuda()
attention_mask_batch = attention_mask[i].unsqueeze(0).cuda()
labels_batch = labels[i].unsqueeze(0).cuda()
outputs = model_engine(input_ids_batch, attention_mask=attention_mask_batch, labels=labels_batch)
loss = outputs.loss
model_engine.backward(loss)
model_engine.step()
print(f"Epoch {epoch}, Step {i}, Loss: {loss.item()}")
# 保存模型
if model_engine.is_main_process: # Only save on the main process
model_engine.module.save_pretrained("./gpt2_deepspeed")
tokenizer.save_pretrained("./gpt2_deepspeed")代码解释：
- 这个示例代码展示了如何使用DeepSpeed训练一个简单的GPT-2模型。
- 代码中使用了ZeRO-3和FP16混合精度训练，可以显著减少显存占用。
- 需要安装
transformers和deepspeed库。
六、总结与展望
DeepSpeed和FairScale为我们提供了强大的工具，可以轻松地训练大规模模型。但是，分布式训练仍然是一个充满挑战的领域，需要不断探索和优化。
未来，我们可以期待：
- 更高效的模型并行技术
- 更智能的资源调度
- 更易用的分布式训练框架
希望今天的分享能帮助大家更好地理解和使用DeepSpeed和FairScale，让大家都能训练出自己的大模型！ 谢谢大家！ 散会！