update script

convert_flax.py

@@ -1,15 +1,18 @@
 import argparse
+import json
 from pathlib import Path
 
+import jax
 import jax.numpy as jnp
 import numpy as np
+import orbax.checkpoint as ocp
 from safetensors.flax import save_file
 from tqdm import tqdm
 
-from gemma import gm
+SIGLIP_PREFIX = "SigLiPFromPatches_0/siglip_encoder"
 
 
-def flatten(x: jnp.ndarray, start: int = 0, end: int = -1):
+def flatten(x: np.ndarray, start: int = 0, end: int = -1):
     if start < 0:
         start += x.ndim
     if end < 0:
@@ -18,27 +21,27 @@ def flatten(x: jnp.ndarray, start: int = 0, end: int = -1):
     return x.reshape(new_shape)
 
 
-def unflatten(x: jnp.ndarray, dim: int, sizes: tuple[int, ...]):
+def unflatten(x: np.ndarray, dim: int, sizes: tuple[int, ...]):
     new_shape = x.shape[:dim] + tuple(sizes) + x.shape[dim + 1 :]
     return x.reshape(new_shape)
 
 
 # correct quantization parameters mean quantization error = 0 (or close to 0)
-def check_groups(groups: jnp.ndarray, scales: jnp.ndarray, dim: int):
+def check_groups(groups: np.ndarray, scales: np.ndarray, dim: int):
     # groups: (a, b, c, 32, d, e, f)
     # scales: (a, b, c,  1, d, e, f)
     inv_scale = 1.0 / scales.clip(1e-12)
-    q_group = jnp.round(groups * inv_scale)
-    max_diff = jnp.abs(q_group * scales - groups).max(dim, keepdims=True)
+    q_group = np.round(groups * inv_scale)
+    max_diff = np.abs(q_group * scales - groups).max(dim, keepdims=True)
     return max_diff < 1e-6, max_diff
 
 
-def find_scales(w: jnp.ndarray, dim: int, pbar: bool = True):
+def find_scales(w: np.ndarray, dim: int):
     w = unflatten(w, dim, (-1, 32))
     group_range = w.max(dim + 1, keepdims=True) - w.min(dim + 1, keepdims=True)
 
     scales = np.zeros_like(group_range)
-    for q in tqdm(range(15, 0, -1), disable=not pbar):
+    for q in range(15, 0, -1):
         try_scale = group_range / q
         ok, _ = check_groups(w, try_scale, dim + 1)
         scales[ok] = try_scale[ok]
@@ -49,10 +52,10 @@ def find_scales(w: jnp.ndarray, dim: int, pbar: bool = True):
     return scales.squeeze(dim + 1)
 
 
-def convert_siglip(params):
+def convert_siglip(params, num_layers: int):
     state_dict = dict()
 
-    def convert_layer(prefix: str, layer: dict[str, jnp.ndarray]):
+    def convert_layer(prefix: str, layer: dict[str, np.ndarray]):
         bias = layer["bias"]
 
         if "kernel" in layer:
@@ -85,96 +88,129 @@ def convert_siglip(params):
         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)
+    convert_layer("embeddings.patch_embedding.", params[f"{SIGLIP_PREFIX}/embedding"])
+    state_dict["embeddings.position_embedding.weight"] = params[SIGLIP_PREFIX]["pos_embedding"].squeeze(0)
+    convert_layer("post_layernorm.", params[f"{SIGLIP_PREFIX}/Transformer/encoder_norm"])
 
-    trunk = params["Transformer"]
-    convert_layer("post_layernorm.", trunk["encoder_norm"])
-
-    layer_idx = 0
-    while f"encoderblock_{layer_idx}" in trunk:
+    for layer_idx in range(num_layers):
         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"])
+        layer_prefix = f"{SIGLIP_PREFIX}/Transformer/encoderblock_{layer_idx}/"
 
-        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"])
+        convert_layer(f"{prefix}layer_norm1.", params[f"{layer_prefix}LayerNorm_0"])
+        convert_layer(f"{prefix}layer_norm2.", params[f"{layer_prefix}LayerNorm_1"])
 
