Foundation Model Engineering

9.3 Parameter-Efficient Fine-Tuning (PEFT)

Parameter-Efficient Fine-Tuning (PEFT) freezes most pretrained weights and learns a small task-specific state. It lowers gradient and optimizer-state cost, enables many adapters per base model, and can simplify rollback. It does not make activation memory disappear, and the base model, tokenizer, template, adapter, and runtime remain one compatibility unit.


Start with a Memory Budget

For full fine-tuning, account separately for model weights, master weights if used, gradients, optimizer moments, activations, temporary kernels, communication buffers, and checkpoints. A statement such as “70B needs multiple terabytes” is only meaningful after specifying precision, optimizer, sharding, sequence length, microbatch, and activation recomputation.

LoRA removes optimizer states and gradients for frozen base weights, but forward activations still scale with batch and sequence length. Gradient checkpointing trades extra compute for activation memory. FSDP or ZeRO shards model-related states across devices; tensor/context parallelism addresses layers or sequences that do not fit locally. Choose from a measured memory worksheet rather than by model size alone.


LoRA: A Low-Rank Update

For a frozen weight W0Rd×kW_0\in\mathbb{R}^{d\times k}, LoRA learns

W=W0+ΔW,ΔW=αrBA,W = W_0 + \Delta W, \qquad \Delta W = \frac{\alpha}{r}BA,

where ARr×kA\in\mathbb{R}^{r\times k}, BRd×rB\in\mathbb{R}^{d\times r}, and rmin(d,k)r\ll\min(d,k). Initializing one factor to zero keeps the initial function unchanged [1].

The decisive choices are not only rank rr. Sweep rank, α\alpha, learning rate, effective batch in supervised tokens, token budget, dropout, and target modules. q_proj/v_proj is a cheap baseline; all-linear often gives more capacity at higher memory and checkpoint size. Train embeddings or lm_head only when new tokens, domain vocabulary, or output calibration justify it—and record tied-weight behavior.

Educational Merge Semantics

The module below illustrates an idempotent merge/unmerge contract. A real injector must wrap the actual pretrained module rather than constructing a random replacement, preserve bias/dtype/device, handle sharding, and test output equivalence.

import torch
import torch.nn as nn

class LoRAWrapper(nn.Module):
    def __init__(self, base: nn.Linear, rank=16, alpha=32):
        super().__init__()
        self.base = base
        self.base.requires_grad_(False)
        self.A = nn.Parameter(torch.empty(rank, base.in_features, device=base.weight.device))
        self.B = nn.Parameter(torch.zeros(base.out_features, rank, device=base.weight.device))
        nn.init.kaiming_uniform_(self.A, a=5 ** 0.5)
        self.scale = alpha / rank
        self.merged = False

    def delta(self):
        return (self.B @ self.A).to(self.base.weight.dtype) * self.scale

    def forward(self, x):
        if self.merged:
            return self.base(x)
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scale

    @torch.no_grad()
    def merge(self):
        if not self.merged:
            self.base.weight.add_(self.delta())
            self.merged = True

    @torch.no_grad()
    def unmerge(self):
        if self.merged:
            self.base.weight.sub_(self.delta())
            self.merged = False

Test repeated merge() calls, merge→unmerge restoration, and outputs before/after merge within a dtype-appropriate tolerance. Do not set frozen base weights trainable merely to mark merge state.


QLoRA: Quantized Storage, Floating-Point Updates

QLoRA stores frozen base weights in 4-bit NF4, optionally uses double quantization for quantization constants, and computes through dequantized values while LoRA parameters remain trainable [2]. The compute dtype—usually BF16 on supported hardware, otherwise carefully tested FP16—is distinct from the storage dtype.

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model_id = "your-pinned-base-checkpoint"
quant = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id, quantization_config=quant, device_map="auto"
)
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
model = get_peft_model(model, LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules="all-linear",
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
))

This configuration is a starting point, not a universal recipe. Verify kernel/backend support, normalization and output-head dtypes, peak memory, trainable parameter names, and that gradient checkpointing interacts correctly with caches. Quantization error can change the best learning rate and rank.


Choosing and Shipping an Adapter

Compare at least a cheap LoRA baseline, a higher-capacity target-module/rank setting, and full fine-tuning only when the budget permits. Keep data order, token budget, evaluation, and tuning effort equal. DoRA, prompt tuning, and other variants can be useful, but evidence from one architecture or vision benchmark does not establish a universal LLM ranking.

Version the deployment as one artifact bundle:

  • base checkpoint revision and hash;
  • tokenizer files, chat template, special-token map, and generation config;
  • adapter method/config/weights and library versions;
  • precision and quantization configuration;
  • dataset manifest, code/config/container, and evaluation report.

Loading the right adapter on the wrong base can produce plausible but invalid output. Add an explicit compatibility guard. For merged deployment, merge into an explicitly selected floating-point dtype, save a new immutable checkpoint, and compare logits/generations against unmerged inference. Do not assume merging directly into 4-bit weights is lossless or reversible. Retain the original base and adapter for rollback.

LoRA Parameter Efficiency Calculator

Full Fine-Tuning

Matrix W (d × d)

16,777,216
Trainable Parameters

LoRA

Matrices A (d × r) + B (r × d)

131,072
Trainable Parameters
Parameter Reduction: 99.2188%

Run and Release Gates

Monitor token-normalized loss, trainable/frozen parameter counts, gradient norms by adapter module, nonfinite values, peak allocated/reserved memory, tokens per second, data wait, and domain/length slices. Save adapter plus optimizer, scheduler, scaler, RNG, sampler, and data cursor for exact continuation.

Abort on unexpected trainable base parameters, missing target modules, persistent overflow, memory beyond the declared headroom, or held-out regression beyond the gate. Release only after paired domain, base-retention, safety, format/tool, latency, and cost evaluation. Exercise both adapter unload and last-known-good base switching.


Quizzes

Quiz 1: Why can a LoRA run still run out of memory even though less than one percent of parameters are trainable? LoRA greatly reduces trainable gradients and optimizer states, but the frozen weights, activations, temporary kernels, communication buffers, and evaluation still consume memory. Long sequences and microbatch size can dominate activation memory.

Quiz 2: What is the difference between QLoRA’s storage dtype and compute dtype? Frozen base weights are stored in 4-bit NF4, while matrix operations use dequantized floating-point values such as BF16. Adapter weights and optimizer state are also trained in floating-point representations.

Quiz 3: Why must merge be idempotent? A second merge must not add the same delta again. Explicit merged state plus tested unmerge behavior prevents silent weight corruption and makes deployment rollback auditable.

Quiz 4: Which files identify a deployable adapter? At minimum: exact base revision, tokenizer and chat template, special-token and generation configuration, adapter config and weights, precision/quantization settings, code/library versions, and the evaluation report.


References

  1. Hu, E. J., et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685.
  2. Dettmers, T., et al. (2023). QLoRA: Efficient Finetuning of Quantized LLMs. arXiv:2305.14314.
  3. Hugging Face. PEFT documentation. Documentation.