fairscale
Meta开源的PyTorch大规模训练扩展库,提供优化器分片、流水线并行、张量并行和激活检查点等高级
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
Meta开源的PyTorch大规模训练扩展库,提供优化器分片、流水线并行、张量并行和激活检查点等高级
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
如果你曾经训练过 GPT、OPT 或任何超过 10 亿参数的大语言模型,一定会遇到那个熟悉的报错:CUDA out of memory。显存不够用,几乎是每个 AI 工程师成长路上的「成年礼」。而当你尝试去解决这个问题时,很快就会发现——数据并行、模型并行、流水线并行、梯度累积……每一个概念都能让你啃上好几天。
FairScale 就是为了解决这个痛点而生的。它是 Meta(原 Facebook)旗下的 facebookresearch 团队开源的 PyTorch 扩展库,专门为大规模、高性能训练场景打造。简单来说,它提供了一系列「显存不够用时的瑞士军刀」,让你用几行代码就能把训练规模从单机单卡扩展到多机多卡,而不用从零手写分布式训练逻辑。
这个项目在 GitHub 上拥有超过 3400 颗星,被广泛应用在 OPT、LLaMA 等知名大模型的训练流程中,是工业级大模型训练基础设施的重要组成部分。
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 提供了开箱即用的封装。
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。毕竟在大模型时代,训练基础设施的迭代速度非常快,选择有活跃社区支持的方案,通常能省去很多「踩坑」的时间。
| 维度 | 评分 | 说明 |
|---|---|---|
| 功能完整度 | ★★★★☆ | 覆盖主流并行策略,但部分模块较基础 |
| 文档质量 | ★★★★★ | Facebook 官方出品,文档详尽 |
| 活跃度 | ★★★☆☆ | 更新放缓,建议关注 PyTorch FSDP |
| 易用性 | ★★★★★ | API 设计优雅,与 PyTorch 无缝集成 |
| 工业应用 | ★★★★☆ | OPT 等大模型验证过,可靠性高 |