Blame for scripts/run_elf.py

#!/usr/bin/env python3
"""Independent, offline ELF-B generation and synchronized CUDA timing.

Uses the locked official model and sampler without editing upstream sources.
Outputs intentionally do not claim compatibility with the legacy DLB registry.
"""
from __future__ import annotations

import argparse
from dataclasses import asdict
import json
import math
from pathlib import Path
import sys
import time

from elf_common import (ROOT, SOURCE_COMMIT, asset, default_sampling, manifest, offline,
                        preflight_conditions, read_jsonl, require_server, sha256, source_path, write_json)


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("mode", choices=["generate", "timing"])
    p.add_argument("--task", choices=["owt", "owt-prefix", "wmt14", "xsum"], required=True)
    p.add_argument("--steps", type=int, required=True)
    p.add_argument("--root", type=Path, default=ROOT)
    p.add_argument("--input", type=Path)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--num-samples", type=int, default=1024,
                   help="Prompt count; 0 means the entire conditional input file")
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--sampler", choices=["ode", "sde"])
    p.add_argument("--cfg", type=float)
    p.add_argument("--sc-cfg", type=float)
    p.add_argument("--gamma", type=float)
    p.add_argument("--compile", action="store_true")
    p.add_argument("--diversity", action="store_true", help="C64: four extra completions for first 256 prompts")
    p.add_argument("--timing-prompt-index", type=int, default=0)
    p.add_argument("--warmups", type=int, default=5)
    p.add_argument("--repeats", type=int, default=32)
    args = p.parse_args(argv)
    if args.steps < 1 or args.batch_size < 1 or args.num_samples < 0:
        p.error("steps/batch-size must be positive and num-samples nonnegative")
    if (args.task != "owt") != (args.input is not None):
        p.error("--input is required exactly for conditional tasks")
    if args.task == "owt" and args.num_samples == 0:
        p.error("unconditional generation needs an explicit positive sample count")
    if args.diversity and (args.task != "owt-prefix" or args.mode != "generate"):
        p.error("--diversity is only for owt-prefix generation")
    if args.mode == "timing" and args.batch_size != 1:
        p.error("primary timing requires --batch-size 1")
    if args.warmups < 1 or args.repeats < 2:
        p.error("use at least one warmup and two measured repeats")
    return args


