export-gpt2.py

#!/usr/bin/env python3
#
# This is a script to export Hugging Face GPT-2 weights to a flat binary format with JSON manifest.
# The output is a single .bin file containing all tensors in a flat format, as well as a .manifest.json file
# that describes the tensors, their shapes, offsets, and other metadata.
#
# main.swift in turn loads this binary and manifest to interpret the model weights.

import argparse, json, math, os, sys, hashlib
from typing import Dict, Any, List, Tuple

import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoConfig

# ---------- Helpers ----------

DTYPE_MAP = {
    "float32": (np.float32, "f4"),
    "float16": (np.float16, "f2"),
    "bfloat16": (np.float16, "f2"),  # convert to f16 for portability unless you really want bf16
}

def classify_param(name: str) -> str:
    """
    Rough categorization for GPT-2 parameter names.
    """
    if name.startswith("transformer.wte."):
        return "token_embedding"
    if name.startswith("transformer.wpe."):
        return "positional_embedding"
    if name.startswith("transformer.ln_f."):
        return "final_layernorm"
    if name.startswith("lm_head."):
        return "lm_head"
    if ".attn." in name:
        if "c_attn." in name:
            return "attn_qkv_proj"
        if "c_proj." in name:
            return "attn_out_proj"
        if "bias" in name:
            return "attn_bias"
        return "attention"
    if ".mlp." in name:
        if "c_fc." in name:
            return "mlp_fc_in"
        if "c_proj." in name:
            return "mlp_fc_out"
        return "mlp"
    if ".ln_1." in name:
        return "pre_attn_layernorm"
    if ".ln_2." in name:
        return "pre_mlp_layernorm"
    return "other"

def sha256_bytes(b: bytes) -> str:
    h = hashlib.sha256()
    h.update(b)
    return h.hexdigest()

def align_offset(offset: int, align: int) -> int:
    if align <= 1:
        return offset
    return (offset + (align - 1)) // align * align

def torch_to_numpy(t: torch.Tensor, np_dtype: np.dtype) -> np.ndarray:
    # Move to CPU, convert dtype, ensure contiguous, and little-endian
    arr = t.detach().to("cpu").contiguous().numpy()
    # If torch dtype doesn't match requested, cast here.
    if arr.dtype == np.float32 and np_dtype == np.float16:
        arr = arr.astype(np.float16, copy=False)
    elif arr.dtype == np.float16 and np_dtype == np.float32:
        arr = arr.astype(np.float32, copy=False)
    elif arr.dtype == np.float32 and np_dtype == np.float32:
        pass
    elif arr.dtype == np.float16 and np_dtype == np.float16:
        pass
    else:
        # For any unexpected dtype, just try to cast to requested
        arr = arr.astype(np_dtype, copy=False)

    # Ensure little-endian explicitly
    if arr.dtype.byteorder not in ('<', '='):
        arr = arr.byteswap().newbyteorder('<')
    return arr

# ---------- Main exporter ----------

def export_model(model_id: str, out_prefix: str, dtype: str, align: int) -> None:
    if dtype not in DTYPE_MAP:
        raise ValueError(f"Unsupported dtype '{dtype}'. Choose from: {', '.join(DTYPE_MAP.keys())}")

    np_dtype, np_code = DTYPE_MAP[dtype]

    print(f"Loading config: {model_id}", file=sys.stderr)
    config = AutoConfig.from_pretrained(model_id)
    print(f"Loading model weights (no tokenizer needed): {model_id}", file=sys.stderr)
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True)
    model.eval()

    state: Dict[str, torch.Tensor] = model.state_dict()

    # Prepare outputs
    bin_path = f"{out_prefix}.bin"
    man_path = f"{out_prefix}.manifest.json"

    os.makedirs(os.path.dirname(os.path.abspath(bin_path)), exist_ok=True)

    manifest: Dict[str, Any] = {
        "format": "gpt2-flatbin-v1",
        "model_id": model_id,
        "dtype": dtype,
        "alignment": align,
        "total_tensors": len(state),
        "config": {
            "vocab_size": getattr(config, "vocab_size", None),
            "n_layer": getattr(config, "n_layer", None),
            "n_head": getattr(config, "n_head", None),
            "n_embd": getattr(config, "n_embd", None),
            "n_positions": getattr(config, "n_positions", None),
            "bos_token_id": getattr(config, "bos_token_id", None),
            "eos_token_id": getattr(config, "eos_token_id", None),
            "tie_word_embeddings": getattr(config, "tie_word_embeddings", None),
        },
        "tensors": [],   # filled below
    }

    byte_offset = 0
    with open(bin_path, "wb") as fbin:
        for name, tensor in state.items():
            layer_type = classify_param(name)
            shape = list(tensor.shape)

            # Align file position if requested
            aligned = align_offset(byte_offset, align)
            if aligned > byte_offset:
                pad = aligned - byte_offset
                fbin.write(b"\x00" * pad)
                byte_offset = aligned

            # Convert to numpy (requested dtype) and write
            arr = torch_to_numpy(tensor, np_dtype=np_dtype)
            raw = arr.tobytes(order="C")
            fbin.write(raw)

            entry = {
                "name": name,
                "layer_type": layer_type,
                "shape": shape,
                "dtype": dtype,           # stored dtype in the bin file
                "byte_offset": byte_offset,
                "nbytes": len(raw),
                "sha256": sha256_bytes(raw),
            }
            manifest["tensors"].append(entry)

            byte_offset += len(raw)

    with open(man_path, "w", encoding="utf-8") as fman:
        json.dump(manifest, fman, ensure_ascii=False, indent=2)

    print(f"Wrote: {bin_path}")
    print(f"Wrote: {man_path}")
    print("Done.")

# ---------- CLI ----------

def main():
    ap = argparse.ArgumentParser(description="Export Hugging Face GPT-2 weights into a single binary + JSON manifest.")
    ap.add_argument("--model", required=True, help="Model ID (e.g., 'gpt2', 'gpt2-medium', or a repo like 'openai-community/gpt2').")
    ap.add_argument("--out", required=True, help="Output prefix (e.g., 'gpt2_export'). Writes .bin and .manifest.json.")
    ap.add_argument("--dtype", default="float32", choices=list(DTYPE_MAP.keys()), help="Storage dtype.")
    ap.add_argument("--align", type=int, default=64, help="Byte alignment for each tensor (use 1 to disable).")
    args = ap.parse_args()
    export_model(args.model, args.out, args.dtype, args.align)

if __name__ == "__main__":
    main()