ScaleRLT: audit IQuest dual-attention replay

benchmarks/audit_iquest_replay.py

@@ -0,0 +1,212 @@
+"""Compare functional replay with unmodified official token-by-token autograd."""
+
+import argparse
+import copy
+import hashlib
+import importlib
+import json
+import sys
+import types
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+ROOT = Path(__file__).resolve().parents[1]
+SOURCE = ROOT / "benchmarks/results/iquest-contract/reference"
+sys.path.insert(0, str(ROOT))
+from vime_plugins.iquest.replay import replay
+
+
+def official_forward(model, tokens: torch.Tensor) -> torch.Tensor:
+    cache, logits = None, []
+    for position in range(tokens.shape[1]):
+        mask = torch.zeros((1, 1, 1, position + 1))
+        result = model(
+            input_ids=tokens[:, position : position + 1],
+            past_key_values=cache,
+            attention_mask={"full_attention": mask},
+            use_cache=True,
+            cache_position=torch.tensor([position]),
+            position_ids=torch.tensor([[position]]),
+        )
+        cache = result.past_key_values
+        logits.append(result.logits)
+    return torch.cat(logits, dim=1)
+
+
+def build_model(seed: int):
+    package = types.ModuleType("pinned_iquest")
+    package.__path__ = [str(SOURCE)]
+    sys.modules[package.__name__] = package
+    config_module = importlib.import_module("pinned_iquest.configuration_iquestloopcoder")
+    model_module = importlib.import_module("pinned_iquest.modeling_iquestloopcoder")
+    torch.manual_seed(seed)
+    config = config_module.IQuestLoopCoderConfig(
+        vocab_size=97,
+        hidden_size=32,
+        intermediate_size=48,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        max_position_embeddings=131072,
+        loop_num=2,
+        loop_window_size=64,
+    )
+    config._attn_implementation = "eager"
+    model = model_module.IQuestLoopCoderForCausalLM(config).float()
+    with torch.no_grad():
+        for gate in model.model.gate_projections:
+            gate.weight.normal_(std=0.1)
+            gate.bias.normal_(std=0.1)
+    return model
+
+
+def main(math_backend: bool, output: Path) -> None:
+    torch.set_num_threads(1)
+    manifest = json.loads((SOURCE / "source_manifest.json").read_text())
+    for entry in manifest["files"]:
+        assert hashlib.sha256((SOURCE / entry["path"]).read_bytes()).hexdigest() == entry["sha256"]
+    assert not output.exists()
+    receipt = {
+        "state": "RUNNING",
+        "revision": manifest["revision"],
+        "torch": torch.__version__,
+        "official_weights": False,
+        "device": "cpu",
+        "sdpa_backend": "MATH" if math_backend else "AUTO",
+        "gradient_relative_l2_gate": 2e-6,
+        "logits_atol": 1e-5,
+        "logits_rtol": 1e-5,
+        "optimizer": "AdamW",
+        "learning_rate": 1e-4,
+        "betas": [0.9, 0.999],
+        "eps": 1e-8,
+        "weight_decay": 0.1,
+        "cases": [],
+        "scope": "Tiny supervised numerical/Adam equivalence; not RL reward convergence or 40B execution.",
+    }
+    for seed in (42, 43):
+        for length in (1, 4, 8, 65):
+            for recompute in (False, True):
+                reference = build_model(seed)
+                candidate = copy.deepcopy(reference)
+                tokens = torch.randint(3, 97, (1, length))
+                targets = torch.randint(0, 97, (length,))
+                params = [dict(m.named_parameters()) for m in (reference, candidate)]
+                assert params[0].keys() == params[1].keys()
+                optimizers = [
+                    torch.optim.AdamW(
+                        m.parameters(),
+                        lr=1e-4,
+                        betas=(0.9, 0.999),
+                        eps=1e-8,
+                        weight_decay=0.1,
+                    )
+                    for m in (reference, candidate)
+                ]
+                case = {
+                    "seed": seed,
+                    "length": length,
+                    "recompute": recompute,
+                    "tokens": tokens.tolist(),
+                    "targets": targets.tolist(),
+                    "steps": [],
+                }
+                receipt["cases"].append(case)
+                for step in (1, 2):
+                    before = [{k: p.detach().clone() for k, p in row.items()} for row in params]
+                    for optimizer in optimizers:
+                        optimizer.zero_grad(set_to_none=True)
+                    expected = official_forward(reference, tokens)
+                    actual = replay(candidate, tokens, recompute=recompute)
+                    record = {
+                        "step": step,
+                        "logits_max_abs": (expected - actual).abs().max().item(),
+                        "failures": [],
+                    }
+                    case["steps"].append(record)
+                    output.write_text(json.dumps(receipt, indent=2) + "\n")
+                    if not torch.allclose(actual, expected, atol=1e-5, rtol=1e-5):
+                        record["failures"].append({"field": "logits"})
+                    losses = [F.cross_entropy(logits.reshape(-1, 97), targets) for logits in (expected, actual)]
+                    for loss in losses:
+                        loss.backward()
+                    for row in params:
+                        assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in row.values())
+                    norm = sum(p.grad.double().square().sum().item() for p in params[0].values())
+                    error = sum((params[0][k].grad.double() - p.grad.double()).square().sum().item() for k, p in params[1].items())
+                    record["gradient_relative_l2"] = (error / norm) ** 0.5
+                    assert norm > 0
+                    if record["gradient_relative_l2"] > 2e-6:
+                        record["failures"].append({"field": "gradient_relative_l2"})
+                    for optimizer in optimizers:
+                        optimizer.step()
+                    delta_max = 0.0
+                    for name in params[0]:
+                        p, q = params[0][name], params[1][name]
+                        dp, dq = p - before[0][name], q - before[1][name]
+                        delta_max = max(delta_max, (dp - dq).abs().max().item())
+                        a, b = optimizers[0].state[p], optimizers[1].state[q]
+                        assert a["step"].item() == b["step"].item() == step
+                        for field, left, right, atol, rtol in (
+                            ("parameter_delta", dq, dp, 1e-7, 1e-5),
+                            ("parameter", q, p, 1e-7, 1e-5),
+                            ("exp_avg", b["exp_avg"], a["exp_avg"], 1e-8, 2e-6),
+                            (
+                                "exp_avg_sq",
+                                b["exp_avg_sq"],
+                                a["exp_avg_sq"],
+                                1e-10,
+                                2e-6,
+                            ),
+                        ):
+                            close = torch.isclose(left, right, atol=atol, rtol=rtol)
+                            if not bool(close.all()):
+                                record["failures"].append(
+                                    {
+                                        "tensor": name,
+                                        "field": field,
+                                        "mismatched_elements": int((~close).sum()),
+                                        "max_abs": (left - right).abs().max().item(),
+                                    }
+                                )
+                    record.update(
+                        loss_reference=losses[0].item(),
+                        loss_candidate=losses[1].item(),
+                        parameter_delta_max_abs=delta_max,
+                        all_parameters_and_moments_checked=True,
+                        nonzero_applied_update=all(any(bool((row[k] - initial[k]).any()) for k in row) for row, initial in zip(params, before, strict=True)),
+                        passed=not record["failures"],
+                    )
+                    assert record["nonzero_applied_update"]
+                    output.write_text(json.dumps(receipt, indent=2) + "\n")
+                print(
+                    json.dumps(
+                        {
+                            "seed": seed,
+                            "length": length,
+                            "recompute": recompute,
+                            "steps": 2,
+                        }
+                    ),
+                    flush=True,
+                )
+    receipt["state"] = "PASS" if all(s["passed"] for c in receipt["cases"] for s in c["steps"]) else "FAILED_NUMERICAL_GATE"
+    receipt["case_count"] = 16
+    receipt["paired_steps"] = 32
+    output.write_text(json.dumps(receipt, indent=2) + "\n")
+    print(json.dumps({"state": receipt["state"], "cases": 16, "paired_steps": 32}))
+    raise SystemExit(0 if receipt["state"] == "PASS" else 1)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--sdpa-math", action="store_true")
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    with sdpa_kernel(SDPBackend.MATH) if args.sdpa_math else nullcontext():
+        main(args.sdpa_math, args.output)

