Flux 项目使用与启动教程

Flux 项目使用与启动教程

flux A fast communication-overlapping library for tensor/expert parallelism on GPUs. 项目地址: https://gitcode.com/gh_mirrors/flux17/flux

1. 项目介绍

Flux 是一个面向 GPU 的细粒度计算通信重叠库，旨在为张量/专家并行提供高性能的可插拔内核支持。该库与 PyTorch 兼容，并能轻松集成到现有框架中，支持多种 Nvidia GPU 架构和数据类型。Flux 由 ByteDance Seed 团队开发，致力于提升模型训练和推理的性能。

2. 项目快速启动

安装依赖

首先，需要安装 Flux 的依赖项。在安装之前，确保你的系统中已安装了 CUDA。

git clone --recursive https://github.com/bytedance/flux.git
cd flux
bash ./install_deps.sh

编译安装

根据你的 GPU 架构（例如：80 为 Ampere，89 为 Ada Lovelace，90 为 Hopper），编译相应的版本：

./build.sh --arch 80 --nvshmem
# 或者
./build.sh --arch 89 --nvshmem
# 或者
./build.sh --arch 90 --nvshmem

如果你希望在虚拟环境中安装，可以按照以下步骤操作：

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 安装

如果你想从 PyPI 安装预编译的 wheel 包，可以使用 pip：

pip install byte-flux

确保你的 PyTorch 版本与提供的 wheel 包兼容。

3. 应用案例和最佳实践

以下是一些使用 Flux 的基础示例：

• 仅 GEMM 操作

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

• all-gather 与 GEMM 融合（密集 MLP 层0）

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

• GEMM 与 reduce-scatter 融合（密集 MLP 层1）

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

• all-gather 与分组 GEMM 融合（MoE MLP 层0）

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

• 分组 GEMM 与 reduce-scatter 融合（MoE MLP 层1）

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

更多详细用法和性能数据，请参考官方文档。

4. 典型生态项目

目前，Flux 主要用于 ByteDance 内部的深度学习项目中，特别是在需要利用 GPU 进行大规模并行计算的场景。由于它的高性能和易用性，它有潜力成为各种深度学习框架的一个重要组成部分。未来，我们期待看到更多开源项目和社区采用 Flux，共同推动 GPU 计算的发展。

flux A fast communication-overlapping library for tensor/expert parallelism on GPUs. 项目地址: https://gitcode.com/gh_mirrors/flux17/flux