NVFP4 MoE LoRA

Fine-tune 4-bit NVFP4 MoE checkpoints with LoRA and merge the adapter back into a plain NVFP4 model

ModelOpt NVFP4 checkpoints (e.g. nvidia/Qwen3-30B-A3B-NVFP4) store weights as packed 4-bit FP4 with per-block scales. Axolotl trains LoRA adapters directly on these checkpoints: the base stays 4-bit-packed in memory while the LoRA A/B matrices train in bf16.

Two expert kernels support the NVFP4 expert path:

Kernel Config Compute Hardware
SonicMoE use_sonicmoe: true W4A4 native (W4A16 fallback) Datacenter Blackwell SM100 (W4A4), Hopper (W4A16); consumer Blackwell sm_120 not yet supported
ScatterMoE use_scattermoe: true W4A16 (Marlin) any CUDA GPU sm80+

Usage

base_model: nvidia/Qwen3-30B-A3B-NVFP4

plugins:
  - axolotl.integrations.kernels.KernelsPlugin
use_sonicmoe: true
nvfp4_merge_aware: true

adapter: lora
lora_r: 16
lora_alpha: 32
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
lora_target_parameters:
  - experts.gate_up_proj
  - experts.down_proj

Full config: examples/qwen3/30b-a3b-nvfp4-lora.yaml.

Native TorchAO checkpoints

For native TorchAO NVFP4 dense weights, supported ordinary LoRA projections use merge-aware training by default. The static-weight path supports materialized adapter factors and the supported input/output tensor-parallel layouts. TP saves gather complete adapter factors before writing the PEFT checkpoint. FSDP2 static dense LoRA materializes the frozen base and adapter factors through their owning module forwards. Two-rank projection tests cover independently sharded base/A/B modules, including CPU shard offload, gradients, one optimizer update and merged quantized-weight parity. Full two-rank Trainer tests also cover normal loading, CPU-RAM-efficient loading, and CPU shard offload: adapter checkpoints resume with matching optimizer state and schedules, and format-preserving merged exports match the trained model’s inference. CPU-offloaded gradients stay on CPU during clipping; adapter saves gather temporary copies without gathering the frozen base.

Static native dense LoRA also supports DeepSpeed ZeRO-1/2/3. Merge-aware forwards are installed after DeepSpeed attaches its parameter-materialization hooks, and adapter checkpoints retain the original quantizer recipe. Two-rank tests cover forward and gradient parity, optimizer updates and quantized export; a full ZeRO-3 Trainer test covers fresh resume and merged-native inference. Native static forwards apply quantization from the first forward; the scheduled warmup below applies to the ModelOpt path.

Native dynamic activation quantization supports dense native TorchAO merge-aware LoRA on SM100+ hardware, including FSDP2 and DeepSpeed ZeRO-1/2/3. Each adapted projection quantizes its effective weight and activation group. Every frozen dynamic native linear also receives an input-gradient straight-through estimator, including layers outside lora_target_modules. Its backward multiplies the output gradient by the dequantized base weight; it does not differentiate the activation quantizer. FSDP2 installs these forwards after weight loading and sharding, while DeepSpeed installs them after engine preparation. ZeRO-3 retains its packed-weight backward path.

Set nvfp4_merge_aware: false to explicitly select ordinary LoRA. The frozen-base input-gradient path remains enabled, but the adapter is no longer trained against the quantized effective weight, so merging can lose its learned update. Unsupported variants or incomplete coverage emit an NVFP4 MERGE WARNING and do not receive a merged-NVFP4 parity guarantee. The multi-LoRA plugin owns routed native integration; its sharded combinations need separate qualification. Native dynamic TP training also needs separate qualification. This is separate from the ModelOpt/SonicMoE expert path described below.

The native merge writer preserves TorchAO’s quantization recipe and serialized weight components. Use a format-preserving merge for deployment of a merge-aware adapter; serving its unsnapped adapter alongside the frozen base computes a different model. Native adapter metadata records the TorchAO encoder and each adapted weight’s quantization recipe. A mismatched encoder or recipe produces an NVFP4 MERGE WARNING; the merger continues using the checkpoint’s recipe without a parity guarantee.

