add back add_linear_to_key

weights.py

@@ -1,10 +1,26 @@
 import safetensors
 import torch
 import torch.nn as nn
+import re
 
 from contextlib import contextmanager
 from typing import Callable, List
 
+from .text import build_text_model
+from .config import TextConfig
+
+
+# Our custom linear has an module named linear, so we add linear to the name
+def add_linear_to_key(k: str) -> str:
+    k = k.replace("model.", "")
+    if k.startswith("text.") and ".linear." not in k:
+        k = re.sub(
+            r"(attn\.(?:qkv|proj)|mlp\.fc[12])\.(weight|bias)$",
+            r"\1.linear.\2",
+            k,
+        )
+    return k
+
 
 @contextmanager
 def safetensors_open(safetensors_file: str):
@@ -27,12 +43,17 @@ def safetensors_open(safetensors_file: str):
         yield get_tensor
 
 
-def _load_weights(get_tensor: Callable[[str], torch.Tensor], model: nn.Module) -> None:
+def _load_weights(
+    get_tensor: Callable[[str], torch.Tensor],
+    model: nn.Module,
+    is_quantized: bool = False,
+) -> None:
     """Internal function to load weights using a tensor getter function."""
     model = model.to(dtype=torch.float16)
 
     vision = model.vision
     region = model.region
+
     weight_map = {
         "vision_encoder.encoder.model.visual.patch_embed.linear.weight": vision[
             "patch_emb"
@@ -90,23 +111,42 @@ def _load_weights(get_tensor: Callable[[str], torch.Tensor], model: nn.Module) -
             }
         )
 
-    for i in range(len(model.text["blocks"])):
-        prefix = f"text_model.transformer.h.{i}"
-        blk = model.text["blocks"][i]
-        weight_map.update(
-            {
-                f"{prefix}.ln.weight": blk["ln"].weight,
-                f"{prefix}.ln.bias": blk["ln"].bias,
-                f"{prefix}.mixer.Wqkv.weight": blk["attn"]["qkv"].weight,
-                f"{prefix}.mixer.Wqkv.bias": blk["attn"]["qkv"].bias,
-                f"{prefix}.mixer.out_proj.weight": blk["attn"]["proj"].weight,
-                f"{prefix}.mixer.out_proj.bias": blk["attn"]["proj"].bias,
-                f"{prefix}.mlp.fc1.weight": blk["mlp"]["fc1"].weight,
-                f"{prefix}.mlp.fc1.bias": blk["mlp"]["fc1"].bias,
-                f"{prefix}.mlp.fc2.weight": blk["mlp"]["fc2"].weight,
-                f"{prefix}.mlp.fc2.bias": blk["mlp"]["fc2"].bias,
-            }
-        )
+    if not is_quantized:
+        for i in range(len(model.text["blocks"])):
+            prefix = f"text_model.transformer.h.{i}"
+            blk = model.text["blocks"][i]
+            weight_map.update(
+                {
+                    f"{prefix}.ln.weight": blk["ln"].weight,
+                    f"{prefix}.ln.bias": blk["ln"].bias,
+                    f"{prefix}.mixer.Wqkv.weight": blk["attn"]["qkv"].weight,
+                    f"{prefix}.mixer.Wqkv.bias": blk["attn"]["qkv"].bias,
+                    f"{prefix}.mixer.out_proj.weight": blk["attn"]["proj"].weight,
+                    f"{prefix}.mixer.out_proj.bias": blk["attn"]["proj"].bias,
+                    f"{prefix}.mlp.fc1.weight": blk["mlp"]["fc1"].weight,
+                    f"{prefix}.mlp.fc1.bias": blk["mlp"]["fc1"].bias,
+                    f"{prefix}.mlp.fc2.weight": blk["mlp"]["fc2"].weight,
+                    f"{prefix}.mlp.fc2.bias": blk["mlp"]["fc2"].bias,
+                }
+            )
+    else:  # add special quantized path. this is specific to how bitblas expects weights to be loaded (.qweight)
+        for i in range(len(model.text["blocks"])):
+            prefix = f"text_model.transformer.h.{i}"
+            blk = model.text["blocks"][i]
+            weight_map.update(
+                {
+                    f"{prefix}.ln.qweight": blk["ln"].weight,
+                    f"{prefix}.ln.bias": blk["ln"].bias,
+                    f"{prefix}.mixer.Wqkv.qweight": blk["attn"]["qkv"].weight,
+                    f"{prefix}.mixer.Wqkv.bias": blk["attn"]["qkv"].bias,
+                    f"{prefix}.mixer.out_proj.qweight": blk["attn"]["proj"].weight,
+                    f"{prefix}.mixer.out_proj.bias": blk["attn"]["proj"].bias,
+                    f"{prefix}.mlp.fc1.qweight": blk["mlp"]["fc1"].weight,
+                    f"{prefix}.mlp.fc1.bias": blk["mlp"]["fc1"].bias,
+                    f"{prefix}.mlp.fc2.qweight": blk["mlp"]["fc2"].weight,
+                    f"{prefix}.mlp.fc2.bias": blk["mlp"]["fc2"].bias,
+                }
+            )
 
     for key, tensor in weight_map.items():
         tensor.data.copy_(get_tensor(key))
