CERIN AMROTH · ML systems

How do I QLoRA-train the fused experts of a MoE (per-expert LoRA on 4-bit experts)?

Solved by experts4bit-qlora · this page in the repository (pinned 605d21e67958, the source of this rendering; latest on main, unpinned; not the source of any fact rendered here)

Install routes

From docs/capabilities.json at the pinned commit.

pip install "experts4bit-qlora[train]"

Primary route: minimum/reference training.

Alternatives:

Environment (from the capability register; one list for all of these routes, so not every item applies to each): OS: Linux · Python: >=3.11 tested in CI (pyproject says >=3.10; 3.10 is not tested) · Accelerator: NVIDIA CUDA GPU (bitsandbytes 4-bit) · Requires: bitsandbytes>=0.43; torch>=2.2; transformers>=5.0 for the streaming loader ([train] extra)

Load with load_moe_4bit_streaming, which installs a frozen NF4 Experts4bit stack wrapped in trainable per-expert ExpertsLoRA adapters for every MoE layer, then train with python -m experts4bit_qlora.train or your own loop. In your own loop, add enable_fast_train(model, dgrad=True) from the [fast] extra to run the step through grouped-nf4-gemm's grouped kernels, and assert its return value; the trainer does the same with FAST_TRAIN=1 and otherwise trains through the per-expert reference loop.

The abstraction that fails, and what replaces it

PEFT expects nn.Linear targets. A fused MoE checkpoint stores each layer's experts as stacked 3-D projections — gate_up_proj [E, 2I, H] and down_proj [E, H, I] — so there is no per-expert Linear for target_modules to match. A delta hung on the stack as a whole would also compute the wrong function: the low-rank term has to be applied per routed expert before that expert's activation (act(W_e x + B[e] A[e] x), not act(W_e x) + d). ExpertsLoRA supplies the adapter abstraction — stacked per-expert A[e] / B[e] over the frozen packed base, the delta injected pre-activation — and grouped-nf4-gemm supplies the fused grouped forward and, with dgrad=True, the optional input-gradient kernel.

Ordinary PEFT:     nn.Linear per projection  <-  LoRA A/B hooked on each Linear
Fused checkpoint:  experts.gate_up_proj [E, 2I, H], experts.down_proj [E, H, I]   (no per-expert Linear)
        |
        v
Experts4bit frozen packed base   +   ExpertsLoRA A[e] / B[e], one pair per expert
        |
        v
per-routed-expert delta, applied BEFORE the activation:  act(W_e x + B[e] A[e] x)
        |
        v
grouped forward (grouped-nf4-gemm)   +   optional dgrad (enable_fast_train(model, dgrad=True))
componentstatewhere it is decided
packed expert base (Experts4bit, NF4 by default)frozenquant_type= on the loader; ExpertsLoRA.__init__ sets requires_grad=False on every base parameter
per-expert A[e] / B[e] (ExpertsLoRA)trainable; B is zero-initialised, so the delta starts at exactly zeror, alpha on the loader; R, ALPHA in the trainer
router (mlp.gate.weight)frozen by default — TRAIN_ROUTER=0; TRAIN_ROUTER=1 trains it at 0.1× the LoRA learning rateits own trainer switch, separate from TRAIN_EXPERTS and TRAIN_ATTENTION; the trainer selects parameter names ending in mlp.gate.weight, ExpertsLoRA itself never touches the router, and the placement ablation found training it hurts (e4b.train.ablation)
attention LoRA (LoRALinear over the frozen q/k/v/o)trainable by default (TRAIN_ATTENTION=1)a trainer switch
expert executionreference per-expert loop (default), batched fallback (enable_batched_train), or fused grouped path (enable_fast_train, optional dgrad=True)your call after loading; train.main() enables neither for you
base identityverified, not assumedverify_moe_4bit(model, strict=True) at load; the separate fused-train gate checked frozen bytes with a nonempty hash and a byte-flip control (see the distinction below); python -m experts4bit_qlora.infer serves the adapter over the same NF4 base it was trained against

Symptoms

Why it happens