-        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"])
+        attn_prefix = f"{layer_prefix}MultiHeadDotProductAttention_0/"
+        convert_layer(f"{prefix}self_attn.q_proj.", params[f"{attn_prefix}query"])
+        convert_layer(f"{prefix}self_attn.k_proj.", params[f"{attn_prefix}key"])
+        convert_layer(f"{prefix}self_attn.v_proj.", params[f"{attn_prefix}value"])
+        convert_layer(f"{prefix}self_attn.out_proj.", params[f"{attn_prefix}out"])
 
-        layer_idx += 1
+        mlp_prefix = f"{layer_prefix}MlpBlock_0/"
+        convert_layer(f"{prefix}mlp.fc1.", params[f"{mlp_prefix}Dense_0"])
+        convert_layer(f"{prefix}mlp.fc2.", params[f"{mlp_prefix}Dense_1"])
 
     return state_dict
 
 
 # convert to HF format first, then apply quantization
-def convert_to_hf(params):
-    state_dict = dict()
+def convert_to_hf(path: Path):
+    path = path.absolute()  # orbax only works with absolute path
+    ckpt = ocp.StandardCheckpointer()
+    metadata = dict(ckpt.metadata(path))
+    metadata = jax.tree.map(ocp.utils.to_shape_dtype_struct, metadata)
+
+    num_layers = num_siglip_layers = 0
+    while f"transformer/layer_{num_layers}/attn/_key_norm" in metadata:
+        num_layers += 1
+    while f"{SIGLIP_PREFIX}/Transformer/encoderblock_{num_siglip_layers}/LayerNorm_0" in metadata:
+        num_siglip_layers += 1
+    print(f"{num_layers=}")
+    print(f"{num_siglip_layers=}")
+
+    def load_params(*keys: tuple[str, ...], prefix: str | None = None):
+        # load params with specific keys and params starts with prefix
+        f1 = lambda k: tuple(subkey.key for subkey in k) in keys
+        f2 = lambda k: k[0].key.startswith(prefix)
+
+        # set to None to not load that weights
+        pytree = jax.tree.map_with_path(lambda k, v: v if f1(k) or f2(k) else None, metadata)
+        return ckpt.restore(path, pytree)
 
     # 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"]
-
-    layer_idx = 0
-    while f"layer_{layer_idx}" in params:
-        prefix = f"model.layers.{layer_idx}."
-        layer_params = params[f"layer_{layer_idx}"]
-        state_dict[f"{prefix}input_layernorm.weight"] = layer_params["pre_attention_norm"]["scale"]
-        state_dict[f"{prefix}post_attention_layernorm.weight"] = layer_params["post_attention_norm"]["scale"]
-        state_dict[f"{prefix}pre_feedforward_layernorm.weight"] = layer_params["pre_ffw_norm"]["scale"]
-        state_dict[f"{prefix}post_feedforward_layernorm.weight"] = layer_params["post_ffw_norm"]["scale"]
-
-        prefix = f"model.layers.{layer_idx}.self_attn."
-        attn_params = layer_params["attn"]
-        state_dict[f"{prefix}q_norm.weight"] = attn_params["_query_norm"]["scale"]
-        state_dict[f"{prefix}k_norm.weight"] = attn_params["_key_norm"]["scale"]
+    params = load_params(
+        ("transformer/final_norm", "scale"),
+        prefix="transformer/embedder",
+    )
+    state_dict = dict()
 
-        # (num_heads, hidden_size, head_dim) -> (num_heads * head_dim, hidden_size)
-        state_dict[f"{prefix}q_proj.weight"] = flatten(attn_params["q_einsum"]["w"].transpose(0, 2, 1), end=1)
-        state_dict[f"{prefix}k_proj.weight"] = flatten(attn_params["kv_einsum"]["w"][0].transpose(0, 2, 1), end=1)
-        state_dict[f"{prefix}v_proj.weight"] = flatten(attn_params["kv_einsum"]["w"][1].transpose(0, 2, 1), end=1)
+    if num_siglip_layers > 0:
+        # HF append unused tokens for no reason???
+        embed = params["transformer/embedder"]["input_embedding"]
+        params["transformer/embedder"]["input_embedding"] = np.pad(embed, ((0, 64), (0, 0)))
+        gemma_prefix = "language_model."
 
