Upload export_apochat_litert.py with huggingface_hub

export_apochat_litert.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+"""Export the Apochat-tuned Gemma 4 E2B MLX model to a LiteRT .litertlm artifact.
+
+This script is intentionally **not** run on the 16 GB local Mac. It is designed for a
+machine with at least 32 GB of CPU RAM or a GPU with 24 GB+ VRAM (e.g. a Hugging Face
+Space/Notebook with GPU upgrade).
+
+Pipeline:
+1. Load the public MLX-q4 fused snapshot from Hugging Face.
+2. Dequantize weights to bfloat16 and save as PyTorch-format safetensors shards.
+3. Patch the config so transformers sees a normal bf16 checkpoint.
+4. Run `litert convert` with weight-only int4 quantization to produce .litertlm.
+5. Upload the resulting artifact to a Hugging Face model repo.
+
+Usage (on a high-memory machine / HF Space):
+    pip install -r scripts/requirements_litert_export.txt
+    python scripts/export_apochat_litert.py \
+        --mlx-repo apoapps/apochat-gemma4-e2b-apochat-tuned-v1 \
+        --output-dir ./apochat-litert-build \
+        --upload-repo apoapps/apochat-gemma4-e2b-apochat-tuned-v1-litert
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+import numpy as np
+from huggingface_hub import HfApi, create_repo, hf_hub_download, upload_file, upload_folder
+from safetensors.torch import save_file
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Export Apochat-tuned Gemma 4 E2B to LiteRT")
+    parser.add_argument(
+        "--mlx-repo",
+        default="apoapps/apochat-gemma4-e2b-apochat-tuned-v1",
+        help="Hugging Face repo id containing the fused MLX-q4 model",
+    )
+    parser.add_argument(
+        "--revision",
+        default=None,
+        help="Optional git revision for the MLX repo",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./apochat-litert-build",
+        help="Local directory for intermediate PyTorch checkpoint and final .litertlm",
+    )
+    parser.add_argument(
+        "--upload-repo",
+        default="apoapps/apochat-gemma4-e2b-apochat-tuned-v1-litert",
+        help="HF repo id where the final .litertlm will be uploaded",
+    )
+    parser.add_argument(
+        "--upload-private",
+        action="store_true",
+        help="Make the upload repo private",
+    )
+    parser.add_argument(
+        "--skip-upload",
+        action="store_true",
+        help="Keep the output local; do not upload to HF",
+    )
+    parser.add_argument(
+        "--prefill-lengths",
+        default="256",
+        help="LiteRT prefill signature lengths (comma separated)",
+    )
+    parser.add_argument(
+        "--cache-length",
+        type=int,
+        default=1024,
+        help="LiteRT KV-cache length",
+    )
+    parser.add_argument(
+        "--quantize-recipe",
+        default="weight_only_wi4_afp32",
+        help="LiteRT quantization recipe",
+    )
+    parser.add_argument(
+        "--shard-size",
+        type=int,
+        default=5_000_000_000,
+        help="Target size in bytes per PyTorch safetensors shard",
+    )
+    return parser.parse_args()
+
+
+def download_repo_files(repo_id: str, revision: str | None, local_dir: Path) -> None:
+    """Download all non-weight files from the MLX repo into local_dir."""
+    print(f"Downloading aux files from {repo_id} ...")
+    local_dir.mkdir(parents=True, exist_ok=True)
+    api = HfApi()
+    files = api.list_repo_files(repo_id, repo_type="model", revision=revision)
+    for fname in files:
+        if fname.endswith(".safetensors"):
+            continue
+        print(f"  {fname}")
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=fname,
+            repo_type="model",
+            revision=revision,
+            local_dir=str(local_dir),
+            local_dir_use_symlinks=False,
+        )
+
+
+def patch_config_for_pytorch(config_path: Path) -> None:
+    """Remove MLX quantization config and ensure torch_dtype is bfloat16."""
+    with open(config_path, "r", encoding="utf-8") as f:
+        config: dict[str, Any] = json.load(f)
+
+    config.pop("quantization_config", None)
+    text_config = config.get("text_config")
+    if isinstance(text_config, dict):
+        text_config.pop("quantization_config", None)
+    config["torch_dtype"] = "bfloat16"
+
+    with open(config_path, "w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2)
+
+
+def dequantize_mlx_to_pytorch(
+    mlx_repo: str,
+    revision: str | None,
+    output_dir: Path,
+    shard_size_bytes: int,
+) -> None:
+    """Load MLX-q4 weights, dequantize, and write PyTorch safetensors shards."""
+    print("Loading MLX-q4 weights ...")
+    api = HfApi()
+    files = api.list_repo_files(mlx_repo, repo_type="model", revision=revision)
+    safetensors_files = [f for f in files if f.endswith(".safetensors")]
+    weights: dict[str, mx.array] = {}
+    for fname in safetensors_files:
+        print(f"  {fname}")
+        local_path = hf_hub_download(
+            repo_id=mlx_repo,
+            filename=fname,
+            repo_type="model",
+            revision=revision,
+        )
+        part = mx.load(local_path)
+        if isinstance(part, dict):
+            weights.update(part)
+        else:
+            raise RuntimeError(f"Unexpected MLX load result for {fname}: {type(part)}")
+    print(f"Total tensors: {len(weights)}")
+
+    # Identify quantized triples: weight + scales + biases.
+    quantized: set[str] = set()
+    for name in list(weights.keys()):
+        if name.endswith(".scales"):
+            base = name[: -len(".scales")]
+            if f"{base}.biases" in weights:
+                quantized.add(base)
+
+    print(f"Quantized groups: {len(quantized)}")
+
+    current_shard: dict[str, Any] = {}
+    current_shard_bytes = 0
+    shard_index = 0
+
+    def flush_shard() -> None:
+        nonlocal current_shard, current_shard_bytes, shard_index
+        if not current_shard:
+            return
+        shard_path = output_dir / f"model-{shard_index:05d}-of-?????.safetensors"
+        save_file(current_shard, str(shard_path))
+        print(f"  Saved {shard_path.name} ({len(current_shard)} tensors, {current_shard_bytes / 1e9:.2f} GB)")
+        current_shard = {}
+        current_shard_bytes = 0
+        shard_index += 1
+
+    for name, arr in weights.items():
+        # Skip scale/bias metadata; we'll consume them with the base weight.
+        if name.endswith(".scales") or name.endswith(".biases"):
+            continue
+
+        base = name
+        is_quantized = base in quantized
+
+        if is_quantized:
+            scales = weights[f"{base}.scales"]
+            biases = weights[f"{base}.biases"]
+            # Dequantize to bfloat16 on the MLX device.
+            arr = mx.dequantize(arr, scales, biases, group_size=64, bits=4).astype(mx.bfloat16)
+        elif arr.dtype != mx.bfloat16:
+            arr = arr.astype(mx.bfloat16)
+
+        torch_tensor = mlx_bfloat16_to_torch(arr)
+        current_shard[name] = torch_tensor
+        current_shard_bytes += torch_tensor.nbytes
+
+        if current_shard_bytes >= shard_size_bytes:
+            flush_shard()
+
+    flush_shard()
+
+    # Rewrite the final shard names with the actual count.
+    shards = sorted(output_dir.glob("model-?????-of-?????.safetensors"))
+    total = len(shards)
+    for i, old in enumerate(shards):
+        new = old.with_name(f"model-{i:05d}-of-{total:05d}.safetensors")
+        old.rename(new)
+
+    print(f"Wrote {total} safetensors shard(s) to {output_dir}")
+
+
+def mlx_bfloat16_to_torch(arr: mx.array) -> Any:
+    """Convert an MLX bfloat16 array to a contiguous torch bfloat16 tensor."""
+    import torch
+
+    # MLX bfloat16 cannot be read directly by numpy; bridge via uint16.
+    u16 = np.array(arr.astype(mx.uint16))
+    if not u16.flags.c_contiguous:
+        u16 = np.ascontiguousarray(u16)
+    return torch.from_numpy(u16).view(torch.bfloat16)
+
+
+def run_litert_convert(
+    checkpoint_dir: Path,
+    output_dir: Path,
+    prefill_lengths: str,
+    cache_length: int,
+    quantize_recipe: str,
+) -> Path:
+    """Run `litert convert` on the dequantized checkpoint."""
+    print("Running litert convert ...")
+    cmd = [
+        "litert",
+        "convert",
+        str(checkpoint_dir),
+        "--output",
+        str(output_dir),
+        "--quantize-recipe",
+        quantize_recipe,
+        "--prefill-lengths",
+        prefill_lengths,
+        "--cache-length",
+        str(cache_length),
+        "--bundle-litert-lm",
+    ]
+    subprocess.run(cmd, check=True)
+
+    litertlm_files = list(output_dir.glob("*.litertlm"))
+    if not litertlm_files:
+        raise RuntimeError(f"No .litertlm file found in {output_dir}")
+    return litertlm_files[0]
+
+
+def upload_litert_model(repo_id: str, litertlm_path: Path, private: bool) -> str:
+    """Upload the .litertlm file to HF and return the git revision."""
+    print(f"Uploading {litertlm_path.name} to {repo_id} ...")
+    create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
+    upload_file(
+        repo_id=repo_id,
+        repo_type="model",
+        path_in_repo=litertlm_path.name,
+        path_or_fileobj=str(litertlm_path),
+    )
+    # Get the new revision.
+    api = HfApi()
+    info = api.repo_info(repo_id, repo_type="model")
+    print(f"Uploaded. Revision: {info.sha}")
+    return info.sha
+
+
+def main() -> int:
+    args = parse_args()
+    output_dir = Path(args.output_dir).resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Stage 1: prepare a transformers-compatible checkpoint.
+    pytorch_dir = output_dir / "pytorch_checkpoint"
+    pytorch_dir.mkdir(parents=True, exist_ok=True)
+
+    download_repo_files(args.mlx_repo, args.revision, pytorch_dir)
+    patch_config_for_pytorch(pytorch_dir / "config.json")
+
+    dequantize_mlx_to_pytorch(
+        args.mlx_repo,
+        args.revision,
+        pytorch_dir,
+        shard_size_bytes=args.shard_size,
+    )
+
+    # Stage 2: convert to LiteRT.
+    litert_dir = output_dir / "litert_out"
+    litert_dir.mkdir(parents=True, exist_ok=True)
+    litertlm_path = run_litert_convert(
+        pytorch_dir,
+        litert_dir,
+        args.prefill_lengths,
+        args.cache_length,
+        args.quantize_recipe,
+    )
+    print(f"LiteRT artifact: {litertlm_path}")
+
+    # Stage 3: upload.
+    if not args.skip_upload:
+        upload_litert_model(args.upload_repo, litertlm_path, args.upload_private)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())