MoE Expert Quantization

Reduce VRAM usage when training MoE model adapters by quantizing expert weights on load

Transformers v5 changed MoE expert layers from nn.Linear to fused nn.Parameter (3D+ tensors). This means bitsandbytes can no longer quantize them during model loading, resulting in all expert weights being loaded in full bf16 precision and causing massive VRAM usage.

quantize_moe_experts solves this by quantizing expert weights during model loading. It intercepts the weight loading process, quantizes each expert tensor on the fly, and immediately frees the original bf16 tensor from VRAM. This dramatically reduces peak memory. For example, GLM-4.7-Flash QLoRA drops from ~127GiB to ~23GiB reserved memory.

Tip

For checkpoints that ship already quantized to NVFP4 (ModelOpt exports), see NVFP4 MoE LoRA instead: the experts stay 4-bit-packed and the adapter can merge back into a plain NVFP4 model.

Usage

Enable expert quantization in your Axolotl config:

quantize_moe_experts: true

This works with both 4-bit (QLoRA) and 8-bit (LoRA) quantization.

Expert LoRA targeting

You can optionally apply LoRA adapters directly to expert weights using lora_target_parameters:

lora_target_parameters:
  - mlp.experts.gate_up_proj
  - mlp.experts.down_proj
  # - mlp.gate.weight  # router
Note

lora_dropout must be 0 when using lora_target_parameters.

Requirements

  • Requires (adapter: lora and load_in_8bit: true) or (adapter: qlora and load_in_4bit: true)
  • CUDA GPUs only (not tested with ROCm or other backends)
  • FSDP2 compatible for distributed training

Limitations

  • lora_target_linear is not compatible with quantize_moe_experts. See Expert LoRA targeting instead.
  • With CPU-efficient loading disabled, the FSDP2 QLoRA loader needs the entire quantized model to fit on each GPU before sharding. Use the CPU-staged loader below to avoid this initial replica.
  • Total model parameter count may display incorrectly (trainable param count is correct).
  • FSDP LoRA (8-bit) may have a large initial VRAM spike at the first 1-2 steps, which then drops. QLoRA does not exhibit this.
  • FSDP2 may use more VRAM per GPU than single GPU training due to not all layers being properly sharded across ranks.
  • Model loading takes longer due to on-demand quantization, even on consecutive runs.
  • DeepSpeed has not been tested.

CPU-staged FSDP2 QLoRA loading

For models whose quantized weights do not fit on one GPU, enable:

adapter: qlora
load_in_4bit: true
quantize_moe_experts: true
nf4_backend: bitsandbytes
fsdp_version: 2
fsdp_config:
  cpu_ram_efficient_loading: true
  state_dict_type: SHARDED_STATE_DICT
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: YourDecoderLayer
  reshard_after_forward: true

qlora_sharded_model_loading defaults to true for FSDP2 with adapter: qlora, load_in_4bit: true and cpu_ram_efficient_loading: true when omitted or null. An explicit false is preserved and requires disabling cpu_ram_efficient_loading. For other configurations, the default remains false.

Replace YourDecoderLayer with the decoder class for your model. Rank zero reads and quantizes the original floating-point checkpoint in bounded GPU chunks, storing packed weights in CPU memory. Other ranks construct the model on the meta device. Trainable adapters receive CPU storage before optimizer preparation so Accelerate can map each adapter to its sharded parameter; frozen base weights remain on meta on peer ranks. After FSDP wrapping, rank zero transfers individual parameter shards; no rank needs a complete GPU replica during loading. Rank zero still needs CPU RAM for the complete packed model and checkpoint conversion temporaries. A fused MoE conversion can still require a complete floating-point expert tensor and its source tensors on CPU. Quantization metadata is replicated across ranks. Frozen base linear operations are checkpointed internally so backward recomputes dense weights instead of retaining them for every layer. Training still needs enough GPU memory to gather and execute an individual wrapped layer, including its dequantized weights. While shards are distributed, rank zero also briefly holds one full packed tensor on its accelerator, bounded by the largest quantized tensor; this peak is smaller than the wrapped-layer requirement above.

offload_params controls FSDP parameter offloading after loading. This path supports a single FSDP2 data-parallel mesh, optionally combined with expert parallelism (expert_parallel_size, whose expert slices are cut from rank zero’s packed weights without requantizing); DeepSpeed, tensor/context parallelism, DoRA, modules_to_save and value-dependent LoRA inits (pissa, olora, loftq, corda, eva) are unsupported. olora remains available on the non-staged bitsandbytes loader with cpu_ram_efficient_loading: false (the torchao backend is always staged); pissa and loftq require an unquantized base.

A base model that is already quantized with bitsandbytes 4-bit (its config.json carries a quantization_config with quant_method: bitsandbytes and load_in_4bit) is adopted as-is: the checkpoint’s packed weights and quantization state are staged without requantizing, so bnb_config_kwargs that disagree with the checkpoint are ignored (with a warning). The compute dtype must match the checkpoint’s bnb_4bit_quant_storage when that storage is a floating-point dtype. The torchao backend and quantize_moe_experts still need an unquantized base checkpoint. axolotl merge-lora unpacks the checkpoint’s bitsandbytes components and folds the adapter into the dequantized weights, so such a base always merges to a bf16 checkpoint with no quantization_config; the packed 4-bit format cannot be preserved through the merge. A bitsandbytes 8-bit base is rejected.

LoRA init compatibility

Init Linear targets, non-staged 4-bit Linear targets, CPU-staged lora_target_parameters (fused experts)
default, gaussian yes yes yes
olora yes rejected rejected
pissa, loftq rejected by PEFT (need an unquantized base) rejected rejected
corda, eva not wired rejected rejected

