File size: 1,176 Bytes
576ab11
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
"""
Apply the 64-column Weight-Space Patch (`gemma4_e4b_weight_space_patch.safetensors`)
to any loaded `Gemma4ForCausalLM` or `Gemma4ForConditionalGeneration` instance in <1 second!
"""

import torch
from safetensors.torch import load_file


@torch.no_grad()
def apply_gemma4_weight_space_patch(model, patch_safetensors_path: str = "gemma4_e4b_weight_space_patch.safetensors"):
    tensors = load_file(patch_safetensors_path, device="cpu")
    # Support both Gemma4ForCausalLM (model.model.layers) and Gemma4ForConditionalGeneration (model.model.language_model.layers)
    layers = getattr(model.model, "layers", None)
    if layers is None:
        layers = model.model.language_model.layers

    total_params = 0
    for li in [38, 39, 40, 41]:
        cols = tensors[f"layers.{li}.cols"].long()
        patched_cols = tensors[f"layers.{li}.patched_columns"]
        W = layers[li].mlp.down_proj.weight.data
        W[:, cols.to(W.device)] = patched_cols.to(device=W.device, dtype=W.dtype)
        total_params += patched_cols.numel()

    print(f"Applied Gemma-4-E4B-it weight-space patch ({total_params:,} parameters across 64 columns in Layers 38-41).")
    return model