benchmarks/probe_iquest_source.py

@@ -0,0 +1,120 @@
+"""Execute pinned, unmodified IQuest code with a tiny CPU configuration."""
+
+import argparse
+import importlib
+import json
+import sys
+import types
+from pathlib import Path
+
+import torch
+import transformers
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--case",
+        choices=("cached", "uncached", "checkpointed", "decode"),
+        required=True,
+    )
+    parser.add_argument("--length", type=int, default=8)
+    parser.add_argument("--explicit-mask", action="store_true")
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    source = Path(__file__).resolve().parents[1] / "benchmarks/results/iquest-contract/reference"
+    package = types.ModuleType("pinned_iquest")
+    package.__path__ = [str(source)]
+    sys.modules[package.__name__] = package
+    config_module = importlib.import_module("pinned_iquest.configuration_iquestloopcoder")
+    model_module = importlib.import_module("pinned_iquest.modeling_iquestloopcoder")
+    torch.set_num_threads(1)
+    torch.manual_seed(42)
+    config = config_module.IQuestLoopCoderConfig(
+        vocab_size=97,
+        hidden_size=32,
+        intermediate_size=48,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        max_position_embeddings=131072,
+        loop_num=2,
+        loop_window_size=64,
+        attention_dropout=0.0,
+        use_cache=True,
+    )
+    config._attn_implementation = "eager"
+    model = model_module.IQuestLoopCoderForCausalLM(config).float()
+    with torch.no_grad():
+        for gate in model.model.gate_projections:
+            gate.weight.normal_(std=0.1)
+            gate.bias.normal_(std=0.1)
+    tokens = torch.randint(3, config.vocab_size, (1, args.length))
+    receipt = {
+        "case": args.case,
+        "length": args.length,
+        "torch": torch.__version__,
+        "transformers": transformers.__version__,
+        "revision": json.loads((source / "source_manifest.json").read_text())["revision"],
+        "device": "cpu",
+        "official_weights": False,
+        "config": config.to_dict(),
+        "unique_parameter_count": sum(p.numel() for p in model.parameters()),
+        "explicit_causal_mask_input": args.explicit_mask,
+    }
+    args.output.with_suffix(".input.json").write_text(json.dumps(receipt, indent=2) + "\n")
+    print(
+        json.dumps({"state": "TINY_INPUT_READY", "case": args.case, "length": args.length}),
+        flush=True,
+    )
+
+    def forward(input_ids: torch.Tensor, cache=None):
+        kwargs = {}
+        if args.explicit_mask:
+            seen = 0 if cache is None else cache.get_seq_length()
+            positions = torch.arange(seen, seen + input_ids.shape[1])
+            allowed = torch.arange(seen + input_ids.shape[1])[None, :] <= positions[:, None]
+            mask = torch.where(allowed, 0.0, torch.finfo(torch.float32).min)[None, None]
+            kwargs = {
+                "attention_mask": {"full_attention": mask},
+                "cache_position": positions,
+                "position_ids": positions[None],
+            }
+        return model(
+            input_ids=input_ids,
+            past_key_values=cache,
+            use_cache=args.case != "uncached",
+            **kwargs,
+        )
+
+    if args.case == "checkpointed":
+        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    if args.case == "decode":
+        model.eval()
+        with torch.no_grad():
+            full = forward(tokens).logits
+            cache = None
+            chunks = []
+            for position in range(args.length):
+                result = forward(tokens[:, position : position + 1], cache)
+                cache = result.past_key_values
+                chunks.append(result.logits)
+            decoded = torch.cat(chunks, dim=1)
+        receipt["prefill_decode_max_abs"] = (full - decoded).abs().max().item()
+        receipt["prefill_decode_allclose_1e-5"] = torch.allclose(full, decoded, atol=1e-5, rtol=1e-5)
+        receipt["prefill_finite"] = bool(torch.isfinite(full).all())
+        receipt["decode_finite"] = bool(torch.isfinite(decoded).all())
+    else:
+        result = forward(tokens)
+        loss = result.logits.square().mean()
+        loss.backward()
+        receipt["loss"] = loss.item()
+        receipt["all_parameters_have_finite_gradients"] = all(p.grad is not None and torch.isfinite(p.grad).all().item() for p in model.parameters())
+        receipt["gradient_norm"] = sum(p.grad.square().sum().item() for p in model.parameters()) ** 0.5
+    args.output.write_text(json.dumps(receipt, indent=2) + "\n")
+    print(json.dumps({k: v for k, v in receipt.items() if k != "config"}))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/results/iquest-contract/README.md

