#!/usr/bin/env python3 """Probe DeepSeek-V2-Lite routing sensitivity to its official chat template. The experiment keeps one source-addressable public cohort fixed and renders three inputs per prompt: 1. regular tokenizer input (BOS + content); 2. the official one-user chat template; 3. the same official template with the generation prompt appended. All three variants execute together through official BF16 layers 0--6. The output separates full-input loads from content-only loads, pairs bootstrap indices by source prompt, aligns content tokens by character span, and checks the causal invariant that appending a suffix cannot change a shared prefix. """ from __future__ import annotations import argparse import gc import hashlib import json import math import platform from datetime import datetime, timezone from pathlib import Path from typing import Any import numpy as np import torch import torch.nn.functional as F from safetensors import safe_open from transformers import AutoTokenizer from v2_lite_routing_corpus import ( DOMAIN_LABELS, DOMAIN_ORDER, bootstrap_domain, distribution, git_revision, gpu_identity, interval, js_divergence, load_candidates, load_official_modules, metric_vector, scoped_seed, sha256, text_sha256, ) CONDITIONS = ("raw", "user", "generation") CONDITION_LABELS = { "raw": "BOS + content", "user": "official user template", "generation": "official user template + Assistant:", } COMPARISONS = ( ("raw_to_user", "raw", "user"), ("user_to_generation", "user", "generation"), ) SAMPLE_SALT = "llm-atlas-deepseek-routing-template-control-v1" CHAT_TEMPLATE_REVISION = "604d5664dddd88a0433dbae533b7fe9472482de0" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--artifact-dir", 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("--per-domain", type=int, default=32) parser.add_argument("--content-tokens", type=int, default=23) parser.add_argument("--batch-prompts", type=int, default=8) parser.add_argument("--layers", type=int, default=7) parser.add_argument("--bootstrap", type=int, default=2000) parser.add_argument("--seed", type=int, default=20260729) parser.add_argument("--sample-salt", default=SAMPLE_SALT) parser.add_argument("--device", default="cuda") 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 tokenize_with_offsets( tokenizer: Any, text: str, *, add_special_tokens: bool, ) -> tuple[list[int], list[tuple[int, int]]]: encoded = tokenizer( text, add_special_tokens=add_special_tokens, return_offsets_mapping=True, padding=False, truncation=False, ) return list(encoded.input_ids), [tuple(pair) for pair in encoded.offset_mapping] def content_positions( token_ids: list[int], offsets: list[tuple[int, int]], content_start: int, content_end: int, ) -> tuple[list[int], list[dict[str, int]], int]: positions = [] records = [] crossing = 0 for index, (token_id, (start, end)) in enumerate(zip(token_ids, offsets, strict=True)): if end <= start: continue overlaps = start < content_end and end > content_start inside = start >= content_start and end <= content_end if inside: positions.append(index) records.append( { "position": index, "token_id": token_id, "start": start - content_start, "end": end - content_start, } ) elif overlaps: crossing += 1 return positions, records, crossing def common_prefix_length(left: list[int], right: list[int]) -> int: length = 0 for left_id, right_id in zip(left, right): if left_id != right_id: break length += 1 return length def render_variant(tokenizer: Any, content: str, condition: str) -> dict[str, Any]: if condition == "raw": rendered = content token_ids, offsets = tokenize_with_offsets( tokenizer, rendered, add_special_tokens=True, ) content_start = 0 else: messages = [{"role": "user", "content": content}] add_generation_prompt = condition == "generation" rendered = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=add_generation_prompt, ) official_ids = list( tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=add_generation_prompt, ) ) token_ids, offsets = tokenize_with_offsets( tokenizer, rendered, add_special_tokens=False, ) if token_ids != official_ids: raise RuntimeError(f"{condition} rendered token IDs differ from apply_chat_template") content_start = rendered.index(content) content_end = content_start + len(content) positions, records, crossing = content_positions( token_ids, offsets, content_start, content_end, ) if not positions: raise RuntimeError(f"{condition} has no content tokens") return { "condition": condition, "rendered_sha256": text_sha256(rendered), "token_ids": token_ids, "tokens": len(token_ids), "token_ids_sha256": canonical_hash(token_ids), "content_positions": positions, "content_records": records, "content_tokens": len(positions), "wrapper_tokens": len(token_ids) - len(positions), "boundary_crossing_tokens": crossing, } def select_samples( tokenizer: Any, candidates: dict[str, list[dict[str, str]]], per_domain: int, content_tokens: int, sample_salt: str, ) -> tuple[list[dict[str, Any]], dict[str, dict[str, int]]]: def fixed_token_prefix( text: str, full_offsets: list[tuple[int, int]], ) -> tuple[str, list[int]]: lower_index = max(0, content_tokens - 5) upper_index = min(len(full_offsets) - 1, content_tokens + 7) lower = max(1, full_offsets[lower_index][0]) upper = max(lower, full_offsets[upper_index][1]) matches = [] for end in range(lower, upper + 1): prefix = text[:end] ids = list(tokenizer(prefix, add_special_tokens=False).input_ids) if len(ids) == content_tokens: matches.append((end, ids)) if not matches: raise RuntimeError("could not construct a fixed-token character prefix") end, ids = matches[-1] return text[:end], ids selected = [] counts = {} for domain in DOMAIN_ORDER: ranked = [] for row in candidates[domain]: full_ids, full_offsets = tokenize_with_offsets( tokenizer, row["text"], add_special_tokens=False, ) if len(full_ids) < content_tokens: continue rank = hashlib.sha256( f"{sample_salt}|{domain}|{row['id']}".encode() ).hexdigest() ranked.append((rank, row, len(full_ids), full_offsets)) ranked.sort(key=lambda item: (item[0], item[1]["id"])) eligible = [] canonicalization_failures = 0 for rank, row, source_tokens, full_offsets in ranked: try: content, prefix_ids = fixed_token_prefix(row["text"], full_offsets) except RuntimeError: canonicalization_failures += 1 continue variants = { condition: render_variant(tokenizer, content, condition) for condition in CONDITIONS } shared_content_keys = set.intersection( *( { (record["start"], record["end"], record["token_id"]) for record in variants[condition]["content_records"] } for condition in CONDITIONS ) ) for condition in CONDITIONS: variants[condition]["aligned_content_positions"] = [ record["position"] for record in variants[condition]["content_records"] if (record["start"], record["end"], record["token_id"]) in shared_content_keys ] variants[condition]["aligned_content_tokens"] = len( variants[condition]["aligned_content_positions"] ) eligible.append( { "id": row["id"], "domain": domain, "label": DOMAIN_LABELS[domain], "text_sha256": text_sha256(row["text"]), "source_characters": len(row["text"]), "source_tokens": source_tokens, "content": content, "content_sha256": text_sha256(content), "content_characters": len(content), "canonical_content_token_ids_sha256": canonical_hash(prefix_ids), "selection_rank": rank, "variants": variants, } ) if len(eligible) == per_domain: break if len(eligible) < per_domain: raise RuntimeError( f"{domain} has only {len(eligible)} eligible prompts; need {per_domain}" ) domain_rows = eligible[:per_domain] for index, row in enumerate(domain_rows): row["within_domain_index"] = index selected.extend(domain_rows) counts[domain] = { "candidate_records_after_text_filter": len(candidates[domain]), "eligible_records": len(ranked), "canonicalization_failures_before_selection_complete": ( canonicalization_failures ), "selected_records": len(domain_rows), "source_tokens_min": min(row["source_tokens"] for row in domain_rows), "source_tokens_mean": float( np.mean([row["source_tokens"] for row in domain_rows]) ), "source_tokens_max": max(row["source_tokens"] for row in domain_rows), } return selected, counts def make_batches( samples: list[dict[str, Any]], batch_prompts: int, pad_token_id: int, ) -> list[dict[str, Any]]: ordered = sorted( samples, key=lambda row: ( max(row["variants"][condition]["tokens"] for condition in CONDITIONS), DOMAIN_ORDER.index(row["domain"]), row["id"], ), ) batches = [] for start in range(0, len(ordered), batch_prompts): prompt_rows = ordered[start : start + batch_prompts] variants = [ { "sample": sample, "condition": condition, **sample["variants"][condition], } for sample in prompt_rows for condition in CONDITIONS ] sequence = max(row["tokens"] for row in variants) input_ids = torch.full( (len(variants), sequence), pad_token_id, dtype=torch.long, ) attention_mask = torch.zeros((len(variants), sequence), dtype=torch.long) for index, row in enumerate(variants): length = row["tokens"] input_ids[index, :length] = torch.tensor(row["token_ids"]) attention_mask[index, :length] = 1 batches.append( { "prompt_rows": prompt_rows, "variants": variants, "input_ids": input_ids, "attention_mask": attention_mask, "padded_sequence": sequence, } ) return batches def bootstrap_distributions( loads: np.ndarray, mode: str, sampled: np.ndarray, ) -> np.ndarray: if mode == "token_weighted": values = loads[sampled].sum(axis=1, dtype=np.float64) return values / values.sum(axis=1, keepdims=True) prompt_distributions = loads / loads.sum(axis=1, keepdims=True) values = prompt_distributions[sampled].mean(axis=1) return values / values.sum(axis=1, keepdims=True) def paired_domain( before: np.ndarray, after: np.ndarray, mode: str, replicates: int, seed: int, scope: str, ) -> dict[str, Any]: if before.shape != after.shape: raise ValueError(f"paired shape mismatch: {before.shape} != {after.shape}") rng = np.random.default_rng(scoped_seed(seed, scope)) sampled = rng.integers( 0, before.shape[0], size=(replicates, before.shape[0]), endpoint=False, ) before_point = distribution(before, mode) after_point = distribution(after, mode) before_boot = bootstrap_distributions(before, mode, sampled) after_boot = bootstrap_distributions(after, mode, sampled) before_metrics = metric_vector(before_point) after_metrics = metric_vector(after_point) before_boot_metrics = metric_vector(before_boot) after_boot_metrics = metric_vector(after_boot) metrics = {} for name in before_metrics: delta_boot = after_boot_metrics[name] - before_boot_metrics[name] metrics[name] = { "before": float(before_metrics[name][0]), "after": float(after_metrics[name][0]), "delta_after_minus_before": float( after_metrics[name][0] - before_metrics[name][0] ), "delta_ci95": interval(delta_boot), } tv_boot = 0.5 * np.abs(after_boot - before_boot).sum(axis=1) jsd_boot = js_divergence(before_boot, after_boot) share_delta = after_point - before_point return { "metrics": metrics, "total_variation": { "point": float(0.5 * np.abs(share_delta).sum()), "ci95": interval(tv_boot), }, "js_divergence": { "point": float(js_divergence(before_point, after_point)[0]), "ci95": interval(jsd_boot), "unit": "nats", "upper_bound": math.log(2), }, "expert_share_delta": share_delta.tolist(), "expert_share_delta_ci95": interval(after_boot - before_boot), } def aligned_pairs( left: dict[str, Any], right: dict[str, Any], ) -> list[tuple[int, int]]: right_by_key = { (row["start"], row["end"], row["token_id"]): row["position"] for row in right["content_records"] } return [ ( row["position"], right_by_key[(row["start"], row["end"], row["token_id"])], ) for row in left["content_records"] if (row["start"], row["end"], row["token_id"]) in right_by_key ] def route_alignment( left_routes: torch.Tensor, right_routes: torch.Tensor, pairs: list[tuple[int, int]], ) -> dict[str, Any]: if not pairs: return { "aligned_tokens": 0, "ordered_topk_exact": 0, "set_topk_exact": 0, "mean_topk_overlap": None, "mean_jaccard": None, } ordered_exact = 0 set_exact = 0 overlaps = [] jaccards = [] for left_position, right_position in pairs: left = left_routes[left_position].tolist() right = right_routes[right_position].tolist() ordered_exact += int(left == right) left_set = set(left) right_set = set(right) intersection = len(left_set & right_set) union = len(left_set | right_set) set_exact += int(left_set == right_set) overlaps.append(intersection) jaccards.append(intersection / union) return { "aligned_tokens": len(pairs), "ordered_topk_exact": ordered_exact, "set_topk_exact": set_exact, "ordered_topk_exact_rate": ordered_exact / len(pairs), "set_topk_exact_rate": set_exact / len(pairs), "mean_topk_overlap": float(np.mean(overlaps)), "mean_jaccard": float(np.mean(jaccards)), } def layer_statistics( prompt_rows: list[dict[str, Any]], replicates: int, seed: int, layer_index: int, ) -> dict[str, Any]: scopes = {} for load_scope, load_key in ( ("full_input", "full_load"), ("content_only", "content_load"), ): modes = {} for mode in ("token_weighted", "prompt_balanced"): conditions = {} for condition in CONDITIONS: conditions[condition] = {} for domain in DOMAIN_ORDER: loads = np.asarray( [ row["conditions"][condition][load_key] for row in prompt_rows if row["domain"] == domain ], dtype=np.int64, ) conditions[condition][domain] = bootstrap_domain( loads, mode, replicates, seed, ( f"layer={layer_index}|scope={load_scope}|mode={mode}|" f"condition={condition}|domain={domain}" ), ) comparisons = {} for comparison, before_condition, after_condition in COMPARISONS: comparisons[comparison] = {} for domain in DOMAIN_ORDER: rows = [row for row in prompt_rows if row["domain"] == domain] before = np.asarray( [ row["conditions"][before_condition][load_key] for row in rows ], dtype=np.int64, ) after = np.asarray( [ row["conditions"][after_condition][load_key] for row in rows ], dtype=np.int64, ) comparisons[comparison][domain] = paired_domain( before, after, mode, replicates, seed, ( f"layer={layer_index}|scope={load_scope}|mode={mode}|" f"comparison={comparison}|domain={domain}" ), ) modes[mode] = { "conditions": conditions, "comparisons": comparisons, } scopes[load_scope] = {"modes": modes} return scopes def main() -> None: args = parse_args() root = args.artifact_dir.resolve() shard = root / "model-00001-of-000004.safetensors" required = [ root / "config.json", root / "configuration_deepseek.py", root / "modeling_deepseek.py", root / "model.safetensors.index.json", root / "tokenizer.json", root / "tokenizer_config.json", shard, args.human_eval, args.gsm8k, args.tnews, args.tnews_archive, args.wikitext, ] missing = [str(path) for path in required if not path.exists()] if missing: raise FileNotFoundError(f"missing artifacts: {missing}") if args.device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable") if not 2 <= args.layers <= 7: raise ValueError("need layer 0 plus at least one MoE layer; shard ends at layer 6") if args.per_domain < 2: raise ValueError("per-domain sample must be at least two") if args.content_tokens < 4: raise ValueError("content prefix must contain at least four tokens") if args.batch_prompts < 1: raise ValueError("batch-prompts must be positive") if args.bootstrap < 100: raise ValueError("bootstrap replicates must be at least 100") if not args.sample_salt.strip(): raise ValueError("sample salt must not be empty") torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) torch.backends.cuda.matmul.allow_tf32 = False configuration, modeling = load_official_modules(root) config = configuration.DeepseekV2Config.from_pretrained(root) config._attn_implementation = "eager" tokenizer = AutoTokenizer.from_pretrained( root, trust_remote_code=True, local_files_only=True, ) if not tokenizer.is_fast: raise RuntimeError("offset alignment requires a fast tokenizer") if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" tokenizer_config = json.loads((root / "tokenizer_config.json").read_text()) if tokenizer.chat_template != tokenizer_config.get("chat_template"): raise RuntimeError("loaded chat template differs from tokenizer_config.json") candidates = load_candidates(args) samples, corpus_counts = select_samples( tokenizer, candidates, args.per_domain, args.content_tokens, args.sample_salt, ) batches = make_batches(samples, args.batch_prompts, tokenizer.pad_token_id) device = torch.device(args.device) with safe_open(shard, framework="pt", device="cpu") as handle: embedding = handle.get_tensor("model.embed_tokens.weight") for batch in batches: batch["hidden"] = F.embedding(batch["input_ids"], embedding) del embedding layer_results = [] for layer_index in range(args.layers): prefix = f"model.layers.{layer_index}." with safe_open(shard, framework="pt", device="cpu") as handle: state = { key[len(prefix) :]: handle.get_tensor(key) for key in handle.keys() if key.startswith(prefix) } state_numel = sum(value.numel() for value in state.values()) state_bytes = sum(value.numel() * value.element_size() for value in state.values()) with torch.device("meta"): layer = modeling.DeepseekV2DecoderLayer(config, layer_index) layer.to_empty(device="cpu") layer.load_state_dict(state, strict=True, assign=True) del state layer = layer.to(device=device, dtype=torch.bfloat16).eval() prompt_rows = [] next_hidden = [] causal_invariant = { "shared_prefix_tokens": 0, "ordered_topk_exact": 0, "violations": 0, } for batch in batches: hidden = batch["hidden"].to(device) attention_mask = batch["attention_mask"].to(device) sequence = batch["padded_sequence"] position_ids = torch.arange(sequence, device=device).unsqueeze(0) causal_mask = modeling._prepare_4d_causal_attention_mask( attention_mask, batch["input_ids"].shape, hidden, 0, ) captures: dict[str, torch.Tensor] = {} hook = None if layer_index > 0: def capture_gate(_module: Any, _inputs: Any, output: Any) -> None: captures["topk_ids"] = output[0].detach().cpu() hook = layer.mlp.gate.register_forward_hook(capture_gate) with torch.inference_mode(): outputs = layer( hidden, attention_mask=causal_mask, position_ids=position_ids, past_key_value=None, use_cache=False, ) next_hidden.append(outputs[0].cpu()) if hook is not None: hook.remove() topk = captures["topk_ids"].view( len(batch["variants"]), sequence, config.num_experts_per_tok, ) by_prompt: dict[str, dict[str, tuple[dict[str, Any], torch.Tensor]]] = {} for variant_index, variant in enumerate(batch["variants"]): sample = variant["sample"] condition = variant["condition"] routes = topk[variant_index, : variant["tokens"]] by_prompt.setdefault(sample["id"], {})[condition] = (variant, routes) for sample in batch["prompt_rows"]: variants = by_prompt[sample["id"]] condition_rows = {} for condition in CONDITIONS: variant, routes = variants[condition] full_load = torch.bincount( routes.flatten(), minlength=config.n_routed_experts, ) content_routes = routes[ variant["aligned_content_positions"] ] content_load = torch.bincount( content_routes.flatten(), minlength=config.n_routed_experts, ) condition_rows[condition] = { "input_tokens": variant["tokens"], "content_tokens": variant["content_tokens"], "aligned_content_tokens": variant[ "aligned_content_tokens" ], "wrapper_tokens": variant["wrapper_tokens"], "boundary_crossing_tokens": variant[ "boundary_crossing_tokens" ], "routes": int(full_load.sum()), "content_routes": int(content_load.sum()), "full_load": full_load.tolist(), "content_load": content_load.tolist(), "topk_sha256": canonical_hash(routes.tolist()), "content_topk_sha256": canonical_hash( content_routes.tolist() ), } raw_variant, raw_routes = variants["raw"] user_variant, user_routes = variants["user"] generation_variant, generation_routes = variants["generation"] raw_user_pairs = aligned_pairs(raw_variant, user_variant) raw_user = route_alignment( raw_routes, user_routes, raw_user_pairs, ) prefix_length = common_prefix_length( user_variant["token_ids"], generation_variant["token_ids"], ) prefix_left = user_routes[:prefix_length] prefix_right = generation_routes[:prefix_length] exact_prefix = int( torch.equal(prefix_left, prefix_right) ) causal_invariant["shared_prefix_tokens"] += prefix_length causal_invariant["ordered_topk_exact"] += ( prefix_length if exact_prefix else int( (prefix_left == prefix_right).all(dim=1).sum() ) ) causal_invariant["violations"] += int(not exact_prefix) prompt_rows.append( { "id": sample["id"], "domain": sample["domain"], "conditions": condition_rows, "alignments": { "raw_to_user_content": { **raw_user, "raw_content_tokens": raw_variant[ "content_tokens" ], "user_content_tokens": user_variant[ "content_tokens" ], "raw_alignment_coverage": ( raw_user["aligned_tokens"] / raw_variant["content_tokens"] ), "user_alignment_coverage": ( raw_user["aligned_tokens"] / user_variant["content_tokens"] ), }, "user_to_generation_prefix": { "shared_prefix_tokens": prefix_length, "user_tokens": user_variant["tokens"], "generation_tokens": generation_variant["tokens"], "ordered_topk_exact": exact_prefix == 1, "ordered_topk_exact_tokens": ( prefix_length if exact_prefix else int( (prefix_left == prefix_right) .all(dim=1) .sum() ) ), }, }, } ) del hidden, attention_mask, causal_mask, outputs for batch, hidden in zip(batches, next_hidden, strict=True): batch["hidden"] = hidden result: dict[str, Any] = { "layer": layer_index, "ffn": "dense" if layer_index == 0 else "moe", "state_numel": state_numel, "state_bytes": state_bytes, } if layer_index > 0: prompt_rows.sort( key=lambda row: ( DOMAIN_ORDER.index(row["domain"]), row["id"], ) ) causal_invariant["exact_rate"] = ( causal_invariant["ordered_topk_exact"] / causal_invariant["shared_prefix_tokens"] ) result["prompts"] = prompt_rows result["causal_suffix_invariant"] = causal_invariant result["statistics"] = layer_statistics( prompt_rows, args.bootstrap, args.seed, layer_index, ) layer_results.append(result) del layer, next_hidden gc.collect() if device.type == "cuda": torch.cuda.empty_cache() captured_at = args.captured_at or datetime.now(timezone.utc).isoformat() model_index = json.loads((root / "model.safetensors.index.json").read_text()) selected_identity = [] for sample in samples: selected_identity.append( { "id": sample["id"], "domain": sample["domain"], "within_domain_index": sample["within_domain_index"], "selection_rank": sample["selection_rank"], "text_sha256": sample["text_sha256"], "source_characters": sample["source_characters"], "source_tokens": sample["source_tokens"], "content_sha256": sample["content_sha256"], "content_characters": sample["content_characters"], "canonical_content_tokens": args.content_tokens, "canonical_content_token_ids_sha256": sample[ "canonical_content_token_ids_sha256" ], "conditions": { condition: { key: sample["variants"][condition][key] for key in ( "rendered_sha256", "tokens", "token_ids_sha256", "content_tokens", "aligned_content_tokens", "wrapper_tokens", "boundary_crossing_tokens", ) } for condition in CONDITIONS }, "user_generation_common_prefix_tokens": common_prefix_length( sample["variants"]["user"]["token_ids"], sample["variants"]["generation"]["token_ids"], ), "raw_user_aligned_content_tokens": len( aligned_pairs( sample["variants"]["raw"], sample["variants"]["user"], ) ), } ) result = { "schema_version": 1, "captured_at": captured_at, "evidence_identity": ( "X / official BF16 weights, official tokenizer chat template, " "paired local truncated forward" ), "boundary": { "model": "DeepSeek-V2-Lite base", "executed_layers": list(range(args.layers)), "measured_moe_layers": list(range(1, args.layers)), "total_model_layers": config.num_hidden_layers, "full_model_generation": False, "task_performance": False, "training_or_online_load": False, "expert_semantics_inferred": False, "causal_claim": ( "the user-to-generation shared-prefix equality is a causal-mask " "implementation invariant; raw-to-user differences are descriptive " "protocol sensitivity, not a capability effect" ), "population": ( f"{len(samples)} fixed public prompts across four domains; " "not training data, online traffic, or a task benchmark" ), "code_execution": False, "answers_used": False, }, "provenance": { "model": { "huggingface_model": "deepseek-ai/DeepSeek-V2-Lite", "huggingface_revision": CHAT_TEMPLATE_REVISION, "sha256": { "config": sha256(root / "config.json"), "modeling_code": sha256(root / "modeling_deepseek.py"), "tokenizer": sha256(root / "tokenizer.json"), "tokenizer_config": sha256(root / "tokenizer_config.json"), "index": sha256(root / "model.safetensors.index.json"), "shard_1": sha256(shard), }, "checkpoint_tensor_bytes": model_index["metadata"]["total_size"], "shard_1_bytes": shard.stat().st_size, }, "corpora": { "english": { "name": "WikiText-2 raw validation", "url": "https://huggingface.co/datasets/Salesforce/wikitext", "revision": "b08601e04326c79dfdd32d625aee71d232d685c3", "file_sha256": sha256(args.wikitext), "field_used": "text", }, "chinese": { "name": "CLUE TNEWS public test", "url": "https://github.com/CLUEbenchmark/CLUE", "download_url": "https://storage.googleapis.com/cluebenchmark/tasks/tnews_public.zip", "revision": "9e61ddd3659ddb57ed82b4d0ba0a8613dfb55a2e", "archive_sha256": sha256(args.tnews_archive), "file_sha256": sha256(args.tnews), "field_used": "sentence", }, "code": { "name": "OpenAI HumanEval", "url": "https://github.com/openai/human-eval", "revision": git_revision(args.human_eval), "file_sha256": sha256(args.human_eval), "field_used": "prompt", }, "math": { "name": "OpenAI GSM8K test", "url": "https://github.com/openai/grade-school-math", "revision": git_revision(args.gsm8k), "file_sha256": sha256(args.gsm8k), "field_used": "question", }, }, }, "environment": { "python": platform.python_version(), "platform": platform.platform(), "torch": torch.__version__, "torch_cuda": torch.version.cuda, "transformers": __import__("transformers").__version__, "jinja2": __import__("jinja2").__version__, "safetensors": __import__("safetensors").__version__, "numpy": np.__version__, "pyarrow": __import__("pyarrow").__version__, "device": str(device), "gpu": gpu_identity(device), "matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32, }, "configuration": { "layers": config.num_hidden_layers, "hidden": config.hidden_size, "routed_experts": config.n_routed_experts, "active_routed_experts": config.num_experts_per_tok, "shared_experts": config.n_shared_experts, "first_dense_layers": config.first_k_dense_replace, "router_scoring": config.scoring_func, "router_topk_method": config.topk_method, "normalize_selected_weights": config.norm_topk_prob, }, "template_contract": { "chat_template_revision": CHAT_TEMPLATE_REVISION, "chat_template": tokenizer.chat_template, "chat_template_sha256": text_sha256(tokenizer.chat_template), "bos_token": tokenizer.bos_token, "bos_token_id": tokenizer.bos_token_id, "eos_token": tokenizer.eos_token, "eos_token_id": tokenizer.eos_token_id, "conditions": { condition: CONDITION_LABELS[condition] for condition in CONDITIONS }, "comparisons": [ { "name": name, "before": before, "after": after, } for name, before, after in COMPARISONS ], "scope_split": { "full_input": "all BOS, wrapper, content, newline, and generation-prompt tokens", "content_only": ( "the exact intersection of (relative character span, token ID) " "inside source content across all three conditions; boundary-" "crossing and unaligned tokens are excluded from every condition" ), }, }, "corpus_contract": { "domains": list(DOMAIN_ORDER), "domain_labels": DOMAIN_LABELS, "sample_salt": args.sample_salt, "selection": "ascending SHA256(salt|domain|source_id), then source_id", "per_domain": args.per_domain, "canonical_content_tokens": args.content_tokens, "content_prefix": ( "source text cut at the end offset of the fixed regular-tokenizer " "content-token prefix" ), "counts": corpus_counts, "selected": selected_identity, }, "inference_contract": { "batch_prompts": args.batch_prompts, "variants_per_prompt": len(CONDITIONS), "rows_per_full_batch": args.batch_prompts * len(CONDITIONS), "batches": len(batches), "batch_grouping": ( "all raw/user/generation variants of one source prompt execute " "in the same padded batch" ), "attention": "official eager causal mask", "dtype": "BF16", "total_source_prompts": len(samples), "total_prompt_variants": len(samples) * len(CONDITIONS), "input_tokens_by_condition": { condition: sum( sample["variants"][condition]["tokens"] for sample in samples ) for condition in CONDITIONS }, "content_span_tokens_by_condition": { condition: sum( sample["variants"][condition]["content_tokens"] for sample in samples ) for condition in CONDITIONS }, "aligned_content_tokens_by_condition": { condition: sum( sample["variants"][condition]["aligned_content_tokens"] for sample in samples ) for condition in CONDITIONS }, "routes_per_condition_all_moe_layers": { condition: sum( sample["variants"][condition]["tokens"] for sample in samples ) * config.num_experts_per_tok * (args.layers - 1) for condition in CONDITIONS }, "total_routes_all_conditions_all_moe_layers": sum( sum(sample["variants"][condition]["tokens"] for sample in samples) * config.num_experts_per_tok * (args.layers - 1) for condition in CONDITIONS ), }, "statistical_contract": { "unit": "source prompt", "bootstrap_replicates": args.bootstrap, "seed": args.seed, "paired_indices": ( "the same resampled source-prompt indices are used for before and " "after within each domain/layer/scope/mode" ), "modes": { "token_weighted": "sum selected routes, then normalize", "prompt_balanced": ( "normalize each prompt load, then average prompts equally" ), }, "interval": "2.5th and 97.5th percentiles", "multiple_comparison_correction": False, "hypothesis_test": False, }, "layers": layer_results, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text( json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) print(json.dumps(result, indent=2, ensure_ascii=False)) if __name__ == "__main__": main()