def main():
    args = parse_args()
    require_server()
    offline()
    root = args.root.resolve()
    upstream = source_path(root)
    sys.dont_write_bytecode = True
    sys.path[:0] = [str(upstream / "src"), str(root / "src")]
    import numpy as np
    import torch
    from transformers import AutoTokenizer
    from configs.config import load_config_from_yaml, SamplingConfig
    from modules.model import ELF_models
    from modules.t5_encoder import get_encoder
    from utils.encoder_utils import encode_text, build_self_attn_cond_masks
    from utils.generation_utils import _generate_samples_single_batch, _dlm_decode_batch, shift_left, mask_after_eos
    from utils.sampling_utils import get_sampling_steps
    from dlb.timing import benchmark, cuda_runtime_metadata

    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
        raise RuntimeError("A CUDA GPU supporting bf16 is required for the author precision policy")
    if args.output.exists():
        raise FileExistsError(f"Use a fresh output directory: {args.output}")
    device = torch.device("cuda")
    torch.set_float32_matmul_precision("high")
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    generator = torch.Generator(device="cpu").manual_seed(args.seed)
    checkpoint_task = "owt" if args.task == "owt-prefix" else args.task
    config_task = "de-en" if checkpoint_task == "wmt14" else checkpoint_task
    config_path = upstream / f"src/configs/training_configs/train_{config_task}_ELF-B.yml"
    # The official loader resolves sampling YAML relative to upstream cwd.
    import os
    original_cwd = Path.cwd()
    os.chdir(upstream)
    try:
        config = load_config_from_yaml(str(config_path))
    finally:
        os.chdir(original_cwd)
    config.use_bf16 = True
    config.use_compile = args.compile
    config.use_wandb = False
    config.hf_repo_id = None
    chosen = default_sampling(args.task, args.steps)
    for key, value in (("sampling_method", args.sampler), ("cfg", args.cfg),
                       ("sc_cfg", args.sc_cfg), ("gamma", args.gamma)):
        if value is not None:
            chosen[key] = value
    if args.task in ("owt", "owt-prefix") and chosen["cfg"] != 1:
        raise ValueError("OWT checkpoint uses SC-CFG only; input CFG must be 1")
    if any(not math.isfinite(chosen[k]) or chosen[k] < 0 for k in ("cfg", "sc_cfg", "gamma")):
        raise ValueError("CFG/SC-CFG/gamma must be finite and nonnegative")
    sc = SamplingConfig(sampling_method=chosen["sampling_method"],
                        num_sampling_steps=[args.steps], cfgs=[chosen["cfg"]],
                        self_cond_cfg_scales=[chosen["sc_cfg"]], time_schedule="logit_normal",
                        sde_gamma=chosen["gamma"])
    t5_path = asset(root, "t5")
    tokenizer = AutoTokenizer.from_pretrained(str(t5_path), local_files_only=True)
    rows = []
    input_manifest = None
    if args.input:
        args.input = args.input.resolve()
        input_manifest = json.loads(args.input.with_suffix(".manifest.json").read_text())
        if input_manifest["sha256"] != sha256(args.input) or input_manifest["task"] != args.task:
            raise ValueError("Input manifest does not match task/file")
        rows = read_jsonl(args.input)
        if not rows or len(rows) != input_manifest["count"] or len({r["id"] for r in rows}) != len(rows):
            raise ValueError("Input counts/IDs do not match manifest or input is empty")
        if args.num_samples > len(rows) and args.mode == "generate":
            raise ValueError("Requested more prompts than input rows")
        if args.mode == "timing":
            if not 0 <= args.timing_prompt_index < len(rows):
                raise ValueError("Timing prompt index out of range")
            selected = [rows[0], rows[args.timing_prompt_index]]
        else:
            selected = rows[:args.num_samples] if args.num_samples else rows
        print(f"Checking {len(selected)} conditions before loading ELF weights: {args.input}", flush=True)
        tokens = preflight_conditions(selected, tokenizer, args.task, config.max_input_length, config.max_length)
        for row, ids in zip(selected, tokens):
            row["_elf_condition_ids"] = ids
    args.output.mkdir(parents=True)
    print(f"ELF {args.mode}: task={args.task}, steps={args.steps}, "
          f"split={(input_manifest or {}).get('split', 'unconditional')}, output={args.output}", flush=True)
    checkpoint_dir = asset(root, checkpoint_task)
    assets = manifest(root)
    checkpoint = checkpoint_dir / assets["assets"][checkpoint_task]["checkpoint"]
    # Inference only: use the exact EMA state without constructing an optimizer.
    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
    if not payload.get("ema_params1"):
        raise ValueError("The released EMA parameters are required")
    model = ELF_models["ELF-B"](
        text_encoder_dim=512, max_length=config.max_length,
        attn_drop=config.attn_dropout, proj_drop=config.proj_dropout,
        num_time_tokens=config.num_time_tokens, num_self_cond_cfg_tokens=config.num_self_cond_cfg_tokens,
        vocab_size=tokenizer.vocab_size, num_model_mode_tokens=config.num_model_mode_tokens,
        bottleneck_dim=config.bottleneck_dim)
    model.load_state_dict(payload["ema_params1"], strict=True)
    model = model.to(device).eval().requires_grad_(False)
    del payload
    encoder = None
    if args.input:
        _, encoder = get_encoder(str(t5_path), torch.float32)
        encoder = encoder.to(device).eval().requires_grad_(False)
    def prepare(batch_rows):
        # Encode only the visible source. References never enter the encoder.
        batch = len(batch_rows)
        ids = torch.full((batch, config.max_length), tokenizer.pad_token_id, device=device, dtype=torch.long)
        lens = []
        for i, row in enumerate(batch_rows):
            condition = row["_elf_condition_ids"]
            ids[i, :len(condition)] = torch.tensor(condition, device=device)
            lens.append(len(condition))
        mask = (np.arange(config.max_length)[None, :] < np.array(lens)[:, None])
        enc_mask, _, _ = build_self_attn_cond_masks(mask, mask, xp=np)
        return (ids, torch.as_tensor(enc_mask, device=device),
                torch.as_tensor(mask, dtype=torch.float32, device=device),
                torch.tensor(lens, device=device))

    def encode(prepared):
        return encode_text(input_ids=prepared[0], attention_mask=prepared[1], encoder=encoder,
                           latent_mean=config.latent_mean, latent_std=config.latent_std).float()

    @torch.inference_mode()
    def sample(batch, prepared=None, cached=None):
        cond = cached if cached is not None else encode(prepared) if prepared is not None else None
        times = get_sampling_steps(args.steps, "logit_normal", config.denoiser_p_mean,
                                   config.denoiser_p_std, device=device, dtype=torch.float32)
        z = torch.randn((batch, config.max_length, 512), device=device) * config.denoiser_noise_scale
        latent = _generate_samples_single_batch(model=model, generator=generator, z=z, t_steps=times,
                    cond_seq=cond, cond_seq_mask=prepared[2] if prepared is not None else None,
                    config=config, sampling_config=sc, cfg_scale=chosen["cfg"],
                    self_cond_cfg_scale=chosen["sc_cfg"])
        predicted = _dlm_decode_batch(latent, model, 1.0, config, chosen["sc_cfg"])
        if prepared is not None:
            # Fixed prefix is observed data, restored also in the discrete artifact.
            predicted = torch.where(prepared[2].bool(), prepared[0], predicted)
            predicted = shift_left(predicted, prepared[3], tokenizer.pad_token_id)
            if args.task != "owt-prefix":
                predicted = predicted[:, :config.max_length - config.max_input_length]
        return predicted

    # A separate untimed preflight observes every real model input. Verify both
    # clean prefix projection and the actual NFE, including the neural decoder.
    pilot = prepare([rows[0]]) if rows else None
    pilot_cond = encode(pilot) if pilot is not None else None
    observed = {"calls": 0, "prefix_checked_calls": 0}
    def check_forward(module, inputs, kwargs):
        observed["calls"] += 1
        x = inputs[0]
        if not torch.isfinite(x).all():
            raise ValueError("Non-finite ELF model input")
        if pilot is not None:
            mask = pilot[2].bool()
            visible = x[..., :512][mask]
            clean = pilot_cond[mask]
            if not torch.equal(visible, clean) and not (chosen["cfg"] != 1 and torch.count_nonzero(visible) == 0):
                raise ValueError("ELF prefix was not clamped before a model forward")
            observed["prefix_checked_calls"] += 1
    hook = model.register_forward_pre_hook(check_forward, with_kwargs=True)
    try:
        sample(1, pilot, pilot_cond)
    finally:
        hook.remove()
    expected_calls = args.steps * (1 if chosen["cfg"] == 1 else 2) + 1
    if observed["calls"] != expected_calls:
        raise ValueError(f"Unexpected ELF NFE: {observed['calls']} != {expected_calls}")
    # Pilot must not alter production sample assignment.
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    generator.manual_seed(args.seed)
    if args.compile:
        model = torch.compile(model)

    metadata = {
        "schema": "dlb-elf-v1", "task": args.task, "model": "ELF-B", "source_commit": SOURCE_COMMIT,
        "checkpoint_task": checkpoint_task, "checkpoint_selection": "ema_params1",
        "checkpoint_sha256": sha256(checkpoint), "assets_manifest_sha256": sha256(root / "data/elf/assets.json"),
        "source_config_sha256": sha256(config_path), "runner_sha256": sha256(Path(__file__)),
        "helper_sha256": sha256(Path(__file__).with_name("elf_common.py")),
        "lock_sha256": sha256(root / "artifacts/elf_lock.json"),
        "steps": args.steps, "sampling": chosen, "seed": args.seed, "batch_size": args.batch_size,
        "precision": "fp32 parameters/latents; official bf16 forward autocast; fp32 output heads",
        "matmul_precision": "high", "compile": args.compile,
        "model_canvas_t5_tokens": config.max_length, "input": input_manifest,
        "protocol": "c64_text_t5_v1" if args.task == "owt-prefix" else "elf_native_v1",
        "conditioning": "zero_shot_ood_native_projection" if args.task == "owt-prefix" else "task_trained" if rows else "none",
        "condition_tokenization_policy": (
            "released_ids_or_raw_plus_eos_then_source_cap_v2" if args.task in ("wmt14", "xsum")
            else "c64_raw_no_special_tokens_v1" if args.task == "owt-prefix" else "none"),
        "elf_forward_calls_per_sample": args.steps * (1 if chosen["cfg"] == 1 else 2) + 1,
        "condition_encoder_calls_per_sample": int(bool(rows)),
        "untimed_preflight": observed,
        **cuda_runtime_metadata(),
    }
    write_json(args.output / "request.json", metadata)
    if args.mode == "timing":
        if rows and not 0 <= args.timing_prompt_index < len(rows):
            raise ValueError("Timing prompt index out of range")
        prepared = prepare([rows[args.timing_prompt_index]]) if rows else None
        cached = encode(prepared) if prepared is not None else None
        results = {}
        for name in (["sampler_cached_condition", "encoder_plus_sampler"] if rows else ["sampler"]):
            def generate():
                return sample(1, prepared, cached if name == "sampler_cached_condition" else None)
            torch.cuda.reset_peak_memory_stats()
            result = benchmark(generate, torch.cuda.synchronize, warmups=args.warmups,
                               repeats=args.repeats, batch_size=1)
            results[name] = {**asdict(result), "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
                             "includes": ["noise_and_schedule", "denoising", "final_neural_decode", "argmax", "conditional_projection"]}
        write_json(args.output / "timing.json", {**metadata, "prompt_index": args.timing_prompt_index,
                   "condition_length": prepared[3].tolist() if prepared is not None else [], "results": results})
        print(json.dumps(results, indent=2))
        return
    count = (args.num_samples or len(rows)) if rows else args.num_samples
    schedule = [(i, 0) for i in range(count)]
    if args.diversity:
        schedule.extend((i, completion) for completion in range(1, 5) for i in range(min(256, count)))
    temporary = args.output / "samples.jsonl.tmp"
    start = time.perf_counter()
    with temporary.open("w") as handle:
        for offset in range(0, len(schedule), args.batch_size):
            shard = schedule[offset:offset + args.batch_size]
            batch_rows = [rows[i] for i, _ in shard] if rows else []
            prepared = prepare(batch_rows) if rows else None
            predicted = sample(len(shard), prepared)
            predicted = mask_after_eos(predicted, tokenizer.eos_token_id, tokenizer.pad_token_id).cpu().tolist()
            for j, ((i, completion), ids) in enumerate(zip(shard, predicted)):
                if args.task == "owt-prefix":
                    ids = ids[:config.max_length - int(prepared[3][j].item())]
                text = tokenizer.decode(ids, skip_special_tokens=True)
                row = {"id": offset + j, "prompt_id": i, "completion_id": completion,
                       "generated": text, "token_ids": ids}
                if rows:
                    row.update(input=rows[i]["input"], reference=rows[i]["output"], source_id=rows[i]["id"],
                               condition_t5_length=int(prepared[3][j].item()))
                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
            print(f"{offset + len(shard)}/{len(schedule)} samples", flush=True)
    temporary.replace(args.output / "samples.jsonl")
    write_json(args.output / "generation.json", {**metadata, "sample_count": len(schedule),
               "prompt_count": count, "diversity": args.diversity,
               "samples_sha256": sha256(args.output / "samples.jsonl"),
               "wall_seconds_including_io_not_primary_latency": time.perf_counter() - start})


if __name__ == "__main__":
    main()