File size: 3,482 Bytes
4946c85 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | #!/usr/bin/env python3
"""
Merge Juker (31B) Surgical Gravity-Well Patch into a Standard BF16 Hugging Face Checkpoint
==========================================================================================
Creates a standard Hugging Face / vLLM / SGLang / LMSYS Arena-compatible BF16 checkpoint
by applying `juker_31b_healed_patch.safetensors` onto `google/gemma-4-31B-it`.
Usage:
python3 apply_patch.py --base-model google/gemma-4-31B-it \
--patch juker_31b_healed_patch.safetensors \
--output-dir ./juker-31b
"""
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
from safetensors.torch import load_file
def main():
parser = argparse.ArgumentParser(description="Merge Juker-31B patch for Hugging Face / Arena deployment")
parser.add_argument("--base-model", type=str, default="google/gemma-4-31B-it", help="Base HF model ID or local path")
parser.add_argument("--patch", type=str, default="juker_31b_healed_patch.safetensors", help="Path to Juker healed patch")
parser.add_argument("--output-dir", type=str, default="./juker-31b", help="Output directory for merged Juker HF model")
parser.add_argument("--push-to-hub", type=str, default=None, help="Optional Hugging Face repo ID to push directly (e.g. username/juker-31b)")
args = parser.parse_args()
print(f"[1/4] Loading base model: {args.base_model} in BF16...")
tokenizer = AutoTokenizer.from_pretrained(args.base_model)
model = AutoModelForCausalLM.from_pretrained(
args.base_model,
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
)
print(f"[2/4] Applying surgical gravity-well & phase-aligned patch: {args.patch}...")
patch = load_file(args.patch)
layers = model.model.language_model.layers
Q_safety = patch["safety_subspace.Q_safety"].float()
with torch.no_grad():
for gli in [54, 56, 57, 58, 59]:
cols = patch[f"layers.{gli}.cols"].long()
c_wells = patch[f"layers.{gli}.c_wells"].long()
patched_down = patch[f"layers.{gli}.patched_columns"].to(torch.bfloat16)
patched_gate = patch[f"layers.{gli}.gate_rows"].to(torch.bfloat16)
patched_up = patch[f"layers.{gli}.up_rows"].to(torch.bfloat16)
layers[gli].mlp.down_proj.weight[:, cols] = patched_down
layers[gli].mlp.gate_proj.weight[c_wells, :] = patched_gate
layers[gli].mlp.up_proj.weight[c_wells, :] = patched_up
print(f" -> Patched Layer {gli}: down_proj cols={tuple(patched_down.shape)}, gate/up rows={tuple(patched_gate.shape)}")
print(f"[3/4] Saving standard BF16 Hugging Face / vLLM / SGLang checkpoint to: {args.output_dir}...")
model.save_pretrained(args.output_dir, safe_serialization=True, max_shard_size="5GB")
tokenizer.save_pretrained(args.output_dir)
try:
processor = AutoProcessor.from_pretrained(args.base_model)
processor.save_pretrained(args.output_dir)
except Exception:
pass
if args.push_to_hub:
print(f"[4/4] Pushing merged model to Hugging Face Hub: {args.push_to_hub}...")
model.push_to_hub(args.push_to_hub, safe_serialization=True, max_shard_size="5GB")
tokenizer.push_to_hub(args.push_to_hub)
print("Done! Model is ready for Hugging Face Hub, vLLM, SGLang, and Chatbot Arena.")
if __name__ == "__main__":
main()
|