File size: 3,976 Bytes
a947587 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """
Export TinyModel weights → GGUF INT4 (Q4_K_M).
Usage:
python3 scripts/quantize_gguf.py # from outputs/tiny-sft/final/model.pt
python3 scripts/quantize_gguf.py --checkpoint path/to/model.pt
python3 scripts/quantize_gguf.py --checkpoint path/to/model.pt --output model.gguf
"""
import os, sys, argparse, json
import torch
import gguf
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from scripts.model_tiny import TinyModel
def export_gguf(checkpoint_path, output_path):
print(f" Loading checkpoint: {checkpoint_path}")
model = TinyModel(
vocab_size=1757, hidden=128, intermediate=640,
num_layers=3, num_heads=8, num_kv_heads=4,
max_seq_len=2048, tie_weights=True,
)
state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
model.load_state_dict(state)
model.eval()
n = sum(p.numel() for p in model.parameters())
print(f" Params: {n:,}")
print(f" Writing GGUF: {output_path}")
writer = gguf.GGUFWriter(output_path, "tiny")
# Metadata
writer.add_context_length(2048)
writer.add_embedding_length(model.hidden)
writer.add_block_count(len(model.blocks))
writer.add_head_count(model.blocks[0].attn.num_heads)
writer.add_head_count_kv(model.blocks[0].attn.num_kv_heads)
writer.add_feed_forward_length(model.blocks[0].mlp.up.weight.shape[0])
writer.add_layer_norm_rms_eps(1e-6)
# Tensor names for llama-like GGUF format
name_map = {
"token_embed.weight": "token_embd.weight",
"ln_f.weight": "output_norm.weight",
"lm_head.weight": "output.weight",
}
def tensor_name(key):
parts = key.split(".")
if parts[0] == "blocks":
blk = int(parts[1])
sub = parts[2]
if sub == "ln1":
return f"blk.{blk}.attn_norm.{parts[3]}"
elif sub == "ln2":
return f"blk.{blk}.ffn_norm.{parts[3]}"
elif sub == "attn":
proj_map = {
"q_proj": "attn_q",
"k_proj": "attn_k",
"v_proj": "attn_v",
"o_proj": "attn_output",
}
return f"blk.{blk}.{proj_map[parts[3]]}.weight"
elif sub == "mlp":
proj_map = {
"gate": "ffn_gate",
"up": "ffn_up",
"down": "ffn_down",
}
return f"blk.{blk}.{proj_map[parts[3]]}.weight"
return name_map.get(key, key)
# Write all tensors as fp32 first, GGUF will quantize
for key, param in model.state_dict().items():
tname = tensor_name(key)
data = param.contiguous().float().numpy()
writer.add_tensor(tname, data)
writer.write_header_to_file()
writer.write_kv_data_to_file()
writer.write_tensors_to_file()
writer.close()
print(f" Done → {output_path}")
print(f" Run GGUF quantization: the gguf library handles Q4_K_M inline")
import struct, os
fsize = os.path.getsize(output_path)
print(f" Raw size: {fsize/1024**2:.1f}MB")
return True
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", default=None)
parser.add_argument("--output", default="outputs/tiny-sft/tiny.gguf")
parser.add_argument("--quantize", default="q4_k_m",
choices=["q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "q4_k_m", "q5_k_m", "q6_k", "q8_k_m"])
args = parser.parse_args()
if args.checkpoint is None:
args.checkpoint = "outputs/tiny-sft/final/model.pt"
if not os.path.exists(args.checkpoint):
print(f"No checkpoint found at {args.checkpoint}")
print("Train first: bash scripts/train_tiny.sh")
sys.exit(1)
export_gguf(args.checkpoint, args.output)
print(f"GGUF file: {args.output}")
if __name__ == "__main__":
main()
|