VibeGame-8B / code /scripts /export_model.py
khursheed's picture
Release pipeline source, tests, notebook, and licenses
ae36cfb verified
Raw History Blame Contribute Delete
16.2 kB
#!/usr/bin/env python3
"""Merge one tensor at a time, then export GGUF without loading a full 7B CPU model."""
from __future__ import annotations
import argparse
import gc
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path
BASE = "Qwen/Qwen2.5-Coder-7B-Instruct"
BASE_REVISION = "c03e6d358207e414f1eca0bb1891e29f1db0e242"
LLAMA_REVISION = "11fe02151f79c41d0d4af7da708755d73b9c0da6"
def sha256(path):
h = hashlib.sha256()
with Path(path).open("rb") as f:
for b in iter(lambda: f.read(1024 * 1024), b""):
h.update(b)
return h.hexdigest()
def atomic_json(path, data):
tmp = path.with_name(path.name + ".tmp")
with tmp.open("w") as f:
json.dump(data, f, indent=2)
f.flush()
os.fsync(f.fileno())
tmp.replace(path)
def lora_pairs(adapter_file):
from safetensors import safe_open
pairs = {}
with safe_open(adapter_file, framework="pt", device="cpu") as f:
for name in f.keys(): # noqa: SIM118 - safetensors reader is not a dict
match = re.fullmatch(r"(?:base_model\.model\.)?(.+)\.lora_([AB])(?:\.default)?\.weight", name)
if not match:
raise ValueError(f"Unsupported adapter tensor: {name}; only plain LoRA A/B supported")
pairs.setdefault(match[1] + ".weight", {})[match[2]] = f.get_tensor(name)
if not pairs or any(set(pair) != {"A", "B"} for pair in pairs.values()):
raise ValueError("Adapter contains missing LoRA A/B tensor pairs")
return pairs
def merge_tensor(weight, pair, scale, row_chunk=128):
"""Keep full output in fp16; calculate float32 deltas in small row blocks."""
import torch
# A dtype conversion already allocates. Chaining clone() would transiently
# duplicate the 1GiB vocabulary tensor on the 12GiB free Colab host.
result = weight.to(dtype=torch.float16, copy=True)
if pair is None:
return result
if weight.ndim != 2 or pair["A"].shape[1] != weight.shape[1] or pair["B"].shape[0] != weight.shape[0]:
raise ValueError("LoRA matrix dimensions do not match original base weights")
a = pair["A"].float()
for start in range(0, weight.shape[0], row_chunk):
stop = min(start + row_chunk, weight.shape[0])
# Read original precision here instead of the rounded fp16 result.
block = weight[start:stop].float()
block.addmm_(pair["B"][start:stop].float(), a, beta=1, alpha=scale)
result[start:stop] = block.to(torch.float16)
if not torch.isfinite(result).all():
raise ValueError("Nonfinite weights after fp16 merge")
return result.contiguous()
def merge(adapter, output, cache):
import torch
from huggingface_hub import snapshot_download
from safetensors import safe_open
from safetensors.torch import save_file
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
config = json.loads((adapter / "adapter_config.json").read_text())
if config.get("use_dora") or config.get("use_rslora") or config.get("fan_in_fan_out") or config.get("modules_to_save") or config.get("rank_pattern") or config.get("alpha_pattern") or config.get("bias", "none") != "none":
raise ValueError("Export supports only the fixed plain LoRA configuration used by train.py")
if config.get("r") != 32 or config.get("lora_alpha") != 64:
raise ValueError("Unexpected rank/alpha: expected r=32, alpha=64")
adapter_file = adapter / "adapter_model.safetensors"
if not (adapter / "TRAINING_COMPLETE.json").exists():
raise ValueError("No completed training marker. Finish training before release export.")
signature = {"adapter_sha256": sha256(adapter_file), "base": BASE, "revision": BASE_REVISION,
"merge": "Original unquantized base plus LoRA delta, fp32 row chunks rounded to fp16"}
output.mkdir(parents=True, exist_ok=True)
progress_path = output / "merge_progress.json"
if progress_path.exists():
progress = json.loads(progress_path.read_text())
if progress["signature"] != signature:
raise ValueError("Existing export belongs to a different adapter or base; choose a new directory")
else:
progress = {"signature": signature, "tensors": {}}
completed = output / "MERGE_COMPLETE.json"
if completed.exists() and json.loads(completed.read_text()).get("adapter_sha256") == signature["adapter_sha256"]:
records = progress["tensors"]
expected_count = json.loads(completed.read_text())["tensor_count"]
if len(records) == expected_count and all((output / row["file"]).exists() and sha256(output / row["file"]) == row["sha256"] for row in records.values()):
return output
base_dir = Path(snapshot_download(BASE, revision=BASE_REVISION, cache_dir=cache,
allow_patterns=["*.safetensors", "*.json", "*.txt", "*.model", "*.jinja", "LICENSE*", "NOTICE*"],
max_workers=1))
index = json.loads((base_dir / "model.safetensors.index.json").read_text())
mapping = index["weight_map"]
pairs = lora_pairs(adapter_file)
if set(pairs) - set(mapping):
raise ValueError(f"Adapter targets missing from original base: {sorted(set(pairs)-set(mapping))}")
weight_map, total_size = {}, 0
# Keep one mmap source and one converted tensor alive at a time.
for i, name in enumerate(sorted(mapping)):
filename = f"model-{i+1:05d}-of-{len(mapping):05d}.safetensors"
destination = output / filename
prior = progress["tensors"].get(name)
if prior and destination.exists() and sha256(destination) == prior["sha256"]:
weight_map[name] = filename
total_size += prior["tensor_bytes"]
continue
with safe_open(base_dir / mapping[name], framework="pt", device="cpu") as f:
weight = f.get_tensor(name)
merged = merge_tensor(weight, pairs.get(name), config["lora_alpha"] / config["r"])
tensor_bytes = merged.numel() * merged.element_size()
temp = destination.with_suffix(".partial")
save_file({name: merged}, str(temp), metadata={"format": "pt"})
temp.replace(destination)
del merged, weight
gc.collect()
progress["tensors"][name] = {"file": filename, "tensor_bytes": tensor_bytes, "sha256": sha256(destination)}
atomic_json(progress_path, progress)
weight_map[name] = filename
total_size += tensor_bytes
print(f"Merged tensor {i+1}/{len(mapping)}", flush=True)
atomic_json(output / "model.safetensors.index.json", {"metadata": {"total_size": total_size}, "weight_map": weight_map})
for path in base_dir.iterdir():
if path.is_file() and path.suffix in {".json", ".txt", ".model", ".jinja"} and path.name != "model.safetensors.index.json" or path.is_file() and path.name.startswith(("LICENSE", "NOTICE")):
shutil.copyfile(path, output / path.name)
# Preserve the exact training tokenizer/chat template in all released formats.
for path in adapter.iterdir():
if path.is_file() and ("token" in path.name or path.name in {"vocab.json", "merges.txt", "chat_template.jinja"}):
shutil.copyfile(path, output / path.name)
model_config = json.loads((output / "config.json").read_text())
model_config.pop("quantization_config", None)
model_config["torch_dtype"] = "float16"
atomic_json(output / "config.json", model_config)
atomic_json(output / "MERGE_COMPLETE.json", {**signature, "tensor_count": len(weight_map), "total_size": total_size,
"numeric_validation": "Finite tensors and dimension checks; full-model logits still require runtime validation."})
return output
def run(command, **kwargs):
subprocess.run([str(x) for x in command], check=True, **kwargs)
def converter_environment(llama, workspace):
"""Overlay converter dependencies without changing the live eval runtime."""
from importlib.metadata import version
# The converter pins torch==2.11.0. Reuse the installed CUDA build instead
# of downloading another multi-gigabyte wheel into this CPU-only export.
torch_version = version("torch")
if torch_version.split("+", 1)[0] != "2.11.0":
raise RuntimeError("Pinned GGUF converter requires installed torch2.11.0; preserve the eval environment and use the provided Colab runtime")
environment = workspace / "converter-venv"
# Colab's Python can omit ensurepip. Its existing system pip is visible
# through system-site-packages, so no bootstrap or Python upgrade is needed.
# Always rerun creation: a failed ensurepip attempt can leave the config
# with system-site-packages disabled even though bin/python already exists.
run([sys.executable, "-m", "venv", "--system-site-packages", "--without-pip", environment])
python = environment / "bin/python"
constraints = workspace / "converter-constraints.txt"
constraints.write_text(f"torch==={torch_version}\n")
env = dict(os.environ, TMPDIR=str(workspace), PYTHONNOUSERSITE="1",
PIP_REQUIRE_VIRTUALENV="true", PIP_DISABLE_PIP_VERSION_CHECK="1",
PIP_CONFIG_FILE=os.devnull,
OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1",
NUMEXPR_NUM_THREADS="1")
env.pop("PYTHONPATH", None)
env.pop("PYTHONHOME", None)
for key in ("PIP_TARGET", "PIP_PREFIX", "PIP_ROOT", "PIP_USER"):
env.pop(key, None)
isolation_probe = """import importlib.util, sys, sysconfig
from pathlib import Path
expected = Path(sys.argv[1]).resolve()
assert sys.prefix != sys.base_prefix, 'Converter must run in an isolated venv'
assert Path(sys.prefix).resolve() == expected, 'Unexpected converter prefix'
for key in ('purelib', 'platlib', 'scripts'):
assert Path(sysconfig.get_path(key)).resolve().is_relative_to(expected), 'Package destination escapes converter venv'
assert importlib.util.find_spec('pip') is not None, 'System pip is not visible inside converter venv'
"""
run([python, "-c", isolation_probe, environment], env=env)
run([python, "-m", "pip", "install", "--constraint", constraints,
"-r", llama / "requirements/requirements-convert_hf_to_gguf.txt"], env=env)
package_probe = """import importlib.metadata as m, json, sys
assert sys.prefix != sys.base_prefix, 'Converter must run in its isolated venv'
names = ['torch', 'numpy', 'transformers', 'sentencepiece', 'protobuf', 'huggingface-hub', 'gguf', 'safetensors', 'tokenizers']
print(json.dumps({'python': sys.version.split()[0], 'packages': {name: m.version(name) for name in names}}))
"""
versions = json.loads(subprocess.check_output([str(python), "-c", package_probe], env=env, text=True))
if versions["packages"]["torch"] != torch_version:
raise RuntimeError("Converter environment unexpectedly replaced the shared torch build")
return python, env, versions
def gguf(merged, output, workspace):
output.mkdir(parents=True, exist_ok=True)
workspace.mkdir(parents=True, exist_ok=True)
llama = workspace / "llama.cpp"
if not llama.exists():
run(["git", "init", llama])
run(["git", "-C", llama, "remote", "add", "origin", "https://github.com/ggml-org/llama.cpp.git"])
run(["git", "-C", llama, "fetch", "--depth", "1", "origin", LLAMA_REVISION])
run(["git", "-C", llama, "checkout", "--detach", LLAMA_REVISION])
converter_python, temp_env, converter_versions = converter_environment(llama, workspace)
run(["cmake", "-S", llama, "-B", llama / "build", "-DGGML_CUDA=OFF", "-DLLAMA_CURL=OFF", "-DLLAMA_BUILD_TESTS=OFF", "-DLLAMA_BUILD_SERVER=OFF"])
run(["cmake", "--build", llama / "build", "--config", "Release", "--target", "llama-quantize", "-j", "1"])
f16 = workspace / "vibegame-f16.gguf"
quantized = output / "VibeGame-8B-Q4_K_M.gguf"
# Converter uses lazy tensor loading and temp-backed writing, so it need not
# retain the whole fp16 model. No --no-lazy option and no GPU merge allocation.
run([converter_python, llama / "convert_hf_to_gguf.py", merged, "--outtype", "f16",
"--outfile", f16, "--use-temp-file", "--model-name", "VibeGame-8B (Qwen2.5-Coder-7B)"], env=temp_env)
# Upstream defaults to an 8GiB row-buffer budget; reserve host RAM for
# Colab, the export supervisor, and any completed-run evidence processing.
run([llama / "build/bin/llama-quantize", "--max-buffer-size", "256",
f16, quantized, "Q4_K_M", "1"])
if not quantized.exists() or quantized.stat().st_size < 100_000_000:
raise RuntimeError("GGUF conversion did not produce a plausible 7B model")
# Ollama's explicit template mirrors the Qwen chat template used in training.
(output / "Modelfile").write_text('''FROM ./VibeGame-8B-Q4_K_M.gguf
TEMPLATE """{{- if .Messages }}{{- range .Messages }}<|im_start|>{{ .Role }}
{{ .Content }}<|im_end|>
{{- end }}<|im_start|>assistant
{{ else }}{{ if .System }}<|im_start|>system
{{ .System }}<|im_end|>
{{ end }}{{ if .Prompt }}<|im_start|>user
{{ .Prompt }}<|im_end|>
{{ end }}<|im_start|>assistant
{{ .Response }}{{ end }}"""
PARAMETER stop "<|im_end|>"
PARAMETER stop "<|im_start|>"
PARAMETER num_ctx 4096
PARAMETER num_predict 3072
PARAMETER temperature 0
''')
atomic_json(output / "GGUF_COMPLETE.json", {"sha256": sha256(quantized), "bytes": quantized.stat().st_size,
"quantization": "Q4_K_M", "llama_cpp_revision": LLAMA_REVISION,
"converter_environment": converter_versions,
"resource_limits": {"merge_threads": 1, "build_jobs": 1, "quantizer_threads": 1,
"quantizer_buffer_mib": 256},
"source_merge": json.loads((merged / "MERGE_COMPLETE.json").read_text()),
"runtime_validation": "Not implied by successful conversion; run the frozen evaluation against this artifact."})
# Only delete our disposable intermediate after the quantized file is complete.
f16.unlink()
return quantized
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--adapter", type=Path, required=True)
p.add_argument("--output", type=Path, default=Path("/content/vibegame-exports"))
p.add_argument("--cache", type=Path, default=Path("/content/vibegame-base-cache"))
p.add_argument("--workspace", type=Path, default=Path("/content/vibegame-export-work"))
p.add_argument("--merge-only", action="store_true")
args = p.parse_args()
import psutil
if psutil.virtual_memory().available < 4 * 1024**3:
p.error("At least 4GiB available CPU RAM required; terminate the training subprocess before export")
# Merge and conversion run in stages. A disposable base cache is removed
# after the verified merge, before GGUF's temporary fp16 files are created.
for path in (args.output, args.cache, args.workspace):
path.mkdir(parents=True, exist_ok=True)
disposable_cache = args.cache / "vibegame-disposable-original-base"
merge_done = args.output / "merged-fp16/MERGE_COMPLETE.json"
if not merge_done.exists():
for path in {os.stat(x).st_dev: x for x in (args.output, args.cache)}.values():
if shutil.disk_usage(path).free < 34 * 1024**3:
p.error(f"Need 34GiB free for original-base and fp16 merge staging on {path}; no paid storage is allocated")
merged = merge(args.adapter, args.output / "merged-fp16", disposable_cache)
# This nested directory belongs exclusively to this export command.
if disposable_cache.exists():
shutil.rmtree(disposable_cache)
if not args.merge_only:
for path in {os.stat(x).st_dev: x for x in (args.output, args.workspace)}.values():
if shutil.disk_usage(path).free < 34 * 1024**3:
p.error(f"Merged fp16 is safe at {merged}; need 34GiB free on {path} for GGUF temporary files")
gguf(merged, args.output / "gguf", args.workspace)
print(f"Exported to {args.output}. Publish/download before the free VM disconnects.")
if __name__ == "__main__":
main()