jstarks's picture
Upload Gemma-4-E4B-it Weight-Space Programming patch (163,840 weights) & before/after benchmarks
576ab11 verified
Raw History Blame Contribute Delete
1.18 kB
"""
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