LoRA on Decoder Only: Half Parameters Match Full Fine-Tuning on DeepSeek-V4.1-Flash

Causal Encoder–Decoder Architecture Changes LoRA Placement

DeepSeek-V4.1-Flash is a causal encoder–decoder model: the first 20 layers (encoder) write a compressed global KV cache via CSA2, and the last 20 layers (decoder) read that cache to generate output. Only four layers (2, 8, 14, 20) write the shared KV; subsequent layers reuse the nearest kv_source layer's KV. This reduces 40 per-layer caches to four shared segments plus a 128-token sliding window per layer. Layer 20 is the pivot: its compressor converts the encoder's final hidden state into KV for all decoder layers. Because adapters on the encoder would alter the shared KV, each adapter would need its own cache; adapters on the decoder only change how the shared KV is read, allowing multiple adapters to share one cache if layer 20's KV projection is frozen.

Experimental Setup

Experiments run on Geometry3K (1,472 training, 224 validation problems) with four LoRA configurations, each trained with three seeds for 100 steps (BF16, 4 nodes × 8 H20s, context parallelism 2, expert parallelism 16, optimizer CPU offload, rank 16, α=32, max length 4096, dynamic batching). Learning rate 5e-5 chosen after an 8-step sweep (1e-5 to 2e-4) because its plateau loss was 0.006–0.044 lower than 2e-4 across all configs. Validation loss averaged over steps 50, 70, 100.

Decoder-Only LoRA Matches Full-Model Performance

Four configs compared (Table 1):

Full linear rank 16 (layers 0–39): 1.87B trainable params, val loss 0.293 (0.266–0.317).

Encoder half rank 16 (layers 0–19): 0.94B params, val loss 0.338 (0.329–0.348).

Decoder half rank 16 (layers 20–39): 0.94B params, val loss 0.296 (0.287–0.302).

Full linear rank 8 (layers 0–39): 0.94B params, val loss 0.285 (0.277–0.293).

Decoder half-edge matches full rank 16 (difference 0.003) with half the parameters. Encoder half-edge is consistently worse by 0.036–0.058 per seed. Rank 8 also matches full rank 16 but decoder half-edge still beats encoder half-edge at equal budget. Training loss curves (Fig. 3) show full rank 16 drops fastest early, but decoder half-edge catches up by the plateau.

Additional Validation

Hold-out test set (601 problems): Decoder half-edge loss 0.190 (no ViT adapter) vs full rank 16's 0.201; with ViT adapter 0.193 vs 0.210.

Longer training (300 steps, one seed, no ViT adapter): Minimum losses: rank 8 0.246, full rank 16 0.249, decoder half-edge 0.253 (all within 0.007), encoder half-edge 0.330.

Different task (A2UI, pure text, ~20× tokens/step): Plateau val loss: full rank 16 0.189, rank 8 0.193, decoder half-edge 0.192.

All three checks confirm decoder half-edge is on par with full-model LoRA.

Why Decoder Layers Work Better

Two factors align with results:

Memory concentration in encoder: Three of four kv_source layers are in the encoder; the fourth (layer 20) writes the KV that all decoder layers read. Encoder adapters perturb the shared memory; decoder adapters only change how it is read.

Deeper layers are more redundant: Pruning studies show deeper pretrained layers have more capacity to absorb new tasks while shallow layers retain general knowledge. Decoder layers are also closer to the loss, so gradients traverse a shorter Jacobian chain.

This ordering is robust to encoder learning-rate scaling (halving or layer-wise decay). Freezing layer 20's K/V projection and compressor projection enables cache sharing with no loss penalty (plateau loss changes −0.001, −0.008, −0.002 across three seeds).

Freezing ViT and Aligner Saves Memory and Stabilizes Training

Removing ViT/aligner adapters (9M params) lowers validation loss for three of four configs, speeds steps by 3–9%, and reduces peak VRAM by 5.2 GB (87.2–87.8 → 82.0–82.6 GB). Savings come from skipping ViT backward and recomputation (Fig. 6). With ViT frozen, encoder activations can be dropped; backward passes only through decoder's 20 layers, halving compute. Pipeline parallelism support added via Megatron-LM #8664. Frozen ViT also improves seed stability (range 0.003–0.012 vs 0.015–0.052). If ViT adapters are desired, use 1/10 learning rate (full rank 16: shared LR 0.289 range 0.054 → 1/10 LR 0.273 range 0.022; frozen ViT 0.274 range 0.014). Model still uses image features: swapping validation images raises loss by 0.51–0.66 vs base model's −0.06.

