DOCSEARCHbyAlgolia
Submit Docs

Tutel MoE Implementation

repository·main·Indexed 21 days ago

https://github.com/microsoft/tutel

An optimized Mixture-of-Experts (MoE) implementation for high-performance training and inference. Tutel provides specialized kernels and parallel solutions for models such as DeepSeek, Qwen, Kimi, and GLM-5.x, supporting PyTorch (>= 2.0) across CUDA, ROCm, and CPU environments. It features support for various precisions including NVFP4, MXFP4, and BlockwiseFP8, and offers dynamic configuration for parallelism, sparsity, and capacity.

Tokens
7.8K
Snippets
17
Records
25
Agent score
77%

What's inside Tutel

  1. Overview of Tutel MoE

    main

    Tutel is an optimized Mixture-of-Experts (MoE) implementation designed for modern training and inference with dynamic behaviors. It features "No-penalty Parallelism/Sparsity/Capacity/.. Switching" to handle dynamic behaviors efficiently.

    Supported Environments:

    • Framework: PyTorch (version >= 2.0 recommended).
    • GPUs: CUDA (fp64, fp32, fp16, bf16) and ROCm (fp64, fp32, fp16, bf16).
    • CPU: fp64, fp32.
    • Inference Support: Direct NVFP4, MXFP4, and BlockwiseFP8 inference for MoE-based models including GLM-5.x, DeepSeek-3.x, Kimi-2.x, Qwen3, and GptOSS on hardware like A100, A800, H100, and MI300.

  2. Configure Megablocks for MoE Decoder Inference

    main

    Megablocks can be used to improve decoder inference on a single GPU when num_local_expert >= 2. You can control the megablocks_size parameter in the forward method of the MoE layer. Setting megablocks_size=0 disables it.

    Available options for megablocks_size:

    • 0: Disabled (uses BatchMatmul)
    • 1: Megablocks with block_size = 1
    • 2: Megablocks with block_size = 2
    # Control the switch of megablocks_size (0 for disabled)
self._moe_layer.forward(x, .., megablocks_size=1)

  3. Implementation steps for MoE training in NanoGPT

    main

    To implement MoE training within a transformer model using Tutel, follow these four core steps:

    1. Define the MoE layer: Integrate the moe_layer into your transformer architecture, ensuring you specify the expected weight dtype (e.g., torch.bfloat16).
    2. Broadcast shared parameters: Implement logic to broadcast shared parameter weights to all other GPUs while explicitly preventing the broadcasting of non-shared parameters.
    3. Synchronize gradients: During a single training step, perform an all-reduce on the gradients of all shared parameters, while ensuring non-shared parameters are excluded from the all-reduce operation.
    4. Apply gradients: Call opt.step() to apply the distributed gradients to the model parameters.

  4. Install Tutel

    main

    You can install Tutel either online via pip or by building from source.

    Prerequisites: Prepare PyTorch

    Ensure you have PyTorch >= 2.0.0 installed. Use the appropriate index URL for your hardware:

    • NVIDIA CUDA >= 11.7 (Windows/Linux): python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    • AMD ROCm >= 6.2.2 (Linux): python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2.2
    • CPU (Windows/Linux): python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

    Installation Options

    Option 1: Install Online

    python3 -m pip uninstall tutel -y
python3 -m pip install -v -U --no-build-isolation git+https://github.com/microsoft/tutel@main

    Option 2: Build from Source

    git clone https://github.com/microsoft/tutel --branch main
python3 -m pip uninstall tutel -y
python3 ./tutel/setup.py install --user
    python3 -m pip install -v -U --no-build-isolation git+https://github.com/microsoft/tutel@main

  5. Run MoE or Dense training with the modded-nanogpt example

    main

    This example demonstrates how to replace Dense layers with MoE (Mixture of Experts) layers in a NanoGPT architecture. You can toggle between MoE training and standard Dense training using the USE_MOE environment variable when executing the run.sh script.

    To run MoE training on 8 local GPUs:

    USE_MOE=1 ./run.sh

    To run standard Dense training on 8 local GPUs:

    USE_MOE=0 ./run.sh
    # run MoE training with 8 local GPUs
USE_MOE=1 ./run.sh

