streamft 0.2.0__py3-none-any.whl

streamft/disk_quant_store.py

@@ -0,0 +1,353 @@
+"""Disk-backed (mmap) NF4 weight store: stream weights from SSD into VRAM.
+
+Idea
+----
+The RAM-resident streaming engine (`stream_quant_lora`) keeps the whole 4-bit
+base in host RAM. That is what caps the model size at what fits in RAM
+(a 14B NF4 base is ~8 GB; a 32B one would need ~18 GB). This module removes
+that cap: the pre-quantized weights stay in the bnb-4bit safetensors shards on
+disk and are memory-mapped. `safetensors.safe_open(...).get_tensor()` returns a
+*zero-copy* CPU view into the mmap, so
+
+  * host RAM is used only as an OS-managed page cache (evictable, LRU), i.e.
+    "use as much RAM as is free, keep the rest on SSD" happens for free, and
+  * the model-size limit becomes disk space + step time instead of host RAM.
+
+A decoder block is materialized directly on the GPU (uint8 data copied H2D,
+`Params4bit.from_prequantized` re-attaches the saved NF4 quant state) and moved
+back to `meta` afterwards, so only one block lives in VRAM at a time.
+
+Nothing here mutates the checkpoint: the base stays frozen 4-bit; only the
+LoRA adapters (fp32, CPU) are trained.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+# Components of a bnb NF4 quant state as written by transformers >= 4.4x:
+# one packed JSON blob (non-tensor items) + four tensors.
+QUANT_SUFFIXES = (
+    "weight.quant_state.bitsandbytes__nf4",
+    "weight.absmax",
+    "weight.quant_map",
+    "weight.nested_absmax",
+    "weight.nested_quant_map",
+)
+
+
+class DiskQuantStore:
+    """Key -> mmap'd tensor lookups over a (possibly sharded) bnb-4bit checkpoint."""
+
+    def __init__(self, model_dir, device: str = "cuda"):
+        self.dir = Path(model_dir)
+        index = json.loads((self.dir / "model.safetensors.index.json").read_text())
+        self.weight_map: dict[str, str] = dict(index["weight_map"])
+        self.device = device
+        self._handles: dict[str, object] = {}
+
+    def _handle(self, filename: str):
+        h = self._handles.get(filename)
+        if h is None:
+            from safetensors import safe_open
+
+            h = safe_open(str(self.dir / filename), framework="pt", device="cpu")
+            self._handles[filename] = h
+        return h
+
+    def has(self, key: str) -> bool:
+        return key in self.weight_map
+
+    def get(self, key: str) -> torch.Tensor:
+        """Zero-copy CPU view backed by the memory-mapped shard."""
+        return self._handle(self.weight_map[key]).get_tensor(key)
+
+
+def build_skeleton(model_dir, lora_r: int = 8, lora_alpha: int = 16,
+                   targets: str = "attn", head: str = "bf16",
+                   low_precision_shared: bool = False, with_lora: bool = True):
+    """Instantiate the model on the `meta` device with bnb Linear4bit + LoRA.
+
+    No weight memory is allocated here; only the (tiny) LoRA adapters are real.
+    Mirrors peft's `prepare_model_for_kbit_training`: every parameter is frozen
+    and all non-4bit weights (norms, biases, embedding/head) are cast to fp32 --
+    except the embedding/head when `low_precision_shared=True`, which keeps them
+    in the checkpoint dtype (needed for very large models).
+    Returns (model, config).
+    """
+    from transformers import AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig
+    from transformers.integrations.bitsandbytes import replace_with_bnb_linear
+
+    from .lora import LoRALinear, apply_lora
+
+    cfg = AutoConfig.from_pretrained(str(model_dir))
+    cfg.use_cache = False
+
+    with torch.device("meta"):
+        model = AutoModelForCausalLM.from_config(cfg)
+
+    qcfg = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_storage=torch.uint8,
+    )
+    replace_with_bnb_linear(
+        model,
+        # `lm_head` stays a plain Linear (bf16) unless we explicitly quantize it;
+        # tied models must keep it plain because it *is* the embedding tensor.
+        modules_to_not_convert=(
+            None if (head == "nf4" and not cfg.tie_word_embeddings) else ["lm_head"]
+        ),
+        quantization_config=qcfg,
+        pre_quantized=True,
+    )
+
+    r_attn, r_mlp = (lora_r, 0) if targets == "attn" else (lora_r, lora_r)
+    if not with_lora:
+        r_attn = r_mlp = 0
+
+    # 1) freeze everything, like prepare_model_for_kbit_training
+    for p in model.parameters():
+        p.requires_grad = False
+
+    # 2) cast all non-4bit params to fp32 (the RAM path does this, so the forward
+    #    matches bit-for-bit). `low_precision_shared` keeps *everything* in the
+    #    checkpoint dtype: mixing an fp32 norm with a bf16 head breaks F.linear.
+    if not low_precision_shared:
+        for name, p in list(model.named_parameters()):
+            if p.__class__.__name__ == "Params4bit":
+                continue
+            if p.dtype not in (torch.float16, torch.bfloat16):
+                continue
+            mod_name, _, attr = name.rpartition(".")
+            target = model.get_submodule(mod_name) if mod_name else model
+            target._parameters[attr] = torch.nn.Parameter(
+                target._parameters[attr].data.to(torch.float32), requires_grad=False)
+
+    if with_lora:
+        apply_lora(model, r_attn=r_attn, r_mlp=r_mlp, alpha_attn=lora_alpha,
+                   alpha_mlp=lora_alpha)
+        # LoRA adapters were created on `meta` (that is where the base lives); give
+        # them real CPU storage so the optimizer can train them.
+        import math
+
+        for mod in model.modules():
+            if isinstance(mod, LoRALinear):
+                out_f, in_f = mod.base.weight.shape
+                mod.A = torch.nn.Parameter(
+                    torch.randn(mod.r, in_f, dtype=torch.float32) / math.sqrt(in_f))
+                mod.B = torch.nn.Parameter(torch.zeros(out_f, mod.r, dtype=torch.float32))
+
+    # Tie the LM head to the input embedding when the config asks for it (the
+    # checkpoint then has no separate `lm_head.weight`).
+    if getattr(cfg, "tie_word_embeddings", False):
+        from .stream_quant_lora import base_model_of
+
+        lm = base_model_of(model)
+        lm.lm_head._parameters["weight"] = lm.model.embed_tokens._parameters["weight"]
+    return model, cfg
+
+
+def _rope_theta(cfg) -> float:
+    """Theta location moved between transformers versions; look in all of them."""
+    rp = getattr(cfg, "rope_parameters", None)
+    if isinstance(rp, dict) and "rope_theta" in rp:
+        return float(rp["rope_theta"])
+    rs = getattr(cfg, "rope_scaling", None)
+    if isinstance(rs, dict) and "rope_theta" in rs:
+        return float(rs["rope_theta"])
+    return float(getattr(cfg, "rope_theta", 1_000_000.0))
+
+
+def fix_rotary(inner, cfg):
+    """`rotary_emb` buffers are computed at init, so they are `meta` here.
+
+    Rebuild the module for real (the constructor knows how to read rope config);
+    fall back to filling `inv_freq` by hand.
+    """
+    rot = inner.rotary_emb
+    try:
+        fresh = type(rot)(cfg)
+        if not fresh.inv_freq.is_meta:
+            inner.rotary_emb = fresh
+            return
+    except Exception:
+        pass
+
+    dim = cfg.hidden_size // cfg.num_attention_heads
+    theta = _rope_theta(cfg)
+    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    rot.inv_freq = inv
+    rot.original_inv_freq = inv.clone()
+    if getattr(rot, "attention_scaling", None) is None:
+        rot.attention_scaling = 1.0
+
+
+def _set_param(module, name: str, tensor, requires_grad: bool = False) -> None:
+    """Replace a parameter's storage.
+
+    Works even when the current parameter lives on `meta`: PyTorch forbids
+    `meta_param.data = cuda_tensor` ("incompatible tensor type"), so we swap the
+    `Parameter` object instead. Base weights are never in an optimizer, so losing
+    object identity is fine.
+    """
+    module._parameters[name] = torch.nn.Parameter(tensor, requires_grad=requires_grad)
+
+
+def _set_param_cast(module, name: str, tensor) -> None:
+    """`_set_param` that keeps the destination dtype (e.g. fp32 norms)."""
+    old = module._parameters.get(name)
+    if old is not None and old.dtype != tensor.dtype:
+        tensor = tensor.to(old.dtype)
+    _set_param(module, name, tensor)
+
+
+def materialize_shared(model, store: DiskQuantStore, cfg, device: str = "cuda"):
+    """Load embed, final norm and lm_head once (they are not streamed)."""
+    from .stream_quant_lora import base_model_of
+
+    lm = base_model_of(model)
+    inner = lm.model
+    emb = inner.embed_tokens
+    tied = lm.lm_head.weight is emb.weight
+
+    weight = store.get("model.embed_tokens.weight")
+    if tied:
+        # the same tensor feeds the LM head, so it must sit on the GPU
+        _set_param_cast(emb, "weight", weight.to(device))
+    else:
+        # zero-copy mmap view on CPU: gathering rows only touches the pages for
+        # the tokens actually used
+        _set_param_cast(emb, "weight", weight)
+
+    if store.has("model.norm.weight"):
+        _set_param_cast(inner.norm, "weight", store.get("model.norm.weight").to(device))
+    if tied:
+        lm.lm_head._parameters["weight"] = emb._parameters["weight"]
+    elif store.has("lm_head.weight"):
+        _set_param_cast(lm.lm_head, "weight", store.get("lm_head.weight").to(device))
+
+    fix_rotary(inner, cfg)
+    inner.rotary_emb.to(device)
+    return lm
+
+
+def _layer_prefix(i: int) -> str:
+    return f"model.layers.{i}"
+
+
+def _key_of(prefix: str, module_name: str) -> str:
+    """Map a module path inside a layer to its checkpoint prefix."""
+    name = module_name
+    if name.endswith(".base"):
+        name = name[: -len(".base")]
+    return f"{prefix}.{name}" if name else prefix
+
+
+def load_layer(layer, i: int, store: DiskQuantStore, device: str = "cuda") -> None:
+    """Materialize decoder block `i` on `device` straight from the mmap'd store."""
+    import bitsandbytes as bnb
+
+    prefix = _layer_prefix(i)
+    for name, child in layer.named_modules():
+        sub = _key_of(prefix, name)
+        if isinstance(child, bnb.nn.Linear4bit):
+            wkey = f"{sub}.weight"
+            if store.has(wkey):
+                qs = {s: store.get(f"{sub}.{s}") for s in QUANT_SUFFIXES
+                      if store.has(f"{sub}.{s}")}
+                child.weight = bnb.nn.Params4bit.from_prequantized(
+                    data=store.get(wkey), quantized_stats=qs,
+                    requires_grad=False, device=device, module=child,
+                )
+                child._is_hf_initialized = True
+            if child.bias is not None and store.has(f"{sub}.bias"):
+                _set_param_cast(child, "bias", store.get(f"{sub}.bias").to(device))
+        else:
+            for pname, _p in list(child.named_parameters(recurse=False)):
+                key = f"{sub}.{pname}"
+                if store.has(key):
+                    _set_param_cast(child, pname, store.get(key).to(device))
+            for bname, b in list(child.named_buffers(recurse=False)):
+                key = f"{sub}.{bname}"
+                if store.has(key):
+                    t = store.get(key).to(device)
+                    if b is not None and b.dtype != t.dtype:
+                        t = t.to(b.dtype)
+                    child._buffers[bname] = t
+
+    _move_lora(layer, device)
+
+
+def release_layer(layer, i: int, store: DiskQuantStore) -> None:
+    """Drop block `i`'s weights (back to `meta`) so VRAM/RAM are freed."""
+    for child in layer.modules():
+        for pname, p in list(child.named_parameters(recurse=False)):
+            if p.requires_grad:  # LoRA adapters are kept (moved to CPU)
+                continue
+            child._parameters[pname] = torch.nn.Parameter(
+                torch.empty(tuple(p.shape), dtype=p.dtype, device="meta"),
+                requires_grad=False,
+            )
+        for bname, b in list(child.named_buffers(recurse=False)):
+            if b is not None:
+                child._buffers[bname] = torch.empty(
+                    tuple(b.shape), dtype=b.dtype, device="meta")
+    _move_lora(layer, "cpu")
+
+
+def _move_lora(layer, device) -> None:
+    """LoRA adapters deliberately stay on CPU.
+
+    Gradients are accumulated in CPU buffers and the optimizer steps on CPU, so
+    moving the adapters with the block would put param and grad on different
+    devices. `LoRALinear.forward` hops them to the activation's device instead.
+    """
+    return
+
+
+class LayerCache:
+    """Keep up to `capacity` decoder blocks resident on the GPU (LRU eviction).
+
+    capacity=1 reproduces the plain "one block at a time" streaming. A larger
+    capacity is what makes a bigger GPU fast: blocks stay resident between the
+    forward and the backward pass, and if capacity >= the number of layers the
+    weights are loaded once and never streamed again.
+    """
+
+    def __init__(self, store: "DiskQuantStore", layers, device: str = "cuda",
+                 capacity: int = 1):
+        self.store = store
+        self.layers = layers
+        self.device = device
+        self.capacity = max(1, int(capacity))
+        self.order: list[int] = []
+
+    def stage(self, layer, i: int) -> None:
+        if i in self.order:  # already resident: refresh its LRU position
+            self.order.remove(i)
+            self.order.append(i)
+            return
+        while len(self.order) >= self.capacity:
+            j = self.order.pop(0)
+            release_layer(self.layers[j], j, self.store)
+        load_layer(layer, i, self.store, self.device)
+        self.order.append(i)
+
+    def release(self, layer, i: int) -> None:
+        # blocks stay resident until the LRU evicts them
+        return
+
+    def set_capacity(self, capacity: int) -> None:
+        """Change the window at runtime (used by the measured autotuner)."""
+        self.capacity = max(1, int(capacity))
+        while len(self.order) > self.capacity:
+            j = self.order.pop(0)
+            release_layer(self.layers[j], j, self.store)
+
+    def stats(self) -> dict:
+        return {"resident": len(self.order), "capacity": self.capacity}

