| |
| """Official checkpoint/implementation teacher-forced logits, on a CUDA host. |
| |
| No model downloads. Input IDs come from the C++ engine. This is a numerical |
| reference, not a held-out language-quality benchmark: continuation IDs may have |
| been generated by the old quantized engine. |
| """ |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from safetensors import safe_open |
| from transformers import AutoModelForCausalLM |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--source", type=Path, required=True) |
| parser.add_argument("--prompt-ids", type=Path, required=True) |
| parser.add_argument("--continuation-ids", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--candidate", action="append", default=[], help="NAME:LOGITS.f32") |
| args = parser.parse_args() |
| torch.set_num_threads(4) |
| torch.backends.cuda.matmul.allow_tf32 = False |
| prompt = [int(x) for x in args.prompt_ids.read_text().split()] |
| continuation = [int(x) for x in args.continuation_ids.read_text().split()] |
| assert prompt and continuation |
| model = AutoModelForCausalLM.from_pretrained( |
| args.source, trust_remote_code=True, local_files_only=True, |
| torch_dtype=torch.bfloat16, device_map={"": "cuda"}, attn_implementation="eager", |
| ).eval() |
| |
| |
| weight_map = json.loads((args.source / "model.safetensors.index.json").read_text())["weight_map"] |
| parameters = dict(model.named_parameters()) |
| for shard in sorted(set(weight_map.values())): |
| with safe_open(args.source / shard, framework="pt", device="cpu") as source: |
| for name in source.keys(): |
| if name in parameters and source.get_slice(name).get_dtype() == "F32": |
| parameters[name].data = source.get_tensor(name).to("cuda") |
| ids = torch.tensor([prompt + continuation[:-1]], device="cuda") |
| with torch.inference_mode(): |
| result = model(input_ids=ids, use_cache=False) |
| logits = result.logits[0, len(prompt) - 1:].float().cpu() |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| logits.numpy().tofile(args.output) |
| reference = logits.double() |
| log_p = torch.log_softmax(reference, dim=-1) |
| p = log_p.exp() |
| summary = {"prompt_tokens": len(prompt), "steps": len(continuation), "vocab": logits.shape[-1], |
| "reference": "official checkpoint and model code, BF16, source FP32 parameters restored", |
| "candidates": {}} |
| for item in args.candidate: |
| name, path = item.split(":", 1) |
| candidate = torch.from_numpy(np.fromfile(path, dtype=np.float32).reshape(logits.shape)).double() |
| if not torch.isfinite(candidate).all(): |
| raise ValueError(f"non-finite candidate {name}") |
| cosine = torch.nn.functional.cosine_similarity(reference, candidate, dim=-1) |
| kl = (p * (log_p - torch.log_softmax(candidate, dim=-1))).sum(-1) |
| summary["candidates"][name] = { |
| "cosine_mean": cosine.mean().item(), "cosine_min": cosine.min().item(), |
| "kl_mean": kl.mean().item(), "kl_max": kl.max().item(), |
| "top1_agreement": (reference.argmax(-1) == candidate.argmax(-1)).sum().item(), |
| "teacher_nll": -torch.log_softmax(candidate, dim=-1)[torch.arange(len(continuation)), continuation].mean().item(), |
| } |
| summary["reference_teacher_nll"] = -log_p[torch.arange(len(continuation)), continuation].mean().item() |
| args.output.with_suffix(".json").write_text(json.dumps(summary, indent=2) + "\n") |
| print(json.dumps(summary, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
