Add AutoEP example

training/expert_parallel/README.md

@@ -0,0 +1,92 @@
+# AutoEP Training Example
+
+AutoEP (Auto Expert Parallelism) automatically partitions MoE expert weights across GPUs and uses AllToAll communication to route tokens to the correct experts.
+This example offers a quick start for AutoEP in DeepSpeed.
+
+## Quick Start
+
+### Prerequisites
+
+- 2+ GPUs (the fast grouped GEMM path works only on Hopper and Blackwell GPUs)
+- Dependencies:
+  - PyTorch `>= 2.9.1`
+  - **DeepSpeed** with AutoEP: AutoEP has not been merged into `main` yet. `requirements.txt` installs the **tip of [PR #7938](https://github.com/deepspeedai/DeepSpeed/pull/7938)**. Manual install: `pip install "git+https://github.com/deepspeedai/DeepSpeed.git@refs/pull/7938/head#egg=deepspeed"`.
+  - **`transformers` `>= 5.2`** (see `requirements.txt`)
+  - Qwen3.5 requires specific kernel dependencies. See [VERIFICATION.md](VERIFICATION.md#qwen35-kernel-requirements) for more details.
+  - See `requirements.txt` for other dependencies.
+
+### Run
+
+The following launches causal LM training with **AutoEP + ZeRO-1** on a randomly initialized model built from the original **Qwen3.5-MoE** Hugging Face text config.
+`--num_layers` overrides only the layer count in the original model config, which is useful when testing with limited GPU resources. `--dataset_name` and `--dataset_percentage` choose the Hugging Face training dataset and the percentage of the train split to use.
+
+```bash
+deepspeed --num_gpus 8 train.py \
+    --mode autoep \
+    --model qwen3_5_moe \
+    --autoep_size 8 \
+    --num_layers 8 \
+    --dataset_name wikitext \
+    --dataset_percentage 10.0 \
+    --steps 1000
+```
+
+For this `--mode autoep` / `--model qwen3_5_moe` run, `train.py` derives the following `expert_parallel` section:
+
+```json
+    "expert_parallel": {
+        "enabled": true,
+        "autoep_size": 8,
+        "preset_model": "qwen3_5_moe"
+    }
+```
+
+Here are the key options in the DeepSpeed config for AutoEP:
+
+- **`enabled`** — Turns on AutoEP.
+- **`autoep_size`** — Expert-parallel size. It must be specified with `--autoep_size` in AutoEP mode and must divide both the GPU count and the model's expert count. The benchmark commands use `8` for Qwen3.5 and `4` for Llama4 and Mixtral.
+- **`preset_model`** — DeepSpeed's structural AutoEP preset id. The example's public `--model` choices intentionally use the same ids when an AutoEP preset exists.
+
+This example exposes three public `--model` choices: `qwen3_5_moe`, `llama4`, and `mixtral`. These match the DeepSpeed `preset_model` ids used by AutoEP for the same structures. The underlying AutoEP PR also defines additional structural preset ids: `qwen3_moe`, `deepseek_v2`, and `deepseek_v3`; those are not exposed as `--model` choices in this example.
+
+## Performance Benchmark
+
+We benchmarked Qwen3.5 and Mixtral with 8 layers, and Llama4 with 7 layers. Each model was run under matching conditions for AutoEP and the ZeRO-3 leaf baseline: 8 H100 GPUs, sequence length 1024, micro batch size 1, gradient accumulation 4, 100 optimizer steps, and steps 50-99 measured. The table below reports the side-by-side comparison. Llama4 uses 7 layers because the 8-layer ZeRO-3 leaf baseline OOMed during backward.
+
+| Model | ZeRO-3 leaf | AutoEP (+ZeRO-1) |
+| --- | --- | --- |
+| Qwen3.5 MoE | 42,128.05 tok/s, 34.99 GB | 87,540.15 tok/s, 25.58 GB (`2.08x` throughput, `0.73x` memory vs ZeRO-3) |
+| Llama4 (7 layers) | 19,144.07 tok/s, 56.95 GB | 60,178.91 tok/s, 60.08 GB (`3.14x` throughput, `1.06x` memory vs 7-layer ZeRO-3) |
+| Mixtral 8x7B | 32,622.11 tok/s, 50.47 GB | 69,052.31 tok/s, 35.03 GB (`2.12x` throughput, `0.69x` memory vs ZeRO-3) |
+
+Qwen3.5 reproduction steps and reference loss curves for the AutoEP and ZeRO-3 leaf comparison are in [VERIFICATION.md](VERIFICATION.md).
+
+
+## Important Constraints
+
+### `autoep_size` requirements
+
+- Must be `<= num_experts`
+- Must evenly divide `num_experts`
+- Must evenly divide `world_size`
+- `autoep_size=1` bypasses EP communication entirely (degenerate case)
+
+### Grouped GEMM backend
+
+`torch._grouped_mm` is required for the default production path. With the default `expert_parallel.use_grouped_mm=true`, DeepSpeed fails fast if `torch._grouped_mm` is unavailable. Set `use_grouped_mm=false` in the DeepSpeed config only for functional/debug runs that intentionally use the sequential for-loop path. On A100 (SM80), verify availability and actual throughput since the Hopper fast path may not activate.
+
+### Qwen3.5 linear-attention kernels
+
+For `--model qwen3_5_moe`, the required linear-attention kernel dependencies and verification checks are documented in [VERIFICATION.md](VERIFICATION.md#qwen35-kernel-requirements).
+
+### bf16 requirement
+
+`bf16` is recommended. `fp16` is functionally correct but not optimized for the Hopper grouped-GEMM fast path used by `torch._grouped_mm`.
+
+### Optimizer wiring
+
+AutoEP runs must let DeepSpeed build the optimizer from the JSON config (no client optimizer). This ensures `configure_moe_param_groups()` is invoked to split expert parameters into expert-data-parallel reduction groups.
+
+### Load balancing status
+
+DeepSeek-style auxiliary-loss-free (expert-bias) load balancing is **not yet implemented** in AutoEP.

training/expert_parallel/VERIFICATION.md

@@ -0,0 +1,125 @@
+# Qwen3.5 AutoEP Reproduction
+
+This document describes how to reproduce the Qwen3.5 AutoEP sample and compare it with the ZeRO-3 leaf baseline. It is intentionally environment-neutral: choose local output directories and caches that fit your machine or cluster.
+
+The commands below use:
+
+- model preset: `qwen3_5_moe`
+- DeepSpeed AutoEP preset: `qwen3_5_moe`
+- layers: `8`
+- dataset: `wikitext`, `dataset_percentage=10.0`
+- tokenizer: `Qwen/Qwen3-0.6B`
+- sequence length: `1024`
+- micro batch size: `1`
+- gradient accumulation: `4`
+- world size: `8`
+- steps: `100`, with steps `50-99` treated as the post-warmup measurement window
+
+AutoEP uses the sample's built-in `--mode autoep` config with `--autoep_size 8`: bf16, AdamW, ZeRO stage 1, `expert_parallel.enabled=true`, `autoep_size=8`, `preset_model=qwen3_5_moe`. The baseline uses `--mode zero3_leaf`: bf16, AdamW, ZeRO stage 3, and the Qwen3.5 MoE block registered as a ZeRO leaf module.
+
+## Install
+
+Run from the repository root:
+
+```bash
+cd training/expert_parallel
+python -m pip install -r requirements.txt
+export TOKENIZERS_PARALLELISM=false
+mkdir -p runs/qwen35/{init,autoep,zero3_leaf,compare}
+```
+
+You may set `HF_HOME` and `HF_DATASETS_CACHE` if your environment requires explicit Hugging Face cache locations.
+
+
+## Qwen3.5 Kernel Requirements
+
+For `--model qwen3_5_moe`, `flash-linear-attention`, `causal-conv1d`, `flash-attn`, and `tilelang` on H100/Triton `>= 3.4` are verification requirements, not optional accelerators. The verification should fail if `transformers.utils.import_utils.is_flash_linear_attention_available()` or `is_causal_conv1d_available()` is false, or if a runtime inspection shows `Qwen3_5MoeGatedDeltaNet` using `torch_causal_conv1d_update` or `torch_chunk_gated_delta_rule`.
+
+`flash-attn` is also required when full-attention layers are configured to use `attn_implementation="flash_attention_2"`.
+
+## Create A Shared Initialization
+
+Use the same randomly initialized weights for AutoEP and ZeRO-3 leaf when comparing loss curves. This makes the metric comparison easier to interpret.
+
+```bash
+python utils/prepare_init_weights.py \
+  --model qwen3_5_moe \
+  --num_layers 8 \
+  --seed 42 \
+  --output runs/qwen35/init/qwen35_l8_seed42.safetensors
+```
+
+## Run AutoEP
+
+```bash
+deepspeed --num_gpus 8 --master_port 29104 train.py \
+  --mode autoep \
+  --model qwen3_5_moe \
+  --autoep_size 8 \
+  --num_layers 8 \
+  --steps 100 \
+  --warmup_steps 50 \
+  --log_interval 1 \
+  --seq_len 1024 \
+  --micro_batch_size 1 \
+  --grad_accum 4 \
+  --seed 42 \
+  --dataset_name wikitext \
+  --dataset_percentage 10.0 \
+  --tokenizer_name Qwen/Qwen3-0.6B \
+  --load_init_weights runs/qwen35/init/qwen35_l8_seed42.safetensors \
+  --metrics_out runs/qwen35/autoep/metrics.csv
+```
+
+## Run ZeRO-3 Leaf Baseline
+
+```bash
+deepspeed --num_gpus 8 --master_port 29105 train.py \
+  --mode zero3_leaf \
+  --model qwen3_5_moe \
+  --num_layers 8 \
+  --steps 100 \
+  --warmup_steps 50 \
+  --log_interval 1 \
+  --seq_len 1024 \
+  --micro_batch_size 1 \
+  --grad_accum 4 \
+  --seed 42 \
+  --dataset_name wikitext \
+  --dataset_percentage 10.0 \
+  --tokenizer_name Qwen/Qwen3-0.6B \
+  --load_init_weights runs/qwen35/init/qwen35_l8_seed42.safetensors \
+  --metrics_out runs/qwen35/zero3_leaf/metrics.csv
+```
+
+## Compare Metrics
+
+`compare_metrics.py` compares the loss, throughput, and peak memory reported by the two metrics CSV files, then generates summary JSON plus plots.
+
+```bash
+python utils/compare_metrics.py \
+  --autoep_csv runs/qwen35/autoep/metrics.csv \
+  --zero3_leaf_csv runs/qwen35/zero3_leaf/metrics.csv \
+  --warmup_steps 50 \
+  --out_dir runs/qwen35/compare \
+  --out_json runs/qwen35/compare/summary.json
+```
+
+Useful outputs:
+
+- `runs/qwen35/compare/summary.json`
+- `runs/qwen35/compare/loss_curve.png`
+- `runs/qwen35/compare/peak_memory_bar.png`
+- `runs/qwen35/compare/throughput_bar.png`
+
+Small numeric differences are expected because AutoEP and ZeRO-3 leaf use different distributed execution paths. The comparison is intended to confirm that loss behavior remains aligned while reporting throughput and memory differences under the same model, data, seed, and initialization.
+
+## Reference Loss Curves
+
+The repository includes reference Qwen3.5 loss-curve images from a longer AutoEP and ZeRO-3 leaf comparison. They are useful for checking the expected curve shape after reproducing the workflow above, but the commands above are the source of truth for a fresh run in your own environment.
+
+![Qwen3.5 CE loss curve](images/qwen35_aux_10k_ce_loss_curve.png)
+
+![Qwen3.5 total loss curve](images/qwen35_aux_10k_total_loss_curve.png)
+
+![Qwen3.5 aux loss curve](images/qwen35_aux_10k_aux_loss_curve.png)