streamft/doctor.py

@@ -0,0 +1,34 @@
+"""`streamft doctor` — what this machine can actually do."""
+from __future__ import annotations
+
+import importlib
+import platform
+import sys
+
+
+def main(argv=None) -> int:
+    print("python", sys.version.split()[0], "|", platform.platform())
+    for pkg in ("torch", "transformers", "bitsandbytes", "safetensors", "datasets", "peft"):
+        try:
+            m = importlib.import_module(pkg)
+            print(f"  {pkg:14s} {getattr(m, '__version__', '?')}")
+        except Exception:
+            print(f"  {pkg:14s} MISSING")
+    try:
+        import torch
+
+        print("cuda:", torch.cuda.is_available(), "|", torch.version.cuda)
+        if torch.cuda.is_available():
+            p = torch.cuda.get_device_properties(0)
+            free, total = torch.cuda.mem_get_info(0)
+            print(f"gpu: {p.name} | {total/1e9:.1f} GB total, {free/1e9:.1f} GB free")
+            print(f"capability: {torch.cuda.get_device_capability(0)} | "
+                  f"sdpa: {hasattr(torch.nn.functional, 'scaled_dot_product_attention')}")
+    except Exception as e:
+        print("torch check failed:", e)
+    print("\ntip: streamft fit --model <dir|hf-id> --gpu <GB> --seq <tokens>")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