-        # (num_heads, head_dim, hidden_size) -> (hidden_size, num_heads * head_dim)
-        state_dict[f"{prefix}o_proj.weight"] = flatten(attn_params["attn_vec_einsum"]["w"], end=1).T
+        prefix = "multi_modal_projector.mm_"
+        jax_prefix = "transformer/embedder/"
+        state_dict[f"{prefix}input_projection_weight"] = params[f"{jax_prefix}mm_input_projection"]["w"]
+        state_dict[f"{prefix}soft_emb_norm.weight"] = params[f"{jax_prefix}mm_soft_embedding_norm"]["scale"]
 
-        prefix = f"model.layers.{layer_idx}.mlp."
-        mlp_params = layer_params["mlp"]
-        state_dict[f"{prefix}gate_proj.weight"] = mlp_params["gating_einsum"][0]  # NOTE: may need to transpose?
-        state_dict[f"{prefix}up_proj.weight"] = mlp_params["gating_einsum"][1]
-        state_dict[f"{prefix}down_proj.weight"] = mlp_params["linear"].T
+    else:
+        gemma_prefix = ""
 
-        layer_idx += 1
+    state_dict[f"{gemma_prefix}model.embed_tokens.weight"] = params["transformer/embedder"]["input_embedding"]
+    state_dict[f"{gemma_prefix}model.norm.weight"] = params["transformer/final_norm"]["scale"]
 
-    # 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)))
+    yield state_dict
 
-        for k in list(state_dict.keys()):
-            state_dict[f"language_model.{k}"] = state_dict.pop(k)
+    for layer_idx in range(num_layers):
+        jax_prefix = f"transformer/layer_{layer_idx}/"
+        params = load_params(prefix=jax_prefix)
 
-        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"]
+        state_dict = dict()
+        prefix = f"{gemma_prefix}model.layers.{layer_idx}."
+        state_dict[f"{prefix}input_layernorm.weight"] = params[f"{jax_prefix}pre_attention_norm"]["scale"]
+        state_dict[f"{prefix}post_attention_layernorm.weight"] = params[f"{jax_prefix}post_attention_norm"]["scale"]
+        state_dict[f"{prefix}pre_feedforward_layernorm.weight"] = params[f"{jax_prefix}pre_ffw_norm"]["scale"]
+        state_dict[f"{prefix}post_feedforward_layernorm.weight"] = params[f"{jax_prefix}post_ffw_norm"]["scale"]
+
+        prefix = f"{gemma_prefix}model.layers.{layer_idx}.self_attn."
+        jax_prefix = f"transformer/layer_{layer_idx}/attn/"
+        state_dict[f"{prefix}q_norm.weight"] = params[f"{jax_prefix}_query_norm"]["scale"]
+        state_dict[f"{prefix}k_norm.weight"] = params[f"{jax_prefix}_key_norm"]["scale"]
+
+        # (num_heads, hidden_size, head_dim) -> (num_heads * head_dim, hidden_size)
+        state_dict[f"{prefix}q_proj.weight"] = flatten(params[f"{jax_prefix}q_einsum"]["w"].transpose(0, 2, 1), end=1)
+        state_dict[f"{prefix}k_proj.weight"] = flatten(
+            params[f"{jax_prefix}kv_einsum"]["w"][0].transpose(0, 2, 1), end=1
+        )
+        state_dict[f"{prefix}v_proj.weight"] = flatten(
+            params[f"{jax_prefix}kv_einsum"]["w"][1].transpose(0, 2, 1), end=1
+        )
+
+        # (num_heads, head_dim, hidden_size) -> (hidden_size, num_heads * head_dim)
+        state_dict[f"{prefix}o_proj.weight"] = flatten(params[f"{jax_prefix}attn_vec_einsum"]["w"], end=1).T
+
+        prefix = f"{gemma_prefix}model.layers.{layer_idx}.mlp."
+        jax_prefix = f"transformer/layer_{layer_idx}/mlp/"
+        state_dict[f"{prefix}gate_proj.weight"] = params[f"{jax_prefix}gating_einsum"]["w"][0]
+        state_dict[f"{prefix}up_proj.weight"] = params[f"{jax_prefix}gating_einsum"]["w"][1]
+        state_dict[f"{prefix}down_proj.weight"] = params[f"{jax_prefix}linear"]["w"].T
+
+        yield state_dict
 
