| """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() |
| |
| 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 |
|
|
