Correct the model card (pooling is at the LABELS suffix's last token); add code/classifier from project commit 4ced7554f
3d7c238 verified | """Tiny-model tests on the pod, before the 9B run. Run as ``python -m classifier.selftest --model DIR``. | |
| The tiny model is the real Qwen3.5 text config cut to four layers (full, DeltaNet, DeltaNet, full attention) | |
| with a small vocabulary and FFN, but the **real head counts and head sizes**, which the kernels constrain. | |
| It is randomly initialised; nothing here needs the 9B weights. Checks (exit 1 if any fails): | |
| 1. gradient parity: logits and every LoRA gradient (heads excluded) with the fla + causal-conv1d kernels against | |
| the PyTorch path, on a right-padded batch: worst per-tensor cosine and norm ratio (fp32 if the kernels | |
| accept it, otherwise bf16 with looser bounds); | |
| 2. padding invariance: a row alone and the same row right-padded in a batch give the same head logits; | |
| 3. loss split: one step as one micro-batch and as two micro-batches give the same gradients; | |
| 4. save/load: an adapter + heads checkpoint reloads to the same logits. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import copy | |
| import importlib | |
| import json | |
| import tempfile | |
| from pathlib import Path | |
| VOCAB = 2048 | |
| LORA = dict(lora_r=8, lora_alpha=8, lora_dropout=0.0, head_dropout=0.0) | |
| def tiny_config(model_dir): | |
| from .model import text_config | |
| cfg = copy.deepcopy(text_config(model_dir)) | |
| # full -> linear -> linear -> full: LoRA weights sit upstream of each Gated DeltaNet layer, so a wrong | |
| # kernel backward (including through g/beta) shows up in trainable gradients. | |
| cfg.num_hidden_layers = 4 | |
| if hasattr(cfg, 'layer_types'): | |
| cfg.layer_types = ['full_attention', 'linear_attention', 'linear_attention', 'full_attention'] | |
| cfg.vocab_size = VOCAB | |
| cfg.hidden_size = 1024 | |
| cfg.intermediate_size = 1024 | |
| if hasattr(cfg, 'mtp_num_hidden_layers'): | |
| cfg.mtp_num_hidden_layers = 0 | |
| cfg.tie_word_embeddings = False | |
| cfg.eos_token_id, cfg.pad_token_id = 1, 0 | |
| return cfg | |
| def make_tiny(model_dir, path: Path, dtype): | |
| import torch | |
| from transformers import Qwen3_5ForCausalLM | |
| torch.manual_seed(0) | |
| lm = Qwen3_5ForCausalLM(tiny_config(model_dir)).to(dtype) | |
| lm.save_pretrained(str(path)) | |
| return path | |
| def batch(lengths, device, seed=1): | |
| import torch | |
| g = torch.Generator().manual_seed(seed) | |
| return [torch.randint(10, VOCAB, (n,), generator=g).tolist() for n in lengths] | |
| def flat_grads(model): | |
| import torch | |
| return torch.cat([p.grad.detach().float().flatten() for n, p in model.named_parameters() | |
| if p.requires_grad and p.grad is not None]) | |
| def lora_grads(model) -> dict: | |
| """LoRA gradients by parameter name (heads excluded: they sit after the backbone and would dominate).""" | |
| return {n: p.grad.detach().float().flatten().clone() for n, p in model.named_parameters() | |
| if p.requires_grad and p.grad is not None and 'lora_' in n} | |
| def per_tensor(fast: dict, ref: dict) -> dict: | |
| """Worst per-tensor cosine and norm ratio over LoRA tensors with a non-negligible reference gradient.""" | |
| import torch | |
| top = max(float(v.norm()) for v in ref.values()) | |
| worst_cos, worst_ratio, n = 1.0, 0.0, 0 | |
| for name, r in ref.items(): | |
| if float(r.norm()) < 1e-6 * top: | |
| continue | |
| f = fast[name] | |
| worst_cos = min(worst_cos, float(torch.nn.functional.cosine_similarity(f, r, dim=0))) | |
| worst_ratio = max(worst_ratio, abs(float(f.norm()) / float(r.norm()) - 1.0)) | |
| n += 1 | |
| return {'tensors': n, 'worst_cosine': worst_cos, 'worst_norm_ratio_error': worst_ratio} | |
| def cosine(a, b) -> float: | |
| import torch | |
| return float(torch.nn.functional.cosine_similarity(a, b, dim=0)) | |
| def rel_err(a, b) -> float: | |
| return float((a - b).abs().max() / b.abs().max().clamp_min(1e-12)) | |
| def build_tiny(path, device, dtype_name): | |
| """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 | |
| cfg = dict(LORA, model=str(path), dtype=dtype_name) | |
| model = build(cfg, [3, 5, 11], device, train=True) | |
| g = torch.Generator(device='cpu').manual_seed(7) | |
| 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.02).to(p.dtype)) | |
| return model | |
| def forward_backward(model, rows, labels, device, pad_id=0, split=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(rows)))] | |
| logits_out = [] | |
| for part in parts: | |
| ids, mask, last = pad([rows[i] for i in part], pad_id, device) | |
| y = torch.tensor([labels[i] for i in part], device=device) | |
| logits = model(ids, mask, last) | |
| loss, _ = step_loss(logits, y, counts) | |
| loss.backward() | |
| logits_out.append(torch.cat([lg.detach().float() for lg in logits], dim=1)) | |
| return torch.cat(logits_out, dim=0), flat_grads(model) | |
| def parity(path, device, dtype_name) -> dict: | |
| import torch | |
| from .check_env import fast_path, torch_path | |
| model = build_tiny(path, device, dtype_name) | |
| model.train() | |
| fp = fast_path(model) | |
| rows = batch([37, 64, 9, 120], device) | |
| labels = [[0, 1, 2], [1, -1, 10], [2, 4, 0], [-1, 0, 3]] | |
| torch.manual_seed(0) | |
| logits_fast, _ = forward_backward(model, rows, labels, device) | |
| lora_fast = lora_grads(model) | |
| with torch_path(): | |
| torch.manual_seed(0) | |
| logits_ref, _ = forward_backward(model, rows, labels, device) | |
| lora_ref = lora_grads(model) | |
| flat_fast = torch.cat(list(lora_fast.values())) | |
| flat_ref = torch.cat([lora_ref[k] for k in lora_fast]) | |
| return dict({'dtype': dtype_name, 'fast_path_before': fp, 'logits_rel_err': rel_err(logits_fast, logits_ref), | |
| 'lora_grad_cosine': cosine(flat_fast, flat_ref), 'lora_grad_rel_err': rel_err(flat_fast, flat_ref)}, | |
| **per_tensor(lora_fast, lora_ref)) | |
| def main(argv=None): | |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| p.add_argument('--model', required=True, help='Real checkpoint directory (only its config is used)') | |
| p.add_argument('--device', default='cuda') | |
| p.add_argument('--skip-parity', action='store_true', help='Local CPU dry run only: no kernels to compare') | |
| p.add_argument('--out') | |
| args = p.parse_args(argv) | |
| import torch | |
| from .model import build, save_checkpoint, pad | |
| device = torch.device(args.device) | |
| out, ok = {}, True | |
| with tempfile.TemporaryDirectory() as tmp: | |
| tmp = Path(tmp) | |
| dtype_name = 'float32' | |
| if args.skip_parity: | |
| out['parity'] = 'skipped (--skip-parity)' | |
| ok = False if device.type == 'cuda' else ok | |
| else: | |
| # 1. gradient parity | |
| try: | |
| r = parity(make_tiny(args.model, tmp / 'tiny32', torch.float32), device, 'float32') | |
| r['pass'] = (r['logits_rel_err'] < 1e-3 and r['lora_grad_rel_err'] < 1e-2 and r['worst_cosine'] >= 0.999 | |
| and r['worst_norm_ratio_error'] <= 0.01) | |
| except Exception as e: # kernels may refuse fp32: retry in bf16 with looser bounds | |
| out['parity_float32_error'] = ('%s: %s' % (type(e).__name__, e))[:500] | |
| r = parity(make_tiny(args.model, tmp / 'tiny16', torch.bfloat16), device, 'bfloat16') | |
| r['pass'] = (r['logits_rel_err'] < 5e-2 and r['worst_cosine'] >= 0.99 | |
| and r['worst_norm_ratio_error'] <= 0.05) | |
| r['fast_path_was_active'] = bool(r['fast_path_before']['ok']) | |
| r['pass'] = r['pass'] and r['fast_path_was_active'] | |
| out['parity'] = r | |
| ok &= r['pass'] | |
| dtype_name = r['dtype'] | |
| # 2-4 on the default path, in fp32 (bf16 if the kernels refused fp32 above) | |
| path = tmp / ('tiny32' if dtype_name == 'float32' else 'tiny16') | |
| if not path.exists(): | |
| make_tiny(args.model, path, torch.float32 if dtype_name == 'float32' else torch.bfloat16) | |
| tol = 1e-4 if dtype_name == 'float32' else 3e-2 | |
| model = build_tiny(path, device, dtype_name) | |
| model.eval() | |
| rows = batch([23, 71, 5], device, seed=2) | |
| with torch.no_grad(): | |
| alone = torch.cat([lg.float() for lg in model(*pad([rows[0]], 0, device))], 1) | |
| together = torch.cat([lg.float() for lg in model(*pad(rows, 0, device))], 1)[0:1] | |
| err = rel_err(alone, together) | |
| out['padding_invariance'] = {'rel_err': err, 'pass': err < tol} | |
| ok &= out['padding_invariance']['pass'] | |
| model.train() | |
| rows = batch([30, 31, 60, 12], device, seed=3) | |
| labels = [[0, 1, 2], [1, -1, 10], [2, 4, -1], [0, 0, 3]] | |
| torch.manual_seed(0) | |
| _, g_one = forward_backward(model, rows, labels, device) | |
| torch.manual_seed(0) | |
| _, g_two = forward_backward(model, rows, labels, device, split=[[0, 1], [2, 3]]) | |
| out['loss_split'] = {'grad_cosine': cosine(g_one, g_two), 'grad_rel_err': rel_err(g_two, g_one), | |
| 'pass': (cosine(g_one, g_two) >= (0.9999 if dtype_name == 'float32' else 0.999) | |
| and rel_err(g_two, g_one) < tol * 10)} | |
| ok &= out['loss_split']['pass'] | |
| model.eval() | |
| with torch.no_grad(): | |
| before = torch.cat([lg.float() for lg in model(*pad(rows, 0, device))], 1) | |
| save_checkpoint(model, tmp / 'ckpt', {'test': True}) | |
| cfg = dict(LORA, model=str(path), dtype=dtype_name) | |
| again = build(cfg, [3, 5, 11], device, train=False, adapter_dir=tmp / 'ckpt' / 'adapter', | |
| heads_path=tmp / 'ckpt' / 'heads.pt') | |
| again.eval() | |
| with torch.no_grad(): | |
| after = torch.cat([lg.float() for lg in again(*pad(rows, 0, device))], 1) | |
| err = rel_err(after, before) | |
| out['save_load'] = {'rel_err': err, 'pass': err < tol} | |
| ok &= out['save_load']['pass'] | |
| out['ok'] = bool(ok) | |
| text = json.dumps(out, indent=1, default=str) | |
| print(text) | |
| if args.out: | |
| Path(args.out).write_text(text + '\n') | |
| return 0 if ok else 1 | |
| if __name__ == '__main__': | |
| raise SystemExit(main()) | |