Add ultron/train.py

ultron/train.py

@@ -0,0 +1,270 @@
+"""
+Ultron Training Script — FineWeb-Edu Pretraining
+
+Based on Parcae's training recipe:
+  - AdamW with cosine decay
+  - Per-sequence depth sampling within micro-batches
+  - Backprop truncation: μ_bwd = ⌈μ_rec/2⌉
+  - Linear warmup (2000 steps)
+  - bfloat16 mixed precision
+
+Usage:
+  Single GPU:  python train.py
+  Multi-GPU:   torchrun --nproc_per_node=N train.py
+"""
+
+import os
+import math
+import time
+import json
+import argparse
+from dataclasses import asdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import autocast, GradScaler
+
+from ultron.model import Ultron, UltronConfig
+from ultron.variants import ultron_small, ultron_base, ultron_medium, ultron_large
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Ultron Pretraining")
+    parser.add_argument("--variant", type=str, default="small",
+                        choices=["small", "base", "medium", "large", "custom"])
+    parser.add_argument("--dataset", type=str, default="HuggingFaceFW/fineweb-edu")
+    parser.add_argument("--dataset_subset", type=str, default="sample-10BT")
+    parser.add_argument("--seq_len", type=int, default=1024)
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--grad_accum", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=3e-4)
+    parser.add_argument("--min_lr", type=float, default=3e-5)
+    parser.add_argument("--warmup_steps", type=int, default=2000)
+    parser.add_argument("--max_steps", type=int, default=50000)
+    parser.add_argument("--weight_decay", type=float, default=0.1)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--log_interval", type=int, default=10)
+    parser.add_argument("--save_interval", type=int, default=5000)
+    parser.add_argument("--eval_interval", type=int, default=500)
+    parser.add_argument("--output_dir", type=str, default="./ultron_checkpoints")
+    parser.add_argument("--push_to_hub", action="store_true")
+    parser.add_argument("--hub_model_id", type=str, default=None)
+    parser.add_argument("--use_moe", action="store_true")
+    parser.add_argument("--gradient_checkpointing", action="store_true")
+    parser.add_argument("--per_seq_depth", action="store_true", default=True,
+                        help="Per-sequence depth sampling (Parcae)")
+    return parser.parse_args()
+
+
+def get_lr(step: int, warmup_steps: int, max_steps: int, lr: float, min_lr: float) -> float:
+    """Cosine decay with linear warmup."""
+    if step < warmup_steps:
+        return lr * step / warmup_steps
+    if step >= max_steps:
+        return min_lr
+    progress = (step - warmup_steps) / (max_steps - warmup_steps)
+    return min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * progress))
+
+
+def sample_depth(mu_rec: int, per_sequence: bool = True, batch_size: int = 1) -> list[int]:
+    """Sample recurrence depths.
+
+    Per-sequence sampling (Parcae): each sequence in the batch gets a
+    different depth, reducing loss spike variance.
+    """
+    if per_sequence:
+        # Geometric distribution centered on mu_rec (Parcae's Λ)
+        depths = []
+        for _ in range(batch_size):
+            # Sample from [1, 2*mu_rec] with peak at mu_rec
+            d = max(1, min(2 * mu_rec, int(torch.distributions.Geometric(
+                probs=1.0 / mu_rec
+            ).sample().item()) + 1))
+            depths.append(d)
+        return depths
+    else:
+        return [mu_rec] * batch_size
+
+
+def main():
+    args = get_args()
+
+    # Setup
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    use_bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported()
+
+    # Model config
+    variant_map = {
+        "small": ultron_small,
+        "base": ultron_base,
+        "medium": ultron_medium,
+        "large": ultron_large,
+    }
+    cfg = variant_map[args.variant]()
+
+    if args.use_moe:
+        cfg.use_moe = True
+    if args.gradient_checkpointing:
+        cfg.gradient_checkpointing = True
+    cfg.max_seq_len = args.seq_len
+
+    # Build model
+    model = Ultron(cfg).to(device)
+    total_params = model.get_num_params(non_embedding=False)
+    print(f"Ultron {args.variant} | {total_params:,} params | device: {device}")
+    print(f"Config: {json.dumps(asdict(cfg), indent=2)}")
+    print(f"Spectral radius ρ(A): {model.get_spectral_radius():.6f}")
+
+    # Optimizer (Parcae recipe)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.lr,
+        betas=(0.9, 0.95),
+        weight_decay=args.weight_decay,
+    )
+
+    # Dataset
+    try:
+        from datasets import load_dataset
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        print(f"Loading dataset: {args.dataset}/{args.dataset_subset}...")
+        ds = load_dataset(args.dataset, args.dataset_subset, split="train", streaming=True)
+
+        def tokenize_batch(examples):
+            return tokenizer(examples["text"], truncation=True, max_length=args.seq_len + 1,
+                           padding="max_length", return_tensors="pt")
+
+    except ImportError:
+        print("WARNING: datasets/transformers not installed. Using random data for testing.")
+        ds = None
+
+    # Training loop
+    os.makedirs(args.output_dir, exist_ok=True)
+    scaler = GradScaler() if device.type == "cuda" and not use_bf16 else None
+
+    model.train()
+    step = 0
+    start_time = time.time()
+    running_loss = 0.0
+
+    print(f"\nTraining for {args.max_steps} steps...")
+    print(f"  Effective batch size: {args.batch_size * args.grad_accum}")
+    print(f"  Sequence length: {args.seq_len}")
+    print(f"  Per-sequence depth sampling: {args.per_seq_depth}")
+    print(f"  Mixed precision: {'bf16' if use_bf16 else 'fp16' if scaler else 'fp32'}")
+    print()
+
+    while step < args.max_steps:
+        # Get batch
+        if ds is not None:
+            # Real data path
+            batch_texts = []
+            for item in ds:
+                batch_texts.append(item["text"])
+                if len(batch_texts) >= args.batch_size:
+                    break
+            tokens = tokenizer(batch_texts, truncation=True, max_length=args.seq_len + 1,
+                             padding="max_length", return_tensors="pt")
+            input_ids = tokens["input_ids"].to(device)
+        else:
+            # Random data for testing
+            input_ids = torch.randint(0, cfg.vocab_size, (args.batch_size, args.seq_len + 1), device=device)
+
+        x = input_ids[:, :-1]
+        y = input_ids[:, 1:]
+
+        # Sample depths (Parcae per-sequence sampling)
+        depths = sample_depth(cfg.max_loop_iters, args.per_seq_depth, x.shape[0])
+
+        # Forward pass
+        lr = get_lr(step, args.warmup_steps, args.max_steps, args.lr, args.min_lr)
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+
+        with autocast(dtype=torch.bfloat16 if use_bf16 else torch.float16, enabled=device.type == "cuda"):
+            # For per-sequence depth, we process each sample with its own depth
+            # In practice, we use the mean depth for the batch
+            # (true per-sequence would need custom loss masking)
+            mean_depth = max(1, sum(depths) // len(depths))
+            logits = model(x, n_loops=mean_depth)
+            loss = F.cross_entropy(logits.reshape(-1, cfg.vocab_size), y.reshape(-1))
+            loss = loss / args.grad_accum
+
+        # Backward
+        if scaler:
+            scaler.scale(loss).backward()
+        else:
+            loss.backward()
+
+        running_loss += loss.item() * args.grad_accum
+
+        # Gradient accumulation step
+        if (step + 1) % args.grad_accum == 0:
+            if scaler:
+                scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            if scaler:
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                optimizer.step()
+            optimizer.zero_grad()
+
+        step += 1
+
+        # Logging
+        if step % args.log_interval == 0:
+            avg_loss = running_loss / args.log_interval
+            elapsed = time.time() - start_time
+            tokens_per_sec = (step * args.batch_size * args.seq_len) / elapsed
+            rho = model.get_spectral_radius()
+            print(f"step {step:>6d} | loss {avg_loss:.4f} | lr {lr:.2e} | "
+                  f"ρ(A) {rho:.4f} | depth {mean_depth} | "
+                  f"tok/s {tokens_per_sec:.0f} | elapsed {elapsed:.0f}s")
+            running_loss = 0.0
+
+        # Save checkpoint
+        if step % args.save_interval == 0:
+            ckpt_path = os.path.join(args.output_dir, f"step_{step}.pt")
+            torch.save({
+                "step": step,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "config": asdict(cfg),
+                "loss": avg_loss if step >= args.log_interval else float("inf"),
+            }, ckpt_path)
+            print(f"  Saved checkpoint: {ckpt_path}")
+
+            if args.push_to_hub and args.hub_model_id:
+                try:
+                    from huggingface_hub import HfApi
+                    api = HfApi()
+                    api.upload_file(
+                        path_or_fileobj=ckpt_path,
+                        path_in_repo=f"checkpoints/step_{step}.pt",
+                        repo_id=args.hub_model_id,
+                    )
+                    print(f"  Pushed to hub: {args.hub_model_id}")
+                except Exception as e:
+                    print(f"  Hub push failed: {e}")
+
+    print(f"\nTraining complete! {step} steps in {time.time() - start_time:.0f}s")
+    print(f"Final ρ(A): {model.get_spectral_radius():.6f}")
+
+    # Save final
+    final_path = os.path.join(args.output_dir, "final.pt")
+    torch.save({
+        "model_state_dict": model.state_dict(),
+        "config": asdict(cfg),
+    }, final_path)
+    print(f"Saved final model: {final_path}")
+
+
+if __name__ == "__main__":
+    main()