File size: 9,091 Bytes
e29f88c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """Stage 0 checks shared by the pretrained CLI and a small integration test."""
from collections.abc import Callable
from time import perf_counter
from typing import Any
import torch
from torch import Tensor
from torch.nn import functional as F
from transformers import Qwen2ForCausalLM
from .lora_utils import attach_recurrent_lora
from .model import RecurrentQwen
def validate_stage0(
base: Qwen2ForCausalLM,
cases: dict[str, dict[str, Tensor]],
*,
recurrent_start: int = 6,
recurrent_end: int = 18,
rank: int = 8,
alpha: int = 16,
loop_counts: tuple[int, ...] = (1, 2, 4),
atol: float = 1e-5,
rtol: float = 1e-5,
progress: Callable[[str], None] = print,
) -> dict[str, Any]:
"""Assert architecture contracts; return serializable measured evidence.
Mutates/freezes base and attaches fresh LoRA. No optimizer step is taken.
The supplied cases must be short, nonempty input_ids/attention_mask pairs;
equivalence compares every vocabulary logit at every sequence position.
"""
if not cases or not loop_counts or min(loop_counts) < 1 or max(loop_counts) < 2:
raise ValueError("Provide input cases and positive loop counts including a depth > 1")
base.eval()
base.config.use_cache = False
originals = tuple(base.parameters())
original_ids = {id(p) for p in originals}
original_layers = tuple(base.model.layers)
progress("Computing ordinary-Qwen reference logits")
with torch.no_grad():
references = {name: base(**inputs, use_cache=False).logits for name, inputs in cases.items()}
model = RecurrentQwen(base, recurrent_start, recurrent_end).eval()
assert tuple(model.recurrent.layers) == original_layers[recurrent_start:recurrent_end]
errors = []
for adapted in (False, True):
if adapted:
attach_recurrent_lora(model, rank=rank, alpha=alpha)
progress(f"Checking T=1 equivalence ({'with' if adapted else 'without'} LoRA)")
with torch.no_grad():
for name, inputs in cases.items():
actual = model(**inputs, logits_mode="all").logits
expected = references[name]
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
difference = (actual - expected).abs()
errors.append({
"case": name,
"lora": adapted,
"max_absolute_error": difference.max().item(),
"mean_absolute_error": difference.mean().item(),
"logits_shape": list(actual.shape),
})
del references
progress("Checking shared weights, loop observability, and configurable depth")
parameters_before = tuple(id(p) for p in model.parameters())
calls: list[tuple[int, tuple[int, ...], Tensor, Tensor]] = []
block_calls: list[str] = []
layer_calls: list[int] = []
handles = [model.recurrent.register_forward_hook(
lambda module, args, output: calls.append(
(id(module), tuple(id(p) for p in module.parameters()), args[0], output)
)
)]
for name in ("prelude", "coda"):
handles.append(getattr(model, name).register_forward_hook(
lambda module, args, output, name=name: block_calls.append(name)
))
for layer in model.recurrent.layers:
handles.append(layer.register_forward_hook(
lambda module, args, output: layer_calls.append(id(module))
))
inputs = next(iter(cases.values()))
batch, sequence = inputs["input_ids"].shape
device = inputs["input_ids"].device
allowed = inputs["input_ids"].unique()
if allowed.numel() < 2:
raise ValueError("Validation inputs must contain at least two distinct token IDs")
recurrent_ids = tuple(id(p) for p in model.recurrent.parameters())
trajectories = []
try:
with torch.no_grad():
for depth in loop_counts:
calls.clear()
block_calls.clear()
layer_calls.clear()
labels = allowed[torch.arange(depth, device=device) % allowed.numel()].expand(batch, -1)
start = perf_counter()
output = model(**inputs, num_loops=depth, return_hidden_states=True,
labels=labels, allowed_token_ids=allowed)
if device.type == "mps":
torch.mps.synchronize()
seconds = perf_counter() - start
assert len(calls) == depth
assert block_calls.count("prelude") == 1 and block_calls.count("coda") == depth
assert layer_calls == [id(layer) for layer in model.recurrent.layers] * depth
assert all(call[0] == id(model.recurrent) and call[1] == recurrent_ids for call in calls)
assert tuple(id(p) for p in model.parameters()) == parameters_before
assert calls[0][2] is output.initial_hidden_state
assert all(calls[t][2] is calls[t - 1][3] for t in range(1, depth))
assert len(output.hidden_states) == len(output.loop_logits) == depth
assert output.margins.shape == (batch, depth)
for hidden, logits in zip(output.hidden_states, output.loop_logits):
assert hidden.shape == (batch, sequence, model.config.hidden_size)
assert logits.shape == (batch, model.config.vocab_size)
assert torch.isfinite(hidden).all() and torch.isfinite(logits).all()
assert torch.isfinite(output.margins).all()
trajectories.append({
"num_loops": depth,
"seconds": seconds,
"recurrent_calls": len(calls),
"recurrent_layer_calls": len(layer_calls),
"prelude_calls": block_calls.count("prelude"),
"coda_calls": block_calls.count("coda"),
"predicted_token_ids": [scores.argmax(-1).tolist() for scores in output.loop_logits],
"diagnostic_labels": labels.tolist(),
"allowed_token_ids": allowed.tolist(),
"margins": output.margins.tolist(),
})
finally:
for handle in handles:
handle.remove()
calls.clear()
del output
progress("Checking two-loop backward through the frozen coda (no optimizer step)")
model.train()
model.zero_grad(set_to_none=True)
start = perf_counter()
output = model(**inputs, num_loops=2, return_hidden_states=True)
for hidden in output.hidden_states:
hidden.retain_grad()
# Arbitrary token labels test graph connectivity, not task competence.
loss = F.cross_entropy(output.logits.float(), inputs["input_ids"][:, 0])
loss.backward()
if device.type == "mps":
torch.mps.synchronize()
backward_seconds = perf_counter() - start
assert all(not p.requires_grad and p.grad is None for p in originals)
trainable = [(name, p) for name, p in model.named_parameters() if p.requires_grad]
assert len(trainable) == (recurrent_end - recurrent_start) * 2 * 2
assert {id(p) for p in model.parameters()} - original_ids == {id(p) for _, p in trainable}
gradients = {}
for name, parameter in trainable:
assert name.startswith("recurrent.") and (".lora_A." in name or ".lora_B." in name)
assert parameter.grad is not None and torch.isfinite(parameter.grad).all(), name
norm = parameter.grad.float().norm().item()
if ".lora_B." in name:
assert norm > 0, name
gradients[name] = {"elements": parameter.numel(), "gradient_l2": norm}
hidden_gradients = [hidden.grad.float().norm().item() for hidden in output.hidden_states]
assert all(value > 0 and torch.isfinite(torch.tensor(value)) for value in hidden_gradients)
result = {
"gate_passed": True,
"checks": {name: "PASS" for name in (
"t1_equivalence", "shared_weights", "gradient_scope", "observability", "configurable_depth"
)},
"split": {"prelude": [0, recurrent_start], "recurrent": [recurrent_start, recurrent_end],
"coda": [recurrent_end, len(original_layers)], "end_indices_exclusive": True},
"lora": {"rank": rank, "alpha": alpha, "targets": ["q_proj", "v_proj"], "dropout": 0.0},
"bridge": None,
"tolerance": {"atol": atol, "rtol": rtol},
"equivalence": errors,
"trajectories": trajectories,
"parameters": {"original": sum(p.numel() for p in originals),
"trainable": sum(p.numel() for _, p in trainable),
"trainable_tensors": len(trainable)},
"gradient_check": {"loss": loss.item(), "forward_backward_seconds": backward_seconds,
"hidden_gradient_l2": hidden_gradients, "parameters": gradients,
"frozen_parameters_with_grad": 0, "optimizer_steps": 0},
}
model.zero_grad(set_to_none=True)
return result
|