# run Dense training with 8 local GPUs
USE_MOE=0 ./run.sh

  6. Run Tutel Quick Tests

    main

    Tutel provides several example scripts to verify installation and test different distribution strategies.

    Single Device / GPU Tests

    Run these to test Tutel-optimized MoE on a single device:

    • python3 -m tutel.examples.helloworld --batch_size=16 (Manual distribution)
    • python3 -m tutel.examples.helloworld_ddp --batch_size=16 (PyTorch DDP distribution)
    • python3 -m tutel.examples.helloworld_ddp_tutel --batch_size=16 (Tutel DDP distribution with ZeRO on optimizers)
    • python3 -m tutel.examples.helloworld_amp --batch_size=16 (AMP data type + manual distribution)
    • python3 -m tutel.examples.helloworld_custom_gate_expert --batch_size=16 (Custom gate/expert layer)
    • python3 -m tutel.examples.moe_mnist (End-to-end MNIST)
    • python3 -m tutel.examples.moe_cifar10 (End-to-end CIFAR10)

    Multi-GPU (Single Machine)

    To test 8 GPUs on one machine:

    python3 -m torch.distributed.run --nproc_per_node=8 -m tutel.examples.helloworld --batch_size=16

    Multi-Node (Multiple Machines)

    Option A: Torch Launcher

    # On Node 0
ssh <node-ip-0> python3 -m torch.distributed.run --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=<node-ip-0> -m tutel.examples.helloworld --batch_size=16

# On Node 1
ssh <node-ip-1> python3 -m torch.distributed.run --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=<node-ip-0> -m tutel.examples.helloworld --batch_size=16

    Option B: Tutel Launcher (requires openmpi-bin)

    # Single Node MPI launch
mpiexec -bind-to none -host localhost -x LOCAL_SIZE=8 python3 -m tutel.launcher.run -m tutel.examples.helloworld_ddp_tutel --batch_size=16

# Multi-Node MPI launch for GPU
mpiexec -bind-to none -host <node-ip-0>,<node-ip-1>,.. -x MASTER_ADDR=<node-ip-0> -x LOCAL_SIZE=8 python3 -m tutel.launcher.run -m tutel.examples.helloworld --batch_size=16

# Multi-Node MPI launch for CPU
mpiexec -bind-to none -host localhost -x LOCAL_SIZE=1 -x OMP_NUM_THREADS=1024 python3 -m tutel.launcher.run -m tutel.examples.helloworld --batch_size=16 --device cpu
    python3 -m torch.distributed.run --nproc_per_node=8 -m tutel.examples.helloworld --batch_size=16

  7. Serve GLM-5.x Models (Claude-Code Mode)

    main

    To serve GLM-5, 5.1, or 5.2 models using Tutel, follow these steps:

    1. Download Models

    Use huggingface_hub to download the required NVFP4 versions:

    hf download --local-dir nvidia/GLM-5.2-NVFP4 nvidia/GLM-5.2-NVFP4
hf download --local-dir nvidia/GLM-5.1-NVFP4 nvidia/GLM-5.1-NVFP4
hf download --local-dir nvidia/GLM-5-NVFP4 nvidia/GLM-5-NVFP4

    2. Run the Server via Docker

    For Azure A100x8/H100x8/B200x8 SXM:

    docker run -e WORKER=1 -e LOCAL_SIZE=8 -p 8000:8000 -it --rm --ipc=host --shm-size=8g \
    --ulimit memlock=-1 --ulimit stack=67108864 -v /:/host -w /host$(pwd) \
    -v /usr/lib/x86_64-linux-gnu/libcuda.so.1:/usr/lib/x86_64-linux-gnu/libcuda.so.1 --privileged \
    tutelgroup/deepseek-671b:a100x8-chat-20260707 --serve=core \
      --try_path nvidia/GLM-5.2-NVFP4 \
      --try_path nvidia/GLM-5.1-NVFP4 \
      --try_path nvidia/GLM-5-NVFP4 \
      --max_seq_len 1000000

    For Azure MI300x8 PCIe:

    docker run -e WORKER=1 -e LOCAL_SIZE=8 -p 8000:8000 -it --rm --ipc=host --shm-size=8g \
    --ulimit memlock=-1 --ulimit stack=67108864 -v /:/host -w /host$(pwd) \
    --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --device=/dev/kfd --device=/dev/dri --group-add=video \
    tutelgroup/deepseek-671b:mi300x8-chat-20260721 --serve=core \
      --try_path nvidia/GLM-5.2-NVFP4 \
      --try_path nvidia/GLM-5.1-NVFP4 \
      --try_path nvidia/GLM-5-NVFP4 \
      --max_seq_len 1000000

  8. Convert checkpoint files for different distributed world sizes

    main

    Tutel provides tools to gather distributed checkpoint files into a single 'All-in-One' checkpoint and scatter them back into distributed formats optimized for different numbers of GPUs. This is useful when you train a model on one GPU configuration (e.g., 2 GPUs) and want to run it on another (e.g., 8 GPUs).

    Workflow for standard checkpoints:

    1. Gather: Combine distributed shards (e.g., from 2 GPUs) into one file using tutel.checkpoint.gather.
    2. Scatter: Split the All-in-One file into new shards for a different number of GPUs using tutel.checkpoint.scatter.

    Mathematical Relationship:

    When converting, if you have $X$ global experts and $Y$ GPUs, the number of local experts per GPU is calculated as $X / Y$. The tool supports a shorthand where setting num_local_experts to a negative value ($-Y/X$) helps adapt to the target configuration.

    # 1. Gather shards from 2 GPUs into one file
