用DeepSpeed训练大模型

这个话题之前一直想整理，今天终于抽出时间来写。本文尽量用实际代码说话，避免空谈概念。

先搞清楚原理

DeepSpeed的底层原理涉及模型评估。用代码演示比画图更直观：

import torch
import torch.nn as nn
from torch.optim import AdamW

model = nn.Sequential(
    nn.Linear(100, 256), nn.ReLU(), nn.Dropout(0.2),
    nn.Linear(256, 128), nn.ReLU(),
    nn.Linear(128, 10),
)

optimizer = AdamW(model.parameters(), lr=0.0005, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
criterion = nn.CrossEntropyLoss()

for epoch in range(20):
    x = torch.randn(32, 100)
    y = torch.randint(0, 10, (32,))
    optimizer.zero_grad()
    loss = criterion(model(x), y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    scheduler.step()
    if epoch % 5 == 0:
        print(f"Epoch {epoch}: loss={loss.item():.4f}, lr={scheduler.get_last_lr()[0]:.6f}")

实际怎么用

实际用的时候有几个注意点。首先是配置问题，不同环境参数可能不一样。其次是性能，数据量大的时候要注意模型评估。看这段实际运行的代码：

import torch
import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU())
        self.mu_head = nn.Linear(hidden_dim, latent_dim)
        self.logvar_head = nn.Linear(hidden_dim, latent_dim)
        self.decoder = nn.Sequential(nn.Linear(latent_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, input_dim), nn.Sigmoid())

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        h = self.encoder(x)
        mu, logvar = self.mu_head(h), self.logvar_head(h)
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar

model = VAE()
x = torch.randn(16, 784)
out, mu, logvar = model(x)
print(f"VAE输出: {out.shape}, mu: {mu.shape}")

进阶一点的用法

有些场景下DeepSpeed的常规写法不够用，需要做模型评估。这个改动涉及的地方不多，但要注意顺序：

容易踩的坑

说几个容易踩的坑。第一个是模型评估的问题，很多人在这上面栽过跟头。第二个是版本兼容性，不同版本行为可能不一样。第三个是边界情况，空值和异常输入要处理好。

写在最后

总结一下，DeepSpeed的关键在于模型参数优化。掌握了这个，其他用法都是在这个基础上扩展。建议多写几个小demo练手，光看不动手容易忘。