| """AngstromE1-Nano: one-click Kaggle setup → train → download""" |
| import os, sys, subprocess, time, re, json, math |
| from pathlib import Path |
| import torch |
| import torch.nn as nn |
| from torch.cuda.amp import GradScaler, autocast |
|
|
| |
| |
| |
| GIT_REPO = "https://github.com/er-del/angstrom.git" |
| BRANCH = "main" |
| MODEL_CONFIG = "medium" |
| TRAIN_STEPS = 50000 |
| BATCH_SIZE = 16 |
| SEQ_LEN = 256 |
| SAVE_EVERY = 5000 |
| SSH_PASSWORD = "ChangeMe123!" |
| KAGGLE_OUT = Path("/kaggle/working/angstrom_output") |
| LOCAL_OUT = Path("checkpoints") |
|
|
| def log(msg): print(f"\033[0;32m[+]\033[0m {msg}") |
| def info(msg): print(f" {msg}") |
| def err(msg): print(f"\033[0;31m[-]\033[0m {msg}") |
|
|
| |
| |
| |
| log("=== 1/7: System setup ===") |
| log("Checking CUDA + GPUs...") |
| n_gpus = torch.cuda.device_count() |
| for i in range(n_gpus): |
| p = torch.cuda.get_device_properties(i) |
| info(f"GPU {i}: {p.name} {p.total_memory/1e9:.1f}GB CUDA {torch.version.cuda}") |
| log(f"{n_gpus} GPU(s) available") |
|
|
| log("Installing system packages...") |
| os.system("apt-get update -qq > /dev/null 2>&1") |
| os.system("apt-get install -y -qq openssh-server curl wget git-lfs > /dev/null 2>&1") |
|
|
| |
| |
| |
| log("=== 2/7: Clone repo ===") |
| if os.path.exists("angstrom"): |
| log("angstrom/ already exists, pulling latest...") |
| os.system("cd angstrom && git pull 2>/dev/null") |
| else: |
| os.system(f"git clone {GIT_REPO} angstrom 2>&1") |
| os.system(f"cd angstrom && git checkout {BRANCH} 2>/dev/null") |
|
|
| os.chdir("angstrom") |
| REPO = Path.cwd() |
| info(f"Working dir: {REPO}") |
|
|
| |
| |
| |
| log("=== 3/7: Python dependencies ===") |
| os.system("pip install -q --upgrade pip 2>/dev/null") |
| os.system("pip install -q datasets>=2.16.0 safetensors>=0.4.0 tqdm>=4.65.0 accelerate>=0.25.0 2>/dev/null") |
| os.system("pip install -q -e . 2>/dev/null || true") |
|
|
| try: |
| from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM |
| from angstrom_nano.tokenizer import AngstromNanoTokenizer |
| log("All imports OK") |
| except Exception as e: |
| err(f"Import failed: {e}") |
| sys.exit(1) |
|
|
| |
| |
| |
| log("=== 4/7: Download training data ===") |
| data_script = REPO / "download_data.py" |
| if data_script.exists(): |
| info("Running download_data.py...") |
| os.system(f"python {data_script} 2>&1") |
|
|
| data_path = REPO / "data" / "corpus.txt" |
| if not data_path.exists(): |
| err("Data download failed") |
| info("Writing fallback mini-corpus...") |
| (REPO / "data").mkdir(exist_ok=True) |
| data_path.write_text("Hello world.\n" * 1000000) |
| text = data_path.read_text() |
| info(f"Data: {len(text):,} chars ({len(text)/1e6:.1f}MB)") |
|
|
| |
| |
| |
| log("=== 5/7: Train model ===") |
|
|
| CONFIGS = { |
| "small": { |
| "vocab_size": 8192, "hidden_size": 256, "intermediate_size": 1024, |
| "num_hidden_layers": 8, "num_attention_heads": 8, "num_key_value_heads": 2, |
| "head_dim": 32, "num_local_experts": 4, "num_experts_per_tok": 2, |
| "max_position_embeddings": 2048, "sliding_window": 512, |
| "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True, |
| "tie_word_embeddings": True, |
| }, |
| "medium": { |
| "vocab_size": 4096, "hidden_size": 192, "intermediate_size": 768, |
| "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": 2048, "sliding_window": 512, |
| "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True, |
| "tie_word_embeddings": True, |
| }, |
| "large": { |
| "vocab_size": 16384, "hidden_size": 512, "intermediate_size": 2048, |
| "num_hidden_layers": 12, "num_attention_heads": 16, "num_key_value_heads": 4, |
| "head_dim": 32, "num_local_experts": 8, "num_experts_per_tok": 2, |
| "max_position_embeddings": 4096, "sliding_window": 1024, |
| "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True, |
| "tie_word_embeddings": True, |
| }, |
| } |
|
|
| config_dict = CONFIGS[MODEL_CONFIG] |
| info(f"Using config: {MODEL_CONFIG}") |
|
|
| |
| tok_path = REPO / "checkpoints/tokenizer.json" |
| if tok_path.exists(): |
| tok = AngstromNanoTokenizer.from_bpe_file(str(tok_path)) |
| config_dict["vocab_size"] = len(tok) |
| info(f"Loaded tokenizer: {len(tok)} vocab") |
| else: |
| info("Training new tokenizer...") |
| tok_path.parent.mkdir(parents=True, exist_ok=True) |
| tok = AngstromNanoTokenizer.train_bpe( |
| [str(data_path)], vocab_size=config_dict["vocab_size"], |
| out_path=str(tok_path)) |
| info(f"Trained tokenizer: {len(tok)} vocab") |
|
|
| |
| class TextDataset(torch.utils.data.Dataset): |
| def __init__(self, token_ids, seq_len): |
| self.token_ids = token_ids |
| self.seq_len = seq_len |
| self.n_samples = len(token_ids) - seq_len - 1 |
| def __len__(self): return self.n_samples |
| def __getitem__(self, idx): |
| return (self.token_ids[idx:idx+self.seq_len], |
| self.token_ids[idx+1:idx+self.seq_len+1]) |
|
|
| info("Tokenizing data...") |
| ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long) |
| info(f"Tokens: {len(ids):,} ({len(ids)/1e6:.1f}M)") |
|
|
| dataset = TextDataset(ids, SEQ_LEN) |
| dataloader = torch.utils.data.DataLoader( |
| dataset, batch_size=BATCH_SIZE, shuffle=True, |
| num_workers=2, pin_memory=True, drop_last=True) |
| info(f"Samples: {len(dataset):,}") |
|
|
| |
| device = torch.device("cuda") |
| cfg = AngstromNanoConfig(**config_dict) |
| model: nn.Module = AngstromNanoForCausalLM(cfg) |
| n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| info(f"Model: {n_params:,} params ({n_params*4/1e6:.1f}MB FP32)") |
| model = model.to(device) |
|
|
| if n_gpus > 1: |
| model = nn.DataParallel(model, device_ids=list(range(n_gpus))) |
| info(f"DataParallel across {n_gpus} GPUs") |
|
|
| |
| optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.1, betas=(0.9, 0.95)) |
| scaler = GradScaler(enabled=True) |
|
|
| WARMUP = 500 |
| MIN_LR = 3e-4 |
| GRAD_CLIP = 1.0 |
|
|
| def get_lr(step): |
| if step < WARMUP: |
| return 3e-3 * step / max(1, WARMUP) |
| progress = (step - WARMUP) / max(1, TRAIN_STEPS - WARMUP) |
| return MIN_LR + 0.5 * (3e-3 - MIN_LR) * (1.0 + math.cos(math.pi * progress)) |
|
|
| log("Starting training loop...") |
| model.train() |
| t0 = time.time() |
| running_loss = 0.0 |
| running_steps = 0 |
|
|
| LOCAL_OUT.mkdir(parents=True, exist_ok=True) |
| KAGGLE_OUT.mkdir(parents=True, exist_ok=True) |
|
|
| data_iter = iter(dataloader) |
| for step in range(1, TRAIN_STEPS + 1): |
| try: |
| x, y = next(data_iter) |
| except StopIteration: |
| data_iter = iter(dataloader) |
| x, y = next(data_iter) |
|
|
| x = x.to(device, non_blocking=True) |
| y = y.to(device, non_blocking=True) |
|
|
| with autocast(enabled=True, dtype=torch.float16): |
| out = model(x, labels=y, output_router_logits=True) |
| loss = out["loss"] |
| aux_loss = out.get("aux_loss", torch.tensor(0.0)) |
|
|
| optimizer.zero_grad() |
| scaler.scale(loss).backward() |
|
|
| if GRAD_CLIP > 0: |
| scaler.unscale_(optimizer) |
| nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) |
|
|
| scaler.step(optimizer) |
| scaler.update() |
|
|
| lr = get_lr(step) |
| for pg in optimizer.param_groups: |
| pg["lr"] = lr |
|
|
| running_loss += loss.item() |
| running_steps += 1 |
|
|
| if step % 100 == 0 or step == 1: |
| avg = running_loss / running_steps |
| ppl = math.exp(min(avg, 20)) |
| elapsed = time.time() - t0 |
| tok_s = (BATCH_SIZE * SEQ_LEN * running_steps) / max(1, elapsed) |
| gpu_mem = torch.cuda.memory_allocated(0) / 1e9 |
| aux_val = aux_loss.item() if isinstance(aux_loss, torch.Tensor) else 0.0 |
| print(f" step {step:>6d}/{TRAIN_STEPS} loss={avg:.4f} ppl={ppl:.2f} " |
| f"aux={aux_val:.6f} lr={lr:.1e} tok/s={tok_s:.0f} " |
| f"gpu={gpu_mem:.1f}GB {elapsed:.0f}s") |
| running_loss = 0.0 |
| running_steps = 0 |
|
|
| if step % SAVE_EVERY == 0: |
| log(f"Saving checkpoint step {step}...") |
| sd = model.state_dict() |
| from safetensors.torch import save_file |
| fname = f"checkpoint-{step}.safetensors" |
| save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(LOCAL_OUT / fname)) |
| save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(KAGGLE_OUT / fname)) |
| (LOCAL_OUT / "config.json").write_text(json.dumps(config_dict, indent=2)) |
| (KAGGLE_OUT / "config.json").write_text(json.dumps(config_dict, indent=2)) |
|
|
| |
| log("Training complete! Saving final model...") |
| sd = model.state_dict() |
| from safetensors.torch import save_file |
| for out_dir in [LOCAL_OUT, KAGGLE_OUT]: |
| save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(out_dir / "model_final.safetensors")) |
| (out_dir / "config.json").write_text(json.dumps(config_dict, indent=2)) |
| if tok_path.exists(): |
| import shutil |
| shutil.copy(tok_path, out_dir / "tokenizer.json") |
|
|
| total_h = (time.time() - t0) / 3600 |
| info(f"Total time: {total_h:.1f} hours") |
| info(f"Model saved to: {KAGGLE_OUT}/ ← download from Kaggle 'Output' tab!") |
|
|
| |
| |
| |
| log("=== 6/7: SSH tunnel ===") |
| os.system("sed -i '/^PermitRootLogin\\|^PasswordAuthentication\\|^UseDNS/d' /etc/ssh/sshd_config") |
| with open("/etc/ssh/sshd_config", "a") as f: |
| f.write("PermitRootLogin yes\nPasswordAuthentication yes\nUseDNS no\n") |
| os.system(f"echo 'root:{SSH_PASSWORD}' | chpasswd") |
| os.system("mkdir -p /var/run/sshd && /usr/sbin/sshd 2>/dev/null") |
|
|
| log("Starting bore tunnel...") |
| os.system("curl -sL https://github.com/ekzhang/bore/releases/download/v0.5.2/bore-v0.5.2-x86_64-unknown-linux-musl.tar.gz | tar xz -C /usr/local/bin/ 2>/dev/null") |
|
|
| port = "" |
| if os.path.exists("/usr/local/bin/bore"): |
| bore_proc = subprocess.Popen( |
| ["bore", "local", "22", "--to", "bore.pub"], |
| stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True |
| ) |
| start = time.time() |
| while time.time() - start < 30: |
| assert bore_proc.stdout is not None |
| line = bore_proc.stdout.readline() |
| if line: |
| print(f" {line.rstrip()}") |
| m = re.search(r"bore\.pub:(\d+)", line) |
| if m: |
| port = m.group(1) |
| break |
|
|
| if port: |
| print() |
| print("=" * 60) |
| print(" SSH TUNNEL READY") |
| print("=" * 60) |
| print(f"\n ssh root@bore.pub -p {port} -o StrictHostKeyChecking=no") |
| print(f" Password: {SSH_PASSWORD}") |
| print(f"\n Model files at: {KAGGLE_OUT}/") |
| print(" KEEP THIS CELL RUNNING\n") |
| print("=" * 60) |
| try: |
| while True: time.sleep(10) |
| except KeyboardInterrupt: |
| pass |
|
|
| |
| |
| |
| log("=== 7/7: Complete ===") |
| info("Download from Kaggle: 'Output' tab → angstrom_output/") |
| for f in sorted(KAGGLE_OUT.glob("*")): |
| info(f" {f.name} ({f.stat().st_size / 1e6:.1f}MB)") |
|
|