Value-dependent inits read and rewrite the base weight while the adapter is built. Fused expert parameters have no weight attribute, so PEFT fails inside the init; Axolotl rejects that combination at validation, before any weights load. Resuming an adapter that was saved with a value-dependent init under CPU-staged loading is also rejected: PEFT re-runs that init on load and the packed base cannot take its residual. Convert such an adapter with PEFT’s path_initial_model_for_weight_conversion first.

Long staging runs and progress

ddp_timeout: 21600  # Seconds; six hours for rank-zero loading and metadata exchange.

NF4 staging creates a dedicated Gloo control group on every rank before loading. Its timeout uses AXOLOTL_NCCL_TIMEOUT when set, otherwise ddp_timeout, otherwise 1800 seconds. This applies even if the default process group was initialized earlier. Status and metadata broadcasts use CPU tensors on this group, so peers do not enter NCCL bootstrap while rank zero is quantizing. This does not change an existing training process group’s timeout or the launcher’s rendezvous timeout.

Logs identify checkpoint conversion, metadata preparation and broadcast, adapter initialization, FSDP wrapping, and shard distribution. Long phases emit elapsed-time messages every minute; shard distribution also reports tensor counts. An elapsed-time message indicates the phase has not finished, not proof that it is making progress. Peer metadata is built on meta without copying CPU tensor data.

Sharding and tensor shapes

The staged loader supports the one-dimensional FSDP2 data-parallel mesh with both reshard_after_forward: true and false. Transformer-based and size-based wrapping are supported, as is wrapping only the root module. Size-based wrapping keeps NF4 parametrization containers and expert adapter internals inside their owning module. A root-only wrapper gathers the whole model for computation, so use decoder-layer wrapping when the model does not fit on a single GPU.

Packed parameters stay uint8 under FSDP2 mixed precision: Axolotl stops FSDP2 casting non-floating parameters to the policy’s param_dtype, since converting packed bytes to bf16 values corrupts the quantized base. They are sharded along their first storage dimension, which need not align with expert or quantization block boundaries. FSDP gathers that storage before NF4 reconstruction; logical weights retain the original dense or fused expert shape. Uneven splits and empty adapter shards are supported. Quantization metadata remains replicated. Full and sharded checkpoint formats preserve the global shapes.

Hybrid replicated/sharded meshes (dp_replicate_size > 1) are currently unsupported and rejected. This restriction is separate from the reshard_after_forward setting.

torchao NF4 alternative

Set nf4_backend: torchao to use torchao’s NF4Tensor quantization for the frozen base weights. It can also be used without FSDP on a single GPU. Linear weights and, when enabled, fused expert weights are converted into bounded NF4Tensor chunks. Packed storage is exposed through parametrizations so FSDP2 can shard it.

Torchao 0.18 is the pinned version; the loader also resolves 0.17’s module path, where the NF4 implementation lived before the move. Torchao uses NF4 codes with double-quantized scales, but its scale quantization is different from bitsandbytes. Switching backends changes the quantized base weights. The torchao path uses fixed 64-element weight blocks and 256-scale blocks, with independent double quantization per chunk. It does not call bitsandbytes’ int32 quantization kernels; bounded chunks also limit torchao’s conversion temporaries.

The bitsandbytes path splits quantization and dequantization into block-aligned chunks to avoid the native kernel’s 2**31 - 1 element limit. It compresses the complete scale vector once, preserving bitsandbytes’ double-quantization semantics. The final packed tensor and requested dequantized output still need to fit in memory.

Resume and merge

Use resume_from_checkpoint with the same base checkpoint, NF4 backend, quantization settings and adapter configuration. The base is reconstructed and sharded before restoring the adapter and optimizer state. Both SHARDED_STATE_DICT and FULL_STATE_DICT adapter checkpoints are supported. Optimizer restoration uses Axolotl’s pinned Accelerate loader.

Keep nf4_backend in the configuration passed to axolotl merge-lora and use merge_method: memory_efficient (the default). The legacy CLI merge path is rejected for staged NF4 adapters because it loads the original dense base. Merging materializes the quantized base and applies the adapter delta, producing ordinary floating-point weights. The memory-efficient merge path reproduces the selected backend’s quantization before adding the delta. Keep the training dtype and bnb_config_kwargs.llm_int8_skip_modules unchanged for merging; both are part of the quantized base definition.

Under staged loading, llm_int8_skip_modules entries match in one of three ways:

  • A plain name or dotted path matches by name: lm_head excludes that module wherever it appears, model.layers.0.mlp excludes that subtree, and proj excludes only a module named exactly proj, not q_proj.
  • A leading or trailing dot requests a plain substring match against the full parameter path: _proj. excludes every *_proj weight.
  • Regex metacharacters request a regular-expression search over the full parameter path: layers\.[0-3]\. excludes the first four layers.

lm_head and embed_out are always excluded.

Because these rules differ from Transformers’ modules_to_not_convert, a key that worked under plain QLoRA can match nothing here. At the end of staging the loader logs, for every exclusion entry, its matching tier and how many quantization candidates it excluded, and warns for each llm_int8_skip_modules entry that excluded none; check that warning before assuming a module stayed unquantized.

Implementation details

The quantization is applied by patching transformers to intercept weight loading. When a 3D+ CUDA tensor with “expert” in its name is detected:

  • 4-bit mode: Uses bitsandbytes NF4 parametrization (configurable via bnb_4bit_quant_type).
  • 8-bit mode: Uses a custom row-wise int8 parametrization with bitsandbytes dequantization.

The original bf16 tensor is freed immediately after quantization. Multiple sub-patches are applied to transformers, PEFT and accelerate FSDP2 to support these parametrized expert modules.

For full implementation details, see PR #3439.