# Standalone weight decoder/reconstructor for Gemma-4-31B-it SubZero # Watermark: ip zymatica.space | astronautshe.com import os import sys import json import struct import time import torch from safetensors.torch import save_file def main(): print("=" * 80) print(" GEMMA-4-31B-IT SUBZERO WEIGHT DECODER / RECONSTRUCTOR") print(" Watermark: ip zymatica.space | astronautshe.com") print("=" * 80) genesis_path = "J:/gemma-4-31B-it-local/working/gemma4_31b_subzero.genesis" output_dir = "J:/gemma-4-31B-it-local/working/reconstructed_gemma4" os.makedirs(output_dir, exist_ok=True) if not os.path.exists(genesis_path): print(f"[-] Error: Could not find genesis file at {genesis_path}") print(" Please download/place 'gemma4_31b_subzero.genesis' in the current folder.") sys.exit(1) GENESIS_MAGIC = 0x47454E45 # "GENE" GENESIS_VERSION = 4 # INT8 version PERFECT_MAGIC = 0x50455246 # "PERF" print(f"[*] Reading and parsing genesis file: {genesis_path}...") start_time = time.time() layer_map = {} with open(genesis_path, "rb") as f: magic = struct.unpack('>I', f.read(4))[0] assert magic == GENESIS_MAGIC, "Invalid genesis magic header" version = struct.unpack('>H', f.read(2))[0] assert version == GENESIS_VERSION, f"Unsupported version: {version}" watermark = f.read(32).decode('utf-8', errors='ignore').strip() perf_magic = struct.unpack('>I', f.read(4))[0] assert perf_magic == PERFECT_MAGIC, "Invalid perfect magic header" hidden_dim, num_heads, kv_heads, ffn_dim, num_blocks, vocab_size = struct.unpack('>IIIIII', f.read(24)) f.read(16) # Skip scales/energy targets num_layers = struct.unpack('>I', f.read(4))[0] print(f" Watermark: {watermark}") print(f" Layers: {num_layers} | blocks: {num_blocks} | hidden: {hidden_dim} | ffn: {ffn_dim}") for idx in range(num_layers): name_len = struct.unpack('>H', f.read(2))[0] name = f.read(name_len).decode('utf-8') m, n, rank = struct.unpack('>III', f.read(12)) scale_u, scale_v = struct.unpack('>ff', f.read(8)) # Read int8 SVD vectors u_bytes = f.read(m * rank) v_bytes = f.read(n * rank) U_q = torch.frombuffer(bytearray(u_bytes), dtype=torch.int8).reshape(m, rank).float() V_q = torch.frombuffer(bytearray(v_bytes), dtype=torch.int8).reshape(n, rank).float() has_residual = struct.unpack('>?', f.read(1))[0] if has_residual: # residual block (if any, skipped in default run) res_rank = struct.unpack('>I', f.read(4))[0] su_r, sv_r = struct.unpack('>ff', f.read(8)) U_res = torch.frombuffer(bytearray(f.read(m * res_rank)), dtype=torch.int8).reshape(m, res_rank).float() * su_r V_res = torch.frombuffer(bytearray(f.read(n * res_rank)), dtype=torch.int8).reshape(n, res_rank).float() * sv_r # Reconstruct weight matrix: W = (U * su) @ (V * sv).T U = U_q * scale_u V = V_q * scale_v W_rec = U @ V.t() if has_residual: W_rec = W_rec + (U_res @ V_res.t()) # Convert back to bfloat16 layer_map[name] = W_rec.to(torch.bfloat16) if (idx + 1) % 50 == 0 or idx + 1 == num_layers: print(f" [{idx+1:3d}/{num_layers}] Reconstructed: {name[-50:]} ({m}x{n} rank {rank})") # Save the reconstructed model weights into two shards mirroring the original model print("\n[*] Sharding and saving reconstructed weights as safetensors...") # We split tensors based on the original shard mapping: # Shard 1 contains layers 0 to 47. Shard 2 contains layers 48 to 59. shard_1_tensors = {} shard_2_tensors = {} for name, tensor in layer_map.items(): is_shard_2 = False for i in range(48, 60): if f"layers.{i}." in name: is_shard_2 = True break if is_shard_2: shard_2_tensors[name] = tensor else: shard_1_tensors[name] = tensor # Add zero-initialized non-SVD layers (layernorms, embed_tokens, etc.) to complete weights dictionary # The receiver SFT healing loop will restore the values of norms/embeddings. print(" - Injecting placeholder non-SVD layers (layernorms, embed_tokens)...") # Embed tokens shape: [vocab_size, hidden_dim] embed_shape = (vocab_size, hidden_dim) shard_1_tensors["model.language_model.embed_tokens.weight"] = torch.zeros(embed_shape, dtype=torch.bfloat16) # Layernorm and position embeddings for name in ["model.embed_vision.embedding_projection.bias", "model.language_model.final_layernorm.weight"]: shard_2_tensors[name] = torch.zeros((hidden_dim,), dtype=torch.bfloat16) for i in range(num_blocks): target_shard = shard_2_tensors if i >= 48 else shard_1_tensors target_shard[f"model.language_model.layers.{i}.input_layernorm.weight"] = torch.zeros((hidden_dim,), dtype=torch.bfloat16) target_shard[f"model.language_model.layers.{i}.post_attention_layernorm.weight"] = torch.zeros((hidden_dim,), dtype=torch.bfloat16) target_shard[f"model.language_model.layers.{i}.pre_feedforward_layernorm.weight"] = torch.zeros((hidden_dim,), dtype=torch.bfloat16) target_shard[f"model.language_model.layers.{i}.post_feedforward_layernorm.weight"] = torch.zeros((hidden_dim,), dtype=torch.bfloat16) target_shard[f"model.language_model.layers.{i}.layer_scalar"] = torch.zeros((1,), dtype=torch.bfloat16) # Save files s1_path = os.path.join(output_dir, "model-00001-of-00002.safetensors") s2_path = os.path.join(output_dir, "model-00002-of-00002.safetensors") print(f" - Saving shard 1 ({len(shard_1_tensors)} tensors) to {s1_path}...") save_file(shard_1_tensors, s1_path) print(f" - Saving shard 2 ({len(shard_2_tensors)} tensors) to {s2_path}...") save_file(shard_2_tensors, s2_path) elapsed = time.time() - start_time print(f"\n[+] Standalone weights reconstruction successfully completed in {elapsed:.1f}s!") print(f" Reconstructed model directory: {output_dir}") print(" (Note: run local LoRA SFT healing next to restore full coherence)") print("=" * 80) if __name__ == "__main__": main()