Upload asvd_llm/heal/llama2-13b-base/lowrank_linear.py with huggingface_hub

asvd_llm/heal/llama2-13b-base/lowrank_linear.py

@@ -0,0 +1,83 @@
+"""Minimal standalone loader for the SVD-baseline compressed LLaMA-2-7B checkpoints.
+
+No dependency beyond torch + transformers. The state_dict stores every compressed
+projection as two factors of a LowRankLinear (A: out x k, B: k x in); ranks.json
+holds the per-layer integer rank of each projection. Untouched parameters
+(embeddings, norms) keep their stock transformers names.
+
+Usage:
+    import json, torch
+    from lowrank_linear import build
+    model = build(
+        base_id="NousResearch/Llama-2-7b-chat-hf",
+        ranks=json.load(open("ranks.json")),
+        state_dict_path="state_dict_fp16.pt",
+        device="cuda", dtype=torch.float16,
+    )
+    tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
+    out = model.generate(tokenizer("2+2=", return_tensors="pt").to("cuda").input_ids,
+                         max_new_tokens=16, do_sample=False)
+    print(tokenizer.decode(out[0]))
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LowRankLinear(nn.Module):
+    """y = A (B x).  A: out x r, B: r x in.  (LLaMA projections carry no bias.)"""
+
+    def __init__(self, A: torch.Tensor, B: torch.Tensor):
+        super().__init__()
+        self.A = nn.Parameter(A)
+        self.B = nn.Parameter(B)
+
+    def forward(self, x):
+        return F.linear(F.linear(x, self.B), self.A)
+
+    @torch.no_grad()
+    def materialize_(self) -> nn.Linear:
+        """Replace this module with the dense W = A @ B (inference only; never call
+        while training -- it cuts the gradient path)."""
+        W = (self.A @ self.B).to(self.A.dtype)
+        dense = nn.Linear(self.B.shape[1], self.A.shape[0], bias=False,
+                          device=W.device, dtype=W.dtype)
+        dense.weight.data = W
+        return dense
+
+
+def build(base_id: str, ranks: dict, state_dict_path: str,
+          device: str = "cuda", dtype=torch.float16):
+    from transformers import AutoModelForCausalLM
+
+    model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=dtype)
+    search = {}
+    for name, mod in model.named_modules():
+        if isinstance(mod, nn.Linear):
+            search[name] = mod
+    for path, k in ranks.items():
+        out_f, in_f = search[path].weight.shape
+        model_mod = model
+        for p in path.split(".")[:-1]:
+            model_mod = getattr(model_mod, p)
+        setattr(model_mod, path.split(".")[-1],
+                LowRankLinear(torch.zeros(out_f, int(k), dtype=dtype),
+                              torch.zeros(int(k), in_f, dtype=dtype)))
+    sd = torch.load(state_dict_path, map_location="cpu", weights_only=True)
+    model.load_state_dict(sd, strict=True)
+    del sd
+    return model.to(device)
+
+
+def materialize_all(model):
+    """Fuse every LowRankLinear back into a dense nn.Linear in place."""
+    for name, mod in list(model.named_modules()):
+        if isinstance(mod, LowRankLinear):
+            parent = model
+            for p in name.split(".")[:-1]:
+                parent = getattr(parent, p)
+            setattr(parent, name.split(".")[-1], mod.materialize_())
+    return model