Merging the adapter

Merging LoRA into a 4-bit base is normally lossy: the NVFP4 grid step is 25-50% of a block’s max weight, while a typical trained LoRA delta is well under 1% of the weight magnitude. Re-quantizing base + delta onto the base grid therefore rounds most of the delta away. axolotl merge-lora detects this and warns NEAR-NO-OP expert merge when the delta is far below the grid step.

Merge-aware training (nvfp4_merge_aware)

For recognized ModelOpt NVFP4 bases with SonicMoE and adapter: lora or adapter: multilora, merge-aware training is enabled automatically unless explicitly disabled. nvfp4_merge_aware: true makes that choice explicit and makes training optimize the merged model directly. The forward pass computes

out = x @ Q(dequant(base) + scaling * (B @ A))^T

where Q is the exact quantizer merge-lora writes with (fresh block scales on the base’s per-tensor scale grid). Gradients flow through the quantizer with a straight-through estimator, so sub-grid-step updates accumulate in A/B until they cross FP4 code boundaries. The merged NVFP4 checkpoint is then bitwise identical to the weights training fake-quantized against: what you trained is what you serve.

use_sonicmoe: true
nvfp4_merge_aware: true
# optional warm-up before the fake-quant kicks in:
# int = absolute optimizer step, float in (0, 1) = fraction of total steps
# nvfp4_merge_aware_start_step: 0.1

Then merge as usual:

axolotl merge-lora config.yaml

The saved adapter’s adapter_config.json carries an nvfp4_merge_aware stamp recording the quantizer identity (scale mode, per-tensor-scale policy, torchao version). merge-lora reads it to select the matching writer mode and errors if the environment’s torchao version differs from the one trained with (--override-quantizer downgrades this to a warning). Tools other than Axolotl ignore the extra key; the merged checkpoint itself is a fully standard NVFP4 model.

Important

With nvfp4_merge_aware, the merged checkpoint is the trained model; the raw adapter is only an intermediate artifact. Training optimizes the snapped weights Q(W_eff) while the un-snapped base + scaling * (B @ A) drifts freely (the straight-through estimator erases sub-boundary drift every forward), so serving base + adapter unmerged, or merging with --dequant, gives a model that was never trained and can score worse than the base. Both are rejected or warned against; always serve the merged output.

Constraints:

  • The ModelOpt merge-aware path uses use_sonicmoe: true and either adapter: lora or adapter: multilora. An unsupported backend or an explicit nvfp4_merge_aware: false produces a prominent NVFP4 MERGE WARNING and continues with ordinary LoRA. That fallback can lose learned updates when merged into NVFP4.
  • Dense LoRA kernels (lora_qkv_kernel, lora_o_kernel, lora_mlp_kernel) are opt-in with merge-aware training. Protected projections use their original forward so fake quantization is preserved; other projections can use optimized kernels.
  • Unsupported projections warn and fall back. If any distributed rank reports a fallback, checkpoints do not carry a merge-aware guarantee; stale merge-aware metadata is removed.
  • FSDP2 reconstruction preserves weight and activation scales, including ranks loaded on the meta device. SonicMoE routes received NVFP4 expert batches through the local EP implementation, preserving local expert IDs and -1 sentinel slots. Local routing has CPU output/gradient parity coverage; end-to-end DeepEP dispatch with FSDP2 still requires qualification on a supported NVLink setup. EP gradient clipping includes every rank-local expert and counts shared replicas once. With CPU-offloaded shards, only the scalar norm moves to CUDA for NCCL reduction; gradient shards stay on CPU. Unsupported adapter variants or sharded paths without owning-forward support warn and fall back without a merge-aware guarantee.
  • Fake quantization adds work to each protected projection. The opt-in kernel dispatch has not yet been benchmarked on a full NVFP4 training run; measure throughput and memory for your configuration.

On Qwen3-30B-A3B-NVFP4 (attention and experts both NVFP4), the merged checkpoint reproduces the trained loss (recovery ~1.0 of the adapter’s improvement over base); without merge-aware training the format-preserving merge retains only ~0.1-0.4.

