| """Tiny random Gemma 4 checks before any 12B run. Run as ``python -m training.selftest [--model DIR] --out FILE``. |
| |
| Classifier-plan §3.5. The tiny checkpoint is written in the real format (``Gemma4UnifiedForConditionalGeneration`` |
| with small image and audio projections, so the text-only loader must map and ignore keys exactly as for the 12B), |
| with a byte-level tokenizer that has Gemma's special-token ids (pad 0, eos 1, bos 2). With ``--model`` the real |
| config's attention shapes (head_dim 256, global head_dim 512, KV heads, sliding window) are kept and only the depth, |
| width and vocabulary shrink; without it everything is small (CPU tests). Nothing here needs the 12B weights. |
| |
| Checks (exit 1 if any fails): |
| 1. text-only load: zero missing text weights, multimodal weights ignored, hidden states equal the full model's; |
| 2. heads: one ``Linear(hidden -> K_f)`` per nsb-v1 field (the real head sizes), LoRA on every target projection; |
| 3. padding invariance: a row alone and the same row right-padded in a batch give the same logits (both poolings, |
| float64 so that only a masking or pooling error can fail it); |
| 4. loss split: one step as one micro-batch and as two give the same gradients; |
| 5. save/load: adapter + heads reload to the same logits; |
| 6. CPU/GPU parity (only with CUDA): the same tiny model gives the same logits on both devices. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import tempfile |
| from pathlib import Path |
|
|
| VOCAB = 512 |
| SPECIALS = ["<pad>", "<eos>", "<bos>", "<unk>", "<|turn>", "<turn|>", "<|channel>", "<channel|>"] |
| LORA = dict(lora_r=8, lora_alpha=8, lora_dropout=0.0, head_dropout=0.0) |
|
|
|
|
| def tiny_text(real=None, layers: int = 6) -> dict: |
| """Text config kwargs: 5 sliding + 1 full attention layer (the 12B pattern repeated once).""" |
| kinds = ["sliding_attention"] * (layers - 1) + ["full_attention"] |
| if real is None: |
| return dict(vocab_size=VOCAB, hidden_size=64, intermediate_size=128, num_hidden_layers=layers, |
| num_attention_heads=4, num_key_value_heads=2, head_dim=16, global_head_dim=32, |
| num_global_key_value_heads=1, attention_k_eq_v=True, sliding_window=16, |
| max_position_embeddings=4096, final_logit_softcapping=30.0, layer_types=kinds, |
| pad_token_id=0, eos_token_id=1, bos_token_id=2) |
| |
| d = real.to_dict() |
| keep = ("num_attention_heads", "num_key_value_heads", "head_dim", "sliding_window", "attention_k_eq_v", |
| "hidden_activation", "rms_norm_eps", "final_logit_softcapping", "rope_parameters") |
| out = {k: d[k] for k in keep if d.get(k) is not None} |
| layer = next(iter((d.get("per_layer_config") or {}).values()), {}) |
| if layer.get("head_dim") is not None: |
| out["global_head_dim"] = layer["head_dim"] |
| if layer.get("num_key_value_heads") is not None: |
| out["num_global_key_value_heads"] = layer["num_key_value_heads"] |
| heads = out.get("num_attention_heads", 4) |
| out.update(vocab_size=VOCAB, hidden_size=heads * 32, intermediate_size=heads * 64, num_hidden_layers=layers, |
| layer_types=kinds, max_position_embeddings=4096, pad_token_id=0, eos_token_id=1, bos_token_id=2) |
| return out |
|
|
|
|
| def make_tokenizer(path: Path): |
| """Byte-level tokenizer (every text encodes) with Gemma's special-token ids; saved next to the model.""" |
| from tokenizers import Tokenizer, decoders, models, pre_tokenizers |
| from transformers import PreTrainedTokenizerFast |
|
|
| alphabet = sorted(pre_tokenizers.ByteLevel.alphabet()) |
| vocab = {t: i for i, t in enumerate(SPECIALS)} |
| for ch in alphabet: |
| vocab.setdefault(ch, len(vocab)) |
| tk = Tokenizer(models.BPE(vocab=vocab, merges=[], unk_token="<unk>")) |
| tk.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False) |
| tk.decoder = decoders.ByteLevel() |
| tok = PreTrainedTokenizerFast(tokenizer_object=tk, pad_token="<pad>", eos_token="<eos>", bos_token="<bos>", |
| unk_token="<unk>", additional_special_tokens=SPECIALS[4:]) |
| tok.save_pretrained(str(path)) |
| return tok |
|
|
|
|
| def make_tiny(path: Path, real_model: str | None = None, layers: int = 6, seed: int = 0, multimodal: bool = True): |
| """Save a random tiny Gemma 4 checkpoint (full multimodal class) plus tokenizer; returns the full model.""" |
| import torch |
| from transformers import Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration |
|
|
| real = None |
| if real_model: |
| from .model import text_config |
|
|
| real = text_config(real_model) |
| text = tiny_text(real, layers) |
| vision = dict(patch_size=2, pooling_kernel_size=2, mm_embed_dim=16, mm_posemb_size=8, |
| output_proj_dims=16) if multimodal else None |
| audio = dict(audio_embed_dim=8) if multimodal else None |
| cfg = Gemma4UnifiedConfig(text_config=text, vision_config=vision, audio_config=audio) |
| torch.manual_seed(seed) |
| full = Gemma4UnifiedForConditionalGeneration(cfg).float() |
| path = Path(path) |
| full.save_pretrained(str(path)) |
| make_tokenizer(path) |
| return full |
|
|
|
|
| def tiny_cfg(path, pool: str = "last", dtype: str = "float32") -> dict: |
| return dict(LORA, model=str(path), dtype=dtype, pool=pool, gradient_checkpointing=True) |
|
|
|
|
| def build_tiny(path, device, pool: str = "last", sizes=None, seed: int = 7, dtype: str = "float32"): |
| """Tiny heads model with non-zero LoRA B weights, so LoRA A gradients and the adapter are non-trivial.""" |
| import torch |
|
|
| from .model import build |
| from .taxonomy import Taxonomy |
|
|
| sizes = sizes or Taxonomy.default().sizes() |
| model = build(tiny_cfg(path, pool, dtype), sizes, device, train=True) |
| g = torch.Generator(device="cpu").manual_seed(seed) |
| with torch.no_grad(): |
| for n, p in model.named_parameters(): |
| if "lora_B" in n: |
| p.copy_((torch.randn(p.shape, generator=g) * 0.05).to(p.dtype)) |
| return model |
|
|
|
|
| def rows(lengths, seed: int = 1) -> list[list[int]]: |
| import torch |
|
|
| g = torch.Generator().manual_seed(seed) |
| return [[2] + torch.randint(10, VOCAB - 1, (n - 1,), generator=g).tolist() for n in lengths] |
|
|
|
|
| def flat_grads(model): |
| import torch |
|
|
| return torch.cat([p.grad.detach().float().flatten() for _, p in model.named_parameters() |
| if p.requires_grad and p.grad is not None]) |
|
|
|
|
| def forward_backward(model, batch, labels, device, split=None, weights=None): |
| import torch |
|
|
| from .data import step_counts |
| from .model import pad, step_loss |
|
|
| model.zero_grad(set_to_none=True) |
| counts = step_counts(labels) |
| parts = split or [list(range(len(batch)))] |
| for part in parts: |
| ids, mask, last = pad([batch[i] for i in part], 0, device) |
| y = torch.tensor([labels[i] for i in part], device=device) |
| loss, _ = step_loss(model(ids, mask, last), y, counts, weights) |
| loss.backward() |
| return flat_grads(model) |
|
|
|
|
| def rel_err(a, b) -> float: |
| return float((a - b).abs().max() / b.abs().max().clamp_min(1e-12)) |
|
|
|
|
| def cat_logits(model, batch, device): |
| import torch |
|
|
| from .model import pad |
|
|
| with torch.no_grad(): |
| return torch.cat([lg.float() for lg in model(*pad(batch, 0, device))], 1) |
|
|
|
|
| def run_checks(tmp: Path, device, real_model: str | None = None) -> dict: |
| import torch |
|
|
| from .model import LORA_TARGETS, build, load_backbone, save_checkpoint |
| from .taxonomy import Taxonomy |
|
|
| tax = Taxonomy.default() |
| out: dict = {} |
| path = tmp / "tiny" |
| full = make_tiny(path, real_model) |
| |
| lm, info = load_backbone(path, torch.device("cpu"), dtype=torch.float32) |
| ids = torch.tensor([[2, 50, 60, 70, 0, 0]]) |
| mask = torch.tensor([[1, 1, 1, 1, 0, 0]]) |
| with torch.no_grad(): |
| err = float((lm.model(input_ids=ids, attention_mask=mask).last_hidden_state |
| - full.model.language_model(input_ids=ids, attention_mask=mask).last_hidden_state).abs().max()) |
| ignored = sorted(info["unexpected_keys"]) |
| out["text_only_load"] = {"missing": sorted(info["missing_keys"]), "ignored_examples": ignored[:4], |
| "ignored": len(ignored), "max_abs_diff_vs_full": err, |
| "pass": not info["missing_keys"] and err < 1e-5} |
| del lm, full |
| |
| model = build_tiny(path, device) |
| shapes = [tuple(h.weight.shape) for h in model.heads] |
| hidden = model.heads[0].in_features |
| wrapped = sorted({n.split(".")[-1] for n, m in model.backbone.named_modules() if hasattr(m, "lora_A")}) |
| out["heads"] = {"shapes": shapes, "lora_modules": wrapped, |
| "pass": shapes == [(k, hidden) for k in tax.sizes()] and set(wrapped) == set(LORA_TARGETS)} |
| tol = 1e-4 |
| |
| |
| |
| pad_ok = {} |
| for pool in ("last", "mean"): |
| m = build_tiny(path, device, pool=pool, dtype="float64") |
| m.eval() |
| batch = rows([23, 71, 5], seed=2) |
| alone = cat_logits(m, [batch[0]], device) |
| together = cat_logits(m, batch, device)[0:1] |
| pad_ok[pool] = rel_err(alone, together) |
| out["padding_invariance"] = {"rel_err": pad_ok, "pass": all(v < tol for v in pad_ok.values())} |
| |
| model.train() |
| batch = rows([30, 31, 60, 12], seed=3) |
| labels = [[0, 1, 2, 0, 1, 2, 3, 1, 0, 2], [1, -1, -1, 1, 0, 0, 11, 3, 5, 0], |
| [2, 4, 4, -1, 2, 1, 0, 0, 1, 1], [0, 0, 0, 3, 3, 3, 5, 2, 2, 3]] |
| torch.manual_seed(0) |
| g_one = forward_backward(model, batch, labels, device, weights=tax.weights()) |
| torch.manual_seed(0) |
| g_two = forward_backward(model, batch, labels, device, split=[[0, 1], [2, 3]], weights=tax.weights()) |
| split_err = rel_err(g_two, g_one) |
| cos = float(torch.nn.functional.cosine_similarity(g_one, g_two, dim=0)) |
| |
| |
| out["loss_split"] = {"grad_rel_err": split_err, "grad_cosine": cos, "pass": split_err < 1e-3 and cos >= 0.9999} |
| |
| model.eval() |
| before = cat_logits(model, batch, device) |
| save_checkpoint(model, tmp / "ckpt", tax.names, {"test": True}) |
| again = build(tiny_cfg(path), tax.sizes(), device, train=False, adapter_dir=tmp / "ckpt" / "adapter", |
| heads_path=tmp / "ckpt" / "heads.safetensors") |
| after = cat_logits(again, batch, device) |
| out["save_load"] = {"rel_err": rel_err(after, before), "pass": rel_err(after, before) < tol} |
| |
| if torch.device(device).type == "cuda": |
| cpu = build(tiny_cfg(path), tax.sizes(), torch.device("cpu"), train=False, |
| adapter_dir=tmp / "ckpt" / "adapter", heads_path=tmp / "ckpt" / "heads.safetensors") |
| e = rel_err(cat_logits(cpu, batch, torch.device("cpu")), after.cpu()) |
| out["cpu_gpu_parity"] = {"rel_err": e, "pass": e < 1e-3} |
| else: |
| out["cpu_gpu_parity"] = "skipped (no CUDA device)" |
| out["ok"] = all(v["pass"] for v in out.values() if isinstance(v, dict)) |
| return out |
|
|
|
|
| def main(argv=None) -> int: |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| p.add_argument("--model", help="Real checkpoint directory (only its config is used)") |
| p.add_argument("--device", default="cuda") |
| p.add_argument("--out") |
| args = p.parse_args(argv) |
| import torch |
|
|
| device = torch.device(args.device if args.device != "cuda" or torch.cuda.is_available() else "cpu") |
| with tempfile.TemporaryDirectory() as tmp: |
| out = run_checks(Path(tmp), device, args.model) |
| out["device"] = str(device) |
| if args.device == "cuda" and device.type != "cuda": |
| out["ok"] = False |
| out["error"] = "CUDA requested but not available" |
| text = json.dumps(out, indent=1, default=str) |
| print(text) |
| if args.out: |
| Path(args.out).write_text(text + "\n") |
| return 0 if out["ok"] else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