Module-Level Ablation: Which Linear Layers Need LoRA?

Under full linear rank 16 (no ViT adapter), tested subsets (Fig. 7, Fig. 8):

Indexer (top-k router): No gradient → adapter stays zero.

Attention side modules: Query up-projection ( wq_b), grouped output projection ( wo_a), and compressor. Each projection is split ( wq_a → wq_b, wo_a → wo_b); adapting one half affects the whole. Compressor exists only in four kv_source layers. Dropping all three (1.6% of params) matches full rank 16 (plateau 0.277/0.279 vs 0.285). On decoder half-edge, slim attention slightly better (0.290 vs 0.296); on rank 8, full modules win (0.272 vs 0.281).

Router expert: Down-projection (33% params, 58% gradient energy) can be removed with no loss (0.280 vs 0.285), saving optimizer state (adapter file 1.77B → 1.20B params). Gate/up projection (63% params, 5% gradient energy) is critical: with it loss 0.274 (range 0.014), without 0.293 (range 0.067).

Attention + shared expert adapters (≈3% params) contribute disproportionately: full modules vs router-only drops loss by 0.05 (both gate/up+down), 0.07 (gate/up only), 0.18 (down only).

Zeroing individual module groups in a trained adapter shows attention output projection (0.9% params) has largest single impact (+0.018), while router down-projection (58% gradient energy) has near-zero impact (−0.001). Gradient energy share divided by param share reveals per-parameter sensitivity: ViT ≈63 (effective step ~8× average), compressor KV projection 17 (but not needed), others ≤7. This ratio guides per-module LR scaling (LoRA+, muP). For language modules, shared LR 5e-5 works; LoRA+ with higher up-projection LR hurts (0.301 at 4×, 0.344 at 16×).

Engineering: Six Upstream PRs Enable Training

DeepSeek-V4.1-Flash breaks training-stack assumptions (single-layer recompute, harmless padding, all params on GPU, loadable checkpoints). Six PRs address this:

Megatron-LM #8018: CSA2 state (shared KV, indices, candidate blocks, mHC) must pass through activation checkpointing. Added stateful CSA2 wrapper in hybrid layer stack; supports packed sequences and context parallelism (boundary exchange, each global KV block owned by one rank).

Megatron-LM #7763: Multimodal router with separate text/image expert biases. Padding rows kept in dispatcher with zero routing probability and excluded from load balancing, preserving tensor shapes.

Megatron-LM #7762: Two Engram tables (3.84B rows × 256 dim, accessed at layers 1 and 14) kept in CPU pinned memory, sparse gather per micro-batch. Saves 22.9 GB VRAM per GPU (106.8 → 83.9 GB) with no step-time increase. DDP sync skips these CPU-resident params.

Megatron-Bridge #6199: Converts released HF checkpoint (FP8 32-block quantized, packed MXFP4 experts, mHC, ViT, aligner, Engram) to Megatron-Core and back. Engram streamed row-wise via SafeTensors without materializing full tables. Also enables LoRA on grouped output projection (bare parameter in model).

vLLM #57245: Registers LoRA for multimodal wrapper, grouped output projection, and compressor (direct GEMM calls). Supports BF16/FP8 base weights. Registers ViT/aligner, isolates adapter state per image.

veRL #7483: Corrects unmerged LoRA sync (base + adapter pushed separately). Detects wrapped base-weight naming, unwraps MXFP4 quantization for repacking. Also fixes router replay for multimodal: vLLM logs routing on compact sequence (image as single placeholder); Megatron runs on expanded sequence with image bias. Replay only text rows; image rows use native routing (Fig. 11).

Production Deployment

Mint Recursive (Mind Lab's serverless training platform) now supports DeepSeek-V4.1-Flash. Post-trained model on a financial client's proprietary benchmark beats official API by 5.96%.

Conclusions & Practical Recipe

Adapt decoder only (layers 20–39), freeze layer 20 KV projection → half params, shared KV cache. Full rank 8 is equally valid half-budget alternative.

Freeze ViT and aligner → 3 of 4 configs match or improve loss, 3–9% faster steps, −5.2 GB VRAM, stable seeds. If needed, use 1/10 LR.

Prune by ablation : Drop compressor, wq_b, wo_a, router down-projection (⅓ params). Keep router gate/up, attention + shared expert adapters.

Budget by host memory : With optimizer offload, adapter size affects host RAM and optimizer step time; peak VRAM nearly unchanged. Keep Engram on CPU.

Together these make adapters cheap enough for experience intelligence: many small decoder-side adapters share one base model and one KV cache, each absorbing new experience at far lower cost than retraining.