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 SIGLIP_PREFIX = "SigLiPFromPatches_0/siglip_encoder" def flatten(x: np.ndarray, start: int = 0, end: int = -1): if start < 0: start += x.ndim if end < 0: end += x.ndim new_shape = x.shape[:start] + (-1,) + x.shape[end + 1 :] return x.reshape(new_shape) 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: 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 = 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: 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 range(15, 0, -1): try_scale = group_range / q ok, _ = check_groups(w, try_scale, dim + 1) scales[ok] = try_scale[ok] ok, _ = check_groups(w, scales, dim + 1) assert ok.all() return scales.squeeze(dim + 1) def convert_siglip(params, num_layers: int): state_dict = dict() def convert_layer(prefix: str, layer: dict[str, np.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[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"]) for layer_idx in range(num_layers): prefix = f"encoder.layers.{layer_idx}." layer_prefix = f"{SIGLIP_PREFIX}/Transformer/encoderblock_{layer_idx}/" convert_layer(f"{prefix}layer_norm1.", params[f"{layer_prefix}LayerNorm_0"]) convert_layer(f"{prefix}layer_norm2.", params[f"{layer_prefix}LayerNorm_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"]) 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(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=}") # NOTE: all gemma3 models use tied embeddings, even for the 27B version. params = ckpt.restore(path) state_dict = dict() 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." 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"] else: gemma_prefix = "" 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"] yield state_dict for layer_idx in range(num_layers): jax_prefix = f"transformer/layer_{layer_idx}/" 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 # vision tower if num_siglip_layers > 0: 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 yield state_dict def convert_awq(state_dict: dict[str, np.ndarray]): awq_state_dict = 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 or v.ndim == 1 ): awq_state_dict[k] = v.astype(jnp.bfloat16) continue assert v.ndim == 2 v = v.T # AWQ transpose the weight K, N = v.shape 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]) # AWQ is actually UINT4 (instead of INT4) # hence, we will shift qweight up by 8 (even though Google AQT only uses [-7,7]) # and set zero_point = 8 qweight = (qweight + 8).astype(np.uint32) # AWQ pack 8 int4 into UINT32 in the following layout (from high bits to low bits) # [7 5 3 1 6 4 2 0] along the 2nd dim qweight = qweight.reshape(K, N // 8, 8) qweight_packed = ( (qweight[..., 7] << (7 * 4)) | (qweight[..., 5] << (6 * 4)) | (qweight[..., 3] << (5 * 4)) | (qweight[..., 1] << (4 * 4)) | (qweight[..., 6] << (3 * 4)) | (qweight[..., 4] << (2 * 4)) | (qweight[..., 2] << (1 * 4)) | (qweight[..., 0] << (0 * 4)) ) qweight_packed = qweight_packed.view(np.int32).reshape(K, N // 8) 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"] = scales.astype(jnp.bfloat16) return awq_state_dict if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--ckpt_dir", required=True, type=Path) parser.add_argument("--save_dir", required=True, type=Path) args = parser.parse_args() 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) new_size = sum(v.nbytes for v in sub_state_dict.values()) if size + new_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" # assume that new_size < 5e9 size += new_size total_size += new_size for k, v in sub_state_dict.items(): state_dict[k] = v weight_map[k] = filename 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"), )