Merging without merge-aware training

For adapters trained without the flag, the format-preserving merge reuses the base scale grid and re-rounds codes, which preserves only the part of the delta that crosses code boundaries. Alternatives:

  • axolotl merge-lora config.yaml --dequant: dequantize the base to bf16 and merge losslessly. The output is a bf16 model (~4x larger).
  • Serve base + adapter unmerged (only valid for adapters trained without nvfp4_merge_aware).

See also

Dense LoRA kernels with merge-aware training

lora_qkv_kernel, lora_o_kernel, and lora_mlp_kernel may be enabled explicitly with nvfp4_merge_aware. They remain opt-in. Projections carrying merge-aware forward hooks retain their fake-quantized forward and STE gradients; ordinary projections still use optimized LoRA kernels. A fused group containing a merge-aware projection dispatches its projections individually. Routed experts continue to use the SonicMoE merge-aware path.

This avoids bypassing fake quantization; it does not turn the snapped effective weight into an ordinary low-rank update or guarantee a speedup. Benchmark the mixture of quantized and unquantized projections in the actual checkpoint.

Native weight-only torchao NVFP4Tensor dense weights retain the base layer’s forward and autograd, while a custom LoRA delta function fuses adapter output addition and computes input/A/B gradients. FP32 adapter parameters retain FP32 reduction outputs without full-size FP32 activation copies. Native SwiGLU/GEGLU MLPs also use fused activation forward/backward kernels. This optimizes the adapter side; it does not replace the NVFP4 base kernel. DoRA, merged/disabled adapters, and multiple active PEFT adapters retain PEFT’s forward. Expert kernels are unchanged.

The fused low-rank LoRA kernel path does not implement dynamic NVFP4 activation quantization. For supported dense native TorchAO LoRA projections, the native merge-aware forward runs the deployed quantized operation and uses a dequantized frozen-weight, quantized-activation straight-through surrogate for gradients. The qualified dense path includes unsharded runs, FSDP2, and DeepSpeed ZeRO-1/2/3 as described above; FSDP2 installs after sharding and DeepSpeed after engine preparation. It does not cover lora_target_parameters expert adapters, EP/DeepEP composition, or dynamic TP training; ordinary-LoRA fallback is available only where the underlying backward is supported. These dense-path results do not establish native EP support.

Which NVFP4 loading path is supported?

  • With plugins: [axolotl.integrations.kernels.KernelsPlugin] and use_sonicmoe: true or use_scattermoe: true, ModelOpt NVFP4/MIXED_PRECISION checkpoints retain packed routed experts, while attention/shared linears load as BF16. Their checkpoint input_scale is deliberately discarded for training. The ordinary dense LoRA kernels apply; nvfp4_merge_aware: true additionally preserves the snapped effective-weight objective on marked projections.
  • Axolotl’s axolotl quantize NVFP4 weight configuration produces TorchAO NVFP4WeightOnlyConfig checkpoints. Load that saved checkpoint as base_model, set adapter: lora, and enable lora_qkv_kernel, lora_o_kernel, and lora_mlp_kernel for the optimized native weight-only adapter path.
  • dsv4_fp4_grouped_mode: nvfp4 controls routed-expert activation quantization inside the existing MoE backend; it is separate from native dense activation quantization and retains that backend’s training gradient convention.
  • Gemma’s nonexpert_quantization: nvfp4 uses frozen Marlin W4A16 non-expert modules for an experts-only LoRA configuration. It does not attach dense PEFT adapters and is outside the native TorchAO dense LoRA path above.

Externally supplied TorchAO checkpoints with dynamic NVFP4 activation quantization are supported for dense PEFT LoRA in unsharded runs, FSDP2, and DeepSpeed ZeRO-1/2/3, including the merge-aware path described above. The forward remains TorchAO’s quantized NVFP4 operation. For base-input backpropagation, Axolotl uses the dequantized frozen weight as a straight-through approximation; it is not the exact derivative through activation quantization. This convention applies to these native dense paths, not the fused low-rank LoRA kernel options.