描述

FairScale 是一个用于高性能和大规模训练的 PyTorch 扩展库。该库在扩展基础 PyTorch 功能的同时，还引入了最新的 SOTA 扩展技术。FairScale 以可组合的模块和易于使用的 API 形式提供了最先进的分布式训练技术。这些 API 是研究人员在资源有限的情况下尝试扩展模型时不可或缺的工具。

FairScale 在设计时秉持以下价值观：

• 易用性 - 用户应能够在最小认知负担下理解和使用 FairScale 的 API。

• 模块化 - 用户应能够无缝地将多个 FairScale API 组合到他们的训练循环中。

• 性能 - FairScale 的 API 在扩展性和效率方面均表现出色。

观看入门视频

安装

要安装 FairScale，请参阅以下说明。您可以通过 pip 或 conda 安装软件包，也可以直接从源代码构建。

入门

完整的文档包含入门指南、深入解析以及关于各种 FairScale API 的教程。

FSDP

FullyShardedDataParallel (FSDP) 是扩展大型神经网络模型的推荐方法。该库已被合并到 PyTorch 主干。此处的 FSDP 版本既可用于历史参考，也可用于研究扩展技术中的新奇想法。请参阅以下博客，了解如何使用 FairScale FSDP 及其工作原理。

测试

我们使用 CircleCI 对 FairScale 进行测试，测试环境包括以下 PyTorch 版本（搭配 CUDA 11.2）：

• 最新的稳定版（例如 1.10.0）
• 最新的长期支持版（例如 1.8.1）
• 最近的每日构建版（例如 1.11.0.dev20211101+cu111）

如果您在安装过程中遇到问题，请创建一个issue。

贡献者

我们欢迎各位贡献！请参阅CONTRIBUTING说明，了解如何为 FairScale 做出贡献。

许可证

FairScale 采用BSD-3-Clause 许可证。

fairscale.nn.pipe 源自 torchgpipe，版权归属 Kakao Brain，2019 年，采用 Apache 许可证。

fairscale.nn.model_parallel 源自 Megatron-LM，版权归属 NVIDIA CORPORATION，2020 年，采用 Apache 许可证。

fairscale.optim.adascale 源自 AdaptDL，版权归属 Petuum, Inc.，2020 年，采用 Apache 许可证。

fairscale.nn.misc.flatten_params_wrapper 源自 PyTorch-Reparam-Module，版权归属 Tongzhou Wang，2018 年，采用 MIT 许可证。

引用 FairScale

如果您在论文或其他出版物中使用 FairScale，请使用以下 BibTeX 条目进行引用。

@Misc{FairScale2021,
  author =       {{FairScale authors}},
  title =        {FairScale:  A general purpose modular PyTorch library for high performance and large scale training},
  howpublished = {\url{https://github.com/facebookresearch/fairscale}},
  year =         {2021}
}

FairScale 快速上手指南

FairScale 是一个用于高性能和大规模训练的 PyTorch 扩展库，提供了最新的分布式训练技术（如 FSDP），帮助研究者在有限资源下扩展模型规模。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.7 或更高版本
• PyTorch: 已安装兼容版本的 PyTorch (建议 1.8+，需包含 CUDA 支持以进行 GPU 训练)
• CUDA: 如需使用 GPU 加速，请确保已正确安装 NVIDIA 驱动和 CUDA Toolkit (FairScale 测试环境通常为 CUDA 11.2+)

安装步骤

您可以通过 pip 直接安装稳定版，或者从源码构建以获取最新功能。

方式一：使用 pip 安装（推荐）

pip install fairscale

国内用户加速建议： 如果您在国内网络环境下安装较慢，推荐使用清华或阿里镜像源：

pip install fairscale -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

如果您需要最新的功能或进行开发贡献：

git clone https://github.com/facebookresearch/fairscale.git
cd fairscale
pip install -e .

基本使用

FairScale 的核心功能是作为可组合的模块嵌入到您的现有 PyTorch 训练循环中。目前最常用且推荐的功能是 FSDP (FullyShardedDataParallel)，用于将大模型分片到多个 GPU 上进行训练。

以下是一个最简单的 FSDP 使用示例，展示如何将普通模型包装为分片并行模型：

import torch
import torch.nn as nn
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP

# 1. 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Linear(1024, 10)
        )

    def forward(self, x):
        return self.net(x)

# 2. 初始化模型和优化器
model = SimpleModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 3. 使用 FSDP 包装模型
# 注意：需要在 CUDA 环境下运行，并设置正确的 device_id
if torch.cuda.is_available():
    model = FSDP(model)
    model.to(torch.cuda.current_device())

# 4. 正常的训练循环 (与普通 PyTorch 代码无异)
dummy_input = torch.randn(32, 1024).cuda()
target = torch.randn(32, 10).cuda()

output = model(dummy_input)
loss = nn.MSELoss()(output, target)

loss.backward()
optimizer.step()
optimizer.zero_grad()

print("FSDP 训练步骤执行成功！")

说明：

• FullyShardedDataParallel 会自动处理参数分片、梯度分片和优化器状态分片，显著降低显存占用。
• 除了 FSDP，FairScale 还提供了 Pipe (流水线并行) 和 AdaScale (自适应批大小缩放) 等模块，具体用法可参考官方完整文档。