fairscale

当你的 GPU 显存告急时，这个 Facebook 开源库能帮你把模型塞进显卡

如果你曾经训练过 GPT、OPT 或任何超过 10 亿参数的大语言模型，一定会遇到那个熟悉的报错：CUDA out of memory。显存不够用，几乎是每个 AI 工程师成长路上的「成年礼」。而当你尝试去解决这个问题时，很快就会发现——数据并行、模型并行、流水线并行、梯度累积……每一个概念都能让你啃上好几天。

FairScale 就是为了解决这个痛点而生的。它是 Meta（原 Facebook）旗下的 facebookresearch 团队开源的 PyTorch 扩展库，专门为大规模、高性能训练场景打造。简单来说，它提供了一系列「显存不够用时的瑞士军刀」，让你用几行代码就能把训练规模从单机单卡扩展到多机多卡，而不用从零手写分布式训练逻辑。

这个项目在 GitHub 上拥有超过 3400 颗星，被广泛应用在 OPT、LLaMA 等知名大模型的训练流程中，是工业级大模型训练基础设施的重要组成部分。

背景：Meta 为什么要自己做一套训练扩展库

2020 年左右，Meta 正在训练 GPT-3 级别的大模型（后来的 OPT 系列），遇到了一个普遍问题：PyTorch 官方提供的数据并行（DataParallel）方案在多卡场景下效率不高，尤其是梯度同步的开销随着卡数增加而急剧上升。与此同时，PyTorch 自己的 DistributedDataParallel（DDP）虽然功能完整，但 API 较为底层，开发者需要手动处理很多细节。

在内部积累了一系列实验性方案后，Meta 决定将其中最成熟、最通用的部分整理成一个独立的开源库，也就是 FairScale。它的目标不是替代 PyTorch DDP，而是在 DDP 之上提供更高级的抽象，让研究人员和工程师能更专注于模型本身，而不是底层并行细节。

FairScale 的设计哲学是「模块化」——每个功能都是独立的，可以单独使用，也可以组合使用。这种设计让它既有足够的灵活性，又不会引入过多的复杂度。

核心功能：四大并行技术，一文读懂

FairScale 的功能可以归纳为四个方向，每个方向都对应一种在大模型训练中常用的并行策略：

1. 优化器并行（Optimizer Sharding / OSS）

fairscale.optim.OSS 是 FairScale 最核心的功能之一。在多卡训练中，传统的做法是每张卡都保存一份完整的优化器状态（比如 Adam 的 m/v 矩阵），这在模型很大时会造成严重的显存浪费。OSS 的思路很简单：每个卡只保存一部分参数的优化器状态，通过通信协调来保证参数更新的一致性。

这相当于把原来「每个人背一整本字典」的模式，改成了「每人背字典的一页，大家共享」的模式。实验数据显示，OSS 可以将优化器显存占用减少约 30-40%，同时基本不影响训练收敛速度。

2. 流水线并行（Pipeline Parallelism / Pipe）

对于超大型模型，单卡甚至多卡都放不下整个模型时，需要把模型的不同层放到不同的 GPU 上。FairScale 的 Pipe 模块实现了流水线并行，将模型切分后按阶段（stage）分配到不同设备上，通过微批次（micro-batch）的流水线调度来提高设备利用率。

不过需要注意的是，FairScale 的 Pipe 实现相对基础（属于 1F1B 即「一个前向、一个反向」的经典实现），对于更复杂的需求（如自动并行搜索），建议结合 Megatron-LM 或 DeepSpeed 使用。

3. 模型并行（Model Parallel）

fairscale.nn.model_parallel 包含了对张量并行（Tensor Parallelism）的支持，底层实现了分片（shard）和 all-gather 操作，可以将单个层的参数矩阵分布到多个 GPU 上。相比流水线并行，张量并行的粒度更细，但通信开销也更大。

4. 激活检查点（Activation Checkpointing）

fairscale.nn.checkpoint 实现了激活重计算技术。这是一种用时间换显存的方法：在反向传播时，不保存所有中间激活值，而是重新计算一部分关键激活值。启用后，显存占用可以显著降低，但每个训练步骤的计算时间会有所增加。这是大模型训练的常见技巧，FairScale 提供了开箱即用的封装。