-        siglip_state_dict = convert_siglip(params["vision_encoder"]["siglip_encoder"])
+    # vision tower
+    if num_siglip_layers > 0:
+        params = load_params(prefix=SIGLIP_PREFIX)
+        siglip_state_dict = convert_siglip(params, num_siglip_layers)
         for k, v in siglip_state_dict.items():
             state_dict[f"vision_tower.vision_model.{k}"] = v
-
-    return state_dict
+        yield state_dict
 
 
-def convert_awq(state_dict: dict[str, jnp.ndarray]):
+def convert_awq(state_dict: dict[str, np.ndarray]):
     awq_state_dict = dict()
 
-    for k, v in tqdm(state_dict.items(), total=len(state_dict)):
+    for k, v in state_dict.items():
         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
@@ -186,10 +222,8 @@ def convert_awq(state_dict: dict[str, jnp.ndarray]):
         assert v.ndim == 2
         v = v.T  # AWQ transpose the weight
 
-        # use numpy since jnp is very slow, likely due to bad memory management on CUDA
-        v = np.asarray(v)
         K, N = v.shape
-        scales = find_scales(v, dim=0, pbar=False)  # (K/32, N)
+        scales = find_scales(v, dim=0)  # (K/32, N)
         inv_scale = 1 / scales.clip(1e-12)
         qweight = np.round(v.reshape(K // 32, 32, N) * inv_scale[:, None])
 
@@ -216,7 +250,7 @@ def convert_awq(state_dict: dict[str, jnp.ndarray]):
         prefix = k.removesuffix(".weight")
         awq_state_dict[f"{prefix}.qweight"] = qweight_packed
         awq_state_dict[f"{prefix}.qzeros"] = np.full((K // 32, N // 8), 0x8888_8888, dtype=np.uint32).view(np.int32)
-        awq_state_dict[f"{prefix}.scales"] = jnp.asarray(scales).astype(jnp.bfloat16)
+        awq_state_dict[f"{prefix}.scales"] = scales.astype(jnp.bfloat16)
 
     return awq_state_dict
 
@@ -224,11 +258,37 @@ def convert_awq(state_dict: dict[str, jnp.ndarray]):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--ckpt_dir", required=True, type=Path)
-    parser.add_argument("--save_path", required=True, type=Path)
+    parser.add_argument("--save_dir", required=True, type=Path)
     args = parser.parse_args()
 
-    params = gm.ckpts.load_params(args.ckpt_dir.absolute())
-    state_dict = convert_to_hf(params)
-    awq_state_dict = convert_awq(state_dict)
-    args.save_path.parent.mkdir(parents=True, exist_ok=True)
-    save_file(awq_state_dict, args.save_path)
+    args.save_dir.mkdir(parents=True, exist_ok=True)
+
+    total_size = 0
+    weight_map = dict()
+
+    state_dict = dict()
+    size = 0
+    shard_idx = 0
+    filename = f"model-{shard_idx + 1:05d}.safetensors"
+    for sub_state_dict in tqdm(convert_to_hf(args.ckpt_dir)):
+        sub_state_dict = convert_awq(sub_state_dict)
+
+        for k, v in sub_state_dict.items():
+            state_dict[k] = v
+            size += v.nbytes
+
+            total_size += v.nbytes
+            weight_map[k] = filename
+
+        if size > 5e9:
+            save_file(state_dict, args.save_dir / filename)
+            state_dict = dict()
+            size = 0
+            shard_idx += 1
+            filename = f"model-{shard_idx + 1:05d}.safetensors"
+
+    save_file(state_dict, args.save_dir / filename)
+    json.dump(
+        dict(metadata=dict(total_size=total_size), weight_map=weight_map),
+        open(args.save_dir / "model.safetensors.index.json", "w"),
+    )