| """ |
| Train angstrom_nano with BPE tokenizer on the full merged corpus. |
| Laptop-friendly medium config: ~8.4M params, fits in <500 MB RAM. |
| Trains on CPU in ~1-2 hours. |
| """ |
| import sys; sys.path.insert(0, '.') |
| import math, time, torch |
| from pathlib import Path |
| from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM |
| from angstrom_nano.tokenizer import AngstromNanoTokenizer |
|
|
| torch.manual_seed(42) |
|
|
| |
| |
| |
| corpus_path = Path("data/corpus.txt") |
| if not corpus_path.exists(): |
| print("corpus.txt not found — run prepare_data.py first") |
| sys.exit(1) |
|
|
| text = corpus_path.read_text(encoding="utf-8") |
| print(f"Data: {len(text):,} chars ({len(text)/1e6:.1f} MB)") |
|
|
| tok = AngstromNanoTokenizer.train_bpe( |
| [str(corpus_path)], |
| vocab_size=4096, |
| out_path="checkpoints/tokenizer.json", |
| ) |
| vocab_size = len(tok) |
| print(f"BPE vocab: {vocab_size} tokens") |
|
|
| ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long) |
| print(f"Tokenized: {len(ids):,} tokens ({len(ids)/1e6:.1f}M)") |
|
|
| |
| |
| |
| cfg = AngstromNanoConfig( |
| vocab_size=vocab_size, |
| hidden_size=192, intermediate_size=512, num_hidden_layers=6, |
| num_attention_heads=6, num_key_value_heads=3, head_dim=32, |
| num_local_experts=4, num_experts_per_tok=2, |
| max_position_embeddings=256, sliding_window=64, |
| scoring_func="sigmoid", use_qk_norm=True, use_routing_bias=True, |
| tie_word_embeddings=True, |
| ) |
| model = AngstromNanoForCausalLM(cfg) |
| p = sum(p.numel() for p in model.parameters()) |
| m = sum(p.numel() * p.element_size() for p in model.parameters()) / 1e6 |
| print(f"Model: {p:,} params ({m:.2f} MB in {'FP32' if m > 0 else 'N/A'})") |
|
|
| |
| |
| |
| seq_len = 64 |
| lr = 3e-3 |
| steps = 5000 |
| print_every = 500 |
| opt = torch.optim.AdamW(model.parameters(), lr=lr) |
| t0 = time.time() |
|
|
| for step in range(1, steps + 1): |
| i = torch.randint(0, len(ids) - seq_len - 1, (1,)).item() |
| x = ids[i : i + seq_len].unsqueeze(0) |
| y = ids[i + 1 : i + seq_len + 1].unsqueeze(0) |
|
|
| out = model(x, labels=y, output_router_logits=True) |
| loss = out["loss"] |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
|
|
| if step % print_every == 0 or step == 1: |
| print(f" step {step:>4d} loss={loss.item():.4f} ppl={math.exp(loss.item()):.2f} " |
| f"aux={out['aux_loss'].item():.6f} {time.time()-t0:.0f}s") |
|
|
| |
| |
| |
| out_dir = Path("checkpoints") |
| out_dir.mkdir(exist_ok=True) |
| path = out_dir / "medium_model.safetensors" |
|
|
| import json |
| from safetensors.torch import save_file |
|
|
| cfg_path = out_dir / "medium_config.json" |
| cfg_path.write_text(json.dumps( |
| {"vocab_size": vocab_size, "hidden_size": 192, "intermediate_size": 512, |
| "num_hidden_layers": 6, "num_attention_heads": 6, "num_key_value_heads": 3, |
| "head_dim": 32, "num_local_experts": 4, "num_experts_per_tok": 2, |
| "max_position_embeddings": 256, "sliding_window": 64, |
| "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True, |
| "tie_word_embeddings": True, "model_type": "angstrom_nano"}, indent=2)) |
|
|
| sd = model.state_dict() |
| if sd["lm_head.weight"].data_ptr() == sd["model.embed_tokens.weight"].data_ptr(): |
| sd.pop("lm_head.weight") |
| save_file({k: v.contiguous() for k, v in sd.items()}, str(path)) |
| print(f"\nSaved: {path} ({path.stat().st_size / 1e6:.2f} MB)") |
| print(f"Saved: {cfg_path}") |
|
|
| |
| |
| |
| from safetensors.torch import load_file as safe_load |
| tok = AngstromNanoTokenizer.from_bpe_file("checkpoints/tokenizer.json") |
| d = json.loads(cfg_path.read_text()) |
| cfg2 = AngstromNanoConfig(**d) |
| model2 = AngstromNanoForCausalLM(cfg2) |
| sd = safe_load(str(path)) |
| if "lm_head.weight" not in sd: |
| model2.lm_head.weight = model2.model.embed_tokens.weight |
| model2.load_state_dict(sd, strict=False) |
| model2.eval() |
|
|
| prompt = "def fibonacci" |
| ids_p = torch.tensor([tok.encode(prompt, add_bos=True, add_eos=False)]) |
| out_ids = model2.generate(ids_p, max_new_tokens=20, temperature=0.8) |
| gen = tok.decode(out_ids[0].tolist(), skip_special_tokens=True) |
| clean = ''.join(c if 32 <= ord(c) <= 126 or c in '\n\t' else ' ' for c in gen) |
| print(f"\nPrompt: {prompt}") |
| print(f"Gen : {clean}") |
| print("TRAINING COMPLETE") |
|
|