代码示例：OSS 优化器使用

FairScale 的 API 设计非常友好，以下是一个典型的使用示例：

import torch
from torch.utils.data import DataLoader
from fairscale.optim import OSS

# 定义模型和标准优化器
model = MyLargeModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 替换为 OSS 优化器，自动实现优化器状态分片
optimizer = OSS(
    params=model.parameters(),
    optim=torch.optim.Adam,
    lr=1e-4,
    # weight_decay=1e-4  # 可传入 Adam 原始参数
)

# 之后就是标准的 DDP 训练循环
model = torch.nn.parallel.DistributedDataParallel(model)
for batch in DataLoader(dataset):
    optimizer.zero_grad()
    loss = model(batch).loss
    loss.backward()
    optimizer.step()

只需要把标准优化器包一层 OSS()，就能自动获得优化器状态分片的能力。这种「渐进式增强」的设计思路贯穿整个 FairScale 库，用户不需要大幅改写现有训练代码。

局限与注意事项

FairScale 并非银弹，使用时有几个关键注意点：

项目活跃度较低：FairScale 的最后一个 release（v0.4.13）发布于 2022 年 12 月，距今已有三年多。虽然代码库本身仍可正常使用，但新功能的开发和 bug 修复已经明显放缓。如果你正在启动一个新的大模型训练项目，建议优先考虑 DeepSpeed 或 PyTorch FSDP（FullyShardedDataParallel）——后者已经在 PyTorch 2.0 中成为官方功能，且得到了持续活跃维护。

Pipe 模块较基础：FairScale 的流水线并行实现比较简单，不支持高级的调度策略（如序列并行、异步执行等）。对于需要精细控制流水线调度的场景，Deepspeed Pipeline Parallelism 或 PyTorch Pipeline Parallel 是更好的选择。

与 PyTorch FSDP 的重叠：PyTorch 2.0 引入了原生 FSDP（Fully Sharded Data Parallel），与 FairScale 的 OSS 功能有较大重叠。Meta 内部也在逐步将依赖从 FairScale 迁移到 PyTorch 原生方案。因此，如果你的 PyTorch 版本 >= 2.0，使用原生 FSDP 通常是更推荐的路径。

CUDA 版本依赖：FairScale 需要 CUDA 11.0 及以上版本，如果你的 GPU 较老（如 Pascal 架构的 GTX 1080），可能需要特殊处理。

行业意义：训练基础设施的「中间件」

尽管活跃度有所下降，FairScale 在大模型训练基础设施的发展历程中仍占有重要地位。它代表了 2019-2022 年间大模型训练从手工作坊走向工程化、自动化的一个重要节点。

OPT（Open Pre-trained Transformers）系列模型——Meta 开源的对标 GPT-3 的模型——就是使用 FairScale 训练的。这让研究社区第一次有机会在不开源训练代码的情况下，仍能复现 Meta 的训练流程并进行二次研究。

FairScale 的另一个重要贡献是教育意义：它的源码质量较高、文档相对完善，是理解分布式训练各个概念的绝佳参考实现。很多后来学习大模型训练工程的工程师，都是通过阅读 FairScale 的源码才真正理解了什么是流水线并行、什么是优化器分片。

增长趋势：FairScale 的 Star 曲线在 2022 年中达到高峰（与 OPT 发布同期），之后随着 DeepSpeed 和 PyTorch FSDP 的崛起而趋于平稳。这条曲线本身就是一个很好的「技术生命周期」观察样本。

总结

FairScale 是一款定位清晰、技术扎实的 PyTorch 训练加速库。它用相对轻量的方式，解决了大模型训练中的几个核心痛点：显存不够、扩展困难、优化器状态臃肿。对于仍在使用 PyTorch < 2.0 或需要快速集成 OSS 功能的团队，它仍然是一个可靠的选择。

但对于新项目，建议优先评估 PyTorch 原生 FSDP 或 DeepSpeed。毕竟在大模型时代，训练基础设施的迭代速度非常快，选择有活跃社区支持的方案，通常能省去很多「踩坑」的时间。