How Expert Balance Works

How Expert Balance Works

English | Русский | 中文

The problem

A large MoE model has hundreds of routed experts per layer. On 4× 96 GB cards (384 GB total, effective ~90 GB per card after reserve) the expert weights of big models do not fit — and KV cache needs room too. Yet most traffic only touches a fraction of experts.

The idea: hot and cold banks

• Hot bank — experts that live in GPU memory. Zero transfer cost per token.
• Cold bank — the rest, pinned in host RAM, read over PCIe on demand (UVA: the GPU can access host memory directly, no explicit copy).

With PCIe 4.0 (≈ 25 GB/s effective per card on this host class) a cold hit is slower but the hot set covers the vast majority of routed tokens — so the average stays fast.

How the hot set is chosen

One shared planner (patches/hotsplit_plan.py, numpy-only, unit-tested) does the work:

1. Routing histogram — counts of how often each expert is routed (from balance.sh on sample traffic — shipped with the DeepSeek/GLM recipes; the MiMo recipe ships a reference expert_hist_mix.json instead — re-rank it for your traffic with patches/hotsplit_rerank.py).
2. Greedy hot set — sort experts by routing share per byte (popularity ÷ size) and fill the per-rank budget HOTSPLIT_HOT_GIB.
3. Validation — NaN/negative/mismatched expert counts raise; budget overflow raises.
4. Oracle check — for small cases a brute-force planner confirms greedy is optimal.

The same planner is byte-identical across all four recipes (enforced by tests).

Static vs dynamic

• Static (HOTSPLIT_DYNAMIC=0) — hot set fixed at load time from the histogram.
• Dynamic (HOTSPLIT_DYNAMIC=1, default) — the hot set adapts at runtime: every HOTSPLIT_DYN_STEPS engine steps the live routing counters are evaluated and the weakest hot experts are swapped for the strongest cold ones (host→GPU copies on a dedicated CUDA stream, published via events so a decode-graph replay never sees a half-applied swap).

Dynamic updates are validated at startup (bad values like HOTSPLIT_DYN_STEPS=0 refuse to start), repeated planner failures abort the run, and each update logs a hot/cold hit-rate line — the number that proves the mode is earning its keep.

The knobs

Knob	Meaning	Default
HOTSPLIT_HOT_GIB	per-rank hot budget (GiB)	empty = stock GPU-resident expert bytes (NOT the full bank)
HOTSPLIT_DYNAMIC	runtime adaptation	1
HOTSPLIT_DYN_STEPS	steps between swaps	2000
HOTSPLIT_DYN_MAX_SWAPS	swaps per layer per update	2
HOTSPLIT_DYN_MARGIN	swap only if in > margin × out	1.25
HOTSPLIT_VRAM_RESERVE_GIB	headroom for CUDA graphs / KV cache	2
HOTSPLIT_NUMA_POLICY	preferred / bind / off	preferred

Why the numbers here are trustworthy

This hardware: cards are PCIe Gen 5, but the host slots are PCIe 4.0 and the RAM is DDR4. All recipes measure on the observed platform (nvidia-smi, numastat), not spec sheets.