feat: trace DeepSeek Chat completion depth
This commit is contained in:
wuyang
@@ -586,3 +586,89 @@ See `research/DEEPSEEK_V2_LITE_CHAT_BEHAVIOR_PROTOCOL.md` and
|
||||
`research/DEEPSEEK_V2_LITE_CHAT_BEHAVIOR_AUDIT.md` for the preregistered
|
||||
contract, output-divergence table, offload device map, truncation boundary,
|
||||
rerun coverage, execution fixes, primary sources, and forbidden conclusions.
|
||||
|
||||
## Completion-aware 512-token evaluation
|
||||
|
||||
The completion run uses the same 16 sources and eight-condition batches but
|
||||
reruns all 128 cells with one uniform 512-token budget. The successful formal
|
||||
placement leaves embeddings and layers 0–23 on CUDA and offloads layers 24–26,
|
||||
final norm, and LM head. `expandable_segments:True` is part of the recorded
|
||||
runtime contract.
|
||||
|
||||
```bash
|
||||
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
|
||||
PYTHONPATH=/path/to/transformers-4.41.2-deps \
|
||||
python -B experiments/deepseek/v2_lite_chat_special_token_behavior_probe.py \
|
||||
--artifact-dir /path/to/deepseek-v2-lite-chat \
|
||||
--reference-routing-json \
|
||||
src/data/deepseek-v2-lite-routing-special-token-family-control.json \
|
||||
--human-eval /path/to/HumanEval.jsonl.gz \
|
||||
--gsm8k /path/to/gsm8k/test.jsonl \
|
||||
--tnews /path/to/tnews/test.json \
|
||||
--tnews-archive /path/to/tnews_public.zip \
|
||||
--wikitext /path/to/wikitext-validation.parquet \
|
||||
--output src/data/deepseek-v2-lite-chat-completion-512.json \
|
||||
--per-domain 4 \
|
||||
--max-new-tokens 512 \
|
||||
--gpu-memory 28GiB \
|
||||
--cpu-memory 80GiB
|
||||
```
|
||||
|
||||
The evaluator separates stopping, semantic terminal state, evaluator coverage,
|
||||
and correctness. Each HumanEval candidate runs in a fresh pinned, networkless,
|
||||
read-only Docker container with no host mounts:
|
||||
|
||||
```bash
|
||||
python -B experiments/deepseek/v2_lite_chat_completion_evaluator.py \
|
||||
--behavior-json \
|
||||
src/data/deepseek-v2-lite-chat-completion-512.json \
|
||||
--baseline-json src/data/deepseek-v2-lite-chat-behavior.json \
|
||||
--human-eval /path/to/HumanEval.jsonl.gz \
|
||||
--gsm8k /path/to/gsm8k/test.jsonl \
|
||||
--sandbox-image \
|
||||
python:3.11-alpine@sha256:25976e9d34a0fab1f278cae931f34c8303d97bf0c0d7f85b6b4dcf641d7702a4 \
|
||||
--output \
|
||||
src/data/deepseek-v2-lite-chat-completion-512-eval.json
|
||||
```
|
||||
|
||||
The formal run reaches natural EOS in 121/128 cells, up from 31/128 at the
|
||||
128-token budget. Strict-complete GSM8K numeric exact is 23/32; HumanEval
|
||||
official-test pass is 24/32. These are four source tasks per domain, not
|
||||
benchmark estimates. A fresh-process one-source-per-domain rerun reproduces
|
||||
all 32 complete generated token sequences.
|
||||
|
||||
## Full 27-layer Chat hidden-state and router trace
|
||||
|
||||
`v2_lite_chat_full_depth_trace.py` performs a prompt-only `use_cache=False`
|
||||
forward on the same checkpoint and inputs. It captures hashes and statistics
|
||||
for embedding, every decoder-layer output, final norm, and every MoE gate.
|
||||
Boundary-crossing tokenizer tokens are excluded from the exact target-content
|
||||
scope.
|
||||
|
||||
```bash
|
||||
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
|
||||
PYTHONPATH=/path/to/transformers-4.41.2-deps \
|
||||
python -B experiments/deepseek/v2_lite_chat_full_depth_trace.py \
|
||||
--artifact-dir /path/to/deepseek-v2-lite-chat \
|
||||
--reference-routing-json \
|
||||
src/data/deepseek-v2-lite-routing-special-token-family-control.json \
|
||||
--human-eval /path/to/HumanEval.jsonl.gz \
|
||||
--gsm8k /path/to/gsm8k/test.jsonl \
|
||||
--tnews /path/to/tnews/test.json \
|
||||
--tnews-archive /path/to/tnews_public.zip \
|
||||
--wikitext /path/to/wikitext-validation.parquet \
|
||||
--output src/data/deepseek-v2-lite-chat-full-depth.json \
|
||||
--per-domain 4 \
|
||||
--gpu-memory 28GiB \
|
||||
--cpu-memory 80GiB
|
||||
```
|
||||
|
||||
The formal trace covers 1,537 exact interior content tokens per condition,
|
||||
29 hidden stages, 26 gates, and 1,918,176 top-6 route decisions. A fresh
|
||||
four-source rerun reproduces 1,856 hidden tensor hashes, 1,664 ordered route
|
||||
hashes, 1,664 route-weight hashes, and every derived comparison.
|
||||
|
||||
See `research/DEEPSEEK_V2_LITE_CHAT_COMPLETION_PROTOCOL.md` and
|
||||
`research/DEEPSEEK_V2_LITE_CHAT_COMPLETION_DEPTH_AUDIT.md` for the OOM
|
||||
amendment, task evaluators, completion table, depth curves, reproduction
|
||||
audit, artifact hashes, and claim boundaries.
|
||||
|
||||
@@ -0,0 +1,701 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Evaluate completion state, GSM8K answers, and HumanEval code safely.
|
||||
|
||||
The input is a raw JSON emitted by
|
||||
v2_lite_chat_special_token_behavior_probe.py. Generated code is executed only
|
||||
inside one fresh, networkless, read-only Docker container per candidate. The
|
||||
result preserves fixed-budget and strict-completion metrics separately.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import gzip
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import subprocess
|
||||
import time
|
||||
from collections import Counter
|
||||
from decimal import Decimal, InvalidOperation
|
||||
from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import Any
|
||||
|
||||
|
||||
NUMBER = r"[-+]?(?:\d[\d,]*\.?\d*|\.\d+)"
|
||||
NUMBER_PATTERN = re.compile(NUMBER)
|
||||
BOXED_PATTERN = re.compile(
|
||||
rf"\\boxed\s*\{{\s*({NUMBER})\s*\}}",
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
HASH_PATTERN = re.compile(
|
||||
rf"####\s*({NUMBER})",
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
ANSWER_PATTERN = re.compile(
|
||||
rf"(?:final\s+answer|answer|result|total|profit)"
|
||||
rf"(?:\s+(?:is|equals|will\s+be))?\s*[:=]?\s*"
|
||||
rf"(?:\$|USD\s*)?({NUMBER})",
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
FENCE_PATTERN = re.compile(
|
||||
r"```(?P<lang>[A-Za-z0-9_+-]*)[ \t]*\n?"
|
||||
r"(?P<body>.*?)(?:```|$)",
|
||||
flags=re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--behavior-json", type=Path, required=True)
|
||||
parser.add_argument("--baseline-json", type=Path)
|
||||
parser.add_argument("--human-eval", type=Path, required=True)
|
||||
parser.add_argument("--gsm8k", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--sandbox-image",
|
||||
required=True,
|
||||
help="Pinned image reference including @sha256 digest.",
|
||||
)
|
||||
parser.add_argument("--timeout-seconds", type=float, default=5.0)
|
||||
parser.add_argument("--skip-code-execution", action="store_true")
|
||||
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 load_jsonl(path: Path) -> list[dict[str, Any]]:
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
return [json.loads(line) for line in handle if line.strip()]
|
||||
|
||||
|
||||
def load_tasks(
|
||||
human_eval_path: Path,
|
||||
gsm8k_path: Path,
|
||||
) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]]]:
|
||||
with gzip.open(human_eval_path, "rt", encoding="utf-8") as handle:
|
||||
human_eval = {
|
||||
row["task_id"]: row
|
||||
for line in handle
|
||||
if line.strip()
|
||||
for row in [json.loads(line)]
|
||||
}
|
||||
gsm8k = {
|
||||
f"gsm8k/test/{index:04d}": row
|
||||
for index, row in enumerate(load_jsonl(gsm8k_path))
|
||||
}
|
||||
return human_eval, gsm8k
|
||||
|
||||
|
||||
def normalize_number(value: str | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
cleaned = value.replace(",", "").strip().rstrip(".")
|
||||
try:
|
||||
number = Decimal(cleaned)
|
||||
except InvalidOperation:
|
||||
return None
|
||||
normalized = format(number.normalize(), "f")
|
||||
if "." in normalized:
|
||||
normalized = normalized.rstrip("0").rstrip(".")
|
||||
return "0" if normalized in {"-0", "+0"} else normalized
|
||||
|
||||
|
||||
def last_match(
|
||||
pattern: re.Pattern[str],
|
||||
text: str,
|
||||
) -> str | None:
|
||||
matches = list(pattern.finditer(text))
|
||||
return matches[-1].group(1) if matches else None
|
||||
|
||||
|
||||
def evaluate_math(
|
||||
text: str,
|
||||
hit_eos: bool,
|
||||
gold_answer: str,
|
||||
) -> dict[str, Any]:
|
||||
candidates = [
|
||||
("boxed", last_match(BOXED_PATTERN, text)),
|
||||
("hash_marker", last_match(HASH_PATTERN, text)),
|
||||
("answer_phrase", last_match(ANSWER_PATTERN, text)),
|
||||
]
|
||||
method = "none"
|
||||
extracted = None
|
||||
for name, value in candidates:
|
||||
if value is not None:
|
||||
method = name
|
||||
extracted = value
|
||||
break
|
||||
explicit_terminal = extracted is not None
|
||||
if extracted is None:
|
||||
values = NUMBER_PATTERN.findall(text)
|
||||
if values:
|
||||
method = "last_number_fallback"
|
||||
extracted = values[-1]
|
||||
|
||||
gold_values = NUMBER_PATTERN.findall(
|
||||
gold_answer.rsplit("####", 1)[-1]
|
||||
)
|
||||
gold = normalize_number(gold_values[-1] if gold_values else None)
|
||||
predicted = normalize_number(extracted)
|
||||
fixed_budget_exact = (
|
||||
gold is not None
|
||||
and predicted is not None
|
||||
and gold == predicted
|
||||
)
|
||||
semantic_terminal = bool(hit_eos or explicit_terminal)
|
||||
return {
|
||||
"gold_final": gold,
|
||||
"predicted_final": predicted,
|
||||
"extraction_method": method,
|
||||
"explicit_final_marker": explicit_terminal,
|
||||
"semantic_terminal": semantic_terminal,
|
||||
"evaluator_covered": predicted is not None,
|
||||
"fixed_budget_numeric_exact": fixed_budget_exact,
|
||||
"strict_complete_numeric_exact": (
|
||||
fixed_budget_exact and semantic_terminal
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def code_fences(text: str) -> list[dict[str, str | bool]]:
|
||||
rows = []
|
||||
for match in FENCE_PATTERN.finditer(text):
|
||||
language = match.group("lang").lower()
|
||||
body = match.group("body").strip()
|
||||
rows.append(
|
||||
{
|
||||
"language": language,
|
||||
"body": body,
|
||||
"closed": text[match.start():match.end()].rstrip().endswith(
|
||||
"```"
|
||||
),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def extract_code(
|
||||
text: str,
|
||||
prompt: str,
|
||||
entry_point: str,
|
||||
) -> dict[str, Any]:
|
||||
target = re.compile(
|
||||
rf"(?m)^\s*(?:async\s+)?def\s+{re.escape(entry_point)}\s*\("
|
||||
)
|
||||
fences = code_fences(text)
|
||||
chosen = next(
|
||||
(
|
||||
row for row in fences
|
||||
if target.search(str(row["body"]))
|
||||
and row["language"] in {"", "py", "python", "python3"}
|
||||
),
|
||||
None,
|
||||
)
|
||||
if chosen is None:
|
||||
chosen = next(
|
||||
(
|
||||
row for row in fences
|
||||
if row["language"] in {"", "py", "python", "python3"}
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
if chosen is not None:
|
||||
raw = str(chosen["body"])
|
||||
mode = "fenced"
|
||||
fence_closed = bool(chosen["closed"])
|
||||
else:
|
||||
match = target.search(text)
|
||||
if match:
|
||||
raw = text[match.start():].strip()
|
||||
mode = "direct_definition"
|
||||
else:
|
||||
raw = text.strip()
|
||||
mode = "prompt_completion"
|
||||
fence_closed = False
|
||||
|
||||
contains_entry_point = target.search(raw) is not None
|
||||
candidate = raw if contains_entry_point else prompt + raw
|
||||
try:
|
||||
ast.parse(candidate)
|
||||
ast_ok = True
|
||||
ast_error = None
|
||||
except SyntaxError as error:
|
||||
ast_ok = False
|
||||
ast_error = {
|
||||
"line": error.lineno,
|
||||
"offset": error.offset,
|
||||
"type": type(error).__name__,
|
||||
}
|
||||
return {
|
||||
"candidate": candidate,
|
||||
"candidate_sha256": sha256_bytes(candidate.encode()),
|
||||
"extraction_mode": mode,
|
||||
"contains_entry_point_definition": contains_entry_point,
|
||||
"closed_code_fence": fence_closed,
|
||||
"python_ast_parse": ast_ok,
|
||||
"parse_error": ast_error,
|
||||
}
|
||||
|
||||
|
||||
def sandbox_harness(
|
||||
candidate: str,
|
||||
test: str,
|
||||
entry_point: str,
|
||||
) -> str:
|
||||
return (
|
||||
"import sys\n"
|
||||
"try:\n"
|
||||
+ "\n".join(
|
||||
f" {line}" if line else ""
|
||||
for line in candidate.splitlines()
|
||||
)
|
||||
+ "\n"
|
||||
+ "\n".join(
|
||||
f" {line}" if line else ""
|
||||
for line in test.splitlines()
|
||||
)
|
||||
+ f"\n check({entry_point})\n"
|
||||
+ "except AssertionError:\n"
|
||||
+ " raise SystemExit(10)\n"
|
||||
+ "except BaseException:\n"
|
||||
+ " raise SystemExit(11)\n"
|
||||
+ "raise SystemExit(0)\n"
|
||||
)
|
||||
|
||||
|
||||
def execute_code(
|
||||
candidate: str,
|
||||
task: dict[str, Any],
|
||||
image: str,
|
||||
timeout_seconds: float,
|
||||
) -> dict[str, Any]:
|
||||
harness = sandbox_harness(
|
||||
candidate,
|
||||
task["test"],
|
||||
task["entry_point"],
|
||||
)
|
||||
command = [
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"-i",
|
||||
"--network",
|
||||
"none",
|
||||
"--read-only",
|
||||
"--tmpfs",
|
||||
"/tmp:rw,noexec,nosuid,size=16m",
|
||||
"--memory",
|
||||
"256m",
|
||||
"--memory-swap",
|
||||
"256m",
|
||||
"--pids-limit",
|
||||
"64",
|
||||
"--cpus",
|
||||
"0.5",
|
||||
"--cap-drop",
|
||||
"ALL",
|
||||
"--security-opt",
|
||||
"no-new-privileges",
|
||||
"--user",
|
||||
"65534:65534",
|
||||
image,
|
||||
"python",
|
||||
"-I",
|
||||
"-",
|
||||
]
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
command,
|
||||
input=harness.encode(),
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=timeout_seconds,
|
||||
check=False,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
return {
|
||||
"status": "timeout",
|
||||
"return_code": None,
|
||||
"runtime_ms": (time.perf_counter() - started) * 1000,
|
||||
"harness_sha256": sha256_bytes(harness.encode()),
|
||||
}
|
||||
runtime_ms = (time.perf_counter() - started) * 1000
|
||||
statuses = {
|
||||
0: "passed",
|
||||
10: "assertion_failed",
|
||||
11: "runtime_error",
|
||||
}
|
||||
return {
|
||||
"status": statuses.get(
|
||||
completed.returncode,
|
||||
"sandbox_error",
|
||||
),
|
||||
"return_code": completed.returncode,
|
||||
"runtime_ms": runtime_ms,
|
||||
"harness_sha256": sha256_bytes(harness.encode()),
|
||||
}
|
||||
|
||||
|
||||
def evaluate_code(
|
||||
text: str,
|
||||
hit_eos: bool,
|
||||
task: dict[str, Any],
|
||||
image: str,
|
||||
timeout_seconds: float,
|
||||
skip_execution: bool,
|
||||
) -> dict[str, Any]:
|
||||
extracted = extract_code(
|
||||
text,
|
||||
task["prompt"],
|
||||
task["entry_point"],
|
||||
)
|
||||
execution = {
|
||||
"status": "not_run",
|
||||
"return_code": None,
|
||||
"runtime_ms": None,
|
||||
"harness_sha256": None,
|
||||
}
|
||||
if extracted["python_ast_parse"] and not skip_execution:
|
||||
execution = execute_code(
|
||||
extracted["candidate"],
|
||||
task,
|
||||
image,
|
||||
timeout_seconds,
|
||||
)
|
||||
passed = execution["status"] == "passed"
|
||||
semantic_terminal = bool(
|
||||
hit_eos
|
||||
or extracted["closed_code_fence"]
|
||||
or passed
|
||||
)
|
||||
return {
|
||||
key: value
|
||||
for key, value in extracted.items()
|
||||
if key != "candidate"
|
||||
} | {
|
||||
"semantic_terminal": semantic_terminal,
|
||||
"evaluator_covered": (
|
||||
extracted["python_ast_parse"]
|
||||
and execution["status"] != "not_run"
|
||||
),
|
||||
"execution": execution,
|
||||
"fixed_budget_tests_pass": passed,
|
||||
"strict_complete_tests_pass": passed and semantic_terminal,
|
||||
}
|
||||
|
||||
|
||||
def completion_class(
|
||||
output: dict[str, Any],
|
||||
task_evaluation: dict[str, Any] | None,
|
||||
) -> str:
|
||||
if output["hit_eos"]:
|
||||
return "NATURAL_EOS"
|
||||
if (
|
||||
task_evaluation is not None
|
||||
and task_evaluation.get("semantic_terminal")
|
||||
):
|
||||
return "TASK_TERMINAL_BEFORE_EOS"
|
||||
if output["stopped_at_max_new_tokens"]:
|
||||
if (
|
||||
task_evaluation is not None
|
||||
and task_evaluation.get("extraction_method")
|
||||
== "last_number_fallback"
|
||||
):
|
||||
return "BUDGET_TRUNCATED_WITH_FALLBACK_ONLY"
|
||||
return "BUDGET_TRUNCATED_UNRESOLVED"
|
||||
return "OTHER_STOP"
|
||||
|
||||
|
||||
def baseline_index(
|
||||
baseline: dict[str, Any] | None,
|
||||
) -> dict[tuple[str, str], dict[str, Any]]:
|
||||
if baseline is None:
|
||||
return {}
|
||||
return {
|
||||
(source["id"], output["condition"]): output
|
||||
for source in baseline["sources"]
|
||||
for output in source["outputs"]
|
||||
}
|
||||
|
||||
|
||||
def prefix_audit(
|
||||
source_id: str,
|
||||
output: dict[str, Any],
|
||||
baseline_rows: dict[tuple[str, str], dict[str, Any]],
|
||||
) -> dict[str, Any] | None:
|
||||
prior = baseline_rows.get((source_id, output["condition"]))
|
||||
if prior is None:
|
||||
return None
|
||||
reference = prior["generated_token_ids"]
|
||||
observed = output["generated_token_ids"][: len(reference)]
|
||||
return {
|
||||
"baseline_generated_tokens": len(reference),
|
||||
"compared_tokens": min(
|
||||
len(reference),
|
||||
len(output["generated_token_ids"]),
|
||||
),
|
||||
"prompt_hash_exact": (
|
||||
prior["prompt_token_ids_sha256"]
|
||||
== output["prompt_token_ids_sha256"]
|
||||
),
|
||||
"generated_prefix_exact": observed == reference,
|
||||
"baseline_hit_eos": prior["hit_eos"],
|
||||
}
|
||||
|
||||
|
||||
def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_condition: dict[str, dict[str, Any]] = {}
|
||||
conditions = sorted({row["condition"] for row in rows})
|
||||
for condition in conditions:
|
||||
subset = [
|
||||
row for row in rows
|
||||
if row["condition"] == condition
|
||||
]
|
||||
math_rows = [
|
||||
row for row in subset
|
||||
if row["domain"] == "math"
|
||||
]
|
||||
code_rows = [
|
||||
row for row in subset
|
||||
if row["domain"] == "code"
|
||||
]
|
||||
by_condition[condition] = {
|
||||
"outputs": len(subset),
|
||||
"natural_eos": sum(row["hit_eos"] for row in subset),
|
||||
"budget_truncated": sum(
|
||||
row["stopped_at_max_new_tokens"]
|
||||
for row in subset
|
||||
),
|
||||
"mean_generated_tokens": mean(
|
||||
row["generated_tokens"] for row in subset
|
||||
),
|
||||
"completion_classes": dict(
|
||||
Counter(row["completion_class"] for row in subset)
|
||||
),
|
||||
"math_fixed_budget_exact": sum(
|
||||
row["task_evaluation"][
|
||||
"fixed_budget_numeric_exact"
|
||||
]
|
||||
for row in math_rows
|
||||
),
|
||||
"math_strict_complete_exact": sum(
|
||||
row["task_evaluation"][
|
||||
"strict_complete_numeric_exact"
|
||||
]
|
||||
for row in math_rows
|
||||
),
|
||||
"math_sources": len(math_rows),
|
||||
"code_ast_parse": sum(
|
||||
row["task_evaluation"]["python_ast_parse"]
|
||||
for row in code_rows
|
||||
),
|
||||
"code_executed": sum(
|
||||
row["task_evaluation"]["execution"]["status"]
|
||||
!= "not_run"
|
||||
for row in code_rows
|
||||
),
|
||||
"code_tests_pass": sum(
|
||||
row["task_evaluation"]["fixed_budget_tests_pass"]
|
||||
for row in code_rows
|
||||
),
|
||||
"code_sources": len(code_rows),
|
||||
}
|
||||
prefix_rows = [
|
||||
row["baseline_prefix_audit"]
|
||||
for row in rows
|
||||
if row["baseline_prefix_audit"] is not None
|
||||
]
|
||||
return {
|
||||
"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
|
||||
),
|
||||
"completion_classes": dict(
|
||||
Counter(row["completion_class"] for row in rows)
|
||||
),
|
||||
"baseline_prefix": {
|
||||
"cells": len(prefix_rows),
|
||||
"prompt_hash_exact": sum(
|
||||
row["prompt_hash_exact"] for row in prefix_rows
|
||||
),
|
||||
"generated_prefix_exact": sum(
|
||||
row["generated_prefix_exact"] for row in prefix_rows
|
||||
),
|
||||
},
|
||||
"by_condition": by_condition,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
for path in (
|
||||
args.behavior_json,
|
||||
args.human_eval,
|
||||
args.gsm8k,
|
||||
):
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(path)
|
||||
if args.baseline_json is not None and not args.baseline_json.is_file():
|
||||
raise FileNotFoundError(args.baseline_json)
|
||||
if "@sha256:" not in args.sandbox_image:
|
||||
raise ValueError("--sandbox-image must include an immutable digest")
|
||||
if args.timeout_seconds <= 0:
|
||||
raise ValueError("--timeout-seconds must be positive")
|
||||
|
||||
behavior = json.loads(
|
||||
args.behavior_json.read_text(encoding="utf-8")
|
||||
)
|
||||
baseline = (
|
||||
json.loads(args.baseline_json.read_text(encoding="utf-8"))
|
||||
if args.baseline_json is not None
|
||||
else None
|
||||
)
|
||||
baseline_rows = baseline_index(baseline)
|
||||
human_eval, gsm8k = load_tasks(args.human_eval, args.gsm8k)
|
||||
|
||||
rows = []
|
||||
for source in behavior["sources"]:
|
||||
for output in source["outputs"]:
|
||||
evaluation = None
|
||||
if source["domain"] == "math":
|
||||
evaluation = evaluate_math(
|
||||
output["text"],
|
||||
output["hit_eos"],
|
||||
gsm8k[source["id"]]["answer"],
|
||||
)
|
||||
elif source["domain"] == "code":
|
||||
evaluation = evaluate_code(
|
||||
output["text"],
|
||||
output["hit_eos"],
|
||||
human_eval[source["id"]],
|
||||
args.sandbox_image,
|
||||
args.timeout_seconds,
|
||||
args.skip_code_execution,
|
||||
)
|
||||
rows.append(
|
||||
{
|
||||
"source_id": source["id"],
|
||||
"domain": source["domain"],
|
||||
"condition": output["condition"],
|
||||
"generated_tokens": output["generated_tokens"],
|
||||
"hit_eos": output["hit_eos"],
|
||||
"stopped_at_max_new_tokens": (
|
||||
output["stopped_at_max_new_tokens"]
|
||||
),
|
||||
"prompt_token_ids_sha256": (
|
||||
output["prompt_token_ids_sha256"]
|
||||
),
|
||||
"generated_token_ids_sha256": (
|
||||
output["generated_token_ids_sha256"]
|
||||
),
|
||||
"text_sha256": output["text_sha256"],
|
||||
"task_evaluation": evaluation,
|
||||
"completion_class": completion_class(
|
||||
output,
|
||||
evaluation,
|
||||
),
|
||||
"baseline_prefix_audit": prefix_audit(
|
||||
source["id"],
|
||||
output,
|
||||
baseline_rows,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
result = {
|
||||
"schema_version": 1,
|
||||
"input": {
|
||||
"behavior_path": str(args.behavior_json),
|
||||
"behavior_sha256": sha256_file(args.behavior_json),
|
||||
"baseline_path": (
|
||||
str(args.baseline_json)
|
||||
if args.baseline_json is not None
|
||||
else None
|
||||
),
|
||||
"baseline_sha256": (
|
||||
sha256_file(args.baseline_json)
|
||||
if args.baseline_json is not None
|
||||
else None
|
||||
),
|
||||
"human_eval_sha256": sha256_file(args.human_eval),
|
||||
"gsm8k_sha256": sha256_file(args.gsm8k),
|
||||
"model_revision": behavior["model"]["revision"],
|
||||
"max_new_tokens": behavior[
|
||||
"generation_contract"
|
||||
]["max_new_tokens"],
|
||||
},
|
||||
"sandbox": {
|
||||
"image": args.sandbox_image,
|
||||
"timeout_seconds": args.timeout_seconds,
|
||||
"code_execution_skipped": args.skip_code_execution,
|
||||
"network": "none",
|
||||
"filesystem": "read-only",
|
||||
"user": "65534:65534",
|
||||
"capabilities": "ALL dropped",
|
||||
"memory": "256m",
|
||||
"memory_swap": "256m",
|
||||
"pids_limit": 64,
|
||||
"cpus": 0.5,
|
||||
"tmpfs": "/tmp:rw,noexec,nosuid,size=16m",
|
||||
"host_mounts": 0,
|
||||
},
|
||||
"rows": rows,
|
||||
"summary": summarize(rows),
|
||||
"claim_boundary": [
|
||||
"Four math and four code sources are not benchmark estimates.",
|
||||
"Completion-conditioned metrics are selection-biased diagnostics.",
|
||||
"A passing HumanEval test is functional evidence, not code-safety evidence.",
|
||||
"Fallback last-number extraction is not strict completion.",
|
||||
"Counterfactual token sequences are not official-valid chats.",
|
||||
],
|
||||
"content_hash": canonical_hash(rows),
|
||||
}
|
||||
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),
|
||||
"summary": result["summary"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,872 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Trace all 27 Chat decoder layers and all 26 MoE gates.
|
||||
|
||||
The probe reuses the fixed 16-source, eight-condition completion cohort but
|
||||
runs prompt-only full-depth forwards. It stores hashes and summaries rather
|
||||
than raw hidden tensors. Target-content comparisons use only byte-identical
|
||||
interior content tokens; boundary-crossing tokens are explicitly excluded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
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 torch.nn.functional as F
|
||||
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
|
||||
|
||||
|
||||
EDGE_PAIRS = {
|
||||
"system_eos": ("s0_eos", "s1_eos"),
|
||||
"system_bos": ("s0_bos", "s1_bos"),
|
||||
"system_x": ("s0_x", "s1_x"),
|
||||
"system_period": ("s0_period", "s1_period"),
|
||||
"bos_at_s0": ("s0_eos", "s0_bos"),
|
||||
"bos_at_s1": ("s1_eos", "s1_bos"),
|
||||
"x_at_s0": ("s0_eos", "s0_x"),
|
||||
"x_at_s1": ("s1_eos", "s1_x"),
|
||||
"period_at_s0": ("s0_eos", "s0_period"),
|
||||
"period_at_s1": ("s1_eos", "s1_period"),
|
||||
}
|
||||
SCOPES = ("target_content", "full_input")
|
||||
|
||||
|
||||
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("--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("--gpu-memory", default="28GiB")
|
||||
parser.add_argument("--cpu-memory", default="80GiB")
|
||||
parser.add_argument("--captured-at", default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def canonical_hash(value: Any) -> str:
|
||||
payload = json.dumps(
|
||||
value,
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
return hashlib.sha256(payload).hexdigest()
|
||||
|
||||
|
||||
def tensor_hash(tensor: torch.Tensor) -> str:
|
||||
value = tensor.detach().cpu().contiguous()
|
||||
header = json.dumps(
|
||||
{
|
||||
"shape": list(value.shape),
|
||||
"dtype": str(value.dtype),
|
||||
},
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
raw = value.view(torch.uint8).numpy().tobytes()
|
||||
return hashlib.sha256(header + b"\0" + raw).hexdigest()
|
||||
|
||||
|
||||
def validate_args(args: argparse.Namespace) -> None:
|
||||
if args.per_domain <= 0:
|
||||
raise ValueError("--per-domain must be positive")
|
||||
for path in (
|
||||
args.artifact_dir,
|
||||
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}"
|
||||
)
|
||||
|
||||
|
||||
def condition_positions(
|
||||
variants: dict[str, dict[str, Any]],
|
||||
batch_tokens: int,
|
||||
) -> dict[str, dict[str, list[int]]]:
|
||||
result = {}
|
||||
for condition in special.CONDITIONS:
|
||||
variant = variants[condition]
|
||||
padding = batch_tokens - variant["tokens"]
|
||||
result[condition] = {
|
||||
"target_content": [
|
||||
padding + int(position)
|
||||
for position in variant["content_positions"]
|
||||
],
|
||||
"full_input": list(range(padding, batch_tokens)),
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def hidden_stats(value: torch.Tensor) -> dict[str, Any]:
|
||||
fp32 = value.float()
|
||||
return {
|
||||
"shape": list(value.shape),
|
||||
"dtype": str(value.dtype),
|
||||
"tensor_sha256": tensor_hash(value),
|
||||
"mean": fp32.mean().item(),
|
||||
"std": fp32.std(unbiased=False).item(),
|
||||
"rms": fp32.square().mean().sqrt().item(),
|
||||
"max_abs": fp32.abs().max().item(),
|
||||
}
|
||||
|
||||
|
||||
def compare_hidden(
|
||||
left: torch.Tensor,
|
||||
right: torch.Tensor,
|
||||
) -> dict[str, Any]:
|
||||
if left.shape != right.shape:
|
||||
raise RuntimeError(
|
||||
f"hidden comparison shape mismatch: {left.shape} != "
|
||||
f"{right.shape}"
|
||||
)
|
||||
left_fp32 = left.float()
|
||||
right_fp32 = right.float()
|
||||
delta = right_fp32 - left_fp32
|
||||
denominator = left_fp32.norm(dim=-1).clamp_min(1e-12)
|
||||
relative = delta.norm(dim=-1) / denominator
|
||||
cosine = F.cosine_similarity(
|
||||
left_fp32,
|
||||
right_fp32,
|
||||
dim=-1,
|
||||
eps=1e-12,
|
||||
).clamp(-1.0, 1.0)
|
||||
exact = (left == right).all(dim=-1)
|
||||
return {
|
||||
"tokens": left.shape[0],
|
||||
"exact_hidden_rows": exact.sum().item(),
|
||||
"mean_cosine_similarity": cosine.mean().item(),
|
||||
"min_cosine_similarity": cosine.min().item(),
|
||||
"mean_relative_l2": relative.mean().item(),
|
||||
"max_abs_delta": delta.abs().max().item(),
|
||||
}
|
||||
|
||||
|
||||
def route_distribution(
|
||||
route_ids: torch.Tensor,
|
||||
route_weights: torch.Tensor,
|
||||
experts: int,
|
||||
) -> dict[str, Any]:
|
||||
loads = torch.bincount(
|
||||
route_ids.reshape(-1),
|
||||
minlength=experts,
|
||||
).to(torch.int64)
|
||||
weighted = torch.zeros(experts, dtype=torch.float64)
|
||||
weighted.scatter_add_(
|
||||
0,
|
||||
route_ids.reshape(-1),
|
||||
route_weights.reshape(-1).to(torch.float64),
|
||||
)
|
||||
total = loads.sum().item()
|
||||
probabilities = loads.to(torch.float64) / max(total, 1)
|
||||
nonzero = probabilities[probabilities > 0]
|
||||
entropy = (
|
||||
-(nonzero * nonzero.log()).sum().item() / math.log(experts)
|
||||
if len(nonzero)
|
||||
else 0.0
|
||||
)
|
||||
mean_load = loads.double().mean()
|
||||
cv = (
|
||||
loads.double().std(unbiased=False) / mean_load
|
||||
if mean_load > 0
|
||||
else torch.tensor(0.0)
|
||||
)
|
||||
return {
|
||||
"tokens": route_ids.shape[0],
|
||||
"top_k": route_ids.shape[1],
|
||||
"decisions": total,
|
||||
"ordered_route_sha256": canonical_hash(route_ids.tolist()),
|
||||
"route_weight_sha256": tensor_hash(route_weights),
|
||||
"loads": loads.tolist(),
|
||||
"weighted_loads": weighted.tolist(),
|
||||
"unique_experts": int((loads > 0).sum().item()),
|
||||
"load_cv": cv.item(),
|
||||
"normalized_load_entropy": entropy,
|
||||
"max_load_share": probabilities.max().item(),
|
||||
}
|
||||
|
||||
|
||||
def normalized_load(loads: torch.Tensor) -> torch.Tensor:
|
||||
value = loads.to(torch.float64)
|
||||
return value / value.sum().clamp_min(1e-12)
|
||||
|
||||
|
||||
def compare_routes(
|
||||
left_ids: torch.Tensor,
|
||||
left_weights: torch.Tensor,
|
||||
right_ids: torch.Tensor,
|
||||
right_weights: torch.Tensor,
|
||||
experts: int,
|
||||
) -> dict[str, Any]:
|
||||
if left_ids.shape != right_ids.shape:
|
||||
raise RuntimeError(
|
||||
f"route comparison shape mismatch: {left_ids.shape} != "
|
||||
f"{right_ids.shape}"
|
||||
)
|
||||
ordered_exact = (left_ids == right_ids).all(dim=-1)
|
||||
left_sorted = left_ids.sort(dim=-1).values
|
||||
right_sorted = right_ids.sort(dim=-1).values
|
||||
set_exact = (left_sorted == right_sorted).all(dim=-1)
|
||||
left_hot = F.one_hot(
|
||||
left_ids,
|
||||
num_classes=experts,
|
||||
).sum(dim=1).bool()
|
||||
right_hot = F.one_hot(
|
||||
right_ids,
|
||||
num_classes=experts,
|
||||
).sum(dim=1).bool()
|
||||
intersection = (left_hot & right_hot).sum(dim=-1)
|
||||
union = (left_hot | right_hot).sum(dim=-1).clamp_min(1)
|
||||
|
||||
left_token_weighted = torch.zeros(
|
||||
left_ids.shape[0],
|
||||
experts,
|
||||
dtype=torch.float64,
|
||||
)
|
||||
right_token_weighted = torch.zeros_like(left_token_weighted)
|
||||
left_token_weighted.scatter_add_(
|
||||
1,
|
||||
left_ids,
|
||||
left_weights.to(torch.float64),
|
||||
)
|
||||
right_token_weighted.scatter_add_(
|
||||
1,
|
||||
right_ids,
|
||||
right_weights.to(torch.float64),
|
||||
)
|
||||
token_weighted_tv = 0.5 * (
|
||||
left_token_weighted - right_token_weighted
|
||||
).abs().sum(dim=-1)
|
||||
|
||||
left_load = torch.bincount(
|
||||
left_ids.reshape(-1),
|
||||
minlength=experts,
|
||||
)
|
||||
right_load = torch.bincount(
|
||||
right_ids.reshape(-1),
|
||||
minlength=experts,
|
||||
)
|
||||
load_tv = 0.5 * (
|
||||
normalized_load(left_load)
|
||||
- normalized_load(right_load)
|
||||
).abs().sum()
|
||||
return {
|
||||
"tokens": left_ids.shape[0],
|
||||
"ordered_topk_exact_tokens": ordered_exact.sum().item(),
|
||||
"set_exact_tokens": set_exact.sum().item(),
|
||||
"mean_set_jaccard": (
|
||||
intersection.to(torch.float64)
|
||||
/ union.to(torch.float64)
|
||||
).mean().item(),
|
||||
"mean_token_weighted_tv": token_weighted_tv.mean().item(),
|
||||
"aggregate_load_tv": load_tv.item(),
|
||||
}
|
||||
|
||||
|
||||
class TraceCapture:
|
||||
def __init__(self, experts: int) -> None:
|
||||
self.experts = experts
|
||||
self.positions: dict[str, dict[str, list[int]]] = {}
|
||||
self.batch_tokens = 0
|
||||
self.hidden: dict[str, Any] = {}
|
||||
self.routes: dict[str, Any] = {}
|
||||
|
||||
def start(
|
||||
self,
|
||||
positions: dict[str, dict[str, list[int]]],
|
||||
batch_tokens: int,
|
||||
) -> None:
|
||||
self.positions = positions
|
||||
self.batch_tokens = batch_tokens
|
||||
self.hidden = {}
|
||||
self.routes = {}
|
||||
|
||||
def hidden_hook(self, stage: str):
|
||||
def hook(
|
||||
_module: torch.nn.Module,
|
||||
_inputs: tuple[Any, ...],
|
||||
output: Any,
|
||||
) -> None:
|
||||
tensor = output[0] if isinstance(output, tuple) else output
|
||||
if tensor.ndim != 3:
|
||||
raise RuntimeError(
|
||||
f"{stage} hidden rank {tensor.ndim}, expected 3"
|
||||
)
|
||||
cpu = tensor.detach().cpu()
|
||||
condition_values: dict[str, dict[str, torch.Tensor]] = {}
|
||||
condition_rows = {}
|
||||
for row, condition in enumerate(special.CONDITIONS):
|
||||
condition_values[condition] = {}
|
||||
condition_rows[condition] = {}
|
||||
for scope in SCOPES:
|
||||
positions = self.positions[condition][scope]
|
||||
value = cpu[row, positions, :].contiguous()
|
||||
condition_values[condition][scope] = value
|
||||
condition_rows[condition][scope] = hidden_stats(value)
|
||||
comparisons = {}
|
||||
for name, (left, right) in EDGE_PAIRS.items():
|
||||
comparisons[name] = compare_hidden(
|
||||
condition_values[left]["target_content"],
|
||||
condition_values[right]["target_content"],
|
||||
)
|
||||
self.hidden[stage] = {
|
||||
"conditions": condition_rows,
|
||||
"target_comparisons": comparisons,
|
||||
}
|
||||
|
||||
return hook
|
||||
|
||||
def gate_hook(self, layer: int):
|
||||
def hook(
|
||||
_module: torch.nn.Module,
|
||||
_inputs: tuple[Any, ...],
|
||||
output: Any,
|
||||
) -> None:
|
||||
route_ids, route_weights, _aux = output
|
||||
top_k = route_ids.shape[-1]
|
||||
ids = route_ids.detach().reshape(
|
||||
len(special.CONDITIONS),
|
||||
self.batch_tokens,
|
||||
top_k,
|
||||
).cpu()
|
||||
weights = route_weights.detach().reshape(
|
||||
len(special.CONDITIONS),
|
||||
self.batch_tokens,
|
||||
top_k,
|
||||
).cpu()
|
||||
values: dict[
|
||||
str,
|
||||
dict[str, tuple[torch.Tensor, torch.Tensor]],
|
||||
] = {}
|
||||
condition_rows = {}
|
||||
for row, condition in enumerate(special.CONDITIONS):
|
||||
values[condition] = {}
|
||||
condition_rows[condition] = {}
|
||||
for scope in SCOPES:
|
||||
positions = self.positions[condition][scope]
|
||||
scoped_ids = ids[row, positions, :].contiguous()
|
||||
scoped_weights = weights[
|
||||
row,
|
||||
positions,
|
||||
:,
|
||||
].contiguous()
|
||||
values[condition][scope] = (
|
||||
scoped_ids,
|
||||
scoped_weights,
|
||||
)
|
||||
condition_rows[condition][scope] = (
|
||||
route_distribution(
|
||||
scoped_ids,
|
||||
scoped_weights,
|
||||
self.experts,
|
||||
)
|
||||
)
|
||||
comparisons = {}
|
||||
for name, (left, right) in EDGE_PAIRS.items():
|
||||
left_ids, left_weights = values[left][
|
||||
"target_content"
|
||||
]
|
||||
right_ids, right_weights = values[right][
|
||||
"target_content"
|
||||
]
|
||||
comparisons[name] = compare_routes(
|
||||
left_ids,
|
||||
left_weights,
|
||||
right_ids,
|
||||
right_weights,
|
||||
self.experts,
|
||||
)
|
||||
self.routes[f"layer_{layer:02d}"] = {
|
||||
"layer": layer,
|
||||
"conditions": condition_rows,
|
||||
"target_comparisons": comparisons,
|
||||
}
|
||||
|
||||
return hook
|
||||
|
||||
|
||||
def aggregate_trace(sources: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
hidden_stages = [
|
||||
"embedding",
|
||||
*(f"layer_{index:02d}" for index in range(27)),
|
||||
"final_norm",
|
||||
]
|
||||
router_layers = [
|
||||
f"layer_{index:02d}"
|
||||
for index in range(1, 27)
|
||||
]
|
||||
hidden = {}
|
||||
for stage in hidden_stages:
|
||||
hidden[stage] = {}
|
||||
for edge in EDGE_PAIRS:
|
||||
rows = [
|
||||
source["hidden_stages"][stage][
|
||||
"target_comparisons"
|
||||
][edge]
|
||||
for source in sources
|
||||
]
|
||||
tokens = sum(row["tokens"] for row in rows)
|
||||
hidden[stage][edge] = {
|
||||
"sources": len(rows),
|
||||
"tokens": tokens,
|
||||
"exact_hidden_rows": sum(
|
||||
row["exact_hidden_rows"] for row in rows
|
||||
),
|
||||
"token_weighted_mean_cosine_similarity": sum(
|
||||
row["mean_cosine_similarity"] * row["tokens"]
|
||||
for row in rows
|
||||
) / tokens,
|
||||
"source_mean_relative_l2": mean(
|
||||
row["mean_relative_l2"] for row in rows
|
||||
),
|
||||
"max_abs_delta": max(
|
||||
row["max_abs_delta"] for row in rows
|
||||
),
|
||||
}
|
||||
routes = {}
|
||||
for layer in router_layers:
|
||||
routes[layer] = {}
|
||||
for edge in EDGE_PAIRS:
|
||||
rows = [
|
||||
source["router_layers"][layer][
|
||||
"target_comparisons"
|
||||
][edge]
|
||||
for source in sources
|
||||
]
|
||||
tokens = sum(row["tokens"] for row in rows)
|
||||
routes[layer][edge] = {
|
||||
"sources": len(rows),
|
||||
"tokens": tokens,
|
||||
"ordered_topk_exact_tokens": sum(
|
||||
row["ordered_topk_exact_tokens"] for row in rows
|
||||
),
|
||||
"set_exact_tokens": sum(
|
||||
row["set_exact_tokens"] for row in rows
|
||||
),
|
||||
"token_weighted_mean_set_jaccard": sum(
|
||||
row["mean_set_jaccard"] * row["tokens"]
|
||||
for row in rows
|
||||
) / tokens,
|
||||
"source_mean_token_weighted_tv": mean(
|
||||
row["mean_token_weighted_tv"] for row in rows
|
||||
),
|
||||
"source_mean_aggregate_load_tv": mean(
|
||||
row["aggregate_load_tv"] for row in rows
|
||||
),
|
||||
}
|
||||
target_tokens = sum(
|
||||
source["content_tokens"] for source in sources
|
||||
)
|
||||
return {
|
||||
"sources": len(sources),
|
||||
"conditions_per_source": len(special.CONDITIONS),
|
||||
"decoder_layers_executed": 27,
|
||||
"moe_gates_executed": 26,
|
||||
"hidden_stages": len(hidden_stages),
|
||||
"target_content_tokens_per_condition": target_tokens,
|
||||
"target_route_decisions": (
|
||||
target_tokens
|
||||
* len(special.CONDITIONS)
|
||||
* 26
|
||||
* 6
|
||||
),
|
||||
"hidden": hidden,
|
||||
"routes": routes,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
validate_args(args)
|
||||
captured_at = args.captured_at or datetime.now(
|
||||
timezone.utc
|
||||
).isoformat()
|
||||
revision_contract = behavior.download_revision_contract(
|
||||
args.artifact_dir
|
||||
)
|
||||
special.install_control_contract()
|
||||
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
|
||||
if special.special_family_token_ids(tokenizer) != {
|
||||
"eos": 100001,
|
||||
"bos": 100000,
|
||||
"x": 87,
|
||||
"period": 13,
|
||||
}:
|
||||
raise RuntimeError("pinned tokenizer ID contract changed")
|
||||
|
||||
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 special.CONDITIONS
|
||||
}
|
||||
content_sequences = {
|
||||
tuple(
|
||||
source["variants"][condition]["token_ids"][position]
|
||||
for position in source["variants"][condition][
|
||||
"content_positions"
|
||||
]
|
||||
)
|
||||
for condition in special.CONDITIONS
|
||||
}
|
||||
if len(content_sequences) != 1:
|
||||
raise RuntimeError(
|
||||
f"{source['id']} target token sequence is not exact "
|
||||
"across conditions"
|
||||
)
|
||||
source["content_token_ids_sha256"] = canonical_hash(
|
||||
next(iter(content_sequences))
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
input_device = model.get_input_embeddings().weight.device
|
||||
|
||||
capture = TraceCapture(model.config.n_routed_experts)
|
||||
handles = [
|
||||
model.model.embed_tokens.register_forward_hook(
|
||||
capture.hidden_hook("embedding")
|
||||
),
|
||||
model.model.norm.register_forward_hook(
|
||||
capture.hidden_hook("final_norm")
|
||||
),
|
||||
]
|
||||
for index, layer in enumerate(model.model.layers):
|
||||
handles.append(
|
||||
layer.register_forward_hook(
|
||||
capture.hidden_hook(f"layer_{index:02d}")
|
||||
)
|
||||
)
|
||||
if hasattr(layer.mlp, "gate"):
|
||||
handles.append(
|
||||
layer.mlp.gate.register_forward_hook(
|
||||
capture.gate_hook(index)
|
||||
)
|
||||
)
|
||||
|
||||
traced_sources = []
|
||||
try:
|
||||
for source in source_rows:
|
||||
unpadded = [
|
||||
source["variants"][condition]["token_ids"]
|
||||
for condition in special.CONDITIONS
|
||||
]
|
||||
prompt_lengths = [len(row) for row in unpadded]
|
||||
batch_tokens = max(prompt_lengths)
|
||||
padded = [
|
||||
[int(tokenizer.pad_token_id)]
|
||||
* (batch_tokens - len(row))
|
||||
+ row
|
||||
for row in unpadded
|
||||
]
|
||||
masks = [
|
||||
[0] * (batch_tokens - len(row))
|
||||
+ [1] * len(row)
|
||||
for row in unpadded
|
||||
]
|
||||
positions = condition_positions(
|
||||
source["variants"],
|
||||
batch_tokens,
|
||||
)
|
||||
capture.start(positions, batch_tokens)
|
||||
input_ids = torch.tensor(
|
||||
padded,
|
||||
dtype=torch.long,
|
||||
device=input_device,
|
||||
)
|
||||
attention_mask = torch.tensor(
|
||||
masks,
|
||||
dtype=torch.long,
|
||||
device=input_device,
|
||||
)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.cuda.synchronize()
|
||||
started = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
output = model.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
use_cache=False,
|
||||
output_attentions=False,
|
||||
output_hidden_states=False,
|
||||
return_dict=True,
|
||||
)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
elapsed = time.perf_counter() - started
|
||||
peak = (
|
||||
torch.cuda.max_memory_allocated()
|
||||
if torch.cuda.is_available()
|
||||
else None
|
||||
)
|
||||
if output.last_hidden_state.shape[:2] != input_ids.shape:
|
||||
raise RuntimeError(
|
||||
"final hidden batch/sequence shape does not match input"
|
||||
)
|
||||
expected_hidden = {
|
||||
"embedding",
|
||||
*(f"layer_{index:02d}" for index in range(27)),
|
||||
"final_norm",
|
||||
}
|
||||
expected_routes = {
|
||||
f"layer_{index:02d}"
|
||||
for index in range(1, 27)
|
||||
}
|
||||
if set(capture.hidden) != expected_hidden:
|
||||
raise RuntimeError(
|
||||
f"hidden stages mismatch: {set(capture.hidden)}"
|
||||
)
|
||||
if set(capture.routes) != expected_routes:
|
||||
raise RuntimeError(
|
||||
f"router layers mismatch: {set(capture.routes)}"
|
||||
)
|
||||
content_tokens = len(
|
||||
positions[special.CONDITIONS[0]][
|
||||
"target_content"
|
||||
]
|
||||
)
|
||||
traced_sources.append(
|
||||
{
|
||||
key: value
|
||||
for key, value in source.items()
|
||||
if key not in {"content", "variants"}
|
||||
}
|
||||
| {
|
||||
"prompt_tokens_by_condition": {
|
||||
condition: prompt_lengths[index]
|
||||
for index, condition in enumerate(
|
||||
special.CONDITIONS
|
||||
)
|
||||
},
|
||||
"batch_prompt_tokens": batch_tokens,
|
||||
"content_tokens": content_tokens,
|
||||
"content_token_ids_sha256": source[
|
||||
"content_token_ids_sha256"
|
||||
],
|
||||
"boundary_crossing_tokens_by_condition": {
|
||||
condition: source["variants"][condition][
|
||||
"boundary_crossing_tokens"
|
||||
]
|
||||
for condition in special.CONDITIONS
|
||||
},
|
||||
"forward_seconds": elapsed,
|
||||
"peak_cuda_memory_allocated_bytes": peak,
|
||||
"hidden_stages": capture.hidden,
|
||||
"router_layers": capture.routes,
|
||||
}
|
||||
)
|
||||
finally:
|
||||
for handle in handles:
|
||||
handle.remove()
|
||||
|
||||
device_map = getattr(model, "hf_device_map", {})
|
||||
peaks = [
|
||||
source["peak_cuda_memory_allocated_bytes"]
|
||||
for source in traced_sources
|
||||
if source["peak_cuda_memory_allocated_bytes"] is not None
|
||||
]
|
||||
result = {
|
||||
"schema_version": 1,
|
||||
"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",
|
||||
"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_ids": tokenizer.all_special_ids,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
"conditions": special.FACTORS,
|
||||
},
|
||||
"source_contract": {
|
||||
**reference,
|
||||
"domains": args.domains,
|
||||
"per_domain": args.per_domain,
|
||||
"sources": len(source_rows),
|
||||
"full_source_text_used": True,
|
||||
"target_scope": (
|
||||
"byte-identical interior content tokens; tokens whose "
|
||||
"offset crosses the content boundary are excluded"
|
||||
),
|
||||
},
|
||||
"trace_contract": {
|
||||
"conditions": list(special.CONDITIONS),
|
||||
"scopes": list(SCOPES),
|
||||
"edge_pairs": EDGE_PAIRS,
|
||||
"decoder_layers": 27,
|
||||
"moe_gate_layers": list(range(1, 27)),
|
||||
"experts": model.config.n_routed_experts,
|
||||
"top_k": model.config.num_experts_per_tok,
|
||||
"use_cache": False,
|
||||
"raw_hidden_tensors_saved": False,
|
||||
"raw_route_ids_saved": False,
|
||||
"load_vectors_saved": True,
|
||||
"hashes_cover_exact_bfloat16_hidden_and_route_values": True,
|
||||
},
|
||||
"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,
|
||||
"load_seconds": load_seconds,
|
||||
"forward_seconds": sum(
|
||||
source["forward_seconds"]
|
||||
for source in traced_sources
|
||||
),
|
||||
"load_peak_cuda_memory_allocated_bytes": load_peak_cuda,
|
||||
"peak_cuda_memory_allocated_bytes": max(
|
||||
[load_peak_cuda, *peaks]
|
||||
),
|
||||
"process_max_rss_kib": resource.getrusage(
|
||||
resource.RUSAGE_SELF
|
||||
).ru_maxrss,
|
||||
},
|
||||
"sources": traced_sources,
|
||||
"summary": aggregate_trace(traced_sources),
|
||||
"claim_boundary": [
|
||||
"The trace observes one fixed SFT Chat checkpoint, not the Base checkpoint.",
|
||||
"Target comparisons exclude boundary-crossing tokens and require exact token IDs.",
|
||||
"Router and hidden associations do not establish mediation or capability causality.",
|
||||
"CPU in hf_device_map is offload residency, not proof of CPU matrix execution.",
|
||||
"Prompt-only use_cache=False traces are not generation-time KV or serving traces.",
|
||||
],
|
||||
}
|
||||
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": hashlib.sha256(payload).hexdigest(),
|
||||
"bytes": len(payload),
|
||||
"sources": len(traced_sources),
|
||||
"decoder_layers": 27,
|
||||
"moe_gates": 26,
|
||||
"target_route_decisions": result["summary"][
|
||||
"target_route_decisions"
|
||||
],
|
||||
"load_seconds": load_seconds,
|
||||
"forward_seconds": result["execution"][
|
||||
"forward_seconds"
|
||||
],
|
||||
"peak_cuda_memory_allocated_bytes": result[
|
||||
"execution"
|
||||
]["peak_cuda_memory_allocated_bytes"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -24,6 +24,7 @@ import ast
|
||||
import gzip
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import resource
|
||||
@@ -933,6 +934,9 @@ def main() -> 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": (
|
||||
|
||||
Reference in New Issue
Block a user