physicsrob/torchwright-doom-e1m1

1"""The sole Doom inference program: portable stock-Hugging-Face generation.2 3This file is copied byte-identical to the root of every published Doom4bundle (``<bundle>/infer.py``) and executed there as a subprocess — by5anyone who downloads a bundle and by production render orchestration6(``run.py``) alike.  It is executed, never imported.  It intentionally7imports no TorchWright or ``torchwright_doom`` code: text enters through8the saved tokenizer, a stock ``Phi3ForCausalLM`` produces rows (a "row" is9one tokenizer id — one row of the tied embedding matrix, one sequence10position), and the only outputs are canonical integer ids plus their raw11tokenizer text.  No pixels are produced here: the bundle's standalone12tools (``tools/txt_to_png.py``; token protocol in ``PROTOCOL.md`` in the13source repo) decode those ids into a frame afterward and do no inference.14The bundle manifest's schema, field meanings, and completeness gate live15in ``torchwright_doom/bundle/manifest.py`` in the source repo.16"""17 18from __future__ import annotations19 20import argparse21import hashlib22import json23import os24import time25from pathlib import Path26 27# Multi-shard checkpoints otherwise load serially.  Keep these defaults in the28# portable program so downloaded bundles and production use the same loader.29os.environ.setdefault("HF_ENABLE_PARALLEL_LOADING", "true")30os.environ.setdefault("HF_PARALLEL_LOADING_WORKERS", "8")31 32import torch33import transformers34from transformers import TextGenerationPipeline, pipeline35 36_PROGRESS_INTERVAL_SECONDS = 15.037 38 39def _canonical_json(value) -> bytes:40    return json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8")41 42 43def _sha(data: bytes) -> str:44    return hashlib.sha256(data).hexdigest()45 46 47def _cuda_devices(model) -> list[torch.device]:48    return sorted(49        {parameter.device for parameter in model.parameters() if parameter.is_cuda},50        key=str,51    )52 53 54class _ProgressStreamer:55    """Report generation throughput without changing or collecting tokens."""56 57    def __init__(self, prompt_rows: int, max_new_tokens: int) -> None:58        self.prompt_rows = prompt_rows59        self.max_new_tokens = max_new_tokens60        self.started = time.monotonic()61        self.last_progress = self.started62        self.prefill_seconds: float | None = None63        self.finished: float | None = None64        self.generated_rows = 065        self._saw_prompt = False66 67    def put(self, value: torch.Tensor) -> None:68        # GenerationMixin streams the complete prompt once before any emitted69        # row.  It is already accounted for separately in ``prompt_rows``.70        if not self._saw_prompt:71            self._saw_prompt = True72            return73        now = time.monotonic()74        if self.prefill_seconds is None:75            self.prefill_seconds = now - self.started76            print(77                f"[infer] prefill complete; rows={self.prompt_rows} "78                f"elapsed={self.prefill_seconds:.1f}s",79                flush=True,80            )81        self.generated_rows += value.numel()82        if now - self.last_progress >= _PROGRESS_INTERVAL_SECONDS:83            elapsed = now - self.started - self.prefill_seconds84            last_row = int(value.reshape(-1)[-1])85            print(86                f"[infer] decode rows={self.generated_rows}/{self.max_new_tokens} "87                f"total_position={self.prompt_rows + self.generated_rows} "88                f"elapsed={elapsed:.1f}s "89                f"rows/s={self.generated_rows / elapsed:.1f} "90                f"last_row={last_row}",91                flush=True,92            )93            self.last_progress = now94 95    def end(self) -> None:96        self.finished = time.monotonic()97 98    @property99    def decode_seconds(self) -> float:100        stopped = self.finished or time.monotonic()101        prefill = self.prefill_seconds or 0.0102        return stopped - self.started - prefill103 104 105def main(argv: list[str] | None = None) -> int:106    parser = argparse.ArgumentParser(description="Run stock Phi-3 Doom inference")107    parser.add_argument("--model", type=Path)108    parser.add_argument("--prompt", type=Path)109    parser.add_argument("--output", type=Path, default=Path("out"))110    parser.add_argument(111        "--device", default="cuda" if torch.cuda.is_available() else "cpu"112    )113    parser.add_argument("--max-new-tokens", type=int)114    args = parser.parse_args(argv)115 116    # This file sits at the bundle root, so its own directory is the bundle.117    model_dir = (args.model or Path(__file__).resolve().parent).resolve()118    prompt_path = args.prompt or model_dir / "examples" / "e1m1_prompt.txt"119    manifest = json.loads((model_dir / "doom_bundle_manifest.json").read_text())120    if not manifest.get("validation", {}).get("complete"):121        raise ValueError("Doom bundle manifest is not complete")122 123    prompt_bytes = prompt_path.read_bytes()124    prompt_sha256 = _sha(prompt_bytes)125    bundled_prompt = prompt_sha256 == manifest["prompt"]["sha256"]126 127    load_t0 = time.monotonic()128    model_kwargs = {129        "attn_implementation": "eager",130        # Read each shard's bytes eagerly: deferring them to mmap page faults131        # stalls badly on network filesystems, and eager reads are harmless on132        # local disks.133        "disable_mmap": True,134    }135    generate: TextGenerationPipeline136    if args.device != "cpu":137        # Accelerate builds the skeleton on meta and dispatches each shard138        # directly to the target device.  This avoids a second full-model139        # ``model.to(cuda)`` pass through CPU-backed mmap pages.140        generate = pipeline(141            "text-generation",142            model=str(model_dir),143            dtype=torch.float32,144            model_kwargs=model_kwargs,145            device_map=args.device,146        )147    else:148        generate = pipeline(149            "text-generation",150            model=str(model_dir),151            dtype=torch.float32,152            model_kwargs=model_kwargs,153        )154    tokenizer = generate.tokenizer155    if tokenizer is None:156        raise RuntimeError("text-generation pipeline loaded without a tokenizer")157    model = generate.model158    model.eval()159    cuda_devices = _cuda_devices(model)160    for cuda_device in cuda_devices:161        # Reset after loading: the current allocation still includes all162        # weights, while the peak will additionally capture generation cache163        # and runtime workspace. This is the consumer-fit measurement.164        torch.cuda.reset_peak_memory_stats(cuda_device)165    attention_implementation = getattr(model.config, "_attn_implementation", None)166    if attention_implementation != "eager":167        # Eager is the implementation the published render was validated168        # under; fused kernels change fp accumulation order, and this check169        # keeps every run on the validated numerics.170        raise RuntimeError(171            "Doom inference requires eager attention, got "172            f"{attention_implementation!r}"173        )174    if (175        model.config.original_max_position_embeddings176        != model.config.max_position_embeddings177    ):178        raise RuntimeError(179            "default-RoPE Doom model has inconsistent original/max position "180            "capacity; GenerationMixin would discard its cache at the boundary"181        )182    load_seconds = time.monotonic() - load_t0183 184    prompt_text = prompt_bytes.decode("utf-8")185    encoded_prompt = tokenizer(186        prompt_text,187        return_tensors="pt",188        add_special_tokens=False,189    )190    input_device = next(model.parameters()).device191    prompt_ids = [int(row) for row in encoded_prompt.input_ids[0].tolist()]192    prompt_ids_sha256 = _sha(_canonical_json(prompt_ids))193    # Only the bundled prompt has a manifest row-id expectation; a custom194    # prompt is permitted, never verified, and recorded in the payload as195    # matches_bundled_prompt=false.196    if bundled_prompt and prompt_ids_sha256 != manifest["prompt"]["row_ids_sha256"]:197        raise ValueError("bundled prompt text does not reproduce its manifest rows")198 199    default_new = int(manifest["generation"]["max_new_tokens"])200    max_new = default_new if args.max_new_tokens is None else int(args.max_new_tokens)201    if max_new < 1:202        raise ValueError("max-new-tokens must be >= 1")203    if len(prompt_ids) + max_new > model.config.max_position_embeddings:204        raise ValueError("requested generation exceeds model position capacity")205 206    print(207        f"[infer] model ready in {load_seconds:.1f}s; prompt={len(prompt_ids)} "208        f"max_new_tokens={max_new} device={input_device}",209        flush=True,210    )211    generate_t0 = time.monotonic()212    progress = _ProgressStreamer(len(prompt_ids), max_new)213    print(f"[infer] generation started; max_new_tokens={max_new}", flush=True)214    with torch.inference_mode():215        records = generate(216            prompt_text,217            add_special_tokens=False,218            return_tensors=True,219            do_sample=False,220            use_cache=True,221            max_new_tokens=max_new,222            eos_token_id=tokenizer.eos_token_id,223            pad_token_id=tokenizer.pad_token_id,224            streamer=progress,225        )226    generate_seconds = time.monotonic() - generate_t0227    sequence = records[0]["generated_token_ids"]228    generated = [int(row) for row in sequence[len(prompt_ids) :]]229    for cuda_device in cuda_devices:230        torch.cuda.synchronize(cuda_device)231    cuda_memory = [232        {233            "device": str(cuda_device),234            "peak_allocated_bytes": torch.cuda.max_memory_allocated(cuda_device),235            "peak_reserved_bytes": torch.cuda.max_memory_reserved(cuda_device),236        }237        for cuda_device in cuda_devices238    ]239    prefill_seconds = progress.prefill_seconds or generate_seconds240    decode_seconds = progress.decode_seconds241    raw_text = tokenizer.decode(242        generated, skip_special_tokens=False, clean_up_tokenization_spaces=False243    )244    if not isinstance(raw_text, str):245        raise TypeError("tokenizer returned a batched decode for one row list")246    if tokenizer(raw_text, add_special_tokens=False)["input_ids"] != generated:247        raise ValueError("decoded output text does not round-trip to generated rows")248 249    args.output.mkdir(parents=True, exist_ok=True)250    emitted_ids_sha256 = _sha(_canonical_json(generated))251    stopped = bool(generated and generated[-1] == tokenizer.eos_token_id)252    payload = {253        "format": "torchwright_doom.output_ids.v1",254        "bundle": manifest.get("bundle_identity"),255        "compile_payload_sha256": manifest.get("compile_payload_sha256"),256        "row_vocab_fingerprint": manifest.get("row_vocab_fingerprint"),257        "prompt": {258            "sha256": prompt_sha256,259            "matches_bundled_prompt": bundled_prompt,260            "row_ids": prompt_ids,261            "row_ids_sha256": prompt_ids_sha256,262        },263        "emitted_row_ids": generated,264        "emitted_row_ids_sha256": emitted_ids_sha256,265        "generation": {266            "mode": "transformers_pipeline",267            "max_new_tokens": max_new,268            "termination_reason": "terminal" if stopped else "cap",269        },270        "timing_seconds": {271            "load": load_seconds,272            "prefill": prefill_seconds,273            "decode": decode_seconds,274            "generate": generate_seconds,275        },276        "attention_implementation": attention_implementation,277        "cuda_memory": cuda_memory,278        "transformers_version": transformers.__version__,279    }280    (args.output / "output.ids.json").write_text(281        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"282    )283    (args.output / "output.txt").write_text(raw_text + "\n", encoding="utf-8")284    print(285        f"[infer] wrote {len(generated)} rows in {generate_seconds:.1f}s; "286        f"stopped={payload['generation']['termination_reason']}",287        flush=True,288    )289    for memory in cuda_memory:290        peak_allocated = int(memory["peak_allocated_bytes"])291        peak_reserved = int(memory["peak_reserved_bytes"])292        print(293            f"[infer] {memory['device']} peak allocated="294            f"{peak_allocated / 1024**3:.2f} GiB "295            f"reserved={peak_reserved / 1024**3:.2f} GiB",296            flush=True,297        )298    return 0299 300 301if __name__ == "__main__":302    raise SystemExit(main())303