import os import sys import json import torch import struct import numpy as np from typing import Dict, List, Optional, Tuple sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) from src.model import FSIEdgeModel, FSIEdgeConfig # ============================================================================ # GGUF Format Specifications # ============================================================================ # GGUF File Structure: # - Magic: "GGUF" (4 bytes) # - Version (int) # - Tensor Count (int) # - Metadata KV pairs # - Alignment padding # - Tensor data (quantized) GGUF_MAGIC = 0x46554747 # "GGUF" as uint32 GGUF_VERSION = 3 # Tensor quantization types GGML_TYPE_F32 = 0 GGML_TYPE_F16 = 1 GGML_TYPE_Q4_0 = 2 GGML_TYPE_Q4_1 = 3 GGML_TYPE_Q5_0 = 6 GGML_TYPE_Q5_1 = 7 GGML_TYPE_Q8_0 = 8 GGML_TYPE_Q6_K = 14 GGML_TYPE_Q4_K_M = 21 GGML_TYPE_Q5_K_M = 23 # Changed from 22 to avoid conflict GGML_TYPE_Q3_K_M = 20 GGML_TYPE_Q2_K = 17 GGML_TYPE_Q8_K_M = 24 # Changed from 23 # Metadata keys KEY_GENERAL_ARCHITECTURE = "general.architecture" KEY_GENERAL_NAME = "general.name" KEY_GENERAL_DESCRIPTION = "general.description" KEY_GENERAL_FILE_TYPE = "general.file_type" KEY_GENERAL_PARAMETER_COUNT = "general.parameter_count" KEY_CONTEXT_LENGTH = "context_length" KEY_EMBEDDING_LENGTH = "embedding_length" KEY_BLOCK_COUNT = "block_count" KEY_FEED_FORWARD_LENGTH = "feed_forward_length" KEY_ATTENTION_HEAD_COUNT = "attention.head_count" KEY_ATTENTION_HEAD_COUNT_KV = "attention.head_count_kv" KEY_ATTENTION_LAYERNORM_EPS = "attention.layer_norm_epsilon" KEY_ROPE_DIMENSION_COUNT = "rope.dimension_count" KEY_ROPE_FREQ_BASE = "rope.freq_base" def quantize_q4_0(tensor: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """Quantize to Q4_0: 4-bit block quantization.""" assert tensor.ndim >= 1 shape = tensor.shape # Flatten flat = tensor.astype(np.float32).ravel() n = flat.shape[0] # Pad to block size block_size = 32 if n % block_size != 0: pad_len = block_size - (n % block_size) flat = np.pad(flat, (0, pad_len)) n_blocks = flat.shape[0] // block_size blocks = flat.reshape(n_blocks, block_size) # Per-block quantization scales = np.zeros(n_blocks, dtype=np.float16) quants = np.zeros(n_blocks * block_size // 2, dtype=np.uint8) for i in range(n_blocks): block = blocks[i] scale = np.max(np.abs(block)) / 7.0 if scale == 0: scale = 1.0 scales[i] = np.float16(scale) quant_block = np.clip(np.round(block / scale), -8, 7).astype(np.int8) # Pack into uint8 pairs for j in range(block_size // 2): lo = quant_block[2*j] & 0x0F hi = (quant_block[2*j + 1] & 0x0F) << 4 quants[i * (block_size // 2) + j] = lo | hi return scales, quants def quantize_q8_0(tensor: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """Quantize to Q8_0: 8-bit block quantization.""" flat = tensor.astype(np.float32).ravel() block_size = 32 if flat.shape[0] % block_size != 0: pad_len = block_size - (flat.shape[0] % block_size) flat = np.pad(flat, (0, pad_len)) n_blocks = flat.shape[0] // block_size blocks = flat.reshape(n_blocks, block_size) scales = np.zeros(n_blocks, dtype=np.float16) quants = np.zeros(n_blocks * block_size, dtype=np.uint8) for i in range(n_blocks): block = blocks[i] scale = np.max(np.abs(block)) / 127.0 if scale == 0: scale = 1.0 scales[i] = np.float16(scale) quant_block = np.clip(np.round(block / scale), -128, 127).astype(np.int8) quants[i * block_size:(i + 1) * block_size] = (quant_block & 0xFF).astype(np.uint8) return scales, quants def write_gguf_tensor(f, name: str, tensor: np.ndarray, quant_type: int): """Write tensor to GGUF file.""" name_bytes = name.encode('utf-8') + b'\x00' if quant_type == GGML_TYPE_F32: data = tensor.astype(np.float32).tobytes() elif quant_type == GGML_TYPE_F16: data = tensor.astype(np.float16).tobytes() elif quant_type == GGML_TYPE_Q4_0: scales, quants = quantize_q4_0(tensor) data = scales.tobytes() + quants.tobytes() elif quant_type == GGML_TYPE_Q8_0: scales, quants = quantize_q8_0(tensor) data = scales.tobytes() + quants.tobytes() else: raise ValueError(f"Unsupported quant type: {quant_type}") # Tensor info n_dims = len(tensor.shape) f.write(struct.pack('I', len(name_bytes))) f.write(name_bytes) f.write(struct.pack('I', n_dims)) f.write(struct.pack('I', quant_type)) for d in tensor.shape: f.write(struct.pack('Q', d)) # Pad to alignment alignment = 32 offset = f.tell() padded_offset = (offset + alignment - 1) // alignment * alignment f.write(b'\x00' * (padded_offset - offset)) f.write(data) def export_to_gguf(model, output_path, quant_type=GGML_TYPE_Q4_0, config=None): """Export FSI_Edge model to GGUF format.""" architecture = "fsi_edge" # Build module name mapping param_map = {} state_dict = model.state_dict() n_params = sum(p.numel() for p in model.parameters()) with open(output_path, 'wb') as f: # Header f.write(struct.pack('I', GGUF_MAGIC)) f.write(struct.pack('I', GGUF_VERSION)) # Count tensors tensor_count = len(state_dict) f.write(struct.pack('Q', tensor_count)) # Metadata count (key-value pairs) metadata = { KEY_GENERAL_ARCHITECTURE: architecture, KEY_GENERAL_NAME: f"fsi_edge-{config.d_model if config else '800m'}", KEY_GENERAL_DESCRIPTION: "FSI_Edge: Novel DNA Helix Memory coding model", KEY_GENERAL_FILE_TYPE: quant_type, KEY_GENERAL_PARAMETER_COUNT: n_params, KEY_CONTEXT_LENGTH: config.max_seq_len if config else 16384, KEY_EMBEDDING_LENGTH: config.d_model if config else 1536, KEY_BLOCK_COUNT: config.n_layers if config else 28, KEY_FEED_FORWARD_LENGTH: config.d_ff if config else 6144, KEY_ATTENTION_HEAD_COUNT: config.n_heads if config else 24, KEY_ATTENTION_HEAD_COUNT_KV: config.kv_heads if config else 6, KEY_ATTENTION_LAYERNORM_EPS: 1e-5, KEY_ROPE_DIMENSION_COUNT: (config.d_model // config.n_heads) if config else 64, KEY_ROPE_FREQ_BASE: 10000.0, } f.write(struct.pack('Q', len(metadata))) for key, value in metadata.items(): key_bytes = key.encode('utf-8') + b'\x00' f.write(struct.pack('I', len(key_bytes))) f.write(key_bytes) if isinstance(value, str): val_bytes = value.encode('utf-8') + b'\x00' f.write(struct.pack('I', 8)) # string type f.write(struct.pack('Q', len(val_bytes))) f.write(val_bytes) elif isinstance(value, (int, np.integer)): f.write(struct.pack('I', 4)) # int32 type f.write(struct.pack('i', int(value))) elif isinstance(value, float): f.write(struct.pack('I', 10)) # float32 type f.write(struct.pack('f', float(value))) else: f.write(struct.pack('I', 4)) f.write(struct.pack('i', int(value))) # Write tensors for name, tensor in state_dict.items(): np_tensor = tensor.cpu().numpy() gguf_name = name.replace('.', '_') # Skip non-essential layers for smaller file if any(skip in name for skip in ['helix_step', 'buffer']): continue write_gguf_tensor(f, gguf_name, np_tensor, quant_type) return output_path def convert_pytorch_to_gguf(model_path, model_size='800M', output_path=None, quant='q4_0'): """Convert saved PyTorch checkpoint to GGUF.""" quant_map = { 'f32': GGML_TYPE_F32, 'f16': GGML_TYPE_F16, 'q4_0': GGML_TYPE_Q4_0, 'q8_0': GGML_TYPE_Q8_0, } quant_type = quant_map.get(quant, GGML_TYPE_Q4_0) size_config = { '360M': {'d_model': 1024, 'n_layers': 24, 'n_heads': 16, 'kv_heads': 4, 'd_ff': 4096, 'max_seq_len': 8192}, '800M': {'d_model': 1536, 'n_layers': 28, 'n_heads': 24, 'kv_heads': 6, 'd_ff': 6144, 'max_seq_len': 16384}, '1.5B': {'d_model': 2048, 'n_layers': 32, 'n_heads': 32, 'kv_heads': 8, 'd_ff': 8192, 'max_seq_len': 32768}, } sc = size_config[model_size] config = FSIEdgeConfig(**sc) model = FSIEdgeModel(config) state = torch.load(model_path, map_location='cpu') if 'model_state_dict' in state: model.load_state_dict(state['model_state_dict']) else: model.load_state_dict(state) if output_path is None: output_path = f'fsi_edge-{model_size}-{quant}.gguf' export_to_gguf(model, output_path, quant_type, config) print(f"Exported GGUF: {output_path}") return output_path if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument('--model-path', type=str, required=True) parser.add_argument('--model-size', type=str, default='800M') parser.add_argument('--output', type=str, default=None) parser.add_argument('--quant', type=str, default='q4_0') args = parser.parse_args() convert_pytorch_to_gguf(args.model_path, args.model_size, args.output, args.quant)