@@ -0,0 +1,27 @@
+# ScaleRLT — IQuest semantic qualification
+
+The fixed official source and sixteen physical checkpoint headers are in `reference/`. No weight tensor body is included. The unique parameter count is 39,794,696,320. These are tiny CPU FP32 tests of the checkpoint's architecture, not a 40B run or RL reward curve.
+
+| Actual comparison | AUTO SDPA | MATH SDPA | Status |
+| --- | ---: | ---: | --- |
+| Paired AdamW steps, two seeds × four lengths × ordinary/recomputed × two steps | 32 | 32 | All executed |
+| Maximum logits error against unmodified official token-by-token autograd | 7.15256×10⁻⁷ | 6.25849×10⁻⁷ | 32/32 pass original 10⁻⁵ absolute/relative |
+| Maximum whole-gradient relative L2 error | 3.21867×10⁻⁷ | 3.04581×10⁻⁷ | 32/32 pass original 2×10⁻⁶ |
+| Maximum parameter-delta difference | 2.39909×10⁻⁶ | 1.98185×10⁻⁶ | Original 10⁻⁷ absolute/10⁻⁵ relative gate fails |
+| Complete update checks passed | 24/32 | 24/32 | **FAILED_NUMERICAL_GATE** |
+| Adam moments and counters | 32/32 | 32/32 | Original checks pass |
+| Speed ratio / GPU memory / RL reward convergence | NOT_MEASURED | NOT_MEASURED | No claim |
+
+The unmodified public full-prefill path differs from decode by 0.41336 at length eight and 0.50415 at length 65. Cache-free and HF checkpointed training lose the first-loop global KV. The default cache/mask API also fails in the existing Transformers 4.54.1 environment; explicit mask inputs isolate the mathematical tests without changing source or packages. All failed records remain alongside the complete matrices.
+
+The experimental replay is not registered as a learner provider. See [the contract report](../../../docs/iquest_semantic_qualification.md) for the complete boundaries. No failed check is promoted to full update qualification.
+
+Run in the existing matching environment:
+
+```sh
+python benchmarks/audit_iquest_replay.py --output /tmp/scalerlt-iquest-auto.json
+python benchmarks/audit_iquest_replay.py --sdpa-math --output /tmp/scalerlt-iquest-math.json
+python benchmarks/probe_iquest_source.py --case decode --length 65 --explicit-mask --output /tmp/scalerlt-iquest-decode.json
+```
+
+Both complete replay audits currently return exit code 1 after writing all sixteen cases and thirty-two paired steps. The decode probe records an agreement failure in its JSON. Its exit code alone is not a numerical pass. Do not replace existing result files. The source manifest and `SHA256SUMS.json` bind the retained artifacts; original source notices and LICENSE are included.