wuyang
1085 lines
36 KiB
Python
1085 lines
36 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the preregistered DeepSeek-V2-Lite-Chat sampling grid.
|
|
|
|
The probe keeps the audited eight-row prompt batch from the preceding greedy
|
|
completion experiment and changes only the decoding policy:
|
|
|
|
do_sample=True, temperature=0.3, top_p=0.95, top_k=0
|
|
|
|
Every source x replicate batch receives a SHA-256-derived run seed. The eight
|
|
rows are batch-seed aligned, not common-random-number pairs: Transformers
|
|
4.41.2 samples the whole batch with one ``torch.multinomial`` call and exposes
|
|
no row-specific generator in ``generate``.
|
|
|
|
This is a four-source mechanism probe. It is not a benchmark, a distribution
|
|
estimate, or a causal test of the counterfactual boundary tokens.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import itertools
|
|
import json
|
|
import os
|
|
import platform
|
|
import resource
|
|
import time
|
|
from collections import Counter
|
|
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
|
|
|
|
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-sampling-v1"
|
|
PREREGISTERED_BASE_SEEDS = (
|
|
19683830,
|
|
1560062173,
|
|
3978401375,
|
|
1280933274,
|
|
1467459869,
|
|
1297359489,
|
|
2722953988,
|
|
3330978061,
|
|
)
|
|
EXPECTED_CONDITIONS = (
|
|
"s0_eos",
|
|
"s1_eos",
|
|
"s0_bos",
|
|
"s1_bos",
|
|
"s0_x",
|
|
"s1_x",
|
|
"s0_period",
|
|
"s1_period",
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--artifact-dir", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--reference-routing-json",
|
|
type=Path,
|
|
required=True,
|
|
)
|
|
parser.add_argument("--greedy-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(
|
|
"--domains",
|
|
nargs="+",
|
|
choices=behavior.DEFAULT_DOMAINS,
|
|
default=list(behavior.DEFAULT_DOMAINS),
|
|
)
|
|
parser.add_argument("--per-domain", type=int, default=1)
|
|
parser.add_argument(
|
|
"--base-seeds",
|
|
nargs="+",
|
|
type=int,
|
|
default=list(PREREGISTERED_BASE_SEEDS),
|
|
)
|
|
parser.add_argument("--max-new-tokens", type=int, default=512)
|
|
parser.add_argument("--temperature", type=float, default=0.3)
|
|
parser.add_argument("--top-p", type=float, default=0.95)
|
|
parser.add_argument("--top-k", type=int, default=0)
|
|
parser.add_argument("--gpu-memory", default="28GiB")
|
|
parser.add_argument("--cpu-memory", default="80GiB")
|
|
parser.add_argument("--captured-at", default=None)
|
|
parser.add_argument(
|
|
"--in-process-replay-first-seed",
|
|
action="store_true",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def sha256_bytes(payload: bytes) -> str:
|
|
return hashlib.sha256(payload).hexdigest()
|
|
|
|
|
|
def derive_base_seed(index: int) -> int:
|
|
payload = f"{PROTOCOL_ID}/seed/{index}".encode()
|
|
return int.from_bytes(hashlib.sha256(payload).digest()[:4], "big")
|
|
|
|
|
|
def derive_run_seed(base_seed: int, source_id: str) -> int:
|
|
payload = (
|
|
f"{PROTOCOL_ID}/run\0{base_seed}\0{source_id}"
|
|
).encode()
|
|
raw = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
|
|
return raw % ((1 << 63) - 1)
|
|
|
|
|
|
def rng_snapshot() -> dict[str, Any]:
|
|
cpu_state = torch.get_rng_state().detach().cpu().numpy().tobytes()
|
|
cuda_states = (
|
|
[
|
|
state.detach().cpu().numpy().tobytes()
|
|
for state in torch.cuda.get_rng_state_all()
|
|
]
|
|
if torch.cuda.is_available()
|
|
else []
|
|
)
|
|
cuda_hashes = [sha256_bytes(state) for state in cuda_states]
|
|
return {
|
|
"cpu_sha256": sha256_bytes(cpu_state),
|
|
"cuda_device_sha256": cuda_hashes,
|
|
"cuda_combined_sha256": behavior.canonical_hash(cuda_hashes),
|
|
}
|
|
|
|
|
|
def set_run_seed(seed: int) -> dict[str, Any]:
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
return rng_snapshot()
|
|
|
|
|
|
def token_similarity(
|
|
left: dict[str, Any],
|
|
right: dict[str, Any],
|
|
) -> float:
|
|
distance = behavior.edit_distance(
|
|
left["generated_token_ids"],
|
|
right["generated_token_ids"],
|
|
)
|
|
denominator = max(
|
|
len(left["generated_token_ids"]),
|
|
len(right["generated_token_ids"]),
|
|
1,
|
|
)
|
|
return 1 - distance / denominator
|
|
|
|
|
|
def summarize_values(values: list[float | int]) -> dict[str, Any]:
|
|
return {
|
|
"count": len(values),
|
|
"mean": mean(values) if values else None,
|
|
"min": min(values) if values else None,
|
|
"max": max(values) if values else None,
|
|
}
|
|
|
|
|
|
def greedy_index(
|
|
greedy: dict[str, Any],
|
|
) -> dict[tuple[str, str], dict[str, Any]]:
|
|
return {
|
|
(source["id"], output["condition"]): output
|
|
for source in greedy["sources"]
|
|
for output in source["outputs"]
|
|
}
|
|
|
|
|
|
def prompt_contract(
|
|
source_rows: list[dict[str, Any]],
|
|
greedy_rows: dict[tuple[str, str], dict[str, Any]],
|
|
) -> dict[str, Any]:
|
|
rows = []
|
|
for source in source_rows:
|
|
for condition in EXPECTED_CONDITIONS:
|
|
observed = source["variants"][condition][
|
|
"token_ids_sha256"
|
|
]
|
|
reference = greedy_rows.get((source["id"], condition))
|
|
if reference is None:
|
|
raise RuntimeError(
|
|
"greedy baseline is missing "
|
|
f"{source['id']}/{condition}"
|
|
)
|
|
expected = reference["prompt_token_ids_sha256"]
|
|
exact = observed == expected
|
|
rows.append(
|
|
{
|
|
"source_id": source["id"],
|
|
"condition": condition,
|
|
"observed_sha256": observed,
|
|
"greedy_sha256": expected,
|
|
"exact": exact,
|
|
}
|
|
)
|
|
mismatches = [row for row in rows if not row["exact"]]
|
|
if mismatches:
|
|
raise RuntimeError(
|
|
f"{len(mismatches)} prompt hashes differ from greedy baseline"
|
|
)
|
|
return {
|
|
"cells": len(rows),
|
|
"exact": sum(row["exact"] for row in rows),
|
|
"rows": rows,
|
|
}
|
|
|
|
|
|
def prepare_source_variants(
|
|
args: argparse.Namespace,
|
|
tokenizer: Any,
|
|
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
|
source_rows, reference = behavior.selected_sources(args, tokenizer)
|
|
for source in source_rows:
|
|
source["variants"] = {
|
|
condition: special.prior.render_boundary_variant(
|
|
tokenizer,
|
|
source["content"],
|
|
condition,
|
|
)
|
|
for condition in EXPECTED_CONDITIONS
|
|
}
|
|
for system in (0, 1):
|
|
lengths = {
|
|
source["variants"][condition]["tokens"]
|
|
for condition in EXPECTED_CONDITIONS
|
|
if special.FACTORS[condition]["system"] == system
|
|
}
|
|
if len(lengths) != 1:
|
|
raise RuntimeError(
|
|
f"{source['id']} has unequal prompt lengths "
|
|
f"within system={system}: {sorted(lengths)}"
|
|
)
|
|
return source_rows, reference
|
|
|
|
|
|
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 EXPECTED_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 generate_run(
|
|
*,
|
|
args: argparse.Namespace,
|
|
source: dict[str, Any],
|
|
batch: dict[str, Any],
|
|
model: Any,
|
|
tokenizer: Any,
|
|
gold: dict[str, dict[str, Any]],
|
|
replicate_index: int,
|
|
base_seed: int,
|
|
) -> dict[str, Any]:
|
|
run_seed = derive_run_seed(base_seed, source["id"])
|
|
rng_before = set_run_seed(run_seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.reset_peak_memory_stats()
|
|
torch.cuda.synchronize()
|
|
started = time.perf_counter()
|
|
with torch.inference_mode():
|
|
generated = model.generate(
|
|
input_ids=batch["input_ids"],
|
|
attention_mask=batch["attention_mask"],
|
|
do_sample=True,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
max_new_tokens=args.max_new_tokens,
|
|
eos_token_id=tokenizer.eos_token_id,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
use_cache=True,
|
|
)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
elapsed = time.perf_counter() - started
|
|
rng_after = rng_snapshot()
|
|
peak_cuda = (
|
|
torch.cuda.max_memory_allocated()
|
|
if torch.cuda.is_available()
|
|
else None
|
|
)
|
|
|
|
outputs = []
|
|
for row_index, condition in enumerate(EXPECTED_CONDITIONS):
|
|
full_ids = generated[row_index].detach().cpu().tolist()
|
|
raw_generated = full_ids[batch["batch_prompt_length"] :]
|
|
generated_ids, hit_eos = behavior.truncate_at_eos(
|
|
raw_generated,
|
|
int(tokenizer.eos_token_id),
|
|
)
|
|
text = tokenizer.decode(
|
|
generated_ids,
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=False,
|
|
)
|
|
prompt_length = batch["prompt_lengths"][row_index]
|
|
outputs.append(
|
|
{
|
|
"condition": condition,
|
|
"factors": special.FACTORS[condition],
|
|
"prompt_tokens": prompt_length,
|
|
"left_padding_tokens": (
|
|
batch["batch_prompt_length"] - prompt_length
|
|
),
|
|
"prompt_token_ids_sha256": source["variants"][
|
|
condition
|
|
]["token_ids_sha256"],
|
|
"generated_tokens": len(generated_ids),
|
|
"generated_token_ids": generated_ids,
|
|
"generated_token_ids_sha256": (
|
|
behavior.canonical_hash(generated_ids)
|
|
),
|
|
"hit_eos": hit_eos,
|
|
"stopped_at_max_new_tokens": (
|
|
not hit_eos
|
|
and len(generated_ids) == args.max_new_tokens
|
|
),
|
|
"text": text,
|
|
"text_sha256": sha256_bytes(text.encode()),
|
|
"task_score": behavior.task_score(
|
|
source["domain"],
|
|
source["id"],
|
|
text,
|
|
gold,
|
|
),
|
|
}
|
|
)
|
|
return {
|
|
"replicate_index": replicate_index,
|
|
"replicate_label": f"R{replicate_index}",
|
|
"base_seed": base_seed,
|
|
"run_seed": run_seed,
|
|
"rng_state_before": rng_before,
|
|
"rng_state_after": rng_after,
|
|
"generation_seconds": elapsed,
|
|
"peak_cuda_memory_allocated_bytes": peak_cuda,
|
|
"generated_tokens_including_batch_padding": (
|
|
generated.shape[0]
|
|
* (generated.shape[1] - batch["batch_prompt_length"])
|
|
),
|
|
"outputs": outputs,
|
|
}
|
|
|
|
|
|
def compare_replay(
|
|
original: dict[str, Any],
|
|
replay: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
original_rows = {
|
|
row["condition"]: row for row in original["outputs"]
|
|
}
|
|
replay_rows = {
|
|
row["condition"]: row for row in replay["outputs"]
|
|
}
|
|
rows = []
|
|
for condition in EXPECTED_CONDITIONS:
|
|
left = original_rows[condition]
|
|
right = replay_rows[condition]
|
|
rows.append(
|
|
{
|
|
"condition": condition,
|
|
"run_seed_exact": (
|
|
original["run_seed"] == replay["run_seed"]
|
|
),
|
|
"rng_pre_state_exact": (
|
|
original["rng_state_before"]
|
|
== replay["rng_state_before"]
|
|
),
|
|
"prompt_hash_exact": (
|
|
left["prompt_token_ids_sha256"]
|
|
== right["prompt_token_ids_sha256"]
|
|
),
|
|
"generated_token_ids_exact": (
|
|
left["generated_token_ids"]
|
|
== right["generated_token_ids"]
|
|
),
|
|
"text_exact": left["text"] == right["text"],
|
|
"stop_state_exact": (
|
|
left["hit_eos"] == right["hit_eos"]
|
|
and left["stopped_at_max_new_tokens"]
|
|
== right["stopped_at_max_new_tokens"]
|
|
),
|
|
}
|
|
)
|
|
exact = sum(
|
|
all(value for key, value in row.items() if key != "condition")
|
|
for row in rows
|
|
)
|
|
return {
|
|
"original_replicate_label": original["replicate_label"],
|
|
"replay_replicate_label": replay["replicate_label"],
|
|
"cells": len(rows),
|
|
"all_contract_fields_exact": exact,
|
|
"rows": rows,
|
|
"replay_run": replay,
|
|
}
|
|
|
|
|
|
def source_condition_summary(
|
|
source: dict[str, Any],
|
|
greedy_rows: dict[tuple[str, str], dict[str, Any]],
|
|
) -> dict[str, Any]:
|
|
result = {}
|
|
for condition in EXPECTED_CONDITIONS:
|
|
rows = [
|
|
next(
|
|
output
|
|
for output in run["outputs"]
|
|
if output["condition"] == condition
|
|
)
|
|
for run in source["runs"]
|
|
]
|
|
pairwise = [
|
|
token_similarity(left, right)
|
|
for left, right in itertools.combinations(rows, 2)
|
|
]
|
|
hashes = {
|
|
row["generated_token_ids_sha256"] for row in rows
|
|
}
|
|
greedy = greedy_rows[(source["id"], condition)]
|
|
greedy_ids = greedy["generated_token_ids"]
|
|
result[condition] = {
|
|
"samples": 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(hashes),
|
|
"greedy_full_trajectory_in_sample_set": (
|
|
greedy["generated_token_ids_sha256"] in hashes
|
|
),
|
|
"sample_trajectory_is_greedy_prefix": sum(
|
|
row["generated_token_ids"]
|
|
== greedy_ids[: len(row["generated_token_ids"])]
|
|
for row in rows
|
|
),
|
|
"pairwise_token_similarity": summarize_values(pairwise),
|
|
"generated_tokens": summarize_values(
|
|
[row["generated_tokens"] for row in rows]
|
|
),
|
|
}
|
|
return result
|
|
|
|
|
|
def source_edge_summary(source: dict[str, Any]) -> dict[str, Any]:
|
|
by_condition = {
|
|
condition: [
|
|
next(
|
|
output
|
|
for output in run["outputs"]
|
|
if output["condition"] == condition
|
|
)
|
|
for run in source["runs"]
|
|
]
|
|
for condition in EXPECTED_CONDITIONS
|
|
}
|
|
result = {}
|
|
for name, left_condition, right_condition in special.COMPARISONS:
|
|
left_rows = by_condition[left_condition]
|
|
right_rows = by_condition[right_condition]
|
|
aligned = [
|
|
token_similarity(left, right)
|
|
for left, right in zip(left_rows, right_rows)
|
|
]
|
|
left_nearest = [
|
|
max(token_similarity(left, right) for right in right_rows)
|
|
for left in left_rows
|
|
]
|
|
right_nearest = [
|
|
max(token_similarity(right, left) for left in left_rows)
|
|
for right in right_rows
|
|
]
|
|
left_hashes = {
|
|
row["generated_token_ids_sha256"] for row in left_rows
|
|
}
|
|
right_hashes = {
|
|
row["generated_token_ids_sha256"] for row in right_rows
|
|
}
|
|
intersection = left_hashes & right_hashes
|
|
union = left_hashes | right_hashes
|
|
result[name] = {
|
|
"left": left_condition,
|
|
"right": right_condition,
|
|
"samples_per_side": len(left_rows),
|
|
"batch_seed_aligned_similarity": summarize_values(aligned),
|
|
"symmetric_mean_nearest_neighbor_similarity": mean(
|
|
[*left_nearest, *right_nearest]
|
|
),
|
|
"left_to_right_nearest_similarity": summarize_values(
|
|
left_nearest
|
|
),
|
|
"right_to_left_nearest_similarity": summarize_values(
|
|
right_nearest
|
|
),
|
|
"generated_hash_set_intersection": len(intersection),
|
|
"generated_hash_set_union": len(union),
|
|
"jaccard_on_exact_trajectory_hashes": (
|
|
len(intersection) / len(union) if union else 1.0
|
|
),
|
|
"natural_eos_count_difference_right_minus_left": (
|
|
sum(row["hit_eos"] for row in right_rows)
|
|
- sum(row["hit_eos"] for row in left_rows)
|
|
),
|
|
}
|
|
return result
|
|
|
|
|
|
def summarize(
|
|
sources: list[dict[str, Any]],
|
|
greedy_rows: dict[tuple[str, str], dict[str, Any]],
|
|
) -> dict[str, Any]:
|
|
output_rows = [
|
|
output
|
|
for source in sources
|
|
for run in source["runs"]
|
|
for output in run["outputs"]
|
|
]
|
|
by_source_condition = {
|
|
source["id"]: source_condition_summary(source, greedy_rows)
|
|
for source in sources
|
|
}
|
|
by_source_edge = {
|
|
source["id"]: source_edge_summary(source)
|
|
for source in sources
|
|
}
|
|
different_seed_cells = 0
|
|
comparable_seed_cells = 0
|
|
for source in sources:
|
|
if len(source["runs"]) < 2:
|
|
continue
|
|
first = {
|
|
row["condition"]: row
|
|
for row in source["runs"][0]["outputs"]
|
|
}
|
|
second = {
|
|
row["condition"]: row
|
|
for row in source["runs"][1]["outputs"]
|
|
}
|
|
for condition in EXPECTED_CONDITIONS:
|
|
comparable_seed_cells += 1
|
|
different_seed_cells += (
|
|
first[condition]["generated_token_ids"]
|
|
!= second[condition]["generated_token_ids"]
|
|
)
|
|
replay_rows = [
|
|
source["in_process_replay"]
|
|
for source in sources
|
|
if source.get("in_process_replay") is not None
|
|
]
|
|
return {
|
|
"sources": len(sources),
|
|
"replicates_per_source": (
|
|
len(sources[0]["runs"]) if sources else 0
|
|
),
|
|
"outputs": len(output_rows),
|
|
"natural_eos": sum(row["hit_eos"] for row in output_rows),
|
|
"budget_truncated": sum(
|
|
row["stopped_at_max_new_tokens"] for row in output_rows
|
|
),
|
|
"unique_generated_token_hashes": len(
|
|
{
|
|
row["generated_token_ids_sha256"]
|
|
for row in output_rows
|
|
}
|
|
),
|
|
"first_two_seeds": {
|
|
"comparable_cells": comparable_seed_cells,
|
|
"different_trajectories": different_seed_cells,
|
|
"sampling_divergence_gate_passed": (
|
|
different_seed_cells > 0
|
|
if comparable_seed_cells
|
|
else None
|
|
),
|
|
},
|
|
"in_process_replay": {
|
|
"sources": len(replay_rows),
|
|
"cells": sum(row["cells"] for row in replay_rows),
|
|
"all_contract_fields_exact": sum(
|
|
row["all_contract_fields_exact"]
|
|
for row in replay_rows
|
|
),
|
|
},
|
|
"by_source_condition": by_source_condition,
|
|
"by_source_edge": by_source_edge,
|
|
}
|
|
|
|
|
|
def validate_args(args: argparse.Namespace) -> None:
|
|
if tuple(special.CONDITIONS) != EXPECTED_CONDITIONS:
|
|
raise RuntimeError(
|
|
f"condition order drifted: {special.CONDITIONS}"
|
|
)
|
|
derived = tuple(
|
|
derive_base_seed(index)
|
|
for index in range(len(PREREGISTERED_BASE_SEEDS))
|
|
)
|
|
if derived != PREREGISTERED_BASE_SEEDS:
|
|
raise RuntimeError(
|
|
f"preregistered base-seed table is invalid: {derived}"
|
|
)
|
|
if args.per_domain <= 0:
|
|
raise ValueError("--per-domain must be positive")
|
|
if args.max_new_tokens <= 0:
|
|
raise ValueError("--max-new-tokens must be positive")
|
|
if not args.base_seeds or len(set(args.base_seeds)) != len(
|
|
args.base_seeds
|
|
):
|
|
raise ValueError("--base-seeds must be non-empty and unique")
|
|
if any(seed < 0 for seed in args.base_seeds):
|
|
raise ValueError("--base-seeds must be non-negative")
|
|
if args.temperature <= 0:
|
|
raise ValueError("--temperature must be positive")
|
|
if not 0 < args.top_p <= 1:
|
|
raise ValueError("--top-p must be in (0, 1]")
|
|
if args.top_k != 0:
|
|
raise ValueError("this protocol requires explicit --top-k 0")
|
|
for path in (
|
|
args.artifact_dir,
|
|
args.reference_routing_json,
|
|
args.greedy_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}"
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
validate_args(args)
|
|
revision_contract = behavior.download_revision_contract(
|
|
args.artifact_dir
|
|
)
|
|
special.install_control_contract()
|
|
captured_at = args.captured_at or datetime.now(
|
|
timezone.utc
|
|
).isoformat()
|
|
|
|
official_generation = json.loads(
|
|
(args.artifact_dir / "generation_config.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
expected_official = {
|
|
"do_sample": True,
|
|
"temperature": 0.3,
|
|
"top_p": 0.95,
|
|
"bos_token_id": 100000,
|
|
"eos_token_id": 100001,
|
|
}
|
|
official_mismatches = {
|
|
key: {
|
|
"expected": expected,
|
|
"observed": official_generation.get(key),
|
|
}
|
|
for key, expected in expected_official.items()
|
|
if official_generation.get(key) != expected
|
|
}
|
|
if official_mismatches:
|
|
raise RuntimeError(
|
|
f"official generation config drifted: {official_mismatches}"
|
|
)
|
|
if (
|
|
args.temperature != official_generation["temperature"]
|
|
or args.top_p != official_generation["top_p"]
|
|
):
|
|
raise RuntimeError(
|
|
"temperature/top-p differ from the pinned official config"
|
|
)
|
|
|
|
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 = prepare_source_variants(args, tokenizer)
|
|
greedy = json.loads(
|
|
args.greedy_json.read_text(encoding="utf-8")
|
|
)
|
|
if greedy["model"]["revision"] != behavior.MODEL_REVISION:
|
|
raise RuntimeError("greedy baseline model revision differs")
|
|
greedy_rows = greedy_index(greedy)
|
|
prompts = prompt_contract(source_rows, greedy_rows)
|
|
gold = behavior.load_gold(args)
|
|
|
|
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
|
|
|
|
generated_sources = []
|
|
for source in source_rows:
|
|
batch = prepare_batch(source, tokenizer, input_device)
|
|
runs = [
|
|
generate_run(
|
|
args=args,
|
|
source=source,
|
|
batch=batch,
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
gold=gold,
|
|
replicate_index=index,
|
|
base_seed=base_seed,
|
|
)
|
|
for index, base_seed in enumerate(args.base_seeds)
|
|
]
|
|
replay = None
|
|
if args.in_process_replay_first_seed:
|
|
replay_run = generate_run(
|
|
args=args,
|
|
source=source,
|
|
batch=batch,
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
gold=gold,
|
|
replicate_index=0,
|
|
base_seed=args.base_seeds[0],
|
|
)
|
|
replay_run["replicate_label"] = "R0_REPLAY"
|
|
replay = compare_replay(runs[0], replay_run)
|
|
generated_sources.append(
|
|
{
|
|
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(
|
|
EXPECTED_CONDITIONS
|
|
)
|
|
},
|
|
"batch_prompt_tokens_after_left_padding": batch[
|
|
"batch_prompt_length"
|
|
],
|
|
"batch_padding_side": "left",
|
|
"batch_conditions": len(EXPECTED_CONDITIONS),
|
|
"runs": runs,
|
|
"in_process_replay": replay,
|
|
}
|
|
)
|
|
|
|
device_map = getattr(model, "hf_device_map", {})
|
|
parameter_bytes = {}
|
|
parameter_bytes_by_dtype = {}
|
|
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
|
|
)
|
|
generation_utils_path = Path(
|
|
transformers.generation.utils.__file__
|
|
)
|
|
|
|
result = {
|
|
"schema_version": 1,
|
|
"protocol_id": PROTOCOL_ID,
|
|
"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,
|
|
"domains": args.domains,
|
|
"per_domain": args.per_domain,
|
|
"sources": len(source_rows),
|
|
"full_source_text_used": True,
|
|
"greedy_baseline_path": str(args.greedy_json),
|
|
"greedy_baseline_sha256": behavior.sha256(
|
|
args.greedy_json
|
|
),
|
|
"prompt_hash_audit": prompts,
|
|
},
|
|
"seed_contract": {
|
|
"protocol_id": PROTOCOL_ID,
|
|
"derivation": (
|
|
"base seed = first four SHA-256 bytes of "
|
|
"protocol/seed/index; run seed = first eight bytes of "
|
|
"SHA-256(protocol/run\\0base\\0source) modulo 2^63-1"
|
|
),
|
|
"preregistered_base_seeds": list(
|
|
PREREGISTERED_BASE_SEEDS
|
|
),
|
|
"executed_base_seeds": args.base_seeds,
|
|
"condition_row_order": list(EXPECTED_CONDITIONS),
|
|
"torch_manual_seed_before_each_source_replicate": True,
|
|
"torch_cuda_manual_seed_all_before_each_source_replicate": (
|
|
torch.cuda.is_available()
|
|
),
|
|
"batch_seed_aligned_not_common_random_numbers": True,
|
|
"row_specific_generator_available": False,
|
|
},
|
|
"generation_contract": {
|
|
"conditions": special.FACTORS,
|
|
"official_serialization": {
|
|
level: level == "eos"
|
|
for level in special.BOUNDARY_LEVELS
|
|
},
|
|
"decode": "nucleus_sampling",
|
|
"do_sample": True,
|
|
"temperature": args.temperature,
|
|
"top_p": args.top_p,
|
|
"top_k": args.top_k,
|
|
"max_new_tokens": args.max_new_tokens,
|
|
"use_cache": True,
|
|
"official_generation_config": official_generation,
|
|
"official_parameters_explicitly_passed_to_generate": True,
|
|
"all_eight_conditions_same_source_batch": True,
|
|
"batch_padding": (
|
|
"left padding with PAD=EOS and attention_mask=0; "
|
|
"generation prefix ends at the same batch column"
|
|
),
|
|
"counterfactual_boundary": (
|
|
"BOS/x/period cells edit one pre-target token ID after "
|
|
"official rendering and are not valid official chats"
|
|
),
|
|
},
|
|
"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_generation_utils_path": str(
|
|
generation_utils_path
|
|
),
|
|
"transformers_generation_utils_sha256": behavior.sha256(
|
|
generation_utils_path
|
|
),
|
|
"official_single_gpu_bf16_requirement": "40GB GPU",
|
|
"local_single_gpu_capacity_mib": 32607,
|
|
"offload_required_by_local_capacity": True,
|
|
},
|
|
"sources": generated_sources,
|
|
"summary": summarize(generated_sources, greedy_rows),
|
|
"claim_boundary": [
|
|
"This is a four-source, eight-seed mechanism probe, not a benchmark.",
|
|
"Eight trajectories do not estimate the full generation distribution.",
|
|
"Batch-seed-aligned rows are not common-random-number pairs.",
|
|
"Counterfactual BOS/x/period token sequences are not official-valid chats.",
|
|
"Unique token sequences are not semantic-diversity measurements.",
|
|
"One GSM8K and one HumanEval source do not estimate task ability.",
|
|
"CPU-offloaded eager latency is not serving throughput.",
|
|
"Sampling differences do not identify a hidden-state or router mediator.",
|
|
],
|
|
}
|
|
result["content_hash"] = behavior.canonical_hash(
|
|
{
|
|
"protocol_id": result["protocol_id"],
|
|
"seed_contract": result["seed_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": result["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": {
|
|
"natural_eos": result["summary"][
|
|
"natural_eos"
|
|
],
|
|
"budget_truncated": result["summary"][
|
|
"budget_truncated"
|
|
],
|
|
"unique_generated_token_hashes": result[
|
|
"summary"
|
|
]["unique_generated_token_hashes"],
|
|
"first_two_seeds": result["summary"][
|
|
"first_two_seeds"
|
|
],
|
|
"in_process_replay": result["summary"][
|
|
"in_process_replay"
|
|
],
|
|
},
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|