#!/usr/bin/env python3 """Run the preregistered task-bootstrap grid with explicit common random numbers. The model and the official temperature/top-p distribution are pinned, but the discrete draw is intentionally not ``torch.multinomial``. For each source+tape+step, a SHA-256-derived uniform variate is shared across the four prompt conditions and mapped through each row's token-ID-ordered CDF. """ from __future__ import annotations import argparse import hashlib import json import os import platform import resource import time from datetime import datetime, timezone from pathlib import Path from statistics import mean from typing import Any import accelerate import safetensors import torch import transformers from transformers import ( AutoTokenizer, TemperatureLogitsWarper, TopPLogitsWarper, ) import v2_lite_chat_sampling_probe as sampling import v2_lite_chat_special_token_behavior_probe as behavior import v2_lite_routing_special_token_family_control as special PROTOCOL_ID = "llm-atlas-deepseek-chat-task-bootstrap-crn-v1" MANIFEST_SHA256 = ( "6313e70536c464fe598a930355767524" "18f08016dfd60ac246437c3b43bf2ae1" ) CONDITIONS = ( "s0_eos", "s1_eos", "s0_period", "s1_period", ) DOMAINS = ("code", "math") TEMPERATURE = 0.3 TOP_P = 0.95 TOP_K = 0 FORMAL_MAX_NEW_TOKENS = 512 SMOKE_MAX_NEW_TOKENS = 16 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument( "--execution-mode", choices=("smoke", "formal", "replay"), required=True, ) parser.add_argument("--artifact-dir", type=Path, required=True) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument( "--reference-routing-json", type=Path, required=True, ) parser.add_argument("--human-eval", type=Path, required=True) parser.add_argument("--gsm8k", type=Path, required=True) parser.add_argument("--tnews", type=Path, required=True) parser.add_argument("--tnews-archive", type=Path, required=True) parser.add_argument("--wikitext", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--journal-dir", type=Path) parser.add_argument("--resume", action="store_true") parser.add_argument("--gpu-memory", default="28GiB") parser.add_argument("--cpu-memory", default="80GiB") parser.add_argument("--captured-at", default=None) return parser.parse_args() def sha256_bytes(payload: bytes) -> str: return hashlib.sha256(payload).hexdigest() def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: while chunk := handle.read(16 * 1024 * 1024): digest.update(chunk) return digest.hexdigest() def canonical_hash(value: Any) -> str: return sha256_bytes( json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ).encode() ) def install_condition_contract() -> None: factors = { condition: special.FACTORS[condition] for condition in CONDITIONS } special.BOUNDARY_LEVELS = ("eos", "period") special.CONDITIONS = CONDITIONS special.FACTORS = factors special.SYSTEM_CELLS = { "eos": ("s0_eos", "s1_eos"), "period": ("s0_period", "s1_period"), } special.SYSTEM_EDGE_CONTRASTS = { "period_minus_eos": ("eos", "period"), } special.COMPARISONS = ( ("system_eos", "s0_eos", "s1_eos"), ("system_period", "s0_period", "s1_period"), ("period_at_s0", "s0_eos", "s0_period"), ("period_at_s1", "s1_eos", "s1_period"), ) special.ALIGNMENT_COMPARISONS = special.COMPARISONS special.install_control_contract() def load_manifest(args: argparse.Namespace) -> dict[str, Any]: if sha256_file(args.manifest) != MANIFEST_SHA256: raise RuntimeError("frozen manifest SHA-256 differs") manifest = json.loads(args.manifest.read_text(encoding="utf-8")) if manifest["protocol_id"] != PROTOCOL_ID: raise RuntimeError("manifest protocol ID differs") if manifest["model"]["revision"] != behavior.MODEL_REVISION: raise RuntimeError("manifest model revision differs") if tuple(manifest["conditions"]) != CONDITIONS: raise RuntimeError("manifest condition order differs") if tuple(manifest["domains"]) != DOMAINS: raise RuntimeError("manifest domain order differs") return manifest def validate_args(args: argparse.Namespace) -> None: expected_tokens = ( SMOKE_MAX_NEW_TOKENS if args.execution_mode == "smoke" else FORMAL_MAX_NEW_TOKENS ) if args.resume and args.journal_dir is None: raise ValueError("--resume requires --journal-dir") if args.execution_mode != "smoke" and args.journal_dir is None: raise ValueError("formal/replay execution requires --journal-dir") for path in ( args.artifact_dir, args.manifest, args.reference_routing_json, args.human_eval, args.gsm8k, args.tnews, args.tnews_archive, args.wikitext, ): if not path.exists(): raise FileNotFoundError(path) missing = [ name for name in behavior.MODEL_FILES if not (args.artifact_dir / name).is_file() ] if missing: raise FileNotFoundError( f"artifact directory is incomplete: {missing}" ) if expected_tokens <= 0: raise AssertionError("max-new-token contract is invalid") def tape_uint64(tape: str, source_id: str, step: int) -> int: payload = ( f"{PROTOCOL_ID}\0uniform\0" f"{tape}\0{source_id}\0{step}" ).encode() return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") def uniform_float32( value: int, *, device: torch.device, ) -> torch.Tensor: if not 0 <= value < 1 << 64: raise ValueError("uniform integer is outside uint64") raw = (value + 0.5) / float(1 << 64) result = torch.tensor(raw, dtype=torch.float32, device=device) zero = torch.tensor(0.0, dtype=torch.float32, device=device) one = torch.tensor(1.0, dtype=torch.float32, device=device) lower = torch.nextafter(zero, one) upper = torch.nextafter(one, zero) return torch.clamp(result, min=lower, max=upper) def inverse_cdf_sample( *, input_ids: torch.Tensor, logits: torch.Tensor, uniform: torch.Tensor, temperature_warper: TemperatureLogitsWarper, top_p_warper: TopPLogitsWarper, ) -> tuple[torch.Tensor, dict[str, Any]]: scores = temperature_warper(input_ids, logits) scores = top_p_warper(input_ids, scores) finite = torch.isfinite(scores) if not bool(finite.any(dim=-1).all()): raise RuntimeError("a sampling row has no finite logits") probabilities = torch.softmax(scores.float(), dim=-1) if not bool(torch.isfinite(probabilities).all()): raise RuntimeError("sampling probabilities contain NaN/Inf") cdf = torch.cumsum(probabilities, dim=-1) cdf[:, -1] = 1.0 row_uniforms = uniform.expand(cdf.shape[0]).contiguous() sampled = torch.searchsorted( cdf.contiguous(), row_uniforms.unsqueeze(-1), right=False, ).squeeze(-1) if bool((sampled >= cdf.shape[-1]).any()): raise RuntimeError("inverse-CDF search exceeded vocabulary") return sampled, { "finite_tokens_per_row": finite.sum(dim=-1).detach().cpu().tolist(), "probability_sum_per_row": probabilities.sum( dim=-1 ).detach().cpu().tolist(), } def synthetic_sampler_test() -> dict[str, Any]: input_ids = torch.tensor([[1], [1]], dtype=torch.long) logits = torch.tensor( [[-3.0, 4.0, 1.0, 0.0], [0.0, 1.0, 4.0, -3.0]], dtype=torch.float32, ) temperature_warper = TemperatureLogitsWarper(TEMPERATURE) top_p_warper = TopPLogitsWarper( TOP_P, min_tokens_to_keep=1, ) observations = [] for name, value in ( ("low", 0), ("middle", 1 << 63), ("high", (1 << 64) - 1), ): uniform = uniform_float32(value, device=logits.device) sampled, audit = inverse_cdf_sample( input_ids=input_ids, logits=logits, uniform=uniform, temperature_warper=temperature_warper, top_p_warper=top_p_warper, ) observations.append( { "name": name, "uint64": value, "uniform_float32": float(uniform.item()), "sampled": sampled.tolist(), "audit": audit, } ) passed = ( 0 < observations[0]["uniform_float32"] < 1 and 0 < observations[-1]["uniform_float32"] < 1 and all( all(count < logits.shape[-1] for count in row["audit"][ "finite_tokens_per_row" ]) for row in observations ) and observations[0]["sampled"] == [1, 2] and observations[-1]["sampled"] == [1, 2] ) if not passed: raise RuntimeError( f"synthetic inverse-CDF sampler test failed: {observations}" ) return { "passed": True, "temperature": TEMPERATURE, "top_p": TOP_P, "observations": observations, } def prepare_sources( args: argparse.Namespace, tokenizer: Any, manifest: dict[str, Any], ) -> tuple[list[dict[str, Any]], dict[str, Any], dict[str, Any]]: args.domains = list(DOMAINS) args.per_domain = 32 source_rows, reference = behavior.selected_sources(args, tokenizer) manifest_rows = { row["id"]: row for row in manifest["sources"] } audit_rows = [] for source in source_rows: source["variants"] = { condition: special.prior.render_boundary_variant( tokenizer, source["content"], condition, ) for condition in CONDITIONS } frozen = manifest_rows.get(source["id"]) if frozen is None: raise RuntimeError( f"{source['id']} is absent from frozen manifest" ) if ( frozen["domain"] != source["domain"] or frozen["domain_index"] != source["within_domain_index"] or frozen["source_text_sha256"] != source["source_text_sha256"] ): raise RuntimeError( f"{source['id']} source contract differs" ) for condition in CONDITIONS: observed = source["variants"][condition][ "token_ids_sha256" ] expected = frozen[ "chat_generation_prompt_token_ids_sha256" ][ condition ] audit_rows.append( { "source_id": source["id"], "domain": source["domain"], "domain_index": source["within_domain_index"], "condition": condition, "observed_sha256": observed, "manifest_sha256": expected, "exact": observed == expected, } ) mismatches = [row for row in audit_rows if not row["exact"]] if mismatches: raise RuntimeError( f"{len(mismatches)} prompt hashes differ from manifest" ) if len(audit_rows) != 256: raise RuntimeError( f"expected 256 prompt audit rows, got {len(audit_rows)}" ) return source_rows, reference, { "cells": len(audit_rows), "exact": sum(row["exact"] for row in audit_rows), "rows": audit_rows, } def execution_assignment( *, mode: str, source_rows: list[dict[str, Any]], manifest: dict[str, Any], ) -> list[tuple[dict[str, Any], list[str]]]: frozen = {row["id"]: row for row in manifest["sources"]} assignment = [] for source in source_rows: row = frozen[source["id"]] if mode == "formal": tapes = [*row["main_tapes"], *row["diagnostic_tapes"]] elif mode == "replay": tapes = ["T0"] if row["independent_replay"] else [] else: tapes = ( ["T0", "T1"] if row["domain_index"] in manifest["tape_contract"][ "diagnostic_source_indices_per_domain" ] else [] ) if tapes: assignment.append((source, tapes)) expected = {"formal": (64, 88), "replay": (16, 16), "smoke": (8, 16)} observed = ( len(assignment), sum(len(tapes) for _, tapes in assignment), ) if observed != expected[mode]: raise RuntimeError( f"{mode} assignment is {observed}, expected {expected[mode]}" ) return assignment def prepare_batch( source: dict[str, Any], tokenizer: Any, input_device: torch.device, ) -> dict[str, Any]: unpadded = [ source["variants"][condition]["token_ids"] for condition in CONDITIONS ] prompt_lengths = [len(row) for row in unpadded] batch_prompt_length = max(prompt_lengths) padded = [ [int(tokenizer.pad_token_id)] * (batch_prompt_length - len(row)) + row for row in unpadded ] masks = [ [0] * (batch_prompt_length - len(row)) + [1] * len(row) for row in unpadded ] return { "input_ids": torch.tensor( padded, dtype=torch.long, device=input_device, ), "attention_mask": torch.tensor( masks, dtype=torch.long, device=input_device, ), "prompt_lengths": prompt_lengths, "batch_prompt_length": batch_prompt_length, } def derive_run_seed(tape: str, source_id: str) -> int: payload = ( f"{PROTOCOL_ID}/run\0{tape}\0{source_id}" ).encode() value = int.from_bytes( hashlib.sha256(payload).digest()[:8], "big", ) return value % ((1 << 63) - 1) def generate_tape( *, source: dict[str, Any], tape: dict[str, Any], batch: dict[str, Any], model: Any, tokenizer: Any, gold: dict[str, dict[str, Any]], max_new_tokens: int, ) -> dict[str, Any]: tape_label = tape["label"] run_seed = derive_run_seed(tape_label, source["id"]) rng_before = sampling.set_run_seed(run_seed) if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() started = time.perf_counter() input_ids = batch["input_ids"] attention_mask = batch["attention_mask"] generated_ids: list[list[int]] = [ [] for _ in CONDITIONS ] active = [True for _ in CONDITIONS] uniform_uint64: list[int] = [] uniform_float32_values: list[float] = [] finite_token_ranges: list[tuple[int, int]] = [] past_key_values = None temperature_warper = TemperatureLogitsWarper(TEMPERATURE) top_p_warper = TopPLogitsWarper( TOP_P, min_tokens_to_keep=1, ) with torch.inference_mode(): for step in range(max_new_tokens): model_inputs = model.prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, attention_mask=attention_mask, use_cache=True, ) outputs = model(**model_inputs, return_dict=True) past_key_values = outputs.past_key_values logits = outputs.logits[:, -1, :] z_value = tape_uint64( tape_label, source["id"], step, ) uniform = uniform_float32( z_value, device=logits.device, ) sampled, step_audit = inverse_cdf_sample( input_ids=input_ids, logits=logits, uniform=uniform, temperature_warper=temperature_warper, top_p_warper=top_p_warper, ) sampled_ids = sampled.detach().cpu().tolist() active_before = list(active) uniform_uint64.append(z_value) uniform_float32_values.append(float(uniform.item())) finite_counts = step_audit["finite_tokens_per_row"] finite_token_ranges.append( (min(finite_counts), max(finite_counts)) ) next_for_model = [] next_mask = [] for row_index, is_active in enumerate(active_before): if is_active: token_id = int(sampled_ids[row_index]) generated_ids[row_index].append(token_id) next_for_model.append(token_id) next_mask.append(1) if token_id == int(tokenizer.eos_token_id): active[row_index] = False else: next_for_model.append(int(tokenizer.pad_token_id)) next_mask.append(0) if not any(active) or step + 1 == max_new_tokens: break input_ids = torch.cat( [ input_ids, torch.tensor( next_for_model, dtype=torch.long, device=input_ids.device, ).unsqueeze(-1), ], dim=-1, ) attention_mask = torch.cat( [ attention_mask, torch.tensor( next_mask, dtype=torch.long, device=attention_mask.device, ).unsqueeze(-1), ], dim=-1, ) del past_key_values del outputs if torch.cuda.is_available(): torch.cuda.synchronize() elapsed = time.perf_counter() - started rng_after = sampling.rng_snapshot() peak_cuda = ( torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None ) if rng_before != rng_after: raise RuntimeError( f"{source['id']}/{tape_label}: inference consumed torch RNG" ) output_rows = [] for row_index, condition in enumerate(CONDITIONS): token_ids = generated_ids[row_index] hit_eos = bool( token_ids and token_ids[-1] == int(tokenizer.eos_token_id) ) text = tokenizer.decode( token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False, ) consumed_uniforms = uniform_uint64[: len(token_ids)] output_rows.append( { "condition": condition, "factors": special.FACTORS[condition], "prompt_tokens": batch["prompt_lengths"][row_index], "left_padding_tokens": ( batch["batch_prompt_length"] - batch["prompt_lengths"][row_index] ), "prompt_token_ids_sha256": source["variants"][ condition ]["token_ids_sha256"], "generated_tokens": len(token_ids), "generated_token_ids": token_ids, "generated_token_ids_sha256": canonical_hash(token_ids), "hit_eos": hit_eos, "stopped_at_max_new_tokens": ( not hit_eos and len(token_ids) == max_new_tokens ), "text": text, "text_sha256": sha256_bytes(text.encode()), "uniform_steps_consumed": len(consumed_uniforms), "uniform_uint64_prefix_sha256": canonical_hash( consumed_uniforms ), "uniform_uint64_first_eight_hex": [ f"{value:016x}" for value in consumed_uniforms[:8] ], "task_score": behavior.task_score( source["domain"], source["id"], text, gold, ), } ) return { "replicate_index": tape["index"], "replicate_label": tape_label, "tape_label": tape_label, "base_seed": tape["display_seed"], "run_seed": run_seed, "run_seed_used_only_for_rng_nonconsumption_audit": True, "rng_state_before": rng_before, "rng_state_after": rng_after, "torch_rng_unchanged": True, "uniform_steps_available": len(uniform_uint64), "uniform_uint64_sha256": canonical_hash(uniform_uint64), "uniform_uint64_first_eight_hex": [ f"{value:016x}" for value in uniform_uint64[:8] ], "uniform_float32_first_eight": uniform_float32_values[:8], "finite_tokens_per_step_min": min( value[0] for value in finite_token_ranges ), "finite_tokens_per_step_max": max( value[1] for value in finite_token_ranges ), "generation_seconds": elapsed, "peak_cuda_memory_allocated_bytes": peak_cuda, "outputs": output_rows, } def compare_runs( left: dict[str, Any], right: dict[str, Any], ) -> dict[str, Any]: rows = [] for left_output, right_output in zip( left["outputs"], right["outputs"], strict=True, ): checks = { "prompt_hash_exact": ( left_output["prompt_token_ids_sha256"] == right_output["prompt_token_ids_sha256"] ), "uniform_hash_exact": ( left_output["uniform_uint64_prefix_sha256"] == right_output["uniform_uint64_prefix_sha256"] ), "generated_token_ids_exact": ( left_output["generated_token_ids"] == right_output["generated_token_ids"] ), "text_exact": ( left_output["text"] == right_output["text"] ), "stop_state_exact": ( left_output["hit_eos"] == right_output["hit_eos"] and left_output["stopped_at_max_new_tokens"] == right_output["stopped_at_max_new_tokens"] ), "rng_pre_state_exact": ( left["rng_state_before"] == right["rng_state_before"] ), } rows.append( { "condition": left_output["condition"], **checks, "all_exact": all(checks.values()), } ) return { "cells": len(rows), "all_exact": sum(row["all_exact"] for row in rows), "rows": rows, "passed": all(row["all_exact"] for row in rows), } def journal_header( *, args: argparse.Namespace, source: dict[str, Any], tapes: list[str], max_new_tokens: int, ) -> dict[str, Any]: return { "protocol_id": PROTOCOL_ID, "manifest_sha256": MANIFEST_SHA256, "model_revision": behavior.MODEL_REVISION, "execution_mode": args.execution_mode, "source_id": source["id"], "domain": source["domain"], "domain_index": source["within_domain_index"], "tapes": tapes, "conditions": list(CONDITIONS), "max_new_tokens": max_new_tokens, "temperature": TEMPERATURE, "top_p": TOP_P, "top_k": TOP_K, } def journal_path( journal_dir: Path, source: dict[str, Any], ) -> Path: safe_id = source["id"].replace("/", "-") return ( journal_dir / f"{source['domain']}-{source['within_domain_index']:02d}-{safe_id}.json" ) def load_journal( path: Path, expected_header: dict[str, Any], ) -> dict[str, Any]: payload = json.loads(path.read_text(encoding="utf-8")) if payload.get("header") != expected_header: raise RuntimeError(f"journal header differs: {path}") source = payload.get("source") if canonical_hash(source) != payload.get("source_sha256"): raise RuntimeError(f"journal content hash differs: {path}") return source def write_journal( path: Path, header: dict[str, Any], source: dict[str, Any], ) -> None: path.parent.mkdir(parents=True, exist_ok=True) payload = { "header": header, "source_sha256": canonical_hash(source), "source": source, } temporary = path.with_suffix(".tmp") temporary.write_text( json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) os.replace(temporary, path) def summarize_sources( sources: list[dict[str, Any]], ) -> dict[str, Any]: rows = [ output for source in sources for run in source["runs"] for output in run["outputs"] ] replay = [ source["in_process_replay"] for source in sources if source.get("in_process_replay") is not None ] by_domain = {} for domain in DOMAINS: domain_rows = [ row for source in sources if source["domain"] == domain for run in source["runs"] for row in run["outputs"] ] by_domain[domain] = { "sources": len( { source["id"] for source in sources if source["domain"] == domain } ), "outputs": len(domain_rows), "natural_eos": sum(row["hit_eos"] for row in domain_rows), "budget_truncated": sum( row["stopped_at_max_new_tokens"] for row in domain_rows ), "mean_generated_tokens": ( mean(row["generated_tokens"] for row in domain_rows) if domain_rows else None ), } return { "sources": len(sources), "runs": sum(len(source["runs"]) for source in sources), "outputs": len(rows), "natural_eos": sum(row["hit_eos"] for row in rows), "budget_truncated": sum( row["stopped_at_max_new_tokens"] for row in rows ), "unique_generated_token_hashes": len( {row["generated_token_ids_sha256"] for row in rows} ), "torch_rng_unchanged_runs": sum( run["torch_rng_unchanged"] for source in sources for run in source["runs"] ), "in_process_replay": { "sources": len(replay), "cells": sum(row["cells"] for row in replay), "all_exact": sum(row["all_exact"] for row in replay), "passed": all(row["passed"] for row in replay) if replay else None, }, "by_domain": by_domain, } def main() -> None: args = parse_args() validate_args(args) manifest = load_manifest(args) install_condition_contract() synthetic_test = synthetic_sampler_test() max_new_tokens = ( SMOKE_MAX_NEW_TOKENS if args.execution_mode == "smoke" else FORMAL_MAX_NEW_TOKENS ) captured_at = args.captured_at or datetime.now( timezone.utc ).isoformat() revision_contract = behavior.download_revision_contract( args.artifact_dir ) official_generation = json.loads( (args.artifact_dir / "generation_config.json").read_text( encoding="utf-8" ) ) expected_official = { "do_sample": True, "temperature": TEMPERATURE, "top_p": TOP_P, "bos_token_id": 100000, "eos_token_id": 100001, } mismatches = { key: { "expected": expected, "observed": official_generation.get(key), } for key, expected in expected_official.items() if official_generation.get(key) != expected } if mismatches: raise RuntimeError( f"official generation config drifted: {mismatches}" ) tokenizer = AutoTokenizer.from_pretrained( args.artifact_dir, trust_remote_code=True, local_files_only=True, use_fast=True, ) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id boundary_ids = special.special_family_token_ids(tokenizer) if boundary_ids != { "eos": 100001, "bos": 100000, "x": 87, "period": 13, }: raise RuntimeError( f"unexpected pinned tokenizer IDs: {boundary_ids}" ) source_rows, reference, prompt_audit = prepare_sources( args, tokenizer, manifest, ) assignment = execution_assignment( mode=args.execution_mode, source_rows=source_rows, manifest=manifest, ) gold = behavior.load_gold(args) tape_index = { row["label"]: row for row in manifest["tape_contract"]["tapes"] } generated_sources = [] pending = [] for source, tapes in assignment: header = journal_header( args=args, source=source, tapes=tapes, max_new_tokens=max_new_tokens, ) path = ( journal_path(args.journal_dir, source) if args.journal_dir is not None else None ) if ( path is not None and path.is_file() and args.resume ): generated_sources.append(load_journal(path, header)) else: pending.append((source, tapes, header, path)) if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() load_started = time.perf_counter() _, official_modeling = special.base.load_official_modules( args.artifact_dir ) model = official_modeling.DeepseekV2ForCausalLM.from_pretrained( args.artifact_dir, local_files_only=True, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto", max_memory={ 0: args.gpu_memory, "cpu": args.cpu_memory, }, ) model.eval() if torch.cuda.is_available(): torch.cuda.synchronize() load_seconds = time.perf_counter() - load_started load_peak_cuda = ( torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None ) cuda_after_load = ( torch.cuda.memory_allocated() if torch.cuda.is_available() else None ) input_device = model.get_input_embeddings().weight.device for pending_index, ( source, tapes, header, path, ) in enumerate(pending, start=1): batch = prepare_batch(source, tokenizer, input_device) runs = [ generate_tape( source=source, tape=tape_index[tape], batch=batch, model=model, tokenizer=tokenizer, gold=gold, max_new_tokens=max_new_tokens, ) for tape in tapes ] replay = None if args.execution_mode == "smoke": replay_run = generate_tape( source=source, tape=tape_index["T0"], batch=batch, model=model, tokenizer=tokenizer, gold=gold, max_new_tokens=max_new_tokens, ) replay = compare_runs(runs[0], replay_run) if not replay["passed"]: raise RuntimeError( f"{source['id']}: smoke replay failed" ) generated = { key: value for key, value in source.items() if key not in {"content", "variants"} } | { "prompt_tokens_by_condition": { condition: batch["prompt_lengths"][index] for index, condition in enumerate(CONDITIONS) }, "batch_prompt_tokens_after_left_padding": batch[ "batch_prompt_length" ], "batch_padding_side": "left", "batch_conditions": len(CONDITIONS), "runs": runs, "in_process_replay": replay, } generated_sources.append(generated) if path is not None: write_journal(path, header, generated) print( json.dumps( { "progress": f"{pending_index}/{len(pending)}", "source_id": source["id"], "tapes": tapes, "outputs": len(runs) * len(CONDITIONS), "seconds": sum( run["generation_seconds"] for run in runs ), "journal": str(path) if path else None, }, ensure_ascii=False, ), flush=True, ) generated_sources.sort( key=lambda source: ( DOMAINS.index(source["domain"]), source["within_domain_index"], ) ) summary = summarize_sources(generated_sources) expected_outputs = { "smoke": 64, "formal": 352, "replay": 64, }[args.execution_mode] if summary["outputs"] != expected_outputs: raise RuntimeError( f"observed {summary['outputs']} outputs, " f"expected {expected_outputs}" ) if args.execution_mode == "smoke": first_by_source_condition = {} for source in generated_sources: for run in source["runs"]: for output in run["outputs"]: first_by_source_condition[ ( source["id"], run["tape_label"], output["condition"], ) ] = output["generated_token_ids"] divergent = sum( first_by_source_condition[ (source["id"], "T0", condition) ] != first_by_source_condition[ (source["id"], "T1", condition) ] for source in generated_sources for condition in CONDITIONS ) summary["smoke_tape_divergence"] = { "cells": len(generated_sources) * len(CONDITIONS), "different_trajectories": divergent, "passed": divergent > 0, } if not summary["smoke_tape_divergence"]["passed"]: raise RuntimeError("T0/T1 smoke divergence gate failed") if ( prompt_audit["exact"] != 256 or summary["torch_rng_unchanged_runs"] != summary["runs"] or not summary["in_process_replay"]["passed"] ): raise RuntimeError("smoke audit gate failed") device_map = getattr(model, "hf_device_map", {}) parameter_bytes: dict[str, int] = {} parameter_bytes_by_dtype: dict[str, int] = {} for parameter in model.parameters(): device = str(parameter.device) parameter_bytes[device] = ( parameter_bytes.get(device, 0) + parameter.numel() * parameter.element_size() ) dtype = str(parameter.dtype) parameter_bytes_by_dtype[dtype] = ( parameter_bytes_by_dtype.get(dtype, 0) + parameter.numel() * parameter.element_size() ) checkpoint_index = json.loads( ( args.artifact_dir / "model.safetensors.index.json" ).read_text(encoding="utf-8") ) checkpoint_tensor_bytes = int( checkpoint_index["metadata"]["total_size"] ) if sum(parameter_bytes.values()) != checkpoint_tensor_bytes: raise RuntimeError( "runtime parameter bytes do not match checkpoint index" ) generation_peaks = [ run["peak_cuda_memory_allocated_bytes"] for source in generated_sources for run in source["runs"] if run["peak_cuda_memory_allocated_bytes"] is not None ] peak_cuda = ( max([load_peak_cuda, *generation_peaks]) if load_peak_cuda is not None else None ) logits_process_path = Path( transformers.generation.logits_process.__file__ ) executed_tapes = sorted( { run["tape_label"] for source in generated_sources for run in source["runs"] } ) result = { "schema_version": 1, "protocol_id": PROTOCOL_ID, "execution_mode": args.execution_mode, "captured_at": captured_at, "model": { "repo": "deepseek-ai/DeepSeek-V2-Lite-Chat", "revision": behavior.MODEL_REVISION, "checkpoint_identity": "SFT Chat", "architecture": type(model).__name__, "dtype": "bfloat16", "checkpoint_tensor_bytes": checkpoint_tensor_bytes, "download_revision_contract": revision_contract, "files": { name: { "bytes": (args.artifact_dir / name).stat().st_size, "sha256": behavior.sha256( args.artifact_dir / name ), } for name in behavior.MODEL_FILES }, }, "tokenizer_contract": { "length": len(tokenizer), "vocab_size": tokenizer.vocab_size, "all_special_tokens": tokenizer.all_special_tokens, "all_special_ids": tokenizer.all_special_ids, "boundary_token_ids": boundary_ids, "pad_token_id": tokenizer.pad_token_id, "pad_aliases_eos": ( tokenizer.pad_token_id == tokenizer.eos_token_id ), }, "source_contract": { **reference, "manifest_path": str(args.manifest), "manifest_sha256": MANIFEST_SHA256, "domains": list(DOMAINS), "frozen_per_domain": 32, "executed_sources": len(generated_sources), "full_source_text_used": True, "prompt_hash_audit": prompt_audit, }, "seed_contract": { "protocol_id": PROTOCOL_ID, "tape_derivation": manifest["tape_contract"], "preregistered_base_seeds": [ tape["display_seed"] for tape in manifest["tape_contract"]["tapes"] ], "executed_base_seeds": [ tape_index[label]["display_seed"] for label in executed_tapes ], "executed_tapes": executed_tapes, "condition_row_order": list(CONDITIONS), "explicit_common_random_numbers": True, "same_source_tape_step_same_uniform_across_conditions": True, "torch_rng_used_for_sampling": False, "torch_rng_normalized_only_for_nonconsumption_audit": True, }, "generation_contract": { "conditions": special.FACTORS, "official_serialization": { "eos": True, "period": False, }, "decode": ( "explicit-uniform inverse-CDF nucleus sampler in " "token-ID order" ), "transformers_generate_called": False, "torch_multinomial_called": False, "distribution_temperature": TEMPERATURE, "distribution_top_p": TOP_P, "distribution_top_k": TOP_K, "max_new_tokens": max_new_tokens, "use_cache": True, "softmax_dtype": "torch.float32", "cdf_dtype": "torch.float32", "uniform_dtype": "torch.float32", "uniform_endpoint_clamp": ( "torch.nextafter(0,1) through torch.nextafter(1,0)" ), "cdf_final_value_forced_to_one": True, "temperature_warper": ( "transformers.TemperatureLogitsWarper(0.3)" ), "top_p_warper": ( "transformers.TopPLogitsWarper(" "0.95,min_tokens_to_keep=1)" ), "all_conditions_same_source_batch": True, "batch_conditions": len(CONDITIONS), "finished_rows_append_pad_eos_with_attention_mask_zero": True, "batch_padding": ( "left padding with PAD=EOS and attention_mask=0; " "all prompt endpoints occupy the final prompt column" ), "official_generation_config": official_generation, "official_sampling_parameters_define_distribution_only": True, "counterfactual_boundary": ( "Period cells edit one pre-target token ID after " "official rendering and are not valid official chats." ), "synthetic_sampler_test": synthetic_test, }, "execution": { "python": platform.python_version(), "torch": torch.__version__, "transformers": transformers.__version__, "accelerate": accelerate.__version__, "safetensors": safetensors.__version__, "platform": platform.platform(), "gpu": ( torch.cuda.get_device_name(0) if torch.cuda.is_available() else None ), "gpu_memory_limit": args.gpu_memory, "cpu_memory_limit": args.cpu_memory, "pytorch_cuda_alloc_conf": os.environ.get( "PYTORCH_CUDA_ALLOC_CONF" ), "input_device": str(input_device), "device_map": device_map, "parameter_bytes_by_runtime_parameter_device": ( parameter_bytes ), "parameter_bytes_by_dtype": parameter_bytes_by_dtype, "load_seconds": load_seconds, "load_peak_cuda_memory_allocated_bytes": load_peak_cuda, "cuda_memory_allocated_after_load_bytes": cuda_after_load, "peak_cuda_memory_allocated_bytes": peak_cuda, "process_max_rss_kib": resource.getrusage( resource.RUSAGE_SELF ).ru_maxrss, "transformers_logits_process_path": str( logits_process_path ), "transformers_logits_process_sha256": behavior.sha256( logits_process_path ), "official_single_gpu_bf16_requirement": "40GB GPU", "local_single_gpu_capacity_mib": 32607, "offload_required_by_local_capacity": True, "journal_dir": ( str(args.journal_dir) if args.journal_dir is not None else None ), "resumed_sources": len(generated_sources) - len(pending), }, "sources": generated_sources, "summary": summary, "claim_boundary": [ "HumanEval and GSM8K are analyzed separately.", ( "The main estimand is scoped to the frozen 32 selected " "tasks per domain under tape T0." ), ( "Selected-task bootstrap bands are not benchmark-" "population, model-ability, or seed-uncertainty intervals." ), ( "The explicit inverse-CDF sampler uses the official " "temperature/top-p distribution but is not the exact " "Transformers torch.multinomial trajectory." ), "Counterfactual period sequences are not official-valid chats.", "Passing HumanEval tests is functional, not safety, evidence.", "CPU-offloaded eager latency is not serving throughput.", ], } result["content_hash"] = canonical_hash( { "protocol_id": result["protocol_id"], "generation_contract": result["generation_contract"], "sources": result["sources"], } ) args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text( json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) payload = args.output.read_bytes() print( json.dumps( { "output": str(args.output), "sha256": sha256_bytes(payload), "bytes": len(payload), "sources": len(generated_sources), "outputs": summary["outputs"], "load_seconds": load_seconds, "generation_seconds": sum( run["generation_seconds"] for source in generated_sources for run in source["runs"] ), "peak_cuda_memory_allocated_bytes": peak_cuda, "summary": summary, }, ensure_ascii=False, indent=2, ) ) if __name__ == "__main__": main()