Adapter libraries attach low-rank deltas to nn.Linear. A fused expert stack has no per-expert Linear to hook, and the adapter must land before the nonlinearity for each routed expert. ExpertsLoRA therefore re-implements the expert math: for each expert e the frozen 4-bit projections get a trainable scaling * (x @ A[e].T) @ B[e].T term, with stacked A/B per expert, while the base's projection re-dequantises in backward so no dequantised-expert activation is held (experts4bit_qlora/lora.py, ../BITSANDBYTES.md). The reference forward is a per-expert loop by design; the speed comes from the kernel side.

Which project solves it

experts4bit-qlora owns the adapter (ExpertsLoRA), the streaming loader, the trainer and the training-side offload. grouped-nf4-gemm (GitHub, PyPI) owns the kernels: enable_fast_train routes the ExpertsLoRA training forward through its nf4_qlora.fused_grouped_lora, and dgrad=True routes the backward through its single-launch dgrad kernel. [fast] is the seam. enable_batched_train is the no-extras fallback when the kernel package will not build.

Install

pip install "experts4bit-qlora[train]"   # minimum/reference training: loader + trainer
pip install "experts4bit-qlora[train,fast]"  # complete install for the accelerated example below

Smallest correct example

Needs: Linux, a CUDA GPU supported by the kernels, enough host RAM for the offloaded checkpoint, network access and the model download. Install [train,fast] above. This is a one-step expert-adapter smoke example; it is not a convergence run or a reproduction of the benchmark.

import torch
from transformers import AutoTokenizer
from experts4bit_qlora import enable_fast_train, load_moe_4bit_streaming, verify_moe_4bit

model, config = load_moe_4bit_streaming(
    "Qwen/Qwen3-30B-A3B", "cuda", torch.bfloat16, r=8, alpha=16,
    quant_type="nf4", offload=True,          # offload=True: experts in pinned host RAM, one layer on the GPU
)
verify_moe_4bit(model, strict=True)
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})  # checkpointing is required with offload
n = enable_fast_train(model, dgrad=True)
assert n > 0, "still on the per-expert loop: grouped-nf4-gemm missing?"

# Match the trainer's expert-only parameter selection; freeze attention,
# router and all base weights. The loader has already attached expert LoRA.
trainable = []
for name, parameter in model.named_parameters():
    is_expert_lora = "experts" in name and "lora" in name
    parameter.requires_grad_(is_expert_lora)
    if is_expert_lora:
        trainable.append(parameter)
assert trainable, "no expert adapters to train"
optimizer = torch.optim.AdamW(trainable, lr=1e-4)
model.config.use_cache = False
model.enable_input_require_grads()
model.train()

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-30B-A3B")
batch = tokenizer("A telescope collects light from distant objects.", return_tensors="pt").to("cuda")
optimizer.zero_grad(set_to_none=True)
loss = model(**batch, labels=batch["input_ids"].clone()).loss
assert torch.isfinite(loss).item(), "non-finite training loss"
loss.backward()
grads = [parameter.grad for parameter in trainable if parameter.grad is not None]
assert grads and all(torch.isfinite(grad).all().item() for grad in grads), "missing or non-finite adapter gradients"
optimizer.step()
print(f"one expert-adapter step completed; loss={loss.item():.4f}")
# This smoke example does not save a trained adapter. Use the trainer below
# for dataset loading, evaluation, schedules and adapter checkpoints.

The assertion is there because acceleration must be observed, not assumed: enable_fast_train returns 0 rather than raising when the kernel package is absent, and from the caller's side 0 looks exactly like the per-expert loop it was meant to replace. Keep the assert in your own loop.

The shipped trainer provides dataset training (tatsu-lab/alpaca, response-only loss) and adapter checkpoints from environment variables (there are no CLI flags; --help prints them). It runs the per-expert reference loop unless FAST_TRAIN=1 (the fused path, enable_fast_train(model, dgrad=True), refused when it patches nothing); LAYER_MAJOR=1 implies it. Without MODEL, it and infer default to allenai/OLMoE-1B-7B-0924, not the model above; set MODEL= on both to change it:

