flux

You can get to know us better through the following channels

# Flux: Fine-grained Computation-communication Overlapping GPU Kernel Library

Flux is a communication-overlapping library for dense/MoE models on GPUs, providing high-performance and pluggable kernels to support various parallelisms in model training/inference. Flux's efficient kernels are compatible with Pytorch and can be integrated into existing frameworks easily, supporting various Nvidia GPU architectures and data types. # News [2025/03/10]We have released **COMET: Computation-communication Overlapping for Mixture-of-Experts**. ## Getting started Install Flux either from source or from PyPI. ### Install from Source ```bash git clone --recursive https://github.com/bytedance/flux.git && cd flux # For Ampere(sm80) GPU ./build.sh --arch 80 --nvshmem # For Ada Lovelace(sm89) GPU ./build.sh --arch 89 --nvshmem # For Hopper(sm90) GPU ./build.sh --arch 90 --nvshmem ``` #### Install in a virtual environment Here is a snippet to install Flux in a virtual environment. Let's finish the installation in an virtual environment with CUDA 12.4, torch 2.6.0 and python 3.11. ```bash 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 ``` Then you would expect a wheel package under `dist/` folder that is suitable for your virtual environment. ### Install from PyPI We also provide some pre-built wheels for Flux, and you can directly install with pip if your wanted version is available. Currently we provide wheels for the following configurations: torch(2.4.0, 2.5.0, 2.6.0), python(3.10, 3.11), cuda(12.4). ```bash # Make sure that PyTorch is installed. pip install byte-flux ``` ### Customized Installation #### Build options for source installation 1. Add `--nvshmem` to build Flux with NVSHMEM support. It is essential for the MoE kernels. 2. If you are tired of the cmake process, you can set environment variable `FLUX_BUILD_SKIP_CMAKE` to 1 to skip cmake if `build/CMakeCache.txt` already exists. 3. If you want to build a wheel package, add `--package` to the build command. find the output wheel file under dist/ #### Dependencies Flux depends on `NCCL` and `CUTLASS`, which are located under `3rdparty/`, and `NVSHMEM`, which you can install by pip. 1. NCCL: Managed by git submodule automatically. 2. NVSHMEM: It's suggested that you install nvshmem by `pip install nvidia-nvshmem-cu12`; If you want to build nvshmem from source, you can download it from https://developer.nvidia.com/nvshmem. Flux is tested with nvshmem 3.2.5/3.3.9 3. CUTLASS: Flux leverages CUTLASS to generate high-performance GEMM kernels. We currently use CUTLASS 4.0.0 ## Quick Start Below are commands to run some basic demos once you have installed Flux successfully. ``` … ``` You can check out the documentations for more details! * For a more detailed usage on MoE kernels, please refer to [Flux MoE Usage](https://github.com/bytedance/flux/blob/main/docs/moe_usage.md). Try some [examples](https://github.com/bytedance/flux/blob/main/examples) as a quick start. A [minimal MoE layer](https://github.com/bytedance/flux/blob/main/examples/moe_flux_only.py) can be implemented within only a few tens of lines of code using Flux! * For some performance numbers, please refer to [Performance Doc](https://github.com/bytedance/flux/blob/main/docs/performance.md). * To learn more about the design principles of Flux, please refer to [Design Doc](https://github.com/bytedance/flux/blob/main/docs/design.md). ## [License](./LICENSE) The Flux Project is under the Apache License v2.0. ## Citation If you use Flux in a scientific publication, we encourage you to add the following reference to the related papers: ``` … ``` ## Reference * [ArXiv Paper (Flux)](http://arxiv.org/abs/2406.06858) * [ArXiv Paper (Comet)](https://arxiv.org/abs/2502.19811) # About [ByteDance Seed Team](https://team.doubao.com/) Founded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society.