您可以通过以下渠道进一步了解我们👇

Flux：细粒度计算-通信重叠的 GPU 内核库

Flux 是一个用于 GPU 上密集模型和 MoE 模型的通信重叠库，提供高性能、可插拔的内核，以支持模型训练和推理中的各种并行策略。

Flux 的高效内核与 PyTorch 兼容，可以轻松集成到现有框架中，支持多种 NVIDIA GPU 架构和数据类型。

新闻

[2025/03/10]🔥我们发布了 COMET：面向专家混合模型的计算-通信重叠。

快速入门

您可以从源代码或 PyPI 安装 Flux。

从源代码安装

git clone --recursive https://github.com/bytedance/flux.git && cd flux

# 对于 Ampere (sm80) GPU
./build.sh --arch 80 --nvshmem
# 对于 Ada Lovelace (sm89) GPU
./build.sh --arch 89 --nvshmem
# 对于 Hopper (sm90) GPU
./build.sh --arch 90 --nvshmem

在虚拟环境中安装

以下是在虚拟环境中安装 Flux 的示例。让我们在一个配备 CUDA 12.4、torch 2.6.0 和 python 3.11 的虚拟环境中完成安装。

conda create -n flux python=3.11
conda activate flux
pip3 install packaging
pip3 install ninja
pip3 install torch==2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

./build.sh --clean-all
./build.sh --arch "80;89;90" --nvshmem --package

随后，您将在 dist/ 文件夹下找到适合您的虚拟环境的 wheel 包。

从 PyPI 安装

我们还提供了 Flux 的预构建 wheel 包，如果所需的版本可用，您可以直接使用 pip 进行安装。目前我们为以下配置提供了 wheel 包：torch（2.4.0、2.5.0、2.6.0）、python（3.10、3.11）、cuda（12.4）。

# 确保已安装 PyTorch。
pip install byte-flux

自定义安装

源码安装的构建选项

1. 添加 --nvshmem 以构建支持 NVSHMEM 的 Flux。这对于 MoE 内核至关重要。
2. 如果您对 cmake 流程感到厌倦，可以将环境变量 FLUX_BUILD_SKIP_CMAKE 设置为 1，以跳过 cmake，前提是 build/CMakeCache.txt 已存在。
3. 如果您想构建 wheel 包，请在构建命令中添加 --package。生成的 wheel 文件将位于 dist/ 目录下。

依赖项

Flux 依赖于 NCCL 和 CUTLASS，它们位于 3rdparty/ 目录下，以及 NVSHMEM，后者可以通过 pip 安装。

1. NCCL：由 git submodule 自动管理。
2. NVSHMEM：建议您通过 pip install nvidia-nvshmem-cu12 安装 nvshmem；如果您想从源代码构建 nvshmem，可以从 https://developer.nvidia.com/nvshmem 下载。Flux 已在 nvshmem 3.2.5/3.3.9 上进行了测试。
3. CUTLASS：Flux 利用 CUTLASS 生成高性能 GEMM 内核。我们目前使用 CUTLASS 4.0.0。

快速开始

以下是成功安装 Flux 后运行一些基本演示的命令。

# 仅 GEMM
python3 test/python/gemm_only/test_gemm_only.py 4096 12288 6144 --dtype=float16

# 全归约与 GEMM 融合（密集 MLP 第 0 层）
./launch.sh test/python/ag_gemm/test_ag_kernel.py 4096 49152 12288 --dtype=float16 --iters=10

# GEMM 与规约-散射融合（密集 MLP 第 1 层）
./launch.sh test/python/gemm_rs/test_gemm_rs.py 4096 12288 49152 --dtype=float16 --iters=10

# 全归约与分组 GEMM 融合（MoE MLP 第 0 层）
./launch.sh test/python/moe_ag_scatter/test_moe_ag.py

# 分组 GEMM 与规约-散射融合（MoE MLP 第 1 层）
./launch.sh test/python/moe_gather_rs/test_moe_gather_rs.py

您还可以查阅文档以获取更多详细信息！

• 关于 MoE 内核的更详细用法，请参阅 Flux MoE 使用指南。您可以尝试一些 示例 作为快速入门。使用 Flux，只需几十行代码即可实现一个 最小化 MoE 层！
• 如需查看性能数据，请参阅 性能文档。
• 若要深入了解 Flux 的设计原则，请参阅 设计文档。

许可证

Flux 项目采用 Apache License v2.0 许可证。

引用

如果您在科学出版物中使用 Flux，我们鼓励您在相关论文中添加以下参考文献：