@@ -120,34 +160,37 @@ def _load_weights(get_tensor: Callable[[str], torch.Tensor], model: nn.Module) -
 def load_weights_from_safetensors(weights_file: str, model: nn.Module) -> None:
     """Load weights from a safetensors file into a MoondreamModel instance."""
     with safetensors_open(weights_file) as get_tensor:
+        all_keys = get_tensor.keys()
         if (
-            "vision.blocks.0.attn.proj.bias" in get_tensor.keys()
-            or "model.vision.blocks.0.attn.proj.bias" in get_tensor.keys()
+            "vision.blocks.0.attn.proj.bias" in all_keys
+            or "model.vision.blocks.0.attn.proj.bias" in all_keys
         ):
             with safetensors_open(weights_file) as get_tensor:
-                tensors = {
-                    k.replace("model.", ""): get_tensor(k) for k in get_tensor.keys()
-                }
+                tensors = {add_linear_to_key(k): get_tensor(k) for k in all_keys}
                 model.load_state_dict(tensors, strict=False)
         else:
             # Wrap the get_tensor function to handle key normalization
-            name_map = {k.replace("._orig_mod", ""): k for k in get_tensor.keys()}
+            name_map = {k.replace("._orig_mod", ""): k for k in all_keys}
             _load_weights(
-                lambda x: get_tensor(name_map[x]).to(dtype=torch.bfloat16), model
+                lambda x: get_tensor(name_map[x]).to(dtype=torch.float16),
+                model
             )
 
 
 def load_weights_from_pt(weights_file: str, model: nn.Module) -> None:
     """Load weights from a PyTorch file into a MoondreamModel instance."""
-    device = str(torch.empty(0).device)
-    tensors = torch.load(weights_file, map_location=device, weights_only=True)
-    if "vision.blocks.0.attn.proj.bias" in tensors.keys():
-        missing_keys, unexpected_keys = model.load_state_dict(tensors, strict=False)
-        print("Missing keys:", missing_keys)
-        print("Unexpected keys:", unexpected_keys)
+    tensors = torch.load(weights_file, map_location="cpu", weights_only=True)
+    all_keys = tensors.keys()
+
+    if (
+        "vision.blocks.0.attn.proj.bias" in all_keys
+        or "model.vision.blocks.0.attn.proj.bias" in all_keys
+    ):
+        tensors = {add_linear_to_key(k): v for k, v in tensors.items()}
+        model.load_state_dict(tensors, strict=False)
     else:
         tensors = {
-            k.replace("._orig_mod", ""): v.to(dtype=torch.bfloat16)
+            k.replace("._orig_mod", ""): v.to(dtype=torch.float16)
             for k, v in tensors.items()
         }
         _load_weights(lambda x: tensors[x], model)
@@ -168,4 +211,4 @@ def load_weights_into_model(weights_file: str, model: nn.Module) -> None:
 
     # Make all parameters contiguous
     for param in model.parameters():
-        param.data = param.data.contiguous()
+        param.data = param.data.contiguous()