| |
| """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" |
| |
| 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"] |
| |
| 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): |
| |
| 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: |
| |
| 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 |
|
|
| |
| |
| 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}") |
| |
| 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() |
|
|