@misc{chang2024flux,
      title={FLUX：通过内核融合实现 GPU 上快速软件级通信重叠},
      author={Li-Wen Chang、Wenlei Bao、Qi Hou、Chengquan Jiang、Ningxin Zheng、Yinmin Zhong、Xuanrun Zhang、Zuquan Song、Ziheng Jiang、Haibin Lin、Xin Jin、Xin Liu},
      year={2024},
      eprint={2406.06858},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

@misc{zhang2025comet,
      title={Comet：面向专家混合模型的细粒度计算-通信重叠},
      author={Shulai Zhang、Ningxin Zheng、Haibin Lin、Ziheng Jiang、Wenlei Bao、Chengquan Jiang、Qi Hou、Weihao Cui、Size Zheng、Li-Wen Chang、Quan Chen 和 Xin Liu},
      year={2025},
      eprint={2502.19811},
      archivePrefix={arXiv},
      primaryClass={cs.DC}
}

参考资料

• ArXiv 论文（Flux）
• ArXiv 论文（Comet）

关于字节跳动 Seed 团队

字节跳动 Seed 团队成立于 2023 年，致力于打造业界最先进的人工智能基础模型。团队立志成为世界一流的研究团队，为科学与社会的进步作出重大贡献。

Flux 快速上手指南

Flux 是由字节跳动 Seed 团队开源的高性能 GPU 内核库，专为密集模型（Dense）和混合专家模型（MoE）设计。它通过细粒度的计算 - 通信重叠技术，显著提升模型训练和推理的并行效率，并兼容 PyTorch 生态。

环境准备

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

• 操作系统: Linux
• GPU 架构: 支持 NVIDIA Ampere (sm80), Ada Lovelace (sm89), 或 Hopper (sm90) 架构。
• 软件依赖:
  • Python 3.10 或 3.11
  • CUDA 12.4 (推荐)
  • PyTorch 2.4.0 / 2.5.0 / 2.6.0
  • CMake, Ninja
• 核心依赖库:
  • NCCL: 自动通过 git submodule 管理。
  • CUTLASS: Flux 使用 CUTLASS 4.0.0 生成高性能 GEMM 内核。
  • NVSHMEM: MoE 内核必需。建议通过 pip 安装：

    pip install nvidia-nvshmem-cu12
    

安装步骤

您可以选择从源码编译或使用预构建的 Wheel 包进行安装。

方式一：从 PyPI 安装（推荐）

如果您使用的 PyTorch 和 Python 版本在支持列表中（torch 2.4.0-2.6.0, python 3.10-3.11, cuda 12.4），可直接使用 pip 安装：

# 确保已安装 PyTorch
pip install byte-flux

方式二：从源码编译

如果需要自定义架构支持或使用最新代码，请从源码编译。

1. 克隆仓库

   git clone --recursive https://github.com/bytedance/flux.git && cd flux
   

2. 执行编译 根据您的 GPU 架构选择对应的编译命令（需添加 --nvshmem 以支持 MoE）：

   # 对于 Ampere (sm80) GPU
   ./build.sh --arch 80 --nvshmem
   
   # 对于 Ada Lovelace (sm89) GPU
   ./build.sh --arch 89 --nvshmem
   
   # 对于 Hopper (sm90) GPU
   ./build.sh --arch 90 --nvshmem
   

   提示: 若需在虚拟环境中构建 Wheel 包，可先配置好 conda 环境（如 Python 3.11 + Torch 2.6.0 + CUDA 12.4），然后运行：

   ./build.sh --clean-all
   ./build.sh --arch "80;89;90" --nvshmem --package
   

   生成的 .whl 文件将位于 dist/ 目录下。

基本使用

安装完成后，您可以运行官方提供的测试脚本来验证功能并体验基础算子。以下是最简单的使用示例：

1. 基础 GEMM 测试

仅测试矩阵乘法内核：

python3 test/python/gemm_only/test_gemm_only.py 4096 12288 6144 --dtype=float16

2. 密集模型算子测试

测试 All-Gather 与 GEMM 融合（对应 Dense MLP 第 0 层）：

./launch.sh test/python/ag_gemm/test_ag_kernel.py 4096 49152 12288 --dtype=float16 --iters=10

测试 GEMM 与 Reduce-Scatter 融合（对应 Dense MLP 第 1 层）：

./launch.sh test/python/gemm_rs/test_gemm_rs.py 4096 12288 49152 --dtype=float16 --iters=10

3. MoE 模型算子测试

测试 All-Gather 与 Grouped GEMM 融合（对应 MoE MLP 第 0 层）：

./launch.sh test/python/moe_ag_scatter/test_moe_ag.py

测试 Grouped GEMM 与 Reduce-Scatter 融合（对应 MoE MLP 第 1 层）：

./launch.sh test/python/moe_gather_rs/test_moe_gather_rs.py

进阶使用: 如需在项目中集成 MoE 层，可参考 examples/moe_flux_only.py，仅需数十行代码即可实现一个最小化的 MoE 层。详细文档请参阅项目仓库中的 docs/ 目录。