wuyang
796 lines
29 KiB
Python
796 lines
29 KiB
Python
#!/usr/bin/env python3
|
|
"""Run one preregistered AttnRes activation-gradient/depth experiment cell."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import platform
|
|
import statistics
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
|
|
ARCHITECTURES = ("baseline", "block")
|
|
DEPTHS = (16, 32)
|
|
EXPECTED_SEEDS = (2026073001, 2026073002, 2026073003)
|
|
DIAGNOSTIC_STEPS = (0, 100, 500, 2000, 4000, 8000)
|
|
FORMAL_STEPS = 8000
|
|
SMOKE_STEPS = 20
|
|
CONTEXT = 256
|
|
VOCABULARY = 256
|
|
BLOCK_GROUPS = 8
|
|
PEAK_LR = 3e-4
|
|
MIN_LR = 3e-5
|
|
WARMUP_STEPS = 400
|
|
WEIGHT_DECAY = 0.1
|
|
BETAS = (0.9, 0.95)
|
|
ADAM_EPS = 1e-8
|
|
GRAD_CLIP = 1.0
|
|
|
|
|
|
def load_round04_module() -> Any:
|
|
path = Path(__file__).resolve().parents[1] / "attnres" / "train.py"
|
|
spec = importlib.util.spec_from_file_location("k3_attnres_round04_train", path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"cannot import Round 04 runner from {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
round04 = load_round04_module()
|
|
|
|
|
|
def configure_round04_globals(depth: int) -> None:
|
|
round04.PROTOCOL_ID = PROTOCOL_ID
|
|
round04.LAYERS = depth
|
|
round04.SUBLAYERS = depth * 2
|
|
round04.BLOCKS = BLOCK_GROUPS
|
|
round04.SUBLAYERS_PER_BLOCK = (depth * 2) // BLOCK_GROUPS
|
|
round04.WARMUP_STEPS = WARMUP_STEPS
|
|
round04.EVAL_STEPS = DIAGNOSTIC_STEPS
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--architecture", choices=ARCHITECTURES, required=True)
|
|
parser.add_argument("--depth", type=int, choices=DEPTHS, required=True)
|
|
parser.add_argument("--seed", type=int, required=True)
|
|
parser.add_argument("--steps", type=int)
|
|
parser.add_argument("--batch-size", type=int, default=32)
|
|
parser.add_argument("--cache-dir", type=Path, required=True)
|
|
parser.add_argument("--manifest", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--validation-windows", type=int, default=64)
|
|
parser.add_argument("--diagnostic-windows", type=int, default=16)
|
|
parser.add_argument("--eval-batch-size", type=int, default=8)
|
|
parser.add_argument("--timing-warmup", type=int, default=20)
|
|
parser.add_argument(
|
|
"--run-kind", choices=("smoke", "formal", "replay"), default="formal"
|
|
)
|
|
args = parser.parse_args()
|
|
expected_steps = SMOKE_STEPS if args.run_kind == "smoke" else FORMAL_STEPS
|
|
if args.steps is None:
|
|
args.steps = expected_steps
|
|
if args.steps != expected_steps:
|
|
raise ValueError(
|
|
f"{args.run_kind} must run exactly {expected_steps} steps, got {args.steps}"
|
|
)
|
|
if args.batch_size != 32:
|
|
raise ValueError("the frozen protocol requires batch size 32")
|
|
if args.validation_windows != 64 or args.diagnostic_windows != 16:
|
|
raise ValueError("the frozen protocol requires 64 validation / 16 diagnostic windows")
|
|
return args
|
|
|
|
|
|
def configure_determinism(seed: int) -> None:
|
|
if os.environ.get("CUBLAS_WORKSPACE_CONFIG") != ":4096:8":
|
|
raise RuntimeError("CUBLAS_WORKSPACE_CONFIG must be :4096:8 before Python starts")
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.use_deterministic_algorithms(True)
|
|
torch.backends.cudnn.benchmark = False
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
torch.set_float32_matmul_precision("highest")
|
|
|
|
|
|
def canonical_sha256(value: Any) -> str:
|
|
payload = json.dumps(
|
|
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
|
).encode()
|
|
return hashlib.sha256(payload).hexdigest()
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def tensor_bytes(tensor: torch.Tensor) -> bytes:
|
|
value = tensor.detach().cpu().contiguous()
|
|
return (
|
|
f"{value.dtype}|{tuple(value.shape)}|".encode()
|
|
+ value.reshape(-1).view(torch.uint8).numpy().tobytes()
|
|
)
|
|
|
|
|
|
def named_state_hash(
|
|
model: nn.Module, *, include_mixers: bool | None
|
|
) -> str:
|
|
digest = hashlib.sha256()
|
|
for name, tensor in sorted(model.state_dict().items()):
|
|
is_mixer = name.startswith("mixers.") or name.startswith("output_mixer.")
|
|
if include_mixers is not None and is_mixer != include_mixers:
|
|
continue
|
|
digest.update(name.encode())
|
|
digest.update(b"\0")
|
|
digest.update(tensor_bytes(tensor))
|
|
return digest.hexdigest()
|
|
|
|
|
|
def state_structure_hash(
|
|
model: nn.Module, *, include_mixers: bool | None
|
|
) -> tuple[str, int, int]:
|
|
digest = hashlib.sha256()
|
|
tensor_count = 0
|
|
element_count = 0
|
|
for name, tensor in sorted(model.state_dict().items()):
|
|
is_mixer = name.startswith("mixers.") or name.startswith("output_mixer.")
|
|
if include_mixers is not None and is_mixer != include_mixers:
|
|
continue
|
|
digest.update(
|
|
f"{name}|{tuple(tensor.shape)}|{tensor.dtype}|{tensor.numel()}\n".encode()
|
|
)
|
|
tensor_count += 1
|
|
element_count += tensor.numel()
|
|
return digest.hexdigest(), tensor_count, element_count
|
|
|
|
|
|
def recursive_state_hash(value: Any) -> str:
|
|
digest = hashlib.sha256()
|
|
|
|
def visit(path: str, item: Any) -> None:
|
|
if torch.is_tensor(item):
|
|
digest.update(f"{path}|tensor|".encode())
|
|
digest.update(tensor_bytes(item))
|
|
elif isinstance(item, dict):
|
|
digest.update(f"{path}|dict|{len(item)}\n".encode())
|
|
for key in sorted(item, key=lambda candidate: str(candidate)):
|
|
visit(f"{path}/{key}", item[key])
|
|
elif isinstance(item, (list, tuple)):
|
|
digest.update(f"{path}|sequence|{len(item)}\n".encode())
|
|
for index, child in enumerate(item):
|
|
visit(f"{path}/{index}", child)
|
|
else:
|
|
digest.update(f"{path}|scalar|{repr(item)}\n".encode())
|
|
|
|
visit("root", value)
|
|
return digest.hexdigest()
|
|
|
|
|
|
@dataclass
|
|
class ActivationTrace:
|
|
block_outputs: list[torch.Tensor]
|
|
layer_input_rms: list[float]
|
|
branch_output_rms: list[float]
|
|
stream_state_rms: list[float]
|
|
depth_weights: list[dict[str, Any]]
|
|
output_weights: dict[str, Any] | None = None
|
|
|
|
|
|
def rms(value: torch.Tensor) -> float:
|
|
return value.float().square().mean().sqrt().detach().cpu().item()
|
|
|
|
|
|
class GradientLanguageModel(round04.ReducedLanguageModel):
|
|
"""Round 04 trunk with aligned post-MLP activation capture."""
|
|
|
|
def forward(
|
|
self, input_ids: torch.Tensor, capture: bool = False
|
|
) -> tuple[torch.Tensor, ActivationTrace | None]:
|
|
embedded = self.embed(input_ids)
|
|
trace = ActivationTrace([], [], [], [], []) if capture else None
|
|
|
|
if self.architecture == "baseline":
|
|
hidden = embedded
|
|
for block in self.blocks:
|
|
attention_input = hidden
|
|
attention_output = block.attention(block.attention_norm(attention_input))
|
|
hidden = hidden + attention_output
|
|
if trace is not None:
|
|
trace.layer_input_rms.append(rms(attention_input))
|
|
trace.branch_output_rms.append(rms(attention_output))
|
|
trace.stream_state_rms.append(rms(hidden))
|
|
mlp_input = hidden
|
|
mlp_output = block.mlp(block.mlp_norm(mlp_input))
|
|
hidden = hidden + mlp_output
|
|
if trace is not None:
|
|
hidden.retain_grad()
|
|
trace.block_outputs.append(hidden)
|
|
trace.layer_input_rms.append(rms(mlp_input))
|
|
trace.branch_output_rms.append(rms(mlp_output))
|
|
trace.stream_state_rms.append(rms(hidden))
|
|
else:
|
|
completed = [embedded]
|
|
partial: torch.Tensor | None = None
|
|
mixer_index = 0
|
|
for block in self.blocks:
|
|
for branch_index in range(2):
|
|
sources = completed + ([] if partial is None else [partial])
|
|
branch_input, weights = self.mixers[mixer_index](
|
|
sources, capture
|
|
)
|
|
mixer_index += 1
|
|
if branch_index == 0:
|
|
branch_output = block.attention(
|
|
block.attention_norm(branch_input)
|
|
)
|
|
else:
|
|
branch_output = block.mlp(block.mlp_norm(branch_input))
|
|
branch_for_residual = branch_output.float()
|
|
partial = (
|
|
branch_for_residual
|
|
if partial is None
|
|
else partial + branch_for_residual
|
|
)
|
|
if trace is not None:
|
|
trace.layer_input_rms.append(rms(branch_input))
|
|
trace.branch_output_rms.append(rms(branch_output))
|
|
trace.stream_state_rms.append(rms(partial))
|
|
trace.depth_weights.append(weights or {})
|
|
if branch_index == 1:
|
|
partial.retain_grad()
|
|
trace.block_outputs.append(partial)
|
|
if mixer_index % round04.SUBLAYERS_PER_BLOCK == 0:
|
|
completed.append(partial)
|
|
partial = None
|
|
if partial is not None or len(completed) != BLOCK_GROUPS + 1:
|
|
raise RuntimeError("Block AttnRes aggregation contract failed")
|
|
if self.output_mixer is None:
|
|
raise RuntimeError("Block AttnRes output mixer missing")
|
|
hidden, output_weights = self.output_mixer(completed, capture)
|
|
if trace is not None:
|
|
trace.output_weights = output_weights
|
|
|
|
normalized = self.final_norm(hidden)
|
|
logits = F.linear(normalized, self.token_embedding.weight)
|
|
return logits, trace
|
|
|
|
|
|
def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
|
return F.cross_entropy(
|
|
logits.float().reshape(-1, VOCABULARY), targets.reshape(-1)
|
|
)
|
|
|
|
|
|
def learning_rate(step: int, total_steps: int) -> float:
|
|
if step <= WARMUP_STEPS:
|
|
return PEAK_LR * step / WARMUP_STEPS
|
|
progress = (step - WARMUP_STEPS) / max(1, total_steps - WARMUP_STEPS)
|
|
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
|
return MIN_LR + (PEAK_LR - MIN_LR) * cosine
|
|
|
|
|
|
@torch.no_grad()
|
|
def evaluate(
|
|
model: GradientLanguageModel,
|
|
corpus: Any,
|
|
window_count: int,
|
|
eval_batch_size: int,
|
|
) -> dict[str, float]:
|
|
model.eval()
|
|
loss_sum = 0.0
|
|
target_count = 0
|
|
for begin in range(0, window_count, eval_batch_size):
|
|
end = min(begin + eval_batch_size, window_count)
|
|
inputs, targets = corpus.fixed_batch(
|
|
corpus.validation_starts, begin, end
|
|
)
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
logits, _ = model(inputs)
|
|
loss = F.cross_entropy(
|
|
logits.float().reshape(-1, VOCABULARY),
|
|
targets.reshape(-1),
|
|
reduction="sum",
|
|
)
|
|
loss_sum += loss.detach().cpu().item()
|
|
target_count += targets.numel()
|
|
nats = loss_sum / target_count
|
|
return {"cross_entropy_nats": nats, "bits_per_byte": nats / math.log(2)}
|
|
|
|
|
|
def mean(values: Iterable[float]) -> float:
|
|
return statistics.fmean(values)
|
|
|
|
|
|
def depth_statistics(values: list[float]) -> dict[str, Any]:
|
|
average = mean(values)
|
|
variance = mean((value - average) ** 2 for value in values)
|
|
quartile = len(values) // 4
|
|
first = mean(values[:quartile])
|
|
last = mean(values[-quartile:])
|
|
ratio = first / last
|
|
return {
|
|
"mean": average,
|
|
"population_cv": math.sqrt(variance) / average,
|
|
"normalized": [value / average for value in values],
|
|
"first_quartile_mean": first,
|
|
"last_quartile_mean": last,
|
|
"first_to_last_ratio": ratio,
|
|
"imbalance_abs_log_ratio": abs(math.log(ratio)),
|
|
}
|
|
|
|
|
|
def core_parameter_gradient_rms(model: GradientLanguageModel) -> list[float]:
|
|
values = []
|
|
for block in model.blocks:
|
|
sum_square = 0.0
|
|
count = 0
|
|
for parameter in block.parameters():
|
|
if parameter.grad is None:
|
|
raise RuntimeError("missing core parameter gradient")
|
|
gradient = parameter.grad.detach().float()
|
|
if not torch.isfinite(gradient).all():
|
|
raise RuntimeError("non-finite core parameter gradient")
|
|
sum_square += gradient.square().sum().detach().cpu().item()
|
|
count += gradient.numel()
|
|
values.append(math.sqrt(sum_square / count))
|
|
return values
|
|
|
|
|
|
def activation_storage_unique(outputs: list[torch.Tensor]) -> bool:
|
|
pointers = [output.untyped_storage().data_ptr() for output in outputs]
|
|
return len(pointers) == len(set(pointers))
|
|
|
|
|
|
def diagnostic(
|
|
model: GradientLanguageModel,
|
|
corpus: Any,
|
|
window_count: int,
|
|
*,
|
|
loss_scale: float = 1.0,
|
|
) -> dict[str, Any]:
|
|
model.eval()
|
|
model.zero_grad(set_to_none=True)
|
|
inputs, targets = corpus.fixed_batch(
|
|
corpus.diagnostic_starts, 0, window_count
|
|
)
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
logits, trace = model(inputs, capture=True)
|
|
unscaled_loss = cross_entropy(logits, targets)
|
|
loss = unscaled_loss * loss_scale
|
|
if trace is None or len(trace.block_outputs) != len(model.blocks):
|
|
raise RuntimeError("aligned activation capture count mismatch")
|
|
expected_shape = (window_count, CONTEXT, round04.D_MODEL)
|
|
if any(tuple(output.shape) != expected_shape for output in trace.block_outputs):
|
|
raise RuntimeError("aligned activation capture shape mismatch")
|
|
if any(output.dtype != torch.float32 for output in trace.block_outputs):
|
|
raise RuntimeError("aligned activation capture must use FP32 residual state")
|
|
if not activation_storage_unique(trace.block_outputs):
|
|
raise RuntimeError("captured block outputs alias storage")
|
|
loss.backward()
|
|
|
|
activation_grad_rms = []
|
|
activation_output_rms = []
|
|
activation_dtypes = []
|
|
for output in trace.block_outputs:
|
|
if output.grad is None:
|
|
raise RuntimeError("captured activation gradient is None")
|
|
gradient = output.grad.detach().float()
|
|
if not torch.isfinite(gradient).all():
|
|
raise RuntimeError("captured activation gradient is non-finite")
|
|
activation_grad_rms.append(
|
|
gradient.square().mean().sqrt().detach().cpu().item()
|
|
)
|
|
activation_output_rms.append(rms(output))
|
|
activation_dtypes.append(str(output.dtype))
|
|
parameter_grad_rms = core_parameter_gradient_rms(model)
|
|
|
|
return {
|
|
"loss_nats": unscaled_loss.detach().cpu().item(),
|
|
"bits_per_byte": unscaled_loss.detach().cpu().item() / math.log(2),
|
|
"loss_scale": loss_scale,
|
|
"capture": {
|
|
"count": len(trace.block_outputs),
|
|
"shape": list(expected_shape),
|
|
"dtypes": activation_dtypes,
|
|
"all_gradients_finite": True,
|
|
"all_gradients_present": True,
|
|
"storage_unique": True,
|
|
"position": (
|
|
"post-MLP Transformer-block output; Block AttnRes is captured "
|
|
"before aggregation-partial reset"
|
|
),
|
|
},
|
|
"activation_grad_rms_by_block": activation_grad_rms,
|
|
"activation_grad_statistics": depth_statistics(activation_grad_rms),
|
|
"activation_output_rms_by_block": activation_output_rms,
|
|
"activation_output_statistics": depth_statistics(activation_output_rms),
|
|
"core_parameter_grad_rms_by_block": parameter_grad_rms,
|
|
"core_parameter_grad_statistics": depth_statistics(parameter_grad_rms),
|
|
"layer_input_rms_by_sublayer": trace.layer_input_rms,
|
|
"branch_output_rms_by_sublayer": trace.branch_output_rms,
|
|
"stream_state_rms_by_sublayer": trace.stream_state_rms,
|
|
"depth_weights": trace.depth_weights,
|
|
"output_weights": trace.output_weights,
|
|
}
|
|
|
|
|
|
def loss_scale_gate(
|
|
model: GradientLanguageModel, corpus: Any, window_count: int
|
|
) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
base = diagnostic(model, corpus, window_count, loss_scale=1.0)
|
|
doubled = diagnostic(model, corpus, window_count, loss_scale=2.0)
|
|
base_values = base["activation_grad_rms_by_block"]
|
|
doubled_values = doubled["activation_grad_rms_by_block"]
|
|
ratios = [
|
|
doubled_value / base_value
|
|
for base_value, doubled_value in zip(base_values, doubled_values)
|
|
]
|
|
base_stats = base["activation_grad_statistics"]
|
|
doubled_stats = doubled["activation_grad_statistics"]
|
|
cv_delta = abs(
|
|
doubled_stats["population_cv"] - base_stats["population_cv"]
|
|
)
|
|
ratio_delta = abs(
|
|
doubled_stats["first_to_last_ratio"]
|
|
- base_stats["first_to_last_ratio"]
|
|
)
|
|
normalized_max_delta = max(
|
|
abs(left - right)
|
|
for left, right in zip(
|
|
base_stats["normalized"], doubled_stats["normalized"]
|
|
)
|
|
)
|
|
passed = (
|
|
all(abs(ratio - 2.0) <= 1e-5 for ratio in ratios)
|
|
and cv_delta <= 1e-6
|
|
and ratio_delta <= 1e-6
|
|
and normalized_max_delta <= 1e-6
|
|
)
|
|
gate = {
|
|
"passed": passed,
|
|
"per_block_scale_ratios": ratios,
|
|
"max_abs_scale_ratio_error": max(abs(ratio - 2.0) for ratio in ratios),
|
|
"population_cv_abs_delta": cv_delta,
|
|
"first_to_last_ratio_abs_delta": ratio_delta,
|
|
"normalized_spectrum_max_abs_delta": normalized_max_delta,
|
|
"thresholds": {
|
|
"scale_ratio_abs": 1e-5,
|
|
"shape_abs": 1e-6,
|
|
},
|
|
}
|
|
if not passed:
|
|
raise RuntimeError(f"loss-scale diagnostic gate failed: {gate}")
|
|
return base, gate
|
|
|
|
|
|
def percentile(values: list[float], quantile: float) -> float:
|
|
return float(np.quantile(np.asarray(values, dtype=np.float64), quantile))
|
|
|
|
|
|
def parameter_inventory(model: GradientLanguageModel) -> dict[str, int]:
|
|
total = sum(parameter.numel() for parameter in model.parameters())
|
|
mixer = sum(
|
|
parameter.numel()
|
|
for name, parameter in model.named_parameters()
|
|
if name.startswith("mixers.") or name.startswith("output_mixer.")
|
|
)
|
|
return {
|
|
"total": total,
|
|
"core": total - mixer,
|
|
"mixer": mixer,
|
|
"embedding": (
|
|
model.token_embedding.weight.numel()
|
|
+ model.position_embedding.weight.numel()
|
|
),
|
|
}
|
|
|
|
|
|
def model_input_gate_hashes(
|
|
corpus: Any, manifest: dict[str, Any], seed: int, batch_size: int
|
|
) -> dict[str, str]:
|
|
values: dict[str, str] = {}
|
|
for step in manifest["windows"]["gate_steps"]:
|
|
raw_digest = hashlib.sha256()
|
|
for row in range(batch_size):
|
|
start = round04.window_start(seed, step, row, len(corpus.train))
|
|
raw_digest.update(
|
|
np.asarray(
|
|
corpus.train[start : start + CONTEXT + 1], dtype=np.uint8
|
|
).tobytes()
|
|
)
|
|
expected_raw_hash = manifest["windows"][
|
|
"gate_training_tensor_sha256"
|
|
][str(seed)][str(step)]
|
|
if raw_digest.hexdigest() != expected_raw_hash:
|
|
raise RuntimeError(f"manifest gate tensor mismatch at step {step}")
|
|
inputs, targets = corpus.training_batch(seed, step, batch_size)
|
|
digest = hashlib.sha256()
|
|
digest.update(tensor_bytes(inputs))
|
|
digest.update(tensor_bytes(targets))
|
|
values[str(step)] = digest.hexdigest()
|
|
return values
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is required by the frozen protocol")
|
|
if args.seed not in EXPECTED_SEEDS:
|
|
raise ValueError(f"seed is not preregistered: {args.seed}")
|
|
configure_round04_globals(args.depth)
|
|
configure_determinism(args.seed)
|
|
device = torch.device("cuda")
|
|
|
|
manifest = json.loads(args.manifest.read_text())
|
|
if manifest["protocol_id"] != PROTOCOL_ID:
|
|
raise ValueError("manifest protocol mismatch")
|
|
if manifest["windows"]["formal_steps"] != FORMAL_STEPS:
|
|
raise ValueError("manifest formal-step mismatch")
|
|
corpus = round04.ByteCorpus(args.cache_dir, manifest, device)
|
|
|
|
model = GradientLanguageModel(args.architecture).to(device)
|
|
public_structure_hash, public_tensors, public_elements = state_structure_hash(
|
|
model, include_mixers=False
|
|
)
|
|
initial_public_hash = named_state_hash(model, include_mixers=False)
|
|
initial_mixer_hash = (
|
|
named_state_hash(model, include_mixers=True)
|
|
if args.architecture == "block"
|
|
else None
|
|
)
|
|
input_gate_hashes = model_input_gate_hashes(
|
|
corpus, manifest, args.seed, args.batch_size
|
|
)
|
|
|
|
decay_parameters: list[nn.Parameter] = []
|
|
no_decay_parameters: list[nn.Parameter] = []
|
|
for parameter in model.parameters():
|
|
(decay_parameters if parameter.ndim >= 2 else no_decay_parameters).append(
|
|
parameter
|
|
)
|
|
optimizer = torch.optim.AdamW(
|
|
[
|
|
{"params": decay_parameters, "weight_decay": WEIGHT_DECAY},
|
|
{"params": no_decay_parameters, "weight_decay": 0.0},
|
|
],
|
|
lr=PEAK_LR,
|
|
betas=BETAS,
|
|
eps=ADAM_EPS,
|
|
)
|
|
|
|
evaluation_steps = sorted(
|
|
set(step for step in DIAGNOSTIC_STEPS if step <= args.steps)
|
|
| {0, args.steps}
|
|
)
|
|
evaluations = [
|
|
{
|
|
"step": 0,
|
|
**evaluate(
|
|
model, corpus, args.validation_windows, args.eval_batch_size
|
|
),
|
|
}
|
|
]
|
|
if args.run_kind == "smoke":
|
|
initial_diagnostic, gradient_gate = loss_scale_gate(
|
|
model, corpus, args.diagnostic_windows
|
|
)
|
|
else:
|
|
initial_diagnostic = diagnostic(
|
|
model, corpus, args.diagnostic_windows
|
|
)
|
|
gradient_gate = None
|
|
diagnostics = [{"step": 0, **initial_diagnostic}]
|
|
model.zero_grad(set_to_none=True)
|
|
|
|
training_history: list[dict[str, float | int]] = []
|
|
step_times: list[float] = []
|
|
model.train()
|
|
for step in range(1, args.steps + 1):
|
|
lr = learning_rate(step, args.steps)
|
|
for group in optimizer.param_groups:
|
|
group["lr"] = lr
|
|
inputs, targets = corpus.training_batch(args.seed, step, args.batch_size)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
|
|
torch.cuda.synchronize()
|
|
started = time.perf_counter()
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
logits, _ = model(inputs)
|
|
loss = cross_entropy(logits, targets)
|
|
if not torch.isfinite(loss):
|
|
raise RuntimeError(f"non-finite loss at step {step}: {loss}")
|
|
loss.backward()
|
|
unclipped_norm = torch.nn.utils.clip_grad_norm_(
|
|
model.parameters(), GRAD_CLIP
|
|
)
|
|
optimizer.step()
|
|
torch.cuda.synchronize()
|
|
elapsed_ms = (time.perf_counter() - started) * 1000
|
|
|
|
if step == args.timing_warmup:
|
|
torch.cuda.reset_peak_memory_stats()
|
|
elif step > args.timing_warmup:
|
|
step_times.append(elapsed_ms)
|
|
if step == 1 or step % 10 == 0 or step == args.steps:
|
|
training_history.append(
|
|
{
|
|
"step": step,
|
|
"loss_nats": loss.detach().cpu().item(),
|
|
"bits_per_byte": loss.detach().cpu().item() / math.log(2),
|
|
"learning_rate": lr,
|
|
"unclipped_grad_norm": float(unclipped_norm.detach().cpu()),
|
|
}
|
|
)
|
|
|
|
if step in evaluation_steps and step != 0:
|
|
evaluations.append(
|
|
{
|
|
"step": step,
|
|
**evaluate(
|
|
model,
|
|
corpus,
|
|
args.validation_windows,
|
|
args.eval_batch_size,
|
|
),
|
|
}
|
|
)
|
|
diagnostics.append(
|
|
{
|
|
"step": step,
|
|
**diagnostic(
|
|
model, corpus, args.diagnostic_windows
|
|
),
|
|
}
|
|
)
|
|
model.zero_grad(set_to_none=True)
|
|
model.train()
|
|
|
|
training_peak_allocated = torch.cuda.max_memory_allocated()
|
|
training_peak_reserved = torch.cuda.max_memory_reserved()
|
|
final_public_hash = named_state_hash(model, include_mixers=False)
|
|
final_mixer_hash = (
|
|
named_state_hash(model, include_mixers=True)
|
|
if args.architecture == "block"
|
|
else None
|
|
)
|
|
final_full_hash = named_state_hash(model, include_mixers=None)
|
|
optimizer_hash = recursive_state_hash(optimizer.state_dict())
|
|
timing = {
|
|
"warmup_steps_excluded": args.timing_warmup,
|
|
"measured_steps": len(step_times),
|
|
"mean_ms": mean(step_times) if step_times else None,
|
|
"median_ms": statistics.median(step_times) if step_times else None,
|
|
"p95_ms": percentile(step_times, 0.95) if step_times else None,
|
|
"peak_allocated_bytes": training_peak_allocated,
|
|
"peak_reserved_bytes": training_peak_reserved,
|
|
}
|
|
result = {
|
|
"schema_version": 1,
|
|
"protocol_id": PROTOCOL_ID,
|
|
"run_kind": args.run_kind,
|
|
"architecture": args.architecture,
|
|
"depth": args.depth,
|
|
"seed": args.seed,
|
|
"steps": args.steps,
|
|
"batch_size": args.batch_size,
|
|
"target_bytes_seen": args.steps * args.batch_size * CONTEXT,
|
|
"manifest": {
|
|
"path": str(args.manifest),
|
|
"file_sha256": file_sha256(args.manifest),
|
|
"formal_schedule_sha256": manifest["windows"][
|
|
"formal_schedule_sha256"
|
|
],
|
|
"validation_tensor_sha256": manifest["windows"][
|
|
"validation_tensor_sha256"
|
|
],
|
|
"diagnostic_tensor_sha256": manifest["windows"][
|
|
"diagnostic_tensor_sha256"
|
|
],
|
|
"input_gate_tensor_hashes": input_gate_hashes,
|
|
},
|
|
"model": {
|
|
"layers": args.depth,
|
|
"sublayers": args.depth * 2,
|
|
"attnres_aggregation_groups": BLOCK_GROUPS,
|
|
"sublayers_per_attnres_group": args.depth * 2 // BLOCK_GROUPS,
|
|
"transformer_blocks_per_attnres_group": args.depth // BLOCK_GROUPS,
|
|
"d_model": round04.D_MODEL,
|
|
"heads": round04.HEADS,
|
|
"d_head": round04.D_HEAD,
|
|
"d_ff": round04.D_FF,
|
|
"context": CONTEXT,
|
|
"vocabulary": VOCABULARY,
|
|
"parameters": parameter_inventory(model),
|
|
},
|
|
"optimizer": {
|
|
"name": "AdamW",
|
|
"betas": list(BETAS),
|
|
"epsilon": ADAM_EPS,
|
|
"weight_decay_ndim_ge_2": WEIGHT_DECAY,
|
|
"peak_lr": PEAK_LR,
|
|
"min_lr": MIN_LR,
|
|
"warmup_steps": WARMUP_STEPS,
|
|
"grad_clip": GRAD_CLIP,
|
|
},
|
|
"hashes": {
|
|
"initial_public_parameter_structure": public_structure_hash,
|
|
"initial_public_parameter_tensors": public_tensors,
|
|
"initial_public_parameter_elements": public_elements,
|
|
"initial_public_parameters": initial_public_hash,
|
|
"initial_mixer_parameters": initial_mixer_hash,
|
|
"final_public_parameters": final_public_hash,
|
|
"final_mixer_parameters": final_mixer_hash,
|
|
"final_model_state": final_full_hash,
|
|
"final_optimizer_state": optimizer_hash,
|
|
},
|
|
"evaluations": evaluations,
|
|
"diagnostics": diagnostics,
|
|
"training_history": training_history,
|
|
"gradient_gate": gradient_gate,
|
|
"timing": timing,
|
|
"environment": {
|
|
"python": platform.python_version(),
|
|
"torch": torch.__version__,
|
|
"cuda": torch.version.cuda,
|
|
"gpu": torch.cuda.get_device_name(0),
|
|
"compute_capability": list(torch.cuda.get_device_capability(0)),
|
|
"cublas_workspace_config": os.environ["CUBLAS_WORKSPACE_CONFIG"],
|
|
"deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
|
|
"autocast": "cuda-bfloat16-forward-fp32-cross-entropy",
|
|
"compile": False,
|
|
},
|
|
}
|
|
result["canonical_sha256_without_self"] = canonical_sha256(result)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = args.output.with_suffix(args.output.suffix + ".tmp")
|
|
temporary.write_text(
|
|
json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
|
)
|
|
os.replace(temporary, args.output)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"output": str(args.output),
|
|
"run_kind": args.run_kind,
|
|
"architecture": args.architecture,
|
|
"depth": args.depth,
|
|
"seed": args.seed,
|
|
"steps": args.steps,
|
|
"final_bpc": evaluations[-1]["bits_per_byte"],
|
|
"final_activation_gradient_cv": diagnostics[-1][
|
|
"activation_grad_statistics"
|
|
]["population_cv"],
|
|
"canonical_sha256": result["canonical_sha256_without_self"],
|
|
"timing": timing,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|