""" train_gpu.py — AngstromE1-Nano GPU Training Optimized for 2x NVIDIA T4 (16GB VRAM each) Features: - Mixed precision (FP16/BF16) - Gradient checkpointing - DataParallel across 2 GPUs - Cosine LR schedule with warmup - Gradient clipping - Periodic checkpointing - Wandb logging (optional) Usage: python train_gpu.py # default config python train_gpu.py --config large # larger model python train_gpu.py --resume checkpoint.pt # resume training """ import sys; sys.path.insert(0, '.') import os import math import time import json import argparse from pathlib import Path import torch import torch.nn as nn from torch.cuda.amp import GradScaler, autocast from torch.nn.parallel import DistributedDataParallel as DDP from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM from angstrom_nano.tokenizer import AngstromNanoTokenizer # ═══════════════════════════════════════════════════════════════════ # Configs # ═══════════════════════════════════════════════════════════════════ 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": 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, }, "large": { "vocab_size": 32000, "hidden_size": 1024, "intermediate_size": 4096, "num_hidden_layers": 24, "num_attention_heads": 16, "num_key_value_heads": 4, "head_dim": 64, "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, }, } # ═══════════════════════════════════════════════════════════════════ # Dataset # ═══════════════════════════════════════════════════════════════════ class TextDataset(torch.utils.data.Dataset): """Memory-mapped token dataset for large corpora.""" def __init__(self, token_ids: torch.Tensor, seq_len: int): 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): x = self.token_ids[idx : idx + self.seq_len] y = self.token_ids[idx + 1 : idx + self.seq_len + 1] return x, y # ═══════════════════════════════════════════════════════════════════ # Training # ═══════════════════════════════════════════════════════════════════ def setup_device(): """Setup multi-GPU or single GPU.""" if not torch.cuda.is_available(): print("WARNING: No GPU found, using CPU (will be slow)") return torch.device("cpu"), 1 n_gpus = torch.cuda.device_count() device = torch.device("cuda:0") print(f"Using {n_gpus} GPU(s):") for i in range(n_gpus): props = torch.cuda.get_device_properties(i) print(f" GPU {i}: {props.name} ({props.total_mem / 1e9:.1f} GB)") return device, n_gpus def get_lr(step, warmup_steps, max_steps, base_lr, min_lr): """Cosine learning rate schedule with warmup.""" if step < warmup_steps: return base_lr * step / max(1, warmup_steps) progress = (step - warmup_steps) / max(1, max_steps - warmup_steps) return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress)) def count_params(model): """Count trainable parameters.""" return sum(p.numel() for p in model.parameters() if p.requires_grad) def save_checkpoint(model, optimizer, scaler, step, loss, config_dict, out_dir): """Save training checkpoint.""" out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) # Save model weights sd = model.module.state_dict() if hasattr(model, 'module') else model.state_dict() if "lm_head.weight" not in sd: sd["lm_head.weight"] = sd["model.embed_tokens.weight"] from safetensors.torch import save_file weights_path = out_dir / f"checkpoint-{step}.safetensors" save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(weights_path)) # Save training state state = { "step": step, "loss": loss, "config": config_dict, "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict() if scaler else None, } torch.save(state, str(out_dir / f"checkpoint-{step}.pt")) # Save config (out_dir / "config.json").write_text(json.dumps(config_dict, indent=2)) print(f" Saved checkpoint at step {step}") def train(args): """Main training loop.""" torch.manual_seed(42) if torch.cuda.is_available(): torch.cuda.manual_seed_all(42) # ── Setup ────────────────────────────────────────────────────── device, n_gpus = setup_device() use_amp = device.type == "cuda" and torch.cuda.is_available() # ── Tokenizer ────────────────────────────────────────────────── tok_path = Path("checkpoints/tokenizer.json") if tok_path.exists(): tok = AngstromNanoTokenizer.from_bpe_file(str(tok_path)) print(f"Loaded tokenizer: {len(tok)} vocab") else: print("Training new tokenizer...") tok = AngstromNanoTokenizer.train_bpe( [str(args.data_path)], vocab_size=args.vocab_size, out_path=str(tok_path), ) print(f"Trained tokenizer: {len(tok)} vocab") # ── Load and tokenize data ───────────────────────────────────── print(f"\nLoading data from {args.data_path}...") text = Path(args.data_path).read_text(encoding="utf-8") print(f" Raw: {len(text):,} chars ({len(text)/1e6:.1f} MB)") ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long) print(f" Tokens: {len(ids):,} ({len(ids)/1e6:.1f}M)") dataset = TextDataset(ids, args.seq_len) print(f" Samples: {len(dataset):,}") dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=2, pin_memory=True, drop_last=True, ) # ── Model ────────────────────────────────────────────────────── config_dict = CONFIGS[args.config] config_dict["vocab_size"] = len(tok) cfg = AngstromNanoConfig(**config_dict) model = AngstromNanoForCausalLM(cfg) n_params = count_params(model) print(f"\nModel: {n_params:,} params ({n_params * 4 / 1e6:.1f} MB FP32)") model = model.to(device) # Multi-GPU if n_gpus > 1: model = nn.DataParallel(model, device_ids=list(range(n_gpus))) print(f" Wrapped in DataParallel across {n_gpus} GPUs") # ── Optimizer ────────────────────────────────────────────────── optimizer = torch.optim.AdamW( model.parameters(), lr=args.lr, weight_decay=0.1, betas=(0.9, 0.95), ) scaler = GradScaler(enabled=use_amp) # ── Resume ───────────────────────────────────────────────────── start_step = 0 if args.resume and Path(args.resume).exists(): print(f"\nResuming from {args.resume}...") ckpt = torch.load(args.resume, map_location=device) if hasattr(model, 'module'): model.module.load_state_dict(ckpt["model"]) else: model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) start_step = ckpt["step"] print(f" Resumed at step {start_step}") # ── Training ─────────────────────────────────────────────────── max_steps = args.steps warmup_steps = args.warmup_steps grad_clip = args.grad_clip log_every = args.log_every save_every = args.save_every print(f"\nTraining for {max_steps} steps, seq_len={args.seq_len}, batch={args.batch_size}") print(f" LR: {args.lr} → {args.min_lr}, warmup: {warmup_steps} steps") print(f" Gradient clipping: {grad_clip}") print(f" Mixed precision: {use_amp}") print(f" Checkpoint every: {save_every} steps") print() model.train() t0 = time.time() running_loss = 0.0 running_steps = 0 for step in range(start_step + 1, max_steps + 1): # Get batch try: x, y = next(dataloader_iter) except (StopIteration, NameError): dataloader_iter = iter(dataloader) x, y = next(dataloader_iter) x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) # Forward pass with AMP with autocast(enabled=use_amp, dtype=torch.float16): out = model(x, labels=y, output_router_logits=True) loss = out["loss"] aux_loss = out["aux_loss"] # Backward pass optimizer.zero_grad() scaler.scale(loss).backward() # Gradient clipping if grad_clip > 0: scaler.unscale_(optimizer) nn.utils.clip_grad_norm_(model.parameters(), grad_clip) scaler.step(optimizer) scaler.update() # Update LR lr = get_lr(step, warmup_steps, max_steps, args.lr, args.min_lr) for param_group in optimizer.param_groups: param_group["lr"] = lr # Logging running_loss += loss.item() running_steps += 1 if step % log_every == 0 or step == 1: avg_loss = running_loss / running_steps ppl = math.exp(min(avg_loss, 20)) # cap to avoid overflow elapsed = time.time() - t0 tokens_per_sec = (args.batch_size * args.seq_len * running_steps) / elapsed gpu_mem = torch.cuda.memory_allocated(0) / 1e9 if device.type == "cuda" else 0 print(f" step {step:>6d}/{max_steps} loss={avg_loss:.4f} ppl={ppl:.2f} " f"aux={aux_loss.item():.6f} lr={lr:.1e} " f"tok/s={tokens_per_sec:.0f} gpu={gpu_mem:.1f}GB " f"{elapsed:.0f}s") running_loss = 0.0 running_steps = 0 # Save checkpoint if step % save_every == 0: save_checkpoint(model, optimizer, scaler, step, avg_loss, config_dict, args.output_dir) # ── Final save ───────────────────────────────────────────────── print("\nTraining complete!") save_checkpoint(model, optimizer, scaler, max_steps, avg_loss, config_dict, args.output_dir) # Save final model as the main model file final_path = Path(args.output_dir) / "model_final.safetensors" sd = model.module.state_dict() if hasattr(model, 'module') else model.state_dict() if "lm_head.weight" not in sd: sd["lm_head.weight"] = sd["model.embed_tokens.weight"] from safetensors.torch import save_file save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(final_path)) print(f"Saved final model: {final_path}") total_time = time.time() - t0 print(f"Total time: {total_time/3600:.1f} hours") # ═══════════════════════════════════════════════════════════════════ # CLI # ═══════════════════════════════════════════════════════════════════ if __name__ == "__main__": parser = argparse.ArgumentParser(description="AngstromE1-Nano GPU Training") parser.add_argument("--config", default="medium", choices=["small", "medium", "large"], help="Model config (default: medium)") parser.add_argument("--data-path", default="data/corpus.txt", help="Path to training corpus") parser.add_argument("--output-dir", default="checkpoints", help="Output directory for checkpoints") parser.add_argument("--vocab-size", type=int, default=16384, help="BPE vocab size (if training tokenizer)") parser.add_argument("--seq-len", type=int, default=512, help="Sequence length (default: 512)") parser.add_argument("--batch-size", type=int, default=4, help="Batch size per GPU (default: 4)") parser.add_argument("--steps", type=int, default=50000, help="Total training steps (default: 50000)") parser.add_argument("--lr", type=float, default=3e-3, help="Peak learning rate (default: 3e-3)") parser.add_argument("--min-lr", type=float, default=3e-4, help="Min learning rate (default: 3e-4)") parser.add_argument("--warmup-steps", type=int, default=500, help="Warmup steps (default: 500)") parser.add_argument("--grad-clip", type=float, default=1.0, help="Gradient clipping (default: 1.0)") parser.add_argument("--log-every", type=int, default=100, help="Log every N steps (default: 100)") parser.add_argument("--save-every", type=int, default=5000, help="Save checkpoint every N steps (default: 5000)") parser.add_argument("--resume", type=str, default=None, help="Resume from checkpoint path") args = parser.parse_args() train(args)