[Question] Any examples for 200B+ MoE training with FSDP?

[cailun01]

Hi XTuner Team,

I noticed in the official documentation that one of the key highlights of Xtuner V1 is its ability to train 200B+ models using FSDP instead of EP.

Breakthrough Performance Bottleneck: First time achieving FSDP training throughput surpassing traditional 3D parallel solutions on MoE models above 200B scale (https://xtuner.readthedocs.io/en/latest/#core-features)

I am currently looking to train the 200B+ MoE Model such as Qwen3-235B-A22B model and would love to leverage this capability. Could you please provide a reference example or a configuration template for a model of this scale?

Thank you for the great work!

[CyCle1024]

@cailun01 Hi, it's a good question. And we are sorry for not providing a best practice for training 200+ MoE model in our documentation. Here is a simple SFT example for Qwen3-235B-A22B, and we would update it into our documentation later:

# /config/config_sft_Qwen3_235B_A22B.py

import os

from xtuner.v1.model.moe.qwen3 import Qwen3MoE235BA22Config
from xtuner.v1.train import TrainerConfig
from xtuner.v1.datasets.sft_tokenize_fn import OpenaiTokenizeFunctionConfig
from xtuner.v1.loss import CELossConfig
from xtuner.v1.float8 import Float8Config
from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
from xtuner.v1.float8.config import Float8Config, ScalingGranularity
from xtuner.v1.config import (
    AdamWConfig,
    FSDPConfig,
    LRConfig,
)

NUM_GPUS = int(os.getenv("WORLD_SIZE", "1"))
EP_SIZE = 1
SP_SIZE = 1
INTRA_LAYER_MICRO_BATCH = 1
ACCUMULATION_STEPS = 1
GLOBAL_BATCH_SIZE = NUM_GPUS * INTRA_LAYER_MICRO_BATCH * ACCUMULATION_STEPS // SP_SIZE

work_dir = os.getenv("WORK_DIR")
data_path = os.getenv("ALPACA_PATH")
model_path = os.getenv("QWEN3_MODEL_PATH")

# Enable FP8 if the environment variable is set
if os.getenv("ENABLE_FP8", "").lower() in ["1", "true", "yes"]:
    float8_cfg = Float8Config(
        scaling_granularity_gemm=ScalingGranularity.TILEWISE,
        scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE,
    )
else:
    float8_cfg = None

# Enable torch compile if the environment variable is set
if os.getenv("ENABLE_TORCH_COMPILE", "").lower() in ["1", "true", "yes"]:
    compile_cfg = True
else:
    compile_cfg = False

model = Qwen3MoE235BA22Config(
    ep_size=EP_SIZE,
    compile_cfg=compile_cfg,
)
optim_cfg = AdamWConfig(lr=6e-05, foreach=True)
lr_cfg = LRConfig(lr_type="cosine", lr_min=1e-6)
fsdp_cfg = FSDPConfig(
    cpu_offload=False,
    ep_size=model.ep_size,
)
loss_cfg = CELossConfig(mode="chunk", chunk_size=1024)
dataset_config = [
    {
        "dataset": DatasetConfig(name="alpaca", anno_path=data_path, sample_ratio=100),
        "tokenize_fn": OpenaiTokenizeFunctionConfig(max_length=32*1024, chat_template="qwen3"),
    },
]

dataloader_config = DataloaderConfig(
    pack_max_length=32*1024,
    pack_extra_buffer_size=5,
    pack_level="soft",
    pack_workers=32,
    num_workers=4,
    # group_by_length=True,
)

trainer = TrainerConfig(
    model_cfg=model,
    load_from=model_path,
    optim_cfg=optim_cfg,
    fsdp_cfg=fsdp_cfg,
    loss_cfg=loss_cfg,
    lr_cfg=lr_cfg,
    dataset_cfg=dataset_config,
    dataloader_cfg=dataloader_config,
    global_batch_size=GLOBAL_BATCH_SIZE,
    total_step=1000,
    checkpoint_interval=None,
    hf_interval=None,
    tokenizer_path=model_path,
    strict_load=False,
    work_dir=work_dir,
    intra_layer_micro_batch=INTRA_LAYER_MICRO_BATCH,
    sp_size=SP_SIZE,
)

#!/bin/bash
# sft_Qwen3_235B_A22B.sh
set -e

export XTUNER_ACTIVATION_OFFLOAD=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

export WORK_DIR=/work_dir/qwen3_sft_235B_a22b
export QWEN3_MODEL_PATH=/model/Qwen3-235B-A22B
export ALPACA_PATH=/data/alpaca
export ENABLE_TORCH_COMPILE=1
export ENABLE_FP8=1
export CONFIG_FILE=/config/config_sft_Qwen3_235B_A22B.py

if [ ! -d ${WORK_DIR}/log ]; then
  mkdir -p ${WORK_DIR}/log
fi

torchrun --nproc-per-node 8 \
  --master_addr=${MASTER_ADDR} \
  --master_port=${MASTER_PORT} \
  --nnodes=${WORLD_SIZE} \
  --node_rank=${RANK} \
  -m xtuner.v1.train.cli.sft \
  --config ${CONFIG_FILE} \
  2>&1 | tee ${WORK_DIR}/log/qwen3_235B_bf16_${RANK}_.log

On each node, we should run bash sft_Qwen3_235B_A22B.sh with $MASTER_ADDR $MASTER_PORT $WORLD_SIZE $RANK environment variables defined. This can be automatically done by using Kubeflow Trainer Operator.

[cailun01]

Hi @CyCle1024 , thanks for you reply! How many H100 GPUs are at least required to run this training task?

[CyCle1024]

Hi @CyCle1024 , thanks for you reply! How many H100 GPUs are at least required to run this training task?

I think about 512 H100 GPUS with 80GB HBM，the setting may not be the most optimized. You may adjust some setting according to your own data.

[CyCle1024]

Hi @CyCle1024 , thanks for you reply! How many H100 GPUs are at least required to run this training task?

@cailun01 We have found a config bug, this bug would disable torch.compile whenever you turn it on or off. Please update FSDPConfig above by:

fsdp_cfg = FSDPConfig(
    torch_compile=True,
    cpu_offload=False,
    ep_size=model.ep_size,
)