streamft/fit.py

@@ -0,0 +1,128 @@
+"""`streamft fit` — will this model fit on my card, and what do we actually know?
+
+Two sources of truth, kept separate on purpose:
+
+  * **memory** is arithmetic from the config: NF4 weight bytes + the shared tensors that
+    stay in fp32 (embeddings / LM head) + a fixed reserve. It answers "whole model, or
+    streaming?" and nothing more.
+  * **speed and block count** are printed only when this exact size+seq was measured on
+    the reference machine (RTX 3050 Laptop 6 GB) and is listed in MEASURED. For anything
+    else the tool refuses to guess and tells you to start at one block and let the
+    trainer's autotune measure the marginal.
+
+Usage:
+  streamft fit --model Qwen/Qwen2.5-14B-Instruct --gpu 6 --seq 512
+  streamft fit --model checkpoints/qwen14b_nf4 --gpu 6 --json
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+# Measured end-to-end on an RTX 3050 Laptop, 6 GB, Windows 11, AC power, NF4 base,
+# LoRA r=8 on q/k/v/o, batch 1, Dolly-15k. step_s = median warm step.
+MEASURED = {
+    (3, 512): {"blocks": 36, "peak_gb": 3.85, "step_s": 0.938, "mode": "all resident"},
+    (7, 512): {"blocks": 28, "peak_gb": 5.51, "step_s": 1.696, "mode": "all resident"},
+    (14, 512): {"blocks": 6, "peak_gb": 3.86, "step_s": 6.726, "mode": "streaming"},
+    (32, 1024): {"blocks": 2, "peak_gb": 3.90, "step_s": 12.541, "mode": "streaming"},
+}
+
+RESERVE_GB = 0.9  # CUDA context + allocator fragmentation + the bf16 recompute buffers
+
+
+def load_config(model: str) -> dict:
+    p = Path(model)
+    if p.is_dir() and (p / "config.json").exists():
+        return json.loads((p / "config.json").read_text(encoding="utf-8"))
+    from huggingface_hub import hf_hub_download
+
+    return json.loads(Path(hf_hub_download(model, "config.json")).read_text(encoding="utf-8"))
+
+
+def model_math(cfg: dict) -> dict:
+    h = cfg.get("hidden_size") or cfg.get("n_embd")
+    layers = cfg.get("num_hidden_layers") or cfg.get("n_layer")
+    inter = cfg.get("intermediate_size") or 4 * h
+    vocab = cfg.get("vocab_size", 0)
+    heads = cfg.get("num_attention_heads", 1)
+    kv = cfg.get("num_key_value_heads", heads)
+    if not (h and layers):
+        raise SystemExit("config.json has no hidden_size / num_hidden_layers")
+
+    attn = h * h + 2 * h * (h * kv // heads) + h * h
+    mlp = 3 * h * inter
+    params = float(layers * (attn + mlp) + vocab * h)
+
+    nf4_gb = params * 0.5 / 1e9 * 1.07          # 4-bit + bitsandbytes overhead
+    tied = bool(cfg.get("tie_word_embeddings", False))
+    # What actually sits on the GPU: the embedding is pinned to CPU by the engine
+    # (only the gathered activation moves over), the head stays in bf16, and the
+    # presets use --shared-dtype as-stored. A tied head has no separate tensor.
+    head_gb = 0.0 if tied else vocab * h * 2 / 1e9
+    return {"params": params, "layers": int(layers), "nf4_gb": nf4_gb,
+            "head_gb": head_gb, "tied": tied}
+
+
+def main(argv: list[str] | None = None) -> int:
+    ap = argparse.ArgumentParser(prog="streamft fit")
+    ap.add_argument("--model", required=True, help="local dir or HF id")
+    ap.add_argument("--gpu", type=float, default=6.0, help="VRAM in GB (default 6)")
+    ap.add_argument("--seq", type=int, default=512)
+    ap.add_argument("--json", action="store_true")
+    args = ap.parse_args(argv)
+
+    m = model_math(load_config(args.model))
+    need = m["nf4_gb"] + m["head_gb"] + RESERVE_GB
+    fits_whole = need <= args.gpu
+    size_b = round(m["params"] / 1e9)
+
+    measured = None
+    for (b, seq), row in MEASURED.items():
+        if abs(size_b - b) <= max(1, round(b * 0.1)) and seq == args.seq:
+            measured = dict(row, ref_size_b=b, ref_seq=seq)
+            break
+
+    out = {
+        "model": args.model,
+        "params_b": round(m["params"] / 1e9, 2),
+        "layers": m["layers"],
+        "nf4_weights_gb": round(m["nf4_gb"], 1),
+        "embedding_on_cpu": True,
+        "lm_head_bf16_gb": round(m["head_gb"], 1),
+        "needed_whole_gb": round(need, 1),
+        "gpu_gb": args.gpu,
+        "seq": args.seq,
+        "fits_whole": fits_whole,
+        "needs_streaming": not fits_whole,
+        "measured_on_rtx3050_6gb": measured,
+        "measured_elsewhere": measured is None,
+    }
+    if args.json:
+        print(json.dumps(out, indent=1))
+        return 0
+
+    print(args.model)
+    print(f"  {out['params_b']}B params, {m['layers']} layers")
+    print(f"  NF4 blocks ~{out['nf4_weights_gb']} GB + bf16 head ~{out['lm_head_bf16_gb']} GB "
+          f"+ reserve {RESERVE_GB} GB (embedding stays on CPU)")
+    print(f"  -> whole model needs ~{out['needed_whole_gb']} GB on {args.gpu:g} GB")
+    if fits_whole:
+        print("  verdict: fits whole at this seq (no streaming needed; >1k can still spill)")
+    else:
+        print("  verdict: needs streaming (one or more blocks on the GPU at a time)")
+    if measured:
+        print(f"  measured on RTX 3050 6 GB: {measured['blocks']} block(s), "
+              f"{measured['peak_gb']} GB peak, {measured['step_s']} s/step at seq {measured['ref_seq']}")
+    else:
+        print("  speed/blocks: not measured for this size at this seq — start with")
+        print("                 --blocks-on-gpu 1 and let autotune grow it (it measures")
+        print("                 the real MB/block on your card instead of guessing)")
+    print("  next: streamft prepare --model <dir> --save-quant checkpoints/<name>")
+    print("        streamft train --model-dir checkpoints/<name> --sft --steps 30")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

streamft/lora.py

@@ -0,0 +1,108 @@
+"""Minimal hand-rolled LoRA for the frozen base model (no peft dependency).
+
+LoRALinear wraps an existing nn.Linear and adds a low-rank update
+    y = base(x) + scale * (x @ A^T @ B^T),   scale = alpha / r
+B starts at zero so the wrapped module is an identity at init.
+A/B are kept in fp32 even when the base is fp16 (fp16 updates would vanish
+below resolution); the low-rank path is cheap.
+
+merge() folds the update into the base weight (in-place) so that inference
+runs at zero extra cost. Call it once, after training.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, r: int = 16, alpha: int = 16, dropout: float = 0.0):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+        self.r = r
+        self.scale = alpha / r
+        w = base.weight
+        self.A = nn.Parameter(torch.randn(r, w.shape[1], device=w.device,
+                                          dtype=torch.float32) / math.sqrt(w.shape[1]))
+        self.B = nn.Parameter(torch.zeros(w.shape[0], r, device=w.device, dtype=torch.float32))
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x):
+        y = self.base(x)
+        h = self.dropout(x).float()
+        # A/B live on CPU (grads are accumulated there and the optimizer steps
+        # there); hop them to the activation's device for the matmul. The low-rank
+        # path is tiny, so this costs a negligible amount of PCIe traffic.
+        A = self.A if self.A.device == h.device else self.A.to(h.device)
+        B = self.B if self.B.device == h.device else self.B.to(h.device)
+        return y + (self.scale * (h @ A.t() @ B.t())).to(y.dtype)
+
+    @torch.no_grad()
+    def merge(self):
+        self.base.weight.data += (self.scale * (self.B @ self.A)).to(self.base.weight.dtype)
+        self.A.data.zero_()
+        self.B.data.zero_()
+
+
+def apply_lora(model, r_attn: int = 16, r_mlp: int = 16, alpha_attn: int = 16,
+               alpha_mlp: int = 16, layers=None):
+    """Replace q/k/v/o and mlp projections with LoRALinear in-place."""
+    base = model.model if hasattr(model, "model") else model
+    n = len(base.layers)
+    if layers is None:
+        layers = list(range(n))
+    for i in layers:
+        layer = base.layers[i]
+        attn = layer.self_attn.orig if hasattr(layer.self_attn, "orig") else layer.self_attn
+        for name, r, a in [("q_proj", r_attn, alpha_attn), ("k_proj", r_attn, alpha_attn),
+                           ("v_proj", r_attn, alpha_attn), ("o_proj", r_attn, alpha_attn)]:
+            mod = getattr(attn, name)
+            if not isinstance(mod, LoRALinear):
+                setattr(attn, name, LoRALinear(mod, r=r, alpha=a))
+        if r_mlp > 0:
+            for name in ("gate_proj", "up_proj", "down_proj"):
+                mod = getattr(layer.mlp, name)
+                if not isinstance(mod, LoRALinear):
+                    setattr(layer.mlp, name, LoRALinear(mod, r=r_mlp, alpha=alpha_mlp))
+
+
+def lora_parameters(model):
+    for mod in model.modules():
+        if isinstance(mod, LoRALinear):
+            yield mod.A
+            yield mod.B
+
+
+def merge_lora(model):
+    for mod in model.modules():
+        if isinstance(mod, LoRALinear):
+            mod.merge()
+
+
+def count_lora_params(model):
+    return sum(p.numel() for p in lora_parameters(model))
+
+
+def lora_state_dict(model):
+    """Adapter-only state dict (A/B tensors), CPU copies — for saving checkpoints."""
+    out = {}
+    for name, p in model.named_parameters():
+        if p.requires_grad and (name.endswith(".A") or name.endswith(".B")):
+            out[name] = p.detach().cpu().clone()
+    return out
+
+
+def set_lora_state_dict(model, sd):
+    """Load an adapter saved with lora_state_dict (module tree must match)."""
+    own = {n for n, p in model.named_parameters() if p.requires_grad}
+    missing = [k for k in sd if k not in own]
+    if missing:
+        raise KeyError(f"adapter keys not found in model: {missing[:4]}")
+    with torch.no_grad():
+        for name, p in model.named_parameters():
+            if name in sd:
+                p.copy_(sd[name].to(device=p.device, dtype=p.dtype))