update conversion script

convert_flax.py

@@ -49,10 +49,76 @@ def find_scales(w: jnp.ndarray, dim: int, pbar: bool = True):
     return scales.squeeze(dim + 1)
 
 
+def convert_siglip(params):
+    state_dict = dict()
+
+    def convert_layer(prefix: str, layer: dict[str, jnp.ndarray]):
+        bias = layer["bias"]
+
+        if "kernel" in layer:
+            w = layer["kernel"]
+            if w.ndim == 2:  # linear layer
+                w = w.T
+
+            elif w.ndim == 3:  # attn projection
+                # qkv projection - (dim, num_heads, head_dim)
+                if bias.ndim == 2:
+                    w = flatten(w, 1, 2).T
+                    bias = bias.reshape(-1)
+
+                # o projection - (num_heads, head_dim, dim)
+                elif bias.ndim == 1:
+                    w = flatten(w, 0, 1).T
+
+            elif w.ndim == 4:  # conv2d layer
+                w = w.transpose(3, 2, 0, 1)
+
+            else:
+                raise RuntimeError(f"Unsupported {w.shape=}")
+
+        elif "scale" in layer:  # layer norm
+            w = layer["scale"]
+
+        else:
+            raise RuntimeError
+
+        state_dict[f"{prefix}weight"] = w
+        state_dict[f"{prefix}bias"] = bias
+
+    convert_layer("embeddings.patch_embedding.", params["embedding"])
+    state_dict["embeddings.position_embedding.weight"] = params["pos_embedding"].squeeze(0)
+
+    trunk = params["Transformer"]
+    convert_layer("post_layernorm.", trunk["encoder_norm"])
+
+    layer_idx = 0
+    while f"encoderblock_{layer_idx}" in trunk:
+        prefix = f"encoder.layers.{layer_idx}."
+        encoder_layer = trunk[f"encoderblock_{layer_idx}"]
+
+        convert_layer(f"{prefix}layer_norm1.", encoder_layer["LayerNorm_0"])
+        convert_layer(f"{prefix}layer_norm2.", encoder_layer["LayerNorm_1"])
+
+        attn_layer = encoder_layer["MultiHeadDotProductAttention_0"]
+        convert_layer(f"{prefix}self_attn.q_proj.", attn_layer["query"])
+        convert_layer(f"{prefix}self_attn.k_proj.", attn_layer["key"])
+        convert_layer(f"{prefix}self_attn.v_proj.", attn_layer["value"])
+        convert_layer(f"{prefix}self_attn.out_proj.", attn_layer["out"])
+
+        mlp_layer = encoder_layer["MlpBlock_0"]
+        convert_layer(f"{prefix}mlp.fc1.", mlp_layer["Dense_0"])
+        convert_layer(f"{prefix}mlp.fc2.", mlp_layer["Dense_1"])
+
+        layer_idx += 1
+
+    return state_dict
+
+
 # convert to HF format first, then apply quantization
 def convert_to_hf(params):
     state_dict = dict()
-    # TODO: output projection
+
+    # NOTE: all gemma3 models use tied embeddings, even for the 27B version.
     state_dict["model.embed_tokens.weight"] = params["embedder"]["input_embedding"]
     state_dict["model.norm.weight"] = params["final_norm"]["scale"]
 
@@ -86,6 +152,22 @@ def convert_to_hf(params):
 
         layer_idx += 1
 
+    # vision tower
+    if "vision_encoder" in params:
+        # HF append unused tokens for no reason???
+        state_dict["model.embed_tokens.weight"] = jnp.pad(state_dict["model.embed_tokens.weight"], ((0, 64), (0, 0)))
+
+        for k in list(state_dict.keys()):
+            state_dict[f"language_model.{k}"] = state_dict.pop(k)
+
+        prefix = "multi_modal_projector.mm_"
+        state_dict[f"{prefix}input_projection_weight"] = params["embedder"]["mm_input_projection"]["w"]
+        state_dict[f"{prefix}soft_emb_norm.weight"] = params["embedder"]["mm_soft_embedding_norm"]["scale"]
+
+        siglip_state_dict = convert_siglip(params["vision_encoder"]["siglip_encoder"])
+        for k, v in siglip_state_dict.items():
+            state_dict[f"vision_tower.vision_model.{k}"] = v
+
     return state_dict
 
 
@@ -93,8 +175,11 @@ def convert_awq(state_dict: dict[str, jnp.ndarray]):
     awq_state_dict = dict()
 
     for k, v in tqdm(state_dict.items(), total=len(state_dict)):
-        # AWQ doesn't support INT4 embeddings
-        if k == "model.embed_tokens.weight" or v.ndim == 1:
+        if (
+            k.endswith("model.embed_tokens.weight")  # AWQ doesn't support INT4 embeddings
+            or k.startswith(("vision_tower", "multi_modal_projector"))  # vision tower is not quantized
+            or v.ndim == 1
+        ):
             awq_state_dict[k] = v.astype(jnp.bfloat16)
             continue