Expert-granularity CPU offload for quantized MoE: 30B-A3B QLoRA in ~7 GB (single resident expert layer, extends the layer_offloading idea)

Use case

quantize_moe_experts (#3439) got fused-expert 4-bit working, and #3747's selective dequant lowered the ceiling further — but the floor for a 30B-class MoE QLoRA is still ~17–24 GiB reserved, so 12–16 GB cards are out. Demand for exactly that exists in the tracker (#3530 OOMs a 32 GB card on gpt-oss-20b; #3402 asks for lower-VRAM MoE training generally).

The missing piece is granularity: layer_offloading (#3512) already proves the pinned-CPU + prefetch-stream pattern in-tree, but at whole-decoder-layer scope its measured savings are modest (~3 GB on Qwen3.5-35B) because attention/router params are small — in a fine-tuned MoE, ~90 % of frozen bytes are the experts. Offloading only the quantized expert tensors, with a single resident expert layer on GPU and prefetch of layer N+1 during layer N's compute, moves almost all of those bytes off-device while touching nothing else.

Evidence it works (measured, single GPU)

I've implemented this out of tree (experts4bit-qlora, MIT) against transformers-v5 fused experts, and measured on a 12 GB RTX A2000:

model	peak VRAM with expert offload	without
Qwen3-30B-A3B QLoRA	7.16 GB	OOM
Gemma-4-26B-A4B QLoRA	8.47 GB	OOM
OLMoE-1B-7B QLoRA	2.57 GB (from 5.97)	5.97 GB

Cost is ~+11 % s/step (OLMoE; pinned-memory H2D copies overlap compute for the larger models). Correctness: the frozen 4-bit forward is bit-identical with offload on/off, including through use_reentrant=False gradient-checkpoint recompute (the recompute re-fetches the same quantized bytes; adapters stay resident and trainable).

Proposed shape

An expert-granularity mode that composes with quantize_moe_experts — conceptually layer_offloading restricted to the expert parametrizations: pin the packed 4-bit expert storage in CPU RAM at load, keep one layer's experts on GPU, prefetch the next layer's on a side stream during compute, evict after use. Scope for a first PR: single-GPU only (explicitly out of scope: FSDP/DDP interaction, expert-parallel), gated behind a config flag, with the checkpointing-recompute path covered by tests.

I'm happy to write the PR against axolotl's parametrization representation (either in-tree as an extension of the layer_offloading mixin, or as an axolotl.integrations plugin — maintainer preference). The numbers above are reproducible from the linked repo; the same primitive is also proposed upstream in bitsandbytes (bitsandbytes-foundation/bitsandbytes#1965, draft until their v0.50), and an in-tree implementation here could switch to that once it lands.