Flux项目安装与配置指南

Flux项目安装与配置指南

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

1. 项目基础介绍

Flux是一个为GPU上的张量/专家并行计算提供高效通信重叠的库。它旨在为密集型或混合专家（MoE）模型提供高性能且可插拔的内核，以支持模型训练/推理中的各种并行策略。Flux的高效内核与PyTorch兼容，并且可以轻松集成到现有的框架中，支持多种NVIDIA GPU架构和数据类型。

主要编程语言包括C++、CUDA和Python。

2. 项目使用的关键技术和框架

• NCCL（NVIDIA Collective Communications Library）：用于在多个GPU之间进行集体通信的库。
• CUTLASS：用于生成高性能GEMM（矩阵乘法）内核的库。
• NVSHMEM：NVIDIA提供的一种用于GPU共享内存的库，对于MoE内核至关重要。

3. 安装和配置准备工作

在开始安装Flux之前，请确保您的系统满足以下要求：

• 操作系统：Linux
• Python版本：3.10或3.11
• CUDA版本：12.4
• PyTorch版本：2.4.0、2.5.0或2.6.0

详细安装步骤

克隆项目

首先，您需要克隆Flux项目到本地：

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

安装依赖

在克隆完项目后，进入项目目录并安装依赖：

cd flux
bash ./install_deps.sh

编译安装

根据您的GPU架构（例如，Ampere(sm80)、Ada Lovelace(sm89)或Hopper(sm90)），运行相应的编译命令：

对于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。以下是在虚拟环境中安装Flux的示例：

conda create -n flux python=3.11
conda activate flux
pip3 install packaging 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安装

如果您的所需版本有可用的预构建wheel包，您也可以直接使用pip进行安装：

pip install byte-flux

确保在安装之前已经安装了PyTorch。

以上就是Flux项目的详细安装和配置指南。按照以上步骤操作，您应该能够成功安装并开始使用Flux库。

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