File size: 8,356 Bytes
81cffe7 | 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 139 140 141 142 143 144 145 146 147 148 149 150 | """Independent V4-family deployment stages; every stage writes a fresh output directory.
`--variant` selects the frozen checkpoint, PyTorch fixtures and CUDA reference
(`vimeml.deployment.v4_variants`); the default keeps the recorded V4.2 commands.
"""
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.v4_variants import DEFAULT, VARIANTS
from vimeml.deployment.bundle import environment, fresh_directory, write_json
from vimeml.training.data import file_sha
TOKENIZER = ROOT / "artifacts/tokenizers/v4-a/ja-unigram-8192"
CHECKPOINT = REFERENCE = CUDA_SCORES = CUDA_LABEL = None
def select(name):
global CHECKPOINT, REFERENCE, CUDA_SCORES, CUDA_LABEL
variant = VARIANTS[name]
CHECKPOINT, REFERENCE = ROOT / variant.checkpoint, ROOT / variant.reference
CUDA_SCORES, CUDA_LABEL = ROOT / variant.cuda_scores, variant.cuda_label
@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, CUDA_LABEL)
return passed, summary
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--threads", type=int, default=4)
parser.add_argument("--variant", choices=sorted(VARIANTS), default=DEFAULT)
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)
select(args.variant)
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()
|