wuyang
716 lines
26 KiB
Python
716 lines
26 KiB
Python
#!/usr/bin/env python3
|
|
"""Train one frozen residual variant for the reduced Attention Residuals study."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import platform
|
|
import statistics
|
|
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-reduced-v1"
|
|
ARCHITECTURES = ("baseline", "full", "block")
|
|
EXPECTED_SEEDS = (2026073001, 2026073002, 2026073003)
|
|
EVAL_STEPS = (0, 100, 250, 500, 1000, 1500, 2000)
|
|
VOCABULARY = 256
|
|
CONTEXT = 256
|
|
LAYERS = 16
|
|
SUBLAYERS = LAYERS * 2
|
|
BLOCKS = 8
|
|
SUBLAYERS_PER_BLOCK = SUBLAYERS // BLOCKS
|
|
D_MODEL = 192
|
|
HEADS = 6
|
|
D_HEAD = D_MODEL // HEADS
|
|
D_FF = 768
|
|
RMS_EPS = 1e-6
|
|
PEAK_LR = 3e-4
|
|
MIN_LR = 3e-5
|
|
WARMUP_STEPS = 100
|
|
WEIGHT_DECAY = 0.1
|
|
BETAS = (0.9, 0.95)
|
|
ADAM_EPS = 1e-8
|
|
GRAD_CLIP = 1.0
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--architecture", choices=ARCHITECTURES, required=True)
|
|
parser.add_argument("--seed", type=int, required=True)
|
|
parser.add_argument("--steps", type=int, default=2000)
|
|
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")
|
|
return parser.parse_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_json_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()
|
|
header = f"{value.dtype}|{tuple(value.shape)}|".encode()
|
|
return header + value.view(torch.uint8).numpy().tobytes()
|
|
|
|
|
|
def 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 window_start(seed: int, step: int, row: int, corpus_length: int) -> int:
|
|
payload = "\0".join(
|
|
[PROTOCOL_ID, "train-window", str(seed), str(step), str(row)]
|
|
).encode()
|
|
value = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
|
|
return value % (corpus_length - (CONTEXT + 1))
|
|
|
|
|
|
class ByteCorpus:
|
|
def __init__(self, cache_dir: Path, manifest: dict[str, Any], device: torch.device):
|
|
self.device = device
|
|
self.train = np.memmap(cache_dir / "train.bin", dtype=np.uint8, mode="r")
|
|
self.validation = np.memmap(
|
|
cache_dir / "validation.bin", dtype=np.uint8, mode="r"
|
|
)
|
|
self.validation_starts = manifest["windows"]["validation_starts"]
|
|
self.diagnostic_starts = manifest["windows"]["diagnostic_starts"]
|
|
|
|
def training_batch(
|
|
self, seed: int, step: int, batch_size: int
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
rows = np.empty((batch_size, CONTEXT + 1), dtype=np.int64)
|
|
for row in range(batch_size):
|
|
start = window_start(seed, step, row, len(self.train))
|
|
rows[row] = self.train[start : start + CONTEXT + 1]
|
|
tensor = torch.from_numpy(rows).to(self.device, non_blocking=False)
|
|
return tensor[:, :-1], tensor[:, 1:]
|
|
|
|
def fixed_batch(
|
|
self, starts: list[int], begin: int, end: int
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
chosen = starts[begin:end]
|
|
rows = np.empty((len(chosen), CONTEXT + 1), dtype=np.int64)
|
|
for row, start in enumerate(chosen):
|
|
rows[row] = self.validation[start : start + CONTEXT + 1]
|
|
tensor = torch.from_numpy(rows).to(self.device, non_blocking=False)
|
|
return tensor[:, :-1], tensor[:, 1:]
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dimension: int):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(dimension))
|
|
|
|
def forward(self, value: torch.Tensor) -> torch.Tensor:
|
|
normalized = value.float() * torch.rsqrt(
|
|
value.float().square().mean(dim=-1, keepdim=True) + RMS_EPS
|
|
)
|
|
return normalized.to(value.dtype) * self.weight
|
|
|
|
|
|
class CausalAttention(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.qkv = nn.Linear(D_MODEL, 3 * D_MODEL, bias=False)
|
|
self.o_proj = nn.Linear(D_MODEL, D_MODEL, bias=False)
|
|
mask = torch.triu(torch.ones(CONTEXT, CONTEXT, dtype=torch.bool), diagonal=1)
|
|
self.register_buffer("causal_mask", mask, persistent=False)
|
|
|
|
def forward(self, value: torch.Tensor) -> torch.Tensor:
|
|
batch, sequence, _ = value.shape
|
|
qkv = self.qkv(value).view(batch, sequence, 3, HEADS, D_HEAD)
|
|
query, key, content = qkv.unbind(dim=2)
|
|
query = query.transpose(1, 2)
|
|
key = key.transpose(1, 2)
|
|
content = content.transpose(1, 2)
|
|
scores = torch.matmul(query, key.transpose(-1, -2)).float() / math.sqrt(D_HEAD)
|
|
scores = scores.masked_fill(
|
|
self.causal_mask[:sequence, :sequence], float("-inf")
|
|
)
|
|
probabilities = torch.softmax(scores, dim=-1).to(query.dtype)
|
|
mixed = torch.matmul(probabilities, content)
|
|
mixed = mixed.transpose(1, 2).contiguous().view(batch, sequence, D_MODEL)
|
|
return self.o_proj(mixed)
|
|
|
|
|
|
class SwiGLU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gate = nn.Linear(D_MODEL, D_FF, bias=False)
|
|
self.up = nn.Linear(D_MODEL, D_FF, bias=False)
|
|
self.down = nn.Linear(D_FF, D_MODEL, bias=False)
|
|
|
|
def forward(self, value: torch.Tensor) -> torch.Tensor:
|
|
return self.down(F.silu(self.gate(value)) * self.up(value))
|
|
|
|
|
|
class TransformerBlock(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.attention_norm = RMSNorm(D_MODEL)
|
|
self.attention = CausalAttention()
|
|
self.mlp_norm = RMSNorm(D_MODEL)
|
|
self.mlp = SwiGLU()
|
|
|
|
|
|
class DepthMixer(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.query = nn.Parameter(torch.zeros(D_MODEL))
|
|
self.key_norm = RMSNorm(D_MODEL)
|
|
|
|
def forward(
|
|
self, sources: list[torch.Tensor], capture: bool = False
|
|
) -> tuple[torch.Tensor, dict[str, Any] | None]:
|
|
values = torch.stack(sources, dim=0)
|
|
keys = self.key_norm(values)
|
|
logits = torch.einsum("d,nbtd->nbt", self.query, keys.float())
|
|
weights = torch.softmax(logits, dim=0)
|
|
output = torch.einsum("nbt,nbtd->btd", weights, values.float()).to(
|
|
values.dtype
|
|
)
|
|
if not capture:
|
|
return output, None
|
|
entropy = -(weights * torch.log(weights.clamp_min(1e-30))).sum(dim=0)
|
|
return output, {
|
|
"mean_weights": weights.mean(dim=(1, 2)).detach().cpu().tolist(),
|
|
"entropy_mean": entropy.mean().detach().cpu().item(),
|
|
"sources": len(sources),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class TraceAccumulator:
|
|
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 ReducedLanguageModel(nn.Module):
|
|
def __init__(self, architecture: str):
|
|
super().__init__()
|
|
self.architecture = architecture
|
|
self.token_embedding = nn.Embedding(VOCABULARY, D_MODEL)
|
|
self.position_embedding = nn.Embedding(CONTEXT, D_MODEL)
|
|
self.blocks = nn.ModuleList([TransformerBlock() for _ in range(LAYERS)])
|
|
self.final_norm = RMSNorm(D_MODEL)
|
|
if architecture == "baseline":
|
|
self.mixers = nn.ModuleList()
|
|
self.output_mixer = None
|
|
else:
|
|
self.mixers = nn.ModuleList([DepthMixer() for _ in range(SUBLAYERS)])
|
|
self.output_mixer = DepthMixer()
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
for module in self.modules():
|
|
if isinstance(module, nn.Embedding):
|
|
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
|
elif isinstance(module, nn.Linear):
|
|
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
|
elif isinstance(module, RMSNorm):
|
|
nn.init.ones_(module.weight)
|
|
scaled = 0.02 / math.sqrt(2 * LAYERS)
|
|
for block in self.blocks:
|
|
nn.init.normal_(block.attention.o_proj.weight, mean=0.0, std=scaled)
|
|
nn.init.normal_(block.mlp.down.weight, mean=0.0, std=scaled)
|
|
for mixer in self.mixers:
|
|
nn.init.zeros_(mixer.query)
|
|
nn.init.ones_(mixer.key_norm.weight)
|
|
if self.output_mixer is not None:
|
|
nn.init.zeros_(self.output_mixer.query)
|
|
nn.init.ones_(self.output_mixer.key_norm.weight)
|
|
|
|
def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
positions = torch.arange(input_ids.shape[1], device=input_ids.device)
|
|
return self.token_embedding(input_ids) + self.position_embedding(positions)[None]
|
|
|
|
def forward(
|
|
self, input_ids: torch.Tensor, capture: bool = False
|
|
) -> tuple[torch.Tensor, TraceAccumulator | None]:
|
|
embedded = self.embed(input_ids)
|
|
trace = (
|
|
TraceAccumulator([], [], [], [])
|
|
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:
|
|
trace.layer_input_rms.append(rms(mlp_input))
|
|
trace.branch_output_rms.append(rms(mlp_output))
|
|
trace.stream_state_rms.append(rms(hidden))
|
|
elif self.architecture == "full":
|
|
sources = [embedded]
|
|
mixer_index = 0
|
|
for block in self.blocks:
|
|
attention_input, weights = self.mixers[mixer_index](sources, capture)
|
|
mixer_index += 1
|
|
attention_output = block.attention(block.attention_norm(attention_input))
|
|
sources.append(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(torch.stack(sources, dim=0))
|
|
)
|
|
trace.depth_weights.append(weights or {})
|
|
mlp_input, weights = self.mixers[mixer_index](sources, capture)
|
|
mixer_index += 1
|
|
mlp_output = block.mlp(block.mlp_norm(mlp_input))
|
|
sources.append(mlp_output)
|
|
if trace is not None:
|
|
trace.layer_input_rms.append(rms(mlp_input))
|
|
trace.branch_output_rms.append(rms(mlp_output))
|
|
trace.stream_state_rms.append(
|
|
rms(torch.stack(sources, dim=0))
|
|
)
|
|
trace.depth_weights.append(weights or {})
|
|
assert self.output_mixer is not None
|
|
hidden, output_weights = self.output_mixer(sources, capture)
|
|
if trace is not None:
|
|
trace.output_weights = output_weights
|
|
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))
|
|
partial = (
|
|
branch_output if partial is None else partial + branch_output
|
|
)
|
|
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 mixer_index % SUBLAYERS_PER_BLOCK == 0:
|
|
completed.append(partial)
|
|
partial = None
|
|
assert partial is None
|
|
assert len(completed) == BLOCKS + 1
|
|
assert self.output_mixer is not None
|
|
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 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
|
|
|
|
|
|
def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
|
return F.cross_entropy(logits.float().view(-1, VOCABULARY), targets.view(-1))
|
|
|
|
|
|
@torch.no_grad()
|
|
def evaluate(
|
|
model: ReducedLanguageModel,
|
|
corpus: ByteCorpus,
|
|
starts: list[int],
|
|
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(starts, begin, end)
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
logits, _ = model(inputs)
|
|
loss = F.cross_entropy(
|
|
logits.float().view(-1, VOCABULARY),
|
|
targets.view(-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 percentile(values: list[float], quantile: float) -> float:
|
|
return float(np.quantile(np.asarray(values, dtype=np.float64), quantile))
|
|
|
|
|
|
def core_parameter_gradient_rms(model: ReducedLanguageModel) -> list[float]:
|
|
values = []
|
|
for block in model.blocks:
|
|
sum_square = 0.0
|
|
count = 0
|
|
for parameter in block.parameters():
|
|
if parameter.grad is None:
|
|
continue
|
|
gradient = parameter.grad.detach().float()
|
|
sum_square += gradient.square().sum().detach().cpu().item()
|
|
count += gradient.numel()
|
|
values.append(math.sqrt(sum_square / count))
|
|
return values
|
|
|
|
|
|
def diagnostic(
|
|
model: ReducedLanguageModel,
|
|
corpus: ByteCorpus,
|
|
window_count: int,
|
|
) -> 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)
|
|
loss = cross_entropy(logits, targets)
|
|
loss.backward()
|
|
gradients = core_parameter_gradient_rms(model)
|
|
assert trace is not None
|
|
return {
|
|
"loss_nats": loss.detach().cpu().item(),
|
|
"bits_per_byte": loss.detach().cpu().item() / math.log(2),
|
|
"layer_input_rms": trace.layer_input_rms,
|
|
"branch_output_rms": trace.branch_output_rms,
|
|
"stream_state_rms": trace.stream_state_rms,
|
|
"core_parameter_grad_rms_by_block": gradients,
|
|
"depth_weights": trace.depth_weights,
|
|
"output_weights": trace.output_weights,
|
|
}
|
|
|
|
|
|
def parameter_inventory(model: ReducedLanguageModel) -> 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.")
|
|
)
|
|
embedding = model.token_embedding.weight.numel() + model.position_embedding.weight.numel()
|
|
return {
|
|
"total": total,
|
|
"core": total - mixer,
|
|
"mixer": mixer,
|
|
"embedding": embedding,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is required by the frozen protocol")
|
|
if args.run_kind != "smoke" and args.seed not in EXPECTED_SEEDS:
|
|
raise ValueError(f"formal/replay seed is not preregistered: {args.seed}")
|
|
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["dataset"]["revision"] != (
|
|
"b08601e04326c79dfdd32d625aee71d232d685c3"
|
|
):
|
|
raise ValueError("dataset revision mismatch")
|
|
corpus = ByteCorpus(args.cache_dir, manifest, device)
|
|
|
|
model = ReducedLanguageModel(args.architecture).to(device)
|
|
initial_common_hash = state_hash(model, include_mixers=False)
|
|
initial_mixer_hash = (
|
|
state_hash(model, include_mixers=True)
|
|
if args.architecture != "baseline"
|
|
else None
|
|
)
|
|
inventory = parameter_inventory(model)
|
|
|
|
decay_parameters: list[nn.Parameter] = []
|
|
no_decay_parameters: list[nn.Parameter] = []
|
|
for parameter in model.parameters():
|
|
if parameter.ndim >= 2:
|
|
decay_parameters.append(parameter)
|
|
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 EVAL_STEPS if step <= args.steps) | {0, args.steps}
|
|
)
|
|
evaluations = [
|
|
{
|
|
"step": 0,
|
|
**evaluate(
|
|
model,
|
|
corpus,
|
|
corpus.validation_starts,
|
|
args.validation_windows,
|
|
args.eval_batch_size,
|
|
),
|
|
}
|
|
]
|
|
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,
|
|
corpus.validation_starts,
|
|
args.validation_windows,
|
|
args.eval_batch_size,
|
|
),
|
|
}
|
|
)
|
|
model.train()
|
|
|
|
training_peak_allocated = torch.cuda.max_memory_allocated()
|
|
training_peak_reserved = torch.cuda.max_memory_reserved()
|
|
diagnostic_result = diagnostic(model, corpus, args.diagnostic_windows)
|
|
final_common_hash = state_hash(model, include_mixers=False)
|
|
final_mixer_hash = (
|
|
state_hash(model, include_mixers=True)
|
|
if args.architecture != "baseline"
|
|
else None
|
|
)
|
|
timing = {
|
|
"warmup_steps_excluded": args.timing_warmup,
|
|
"measured_steps": len(step_times),
|
|
"mean_ms": statistics.fmean(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,
|
|
"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"
|
|
],
|
|
},
|
|
"model": {
|
|
"layers": LAYERS,
|
|
"sublayers": SUBLAYERS,
|
|
"blocks_for_block_attnres": BLOCKS,
|
|
"sublayers_per_attnres_block": SUBLAYERS_PER_BLOCK,
|
|
"d_model": D_MODEL,
|
|
"heads": HEADS,
|
|
"d_head": D_HEAD,
|
|
"d_ff": D_FF,
|
|
"context": CONTEXT,
|
|
"vocabulary": VOCABULARY,
|
|
"parameters": inventory,
|
|
},
|
|
"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_common_parameters": initial_common_hash,
|
|
"initial_mixer_parameters": initial_mixer_hash,
|
|
"final_common_parameters": final_common_hash,
|
|
"final_mixer_parameters": final_mixer_hash,
|
|
},
|
|
"evaluations": evaluations,
|
|
"training_history": training_history,
|
|
"diagnostic": diagnostic_result,
|
|
"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",
|
|
"compile": False,
|
|
},
|
|
}
|
|
result["canonical_sha256_without_self"] = canonical_json_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),
|
|
"architecture": args.architecture,
|
|
"seed": args.seed,
|
|
"steps": args.steps,
|
|
"final_bpc": evaluations[-1]["bits_per_byte"],
|
|
"initial_common_hash": initial_common_hash,
|
|
"final_common_hash": final_common_hash,
|
|
"canonical_sha256": result["canonical_sha256_without_self"],
|
|
"timing": timing,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|