| """ |
| 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 = { |
| "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, |
| }, |
| } |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
| |
| 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)) |
|
|
| |
| 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")) |
|
|
| |
| (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) |
|
|
| |
| device, n_gpus = setup_device() |
| use_amp = device.type == "cuda" and torch.cuda.is_available() |
|
|
| |
| 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") |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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) |
|
|
| |
| if n_gpus > 1: |
| model = nn.DataParallel(model, device_ids=list(range(n_gpus))) |
| print(f" Wrapped in DataParallel across {n_gpus} GPUs") |
|
|
| |
| optimizer = torch.optim.AdamW( |
| model.parameters(), lr=args.lr, |
| weight_decay=0.1, betas=(0.9, 0.95), |
| ) |
| scaler = GradScaler(enabled=use_amp) |
|
|
| |
| 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}") |
|
|
| |
| 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): |
| |
| 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) |
|
|
| |
| 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"] |
|
|
| |
| 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, warmup_steps, max_steps, args.lr, args.min_lr) |
| for param_group in optimizer.param_groups: |
| param_group["lr"] = lr |
|
|
| |
| 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)) |
| 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 |
|
|
| |
| if step % save_every == 0: |
| save_checkpoint(model, optimizer, scaler, step, avg_loss, |
| config_dict, args.output_dir) |
|
|
| |
| print("\nTraining complete!") |
| save_checkpoint(model, optimizer, scaler, max_steps, avg_loss, |
| config_dict, args.output_dir) |
|
|
| |
| 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") |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