STEPS=150 R=8 TRAIN_EXPERTS=1 OFFLOAD_EXPERTS=1 OUT=./out python -m experts4bit_qlora.train

Expected result

verify_moe_4bit(model, strict=True) returns without raising. enable_fast_train returns the number of ExpertsLoRA modules patched — one per MoE layer — and 0 means you are silently on the reference loop (../../CONTRIBUTING.md explains why every enabler returns a count). The trainer writes the adapter (adapter_best.pt) under OUT; python -m experts4bit_qlora.infer serves it over the same NF4 base with ADAPTER=./out/adapter_best.pt.

Supported scope

model_typequantizereference_trainfast_train (the headline path)batched_trainnvme_trainnative_mxfp4_train
olmoesupported (tp1)supported (tp1; e4b.train.olmoe-converges)supported — tp1 OK · PASS on the registered text with real weights (e4b.train.parity.tp1.olmoe.fused.2026-09-05)void — tp1 OK · VOID: the _PAD_WASTE_LIMIT fallback engaged without a counter (…olmoe.batched…)not_tested (outside tp1; the arena ladder, e4b.offload.arena-vs-host-ram, is a separate receipt)n/a
qwen3_moesupported (tp1, resident on a 32 GB card; flagship)supported (tp1; flagship)supported — tp1 OK · PASS resident on one 5090 (…qwen3.fused…), beside the flagship's five datasets (e4b.train.flagship-matrix)void — tp1 OK · VOID: the kernel reached on a fraction of the layers every step (…qwen3.batched…); the dgrad-gate trajectory stands on its own fixturenot_tested (outside tp1; the arena ladder, e4b.offload.arena-vs-host-ram, is a separate receipt)n/a
gemma4_textsupported (tp1: the -it checkpoint loaded on this host, no #344; flagship: base)supported (tp1; flagship)supported — tp1 OK · PASS with the step-wise median inside the band by a small margin (…gemma4.fused…), beside the model-2 flagship (e4b.train.flagship-matrix)void — tp1 OK · VOID: on some steps no layer reached the kernel (…gemma4.batched…)not_tested (outside tp1; the arena ladder, e4b.offload.arena-vs-host-ram, is a separate receipt)n/a
granitemoesupported (tp1: the first direct real-weight load)supported (tp1 OK)supported — tp1 OK · PASS on the corrected-counter re-run (…granite.fused…; attempt 1 a kept HARNESS_ERROR of the harness's counter, …granite.fused.attempt1…, amendment 3) — entered model_families on itsupported — tp1 OK · PASS (…granite.batched…)not_testedn/a
gpt_osssupported (bare GptOssExperts4bit; tp1)refused — no ExpertsLoRA; attention-only QLoRA trains (…gptoss.attn_only…, OK · no pair)refused — enable_fast_train returns 0 (…gptoss.fused…, REFUSED)refused — enable_batched_train returns 0 (…gptoss.batched…, REFUSED)refused — enable_mxfp4_nvme_residency refuses bias-carrying modules (#402; it had defaulted to the V4 epilogue, #397), enable_nvme_train_residency refuses bare modules, and the arena_train=True wrap is refused on structureexperimental — grouped-nf4-gemm's ExpertsMxfp4LoRA; tp1 canary and provenance passed on its own text (…gptoss.mxfp4…, EXPERIMENTAL); never licensed
mixtralsupported (tp1: the first real-weight pass through the w1/w3/w2 fusion, offload=True)supported (tp1, offload)supported — tp1 OK · PASS under offload at half the reference loop's peak VRAM (…mixtral.fused…); entered model_families on this rowsupported — tp1 OK · PASS, the kernel reached everywhere (the 8-expert shape; …mixtral.batched…)not_testedn/a

Each cell is one of supported (completed under the registered protocol with a PASS/OK receipt), refused (with the reason), void (ran, unreadable), harness_error, not_tested, experimental, n/a — per path, never a flat flag; the machine-readable form, with the claim id behind every supported / void / refused cell, is training_support in ../capabilities.json, validated by scripts/check_capabilities.py, and model_families is exactly the families whose fast_train is supported. Row statuses in the tp1 receipt are one of OK / REFUSED / HARNESS_ERROR / ALARM / OOM / NOT_RUN / EXPERIMENTAL with the parity verdict (PASS / FAIL / VOID) as a separate column. On top of that table, the attention-4-bit configuration (TRAIN_ATTN_4BIT, reference_attn4/fused_attn4) has its own receipts since tp2/P40 (2026-09-06): supported on granitemoe, olmoe, qwen3_moe and mixtral (e4b.train.h2h.unsloth.<family>.5090.2026-09-06.arm.e4b.*, internal parity PASS on each); supported on gemma4_text since lane P67 (2026-09-24): its arms run since #435 (P56, tp4, P67; all 115 projections), and P67 measured the per-family floor band that tp1's constant band lacked (#713, e4b.train.p67.gemma4.fused-attn4.floor-band.5090.2026-09-24; tp2/P40's arms had died before a step on quantize_attention_projections_4bit converted 100 projections, expected 120; the bf16-attention fast_train cell above is unchanged); refused on gpt_oss (96 of 96 attention projections carry a bias, and the converter refuses bias-carrying projections). The per-family reasons live in training_support's path entries. DeepSeek-V4 is admitted by structure (the base supplies its clamped epilogue via _apply_gate, ../DEEPSEEK-V4.md) and has no training receipt. The per-family training table with every row's status is the tp1 section of ../ARCHITECTURE_SUPPORT.md; the loader's families are in bitsandbytes-moe-load-in-4bit-still-ooms.md.

Limitations

Measured result

Claim e4b.train.flagship-matrix (measured, receipts in this repository): across two 30B-class MoEs — Qwen3-30B-A3B and Gemma-4-26B-A4B — five datasets each, 200 steps per cell, the fused training path runs 1.52–1.81× faster per step at 0.75–0.81× peak VRAM and 0.80–0.92× energy per step, with parity on both registered training-loss criteria: absolute final training-loss delta and median absolute step-wise training-loss delta, each at most 0.05. These are not held-out evaluation criteria. That ratio is fused-vs-reference with the same offload and checkpointing in both arms. A per-step multiple quoted against the per-expert Python loop on a smaller model (the OLMoE per-expert-loop comparison) has a different denominator and different models; it is not this number, and the two do not compose.

The frozen-byte result is a separate piece of evidence. The Qwen3 matrix receipt withdraws its B1 check: it compared empty placeholders, so it did not establish byte preservation in those cells. The independent fused-train gate hashed 16.31 GB, skipped no empty tensors and detected a deliberate byte flip; that result belongs to that run. The second model's C1 protocol explicitly requires the corrected checks. The matrix's speed, memory, energy and training-loss measurements stand; the invalid B1 does not become valid because another run checked the bytes.

Claim family e4b.train.parity.tp1.* (measured, receipts in ../../bench/train-parity-20260905/tp1/; lane tp1, 2026-09-05, the shipped code on one rented RTX 5090, real weights through the direct loader, 60 steps on the registered clinical text, verdicts in the same registered units with every row classified OK / REFUSED / HARNESS_ERROR / ALARM / OOM / NOT_RUN / EXPERIMENTAL, cost reported and never gated): enable_fast_train(dgrad=True) PASSES on OLMoE-1B-7B-Instruct (e4b.train.parity.tp1.olmoe.fused.2026-09-05, the first fused-vs-reference reading on a registered text with real weights for that family), on Qwen3-30B-A3B resident on the 32 GB card (…qwen3.fused…), on Gemma-4-26B-A4B-it — the -it checkpoint, loaded without #344 on that host — with the step-wise median inside the band by a small margin (…gemma4.fused…), and on Mixtral-8x7B-Instruct under offload=True at half the reference loop's peak VRAM (…mixtral.fused…; the family entered model_families on it); and on Granite-3.1-3B-A800M — the family's first direct real-weight load — on the corrected-counter re-run (…granite.fused…; the first attempt is a kept HARNESS_ERROR row, …granite.fused.attempt1…, a closure bug in the harness's kernel counter, not the shipped code; the family entered model_families on the re-run); enable_batched_train PASSES on Granite (…granite.batched…) and Mixtral (…mixtral.batched…) and is VOID on OLMoE, Qwen3 and Gemma-4 — the kernel was not reached on every layer (below); gpt-oss's fused and batched arms are REFUSED (below). A PASS there is a PASS on one text; no convergence claim and no cross-family ratio is made.

Claim family e4b.train.h2h.unsloth.<family>.5090.2026-09-06 (measured, receipts in ../../bench/h2h-20260906/tp2/; lane tp2 / P40, 2026-09-06, one rented RTX 5090, the pre-registration verbatim in the bundle): the head-to-head against Unsloth, per family — P38's fixture exactly on all six families, one box, one session, e4b at the shipped cut a user installs today (0.35.1 + grouped-nf4-gemm 0.30.2 from PyPI, NF4 attention via TRAIN_ATTN_4BIT) against Unsloth 2026.9.2; every attempt a row, a position only where both primary arms are VALID, the quality reading beside it. Positions on two families: Qwen3-30B-A3B, ratio of median s/step Unsloth/e4b 1.457 (5.986 vs 4.108 s — e4b faster per step; peak 21.371 vs 23.141 GB; 383.9 vs 561.4 J/step; held-out COMPARABLE, Δ +0.0152 ≤ 0.05) — +3.1% from P38's 1.413, inside the pre-registered ±10% anchor band: the sign and size of the per-step position reproduce on a second box (…qwen3.5090.2026-09-06, ….quality-n60, ….e4b-internal-parity); and Mixtral-8x7B, where the footprint trade leads the row: the e4b arm trained its experts under CPU offload at a 3.223 GB peak (the registered design for this family, tp1's offload=True; a trainable-on-smaller-cards capability, its own row …mixtral.5090.2026-09-06.footprint) while Unsloth ran resident (its only mode) at 29.163 GB, and what that VRAM buys it is speed per step — ratio 0.361, a footprint-vs-speed trade, not a kernel deficit; e4b's energy is lower (298.3 vs 350.2 J/step); held-out COMPARABLE, Δ −0.0087; Unsloth's predicted resident OOM did not happen (…mixtral.5090.2026-09-06, ….quality-n60, ….e4b-internal-parity). Statuses and coverage rows are the result on the other four: Granite — the comparator could not train the experts: Unsloth's arm completed but is VOID (it attached LoRA to the attention only, 2,621,440 trainable against e4b's 49,807,360 with ExpertsLoRA on all 32 MoE layers; its MoE-LoRA path never engaged on granitemoe; …granite.5090.2026-09-06.coverage, ….arm.unsloth.ckpt_unsloth), e4b's arms VALID (fused ×5.59 its reference); OLMoE — the comparator could not train the family at all: Unsloth's process died at MoE-LoRA engage before its first receipt write (HARNESS_ERROR; …olmoe.5090.2026-09-06.coverage, ….arm.unsloth.ckpt_unsloth), e4b's arms VALID (fused ×3.73); gpt-oss — three refusals as pre-registered (e4b attention-4-bit refuses on 96/96 bias-carrying projections, the fused path patches nothing, Unsloth's MXFP4 load conversion fails; …gptoss….arm.*); Gemma-4 — both e4b attention-4-bit arms died on the harness's projection-count check, 100 converted where it expected 4 · n_layers = 120 (#412, fixed since by #435; run since in P56 and tp4 and not licensed, #713; the bf16-attention fast_train path stays as tp1 left it) while Unsloth's arm is OK · VALID at 3.510 s/step (…gemma4….arm.unsloth.ckpt_unsloth). e4b's internal fused-vs-reference parity PASSES on all four families that ran both arms (×5.59 / ×3.73 / ×2.69 / ×1.23 per step; informational, tp1 owns the licence). Nothing supersedes P38 — two boxes, two measurements, never averaged — and P38's 200-step curve row still stands beside any Qwen3 position.

Claim family e4b.train.h2h.unsloth.qwen3.5090.2026-09-05 (measured, receipts in ../../bench/h2h-20260905/p38/; lane p38, 2026-09-05, one rented RTX 5090, the pre-registration verbatim in the bundle): the head-to-head against Unsloth, end-to-end, on one identical training problem — Qwen3-30B-A3B at one pinned revision, the registered clinical fixture tokenised once and asserted by sha in every arm, seq 512, r 8 / α 16 on attention q/k/v/o and every expert, the router frozen, the same optimizer, LR, batch, steps, precision and held-out eval on both sides; this package's fused dgrad path with NF4 attention (enable_fast_train(dgrad=True) + the shipped TRAIN_ATTN_4BIT mechanism) against Unsloth 2026.9.2's 4-bit MoE QLoRA path. At 60 steps e4b is faster per step (ratio of median s/step, Unsloth/e4b 1.413: 2.151 vs 1.522 s), lower in peak VRAM (21.371 vs 23.141 GB) and energy per step (157.1 vs 224.7 J), reaches a held-out loss of 0.32 sooner (92.5 vs 130.3 s), and the two are comparable on held-out loss (0.2923 vs 0.2975, |Δ| 0.0052 ≤ the 0.05 reading threshold; …quality-n60). At 200 steps the curves separate in Unsloth's favour (0.2713 vs 0.2881, …curve-n200) — e4b's flattens near 0.29 from step 60 while Unsloth's keeps falling; that row is quoted beside the position wherever the position is quoted, and its causes (the eval schedule, the checkpointing mode, the two stacks' transformers/peft versions, the expert adapter's precision — bf16 here because the loader passes the model dtype to ExpertsLoRA, fp32 on Unsloth's side) are candidates, not established. One workload (≈86 tokens per step, batch 1, resident), one box, one family; no general speed claim; nothing licensed (…e4b-internal-parity is the informational fused-vs-reference PASS on that box, tp1 owns the licence).

Claim e4b.train.h2h.unsloth.qwen3.5090.2026-10-02 (measured, receipts in ../../bench/h2h-2026-10-02/tc1/; lane TC1, 2026-10-02, three rented RTX 5090 hosts, the pre-registration with its amendments in bench/tc1/): the field-recipe position at matched work — Qwen3-30B-A3B at the Unsloth notebooks' recipe with the same 642,514,944 parameters, one per-slot LoRA init and fp32 adapters in both frameworks, Unsloth 2026.9.14 on torch 2.12.1+cu130 with its grouped_mm backend engaged — reads Unsloth/e4b 1.437 [1.434, 1.440] (5.688 vs 8.171 s/step over two draws each) with held-out EQUIVALENT at N = 20 and over 200 steps (.curve-n200), the parity control PASSING, and Unsloth's peak VRAM 3.57 GB lower and energy ×0.72 at that configuration. It supersedes the 2026-09-19 row (4.490), whose Unsloth arm ran on torch 2.8 and therefore on Unsloth's per-expert loop (.loop-fallback-t28: still 6.565 × on a torch-2.8 install). Both steps are host-launch-bound on this card; absolute s/step does not travel between hosts, the within-box ratio does (.host-variance). On an H100 NVL the sign reverses (e4b.train.h2h.unsloth.qwen3.h100.2026-10-02, receipts in ../../bench/h2h-2026-10-02/tc1c/): Unsloth/e4b 0.621 [0.615, 0.628] (medians of two draws) — Unsloth faster by 1.61 × at lower VRAM and energy with the same loss — because the grouped GEMM its path routes through is one launch per call on Hopper and not on Blackwell (.dispatch-profile); the 1.437 is a 5090 position and no "e4b faster" position is quoted for the H100 class.

Evidence

Register: ../claims.json; all of these are measured with receipts in this repository.

Common wrong approaches

Source and freshness

This page is a rendering of docs/solutions/qlora-fused-moe-experts.md at commit 605d21e67958 (sha256 7f877bad422d87da…). Numbers are never copied here: every measured statement cites a claim ID in docs/claims.json at that commit. Repository-relative links resolve to this site's pages where the document is published and to the pinned commit otherwise.

rendered from experts4bit-qlora 0.58.1 · source: pjordanandrsn/experts4bit-qlora@605d21e67958 · checked 2026-10-25