File size: 7,925 Bytes
c83c617 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """Independent V4.2 deployment stages; every stage writes a fresh output directory."""
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
import numpy as np
import torch
from vimeml.deployment import v4_chunk, v4_coreml as v4
from vimeml.deployment.bundle import environment, fresh_directory, write_json
from vimeml.training.data import file_sha
CHECKPOINT = ROOT / "artifacts/models/tiny-ja-v42-e8k-d384-l7/deployment.pt"
TOKENIZER = ROOT / "artifacts/tokenizers/v4-a/ja-unigram-8192"
REFERENCE = ROOT / "outputs/deployment/tiny-ja-v42-epoch20-fp32-input"
CUDA_SCORES = ROOT / "outputs/evaluation-v42-20261020"
@torch.inference_mode()
def baseline(output):
"""Task 1: Mac PyTorch CPU FP32 fixtures and four frozen tracks vs the CUDA run."""
from vimeml.training.infer_v4 import JapaneseLMV4
saved = torch.load(CHECKPOINT, map_location="cpu", weights_only=True)
lm = JapaneseLMV4(CHECKPOINT, TOKENIZER, "cpu")
lm.metadata.update(deployment_checkpoint_sha256=lm.metadata["checkpoint_sha256"],
checkpoint_sha256=saved["source_checkpoint_sha256"], device="cpu",
environment=environment(), threads=torch.get_num_threads())
info, arrays = v4.read_reference(REFERENCE, lm)
errors = []
for name, case in info["cases"].items():
actual = lm.model(torch.from_numpy(arrays[name + "_input_ids"].copy()).long()).numpy()
valid = case["valid_prefix"]
errors.append({"case": name, "check": "stateless", "valid_prefix": valid,
**v4.error_distribution(arrays[name + "_logits"][:, :valid], actual[:, :valid])})
cache, steps = lm.model.new_cache(), []
ids = torch.from_numpy(arrays["long_input_ids"].copy())
for position in range(ids.shape[1]):
steps.append(lm.model.forward_cached(ids[:, position:position + 1], cache))
errors.append({"case": "long", "check": "forward_cached_steps",
**v4.error_distribution(arrays["long_kv_step_logits"], torch.cat(steps, 1).numpy())})
fresh_directory(output)
passed = all(e["outside_tolerance"] == 0 for e in errors)
write_json(output / "fixtures.json", {"format": "vimeml_v4_mac_pytorch_fixtures_v1", "passed": passed,
"model": lm.metadata, "reference": info, "tolerances": {"atol": v4.TOLERANCE, "rtol": v4.TOLERANCE},
"errors": errors})
print(json.dumps({"fixtures_passed": passed, "max_abs": max(e["max_abs"] for e in errors)}), flush=True)
summary = v4.evaluate_tracks(lm, output / "tracks", CUDA_SCORES, "V4.2 FP32 CUDA (2026-10-20)")
return passed, summary
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--threads", type=int, default=4)
stages = parser.add_subparsers(dest="stage", required=True)
commands = {name: stages.add_parser(name) for name in (
"baseline", "export", "convert", "convert-kv", "align", "align-kv", "compress", "compress-kv",
"evaluate", "evaluate-kv", "evaluate-v21",
"convert-kvchunk", "align-kvchunk", "compress-kvchunk", "evaluate-kvchunk")}
for name, command in commands.items():
command.add_argument("--output", type=Path, required=True)
if name not in ("baseline", "export", "compress", "compress-kv", "evaluate-v21", "compress-kvchunk"):
command.add_argument("--bundle", type=Path, required=True)
if name in ("align", "align-kv", "evaluate", "evaluate-kv", "align-kvchunk", "evaluate-kvchunk"):
command.add_argument("--model", type=Path, required=name.endswith(("-kv", "-kvchunk")))
if name in ("compress", "compress-kv", "compress-kvchunk"):
command.add_argument("--source", type=Path, required=True)
command.add_argument("--alignment", type=Path, required=True)
if name == "evaluate-v21":
command.add_argument("--model", type=Path, required=True, help="Deployed V2.1 stateless package")
if name in ("evaluate", "evaluate-kv", "evaluate-v21", "evaluate-kvchunk"):
command.add_argument("--reference", type=Path)
command.add_argument("--reference-label")
commands["align-kv"].add_argument("--stateless", type=Path,
help="Same-precision stateless package; default compares with PyTorch fixtures")
commands["align-kvchunk"].add_argument("--stateless", type=Path,
help="Same-precision stateless package; default compares with PyTorch fixtures")
commands["align-kvchunk"].add_argument("--single", type=Path, help="Same-precision single-token KV package")
args = parser.parse_args()
if args.threads < 1 or args.output.exists():
parser.error("Positive threads and a fresh output path are required.")
torch.set_num_threads(args.threads)
if args.stage == "baseline":
passed, _ = baseline(args.output)
if not passed:
raise SystemExit("Mac PyTorch fixtures failed; evidence preserved.")
elif args.stage == "export":
v4.export_bundle(CHECKPOINT, TOKENIZER, args.output)
elif args.stage == "convert":
v4.convert(args.bundle, args.output)
elif args.stage == "convert-kv":
v4.convert_kv(args.bundle, args.output)
elif args.stage in ("compress", "compress-kv"):
v4.compress(args.source, args.output, args.alignment, kv=args.stage == "compress-kv")
elif args.stage == "align":
lm = v4.CoreMLLMV4(args.bundle, args.model) if args.model else v4.BundleLMV4(args.bundle)
report = v4.align(lm, REFERENCE, args.output)
print(json.dumps({"passed": report["passed"], **report["summary"]}, indent=2))
if not report["passed"]:
raise SystemExit("Strict alignment failed; report preserved. Review ranking quality separately.")
elif args.stage == "align-kv":
runtime = v4.KVRuntimeV4(args.bundle, args.model)
report = v4.align_kv(runtime, args.bundle, REFERENCE, args.output, args.stateless)
print(json.dumps({"passed": report["passed"], "comparison": report["comparison"], **report["summary"]}, indent=2))
if not report["passed"]:
raise SystemExit("KV alignment failed; evidence preserved.")
elif args.stage == "convert-kvchunk":
v4_chunk.convert(args.bundle, args.output)
elif args.stage == "compress-kvchunk":
v4_chunk.compress(args.source, args.output, args.alignment)
elif args.stage == "align-kvchunk":
runtime = v4_chunk.KVChunkRuntimeV4(args.bundle, args.model)
report = v4_chunk.align(runtime, args.bundle, REFERENCE, args.output, args.stateless, args.single)
print(json.dumps({"passed": report["passed"], "comparison": report["comparison"], **report["summary"]}, indent=2))
if not report["passed"]:
raise SystemExit("KV chunk alignment failed; evidence preserved.")
elif args.stage == "evaluate-kvchunk":
v4.evaluate_tracks(v4_chunk.CachedChunkLMV4(args.bundle, args.model), args.output, args.reference,
args.reference_label)
elif args.stage == "evaluate-v21":
# Same frozen pools and scorer for the deployed V2.1 packages (same Mac harness).
from vimeml.deployment import v2
lm = v2.CoreMLLM(ROOT / "artifacts/deployment/tiny-ja-v2.1-extend5-bundle-v1", args.model)
v4.evaluate_tracks(lm, args.output, args.reference, args.reference_label)
else:
if args.stage == "evaluate-kv":
lm = v4.CachedLMV4(args.bundle, args.model)
else:
lm = v4.CoreMLLMV4(args.bundle, args.model) if args.model else v4.BundleLMV4(args.bundle)
v4.evaluate_tracks(lm, args.output, args.reference, args.reference_label)
print(f"Output: {args.output.resolve()}")
if __name__ == "__main__":
main()
|