LoRA

For MoE-specific LoRA (expert weight layout, EP sharding), see MoE LoRA.

LoRA (Low-Rank Adaptation) freezes the base model weights and adds trainable low-rank matrices to selected linear layers. xorl implements FSDP2-compatible LoRA for both local and server training.

How LoRA works

For a linear layer with weight matrix W₀ ∈ ℝ^(d×k), LoRA adds a trainable bypass:

W = W₀ + B · A where A ∈ ℝ^(r×k), B ∈ ℝ^(d×r), and r ≪ min(d, k)

W₀ is frozen. Only A and B are updated during training. The output scaling is lora_alpha / r.

By default, A uses Kaiming uniform initialization and B is zero, so the initial LoRA output is exactly zero and training starts from the pretrained model’s behavior. Set lora_b_init_std above zero to opt into deterministic normal initialization of B (using lora_b_init_seed); that mode intentionally starts with a nonzero adapter delta.

Basic Configuration

Add a lora section to your config:

lora:
  enable_lora: true
  lora_rank: 16
  lora_alpha: 32             # scaling = lora_alpha / lora_rank = 2.0
  lora_target_modules:
    - q_proj
    - k_proj
    - v_proj
    - o_proj
    - gate_proj
    - up_proj
    - down_proj
  save_lora_only: true       # checkpoint only LoRA weights, not base model

Fused projections

Many architectures store input-side projections fused: attention as qkv_proj and the dense or shared MLP as gate_up_proj. xorl audits the requested target set and fails injection if any target is unmatched.

For supported fused modules, split target names automatically create independent logical adapters while retaining the fused base projection. Dynamic LoRA adds each adapter to its corresponding output slice. Exact merged-forward mode canonically folds each adapter into that slice and still runs one fused GEMM. Checkpoint and weight-sync keys remain the standard split names:

  lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]

This is the default path for Qwen2, Qwen3, Qwen3 MoE, Qwen3.5, Llama, OLMo2, and GLM-4 MoE projections implemented by xorl.

Target the fused names. Adapts each fused module as one target, at no throughput cost. One lora_A is shared across the halves, so this is a more constrained parameterization than two independent adapters, and the exported adapter carries fused key names — which do not line up with a HuggingFace base that stores them split:

  lora_target_modules: [qkv_proj, gate_up_proj, down_proj, o_proj]

Unfuse before injection. This explicit fallback is for an architecture whose fused modules do not implement logical adapters, such as GPT-OSS attention. It splits the modules into real q_proj/k_proj/v_proj and gate_proj/up_proj before injection. It requires enable_lora and is rejected with enable_qlora:

lora:
  enable_lora: true
  unfuse_for_lora: true
  lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]

Unfusing changes both performance and floating-point operation order. Requalify trainer-to-sampler equality before using it in a workflow that depends on exact logprobs.

The layout is fixed for a run’s lifetime. Toggling the flag between runs renames mlp.gate_up_proj.* to mlp.gate_proj.* / mlp.up_proj.*, which checkpoint validation rejects with “Checkpoint incompatible with model” — resuming requires the same setting the checkpoint was written with.

Audited architecture defaults

Family	Default plain-LoRA scope	Projection handling
Qwen2, Qwen3, Qwen3 MoE, Llama, OLMo2, GLM-4 MoE	Attention and MLP	Independent logical adapters retain fused qkv and gate/up bases
Qwen3.5 dense and MoE	Attention, GDN g_proj, MLP, routed experts	Separate GDN/full-attention projections plus fused-base logical gate/up adapters
DeepSeek V3, Kimi K2/K2.5, GLM-5	MLA attention and MLP	Architecture-specific separate projections and expert adapters
DeepSeek V4	Attention only	Audited MLA/DSA projection names; routed-expert LoRA is not a generic default
GPT-OSS	Attention only	Set unfuse_for_lora: true for split q/k/v targets
MiniMax M3, Nemotron-H	Attention only	Separate attention projections; expert LoRA is not a generic default

Unknown model families have no guessed default. Set lora_target_modules explicitly after auditing their projection and expert semantics.

Key Parameters

Parameter	Default	Description
enable_lora	false	Enable LoRA injection
lora_rank	16	Rank r of the low-rank decomposition
lora_alpha	16	Scaling factor; effective scale = alpha/rank. Examples may override it to 32.
lora_b_init_std	0.0	Standard deviation for optional deterministic normal LoRA-B initialization. 0.0 preserves the standard no-op start.
lora_b_init_seed	0	Seed used by opt-in nonzero LoRA-B initialization.
lora_target_modules	null	List of module name patterns to inject LoRA into
unfuse_for_lora	false	Explicitly split supported fused projections before injection; use only when fused-base logical adapters are unavailable
save_lora_only	false	Save only LoRA weights in checkpoints

LoRA with MoE Models

For MoE models, you can target expert layers as well. Expert LoRA uses fused group GEMM for efficiency:

lora:
  enable_lora: true
  lora_rank: 16
  lora_alpha: 32
  lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]
  # gate/up/down also select routed experts on supported MoE families

Checkpoint Behavior

With save_lora_only: true, only the LoRA adapter weights are saved. The base model must be available separately to reconstruct the full model.

Checkpoint structure:

outputs/my_run/weights/{run_id}/step_{N}/
├── adapter_config.json
├── adapter_model.safetensors   # LoRA weights only
└── training_state/             # optimizer, scheduler, rng state

Compatible with the HF PEFT adapter format for easy loading:

from peft import PeftModel
model = PeftModel.from_pretrained(base_model, "outputs/my_run/weights/.../step_42")

Periodic Merge (ReLoRA)

For long training runs, periodically merge LoRA into the base weights and restart with fresh LoRA parameters. This avoids rank saturation and allows the effective rank to grow over training.

For QLoRA, see QLoRA — Periodic Merge.

LoRA in Server Training

In server training, LoRA is configured in the server YAML:

enable_lora: true
lora_rank: 16
lora_alpha: 32

The server supports multiple named adapters for multi-task or multi-policy training. Specify a model_id per request to route to a specific adapter:

{
  "batches": [...],
  "loss_fn": "causallm_loss",
  "model_id": "policy_v1"
}

Save a specific adapter:

POST /api/v1/save_weights
{"path": "outputs/adapters/policy_v1", "model_id": "policy_v1"}

Example Configs

See examples/local/dummy/configs/lora/ and examples/server/configs/lora/ for complete LoRA configs across model sizes.

Source

File	Description
src/xorl/lora/modules/base.py	LoraModule abstract base — defines r, lora_alpha, scaling, and from_module() factory
src/xorl/lora/modules/linear.py	LoraLinear — LoRA for nn.Linear; A (Kaiming), B (zeros), forward with scaling
src/xorl/lora/mapping.py	LORA_MAPPING registry and get_lora_class_for_module() lookup
src/xorl/models/layers/moe/lora.py	MoEExpertsLoRA — LoRA for fused expert tensors [E, I, H]