python3 -m tutel.checkpoint.gather --inputs=./states/{rank}-of-{size}.ckpt --input_size=2 --output ./model-all-in-one.ckpt

# 2. Scatter the All-in-One file for 8 GPUs
python3 -m tutel.checkpoint.scatter --input=./model-all-in-one.ckpt --output_size=8 --outputs=./adapted-for-8-gpus/{rank}-of-{size}.ckpt

  9. Serve DeepSeek R1 Chat using Docker on MI300x

    main

    To serve the DeepSeek R1 671B model on 8 MI300 GPUs, follow these three steps:

    1. Download the model: Use huggingface-cli to download the model to a local directory.
    2. Run the Docker container: Use the tutelgroup/deepseek-671b:mi300x8-chat-20250224 image. Ensure you mount the host path and pass the --model_path.
    3. Request a prompt: Use curl to send a POST request to the local port :8000/chat.
    # Step-1: Download Deepseek R1 671B Model
huggingface-cli download deepseek-ai/DeepSeek-R1 --local-dir ./deepseek-ai/DeepSeek-R1

# Step-2: Using 8 MI300 GPUs to Serve Deepseek R1 Chat on Local Port :8000
docker run -it --rm --ipc=host --privileged -p 8000:8000 \
    -v /:/host -w /host$(pwd) tutelgroup/deepseek-671b:mi300x8-chat-20250224 \
    --model_path ./deepseek-ai/DeepSeek-R1

# Step-3: Issue a Prompt Request with curl
curl -X POST http://0.0.0.0:8000/chat -d '{"text": "Calculate the result of: 1 / (sqrt(5) - sqrt(3))"}'

  10. Serve Kimi-K2.6/2.7 and DeepSeek V3.2 (Long-Context Mode)

    main

    To serve long-context models like Kimi or DeepSeek V3.2 using Tutel, follow these steps:

    1. Download Models

    pip3 install -U "huggingface_hub[cli]" --upgrade
hf download moonshotai/Kimi-K2.7-Code --local-dir moonshotai/Kimi-K2.7-Code
hf download moonshotai/Kimi-K2.6 --local-dir moonshotai/Kimi-K2.6
hf download nvidia/Kimi-K2.5-NVFP4 --local-dir nvidia/Kimi-K2.5-NVFP4
hf download nvidia/Kimi-K2-Thinking-NVFP4 --local-dir nvidia/Kimi-K2-Thinking-NVFP4
hf download nvidia/DeepSeek-V3.2-NVFP4 --local-dir nvidia/DeepSeek-V3.2-NVFP4

    2. Run the Server via Docker

    For Azure A100x8/H100x8/B200x8 SXM:

    docker run -e LOCAL_SIZE=8 -e WORKER=1 -it --rm --ipc=host --net=host --shm-size=8g \
    --ulimit memlock=-1 --ulimit stack=67108864 -v /:/host -w /host$(pwd) -v /tmp:/tmp \
    -v /usr/lib/x86_64-linux-gnu/libcuda.so.1:/usr/lib/x86_64-linux-gnu/libcuda.so.1 --privileged \
    tutelgroup/deepseek-671b:a100x8-chat-20260707 --serve=webui --listen_port 8000 \
      --try_path moonshotai/Kimi-K2.7-Code \
      --try_path moonshotai/Kimi-K2.6 \
      --try_path nvidia/Kimi-K2.5-NVFP4 \
      --try_path nvidia/Kimi-K2-Thinking-NVFP4 \
      --try_path nvidia/DeepSeek-V3.2-NVFP4 \
      --try_path nvidia/DeepSeek-R1-NVFP4 \
      --max_seq_len 16384

    For Azure MI300x8 PCIe:

    docker run -e LOCAL_SIZE=8 -e WORKER=1 -it --rm --ipc=host --net=host --shm-size=8g \
    --ulimit memlock=-1 --ulimit stack=67108864 --device=/dev/kfd --device=/dev/dri --group-add=video \
    --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /:/host -w /host$(pwd) -v /tmp:/tmp \
    tutelgroup/deepseek-671b:mi300x8-chat-20260721 --serve=webui --listen_port 8000 \
      --try_path moonshotai/Kimi-K2.7-Code \
      --try_path moonshotai/Kimi-K2.6 \
      --try_path nvidia/Kimi-K2.5-NVFP4 \
      --try_path nvidia/Kimi-K2-Thinking-NVFP4 \
      --try_path nvidia/DeepSeek-V3.2-NVFP4 \
      --try_path nvidia/DeepSeek-R1-NVFP4 \
      --max_seq_len 1000000

    3. Accessing the Server

    • Direct Request (curl): curl -N -X POST http://0.0.0.0:8000/chat -d '{"text": "Your prompt here"}'
    • Python Request: python3 -m tutel.examples.oai_request_stream --url '0.0.0.0:8000' --prompt 'Your prompt here'
    • Web UI: Open http://0.0.0.0:8000 in your browser.

  11. Setup Claude Code for Linux / WSL (Ubuntu >= 24.04)

    main

    To use Claude Code with a local Tutel server on Linux or WSL, follow these steps to install the CLI and configure the environment variables to point to your local endpoint (http://0.0.0.0:8000).

    sudo apt-get install -y npm
sudo npm install -g @anthropic-ai/claude-code@2.1.197
cat > run_claude.sh <<EOF && chmod a+x run_claude.sh
mkdir -p config/
export ANTHROPIC_BASE_URL="http://0.0.0.0:8000"
export ANTHROPIC_API_KEY="sk-ant-api00-local-mock-key"
export CLAUDE_CONFIG_DIR="config"
export DISABLE_AUTOUPDATER=1
echo '{"customApiKeyResponses": {"approved": ["api00-local-mock-key"]}}' > config/.claude.json
claude
EOF

./run_claude.sh

  12. Serve Microsoft VibeVoice (Multimodality Mode)

    main

    To serve the Microsoft VibeVoice framework for expressive, multi-speaker conversational audio, follow these steps:

    1. Download Models

    pip3 install -U "huggingface_hub[cli]" --upgrade
hf download microsoft/VibeVoice-1.5B --local-dir microsoft/VibeVoice-1.5B
hf download Qwen/Qwen2.5-1.5B --local-dir Qwen/Qwen2.5-1.5B
hf download aoi-ot/VibeVoice-Large --local-dir microsoft/VibeVoice-Large
hf download Qwen/Qwen2.5-7B --local-dir Qwen/Qwen2.5-7B

    2. Run the Server via Docker

    For ND_A100/H100/B200 only:

    docker run -e LOCAL_SIZE=1 -it --rm -p 8001:8000 --shm-size=8g \
    --ulimit memlock=-1 --ulimit stack=67108864 -v /:/host -w /host$(pwd) -v /tmp:/tmp \
    -v /usr/lib/x86_64-linux-gnu/libcuda.so.1:/usr/lib/x86_64-linux-gnu/libcuda.so.1 --privileged \
    -e VOICES="https://homepages.inf.ed.ac.uk/htang2/notes/speech-samples/103-1240-0000.wav" \
    tutelgroup/deepseek-671b:a100x8-chat-20251222 --serve=core \
      --try_path ./microsoft/VibeVoice-1.5B \
      --try_path ./microsoft/VibeVoice-Large

    For ND_MI300_192G_v5 only:

    docker run -e LOCAL_SIZE=1 -it --rm -p 8001:8000 --shm-size=8g \
    --ulimit memlock=-1 --ulimit stack=67108864 --device=/dev/kfd --device=/dev/dri --group-add=video \
    --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /:/host -w /host$(pwd) -v /tmp:/tmp \
    -e VOICES="https://homepages.inf.ed.ac.uk/htang2/notes/speech-samples/103-1240-0000.wav" \
    tutelgroup/deepseek-671b:mi300x8-chat-20251222 --serve=core \
      --try_path ./microsoft/VibeVoice-1.5B \
      --try_path ./microsoft/VibeVoice-Large

    3. Audio Generation Request

    Use curl to send a text prompt and receive an MP3 output:

    curl -X POST http://0.0.0.0:8001/chat -d '{"text": "VibeVoice is a novel framework designed for generating expressive, long-form, multi-speaker conversational audio, such as podcasts, from text."}' > sound_output.mp3