Upload smoke_test_300m.py with huggingface_hub

smoke_test_300m.py

@@ -0,0 +1,579 @@
+"""
+LUNA 300M — Smoke Test (2 training steps)
+==========================================
+Downloads the dataset, builds the model, runs 2 training steps,
+saves a checkpoint, runs validation generation, then exits.
+
+If everything prints "SMOKE TEST PASSED" at the end, the full
+training pipeline is confirmed working.
+
+Usage (cloud GPU):
+    python smoke_test_300m.py
+
+Usage (local / CPU):
+    python smoke_test_300m.py --data_path Base/data/litdata_pretrain_final
+
+Options:
+    --data_path   : Path to local litdata dataset (skip download)
+    --hf_repo     : HF dataset repo (default: ASTERIZER/Luna_Dataset)
+    --steps       : Number of test steps (default: 2)
+"""
+
+import os
+import sys
+import gc
+import json
+import time
+import shutil
+import struct
+import argparse
+import traceback
+from pathlib import Path
+
+# ─── Ensure deps ──────────────────────────────────────────────────────────────
+
+def ensure_packages():
+    """Install missing packages quietly."""
+    required = ["torch", "psutil", "yaml", "transformers", "huggingface_hub"]
+    pip_names = {"yaml": "pyyaml"}
+    missing = []
+    for pkg in required:
+        try:
+            __import__(pkg)
+        except ImportError:
+            missing.append(pip_names.get(pkg, pkg))
+    if missing:
+        print(f"  Installing missing packages: {missing}")
+        os.system(f"{sys.executable} -m pip install -q " + " ".join(missing))
+
+ensure_packages()
+
+import yaml
+import psutil
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch.amp import autocast, GradScaler
+
+
+# ─── Stage tracker ────────────────────────────────────────────────────────────
+
+class StageTracker:
+    def __init__(self):
+        self.stages = []
+        self.current = None
+
+    def start(self, name):
+        self.current = name
+        print(f"\n{'─' * 60}")
+        print(f"  STAGE: {name}")
+        print(f"{'─' * 60}")
+
+    def ok(self, detail=""):
+        msg = f"  ✓ {self.current}"
+        if detail:
+            msg += f" — {detail}"
+        print(msg)
+        self.stages.append((self.current, True, detail))
+
+    def fail(self, detail=""):
+        msg = f"  ✗ {self.current}"
+        if detail:
+            msg += f" — {detail}"
+        print(msg)
+        self.stages.append((self.current, False, detail))
+
+    def summary(self):
+        print(f"\n{'=' * 60}")
+        print(f"  SMOKE TEST SUMMARY")
+        print(f"{'=' * 60}")
+        all_pass = True
+        for name, passed, detail in self.stages:
+            icon = "PASS" if passed else "FAIL"
+            line = f"  [{icon}] {name}"
+            if detail:
+                line += f"  ({detail})"
+            print(line)
+            if not passed:
+                all_pass = False
+        print(f"{'=' * 60}")
+        if all_pass:
+            print(f"  >>> SMOKE TEST PASSED — pipeline is ready for full training <<<")
+        else:
+            print(f"  >>> SMOKE TEST FAILED — see above for details <<<")
+        print(f"{'=' * 60}\n")
+        return all_pass
+
+
+# ─── Model (same architecture as train_300m.py) ──────────────────────────────
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len=1024):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        t = torch.arange(max_seq_len).float()
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("cos_cached", emb.cos())
+        self.register_buffer("sin_cached", emb.sin())
+
+    def forward(self, seq_len):
+        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rotary(x, cos, sin):
+    c = cos.unsqueeze(0).unsqueeze(0)
+    s = sin.unsqueeze(0).unsqueeze(0)
+    return x * c + rotate_half(x) * s
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
+        super().__init__()
+        self.n_head = n_head
+        self.head_dim = n_embd // n_head
+        self.rot_dim = int(self.head_dim * rotary_pct)
+        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
+        self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
+        self.rotary = RotaryEmbedding(self.rot_dim, block_size)
+
+    def forward(self, x):
+        B, T, C = x.size()
+        qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        cos, sin = self.rotary(T)
+        q = torch.cat([apply_rotary(q[..., :self.rot_dim], cos, sin), q[..., self.rot_dim:]], dim=-1)
+        k = torch.cat([apply_rotary(k[..., :self.rot_dim], cos, sin), k[..., self.rot_dim:]], dim=-1)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))
+
+
+class MLP(nn.Module):
+    def __init__(self, n_embd):
+        super().__init__()
+        self.fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
+        self.gelu = nn.GELU()
+        self.proj = nn.Linear(4 * n_embd, n_embd, bias=True)
+
+    def forward(self, x):
+        return self.proj(self.gelu(self.fc(x)))
+
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_head, block_size):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
+        self.ln2 = nn.LayerNorm(n_embd)
+        self.mlp = MLP(n_embd)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class LUNAModel(nn.Module):
+    def __init__(self, vocab_size, block_size, n_layer, n_embd, n_head):
+        super().__init__()
+        self.wte = nn.Embedding(vocab_size, n_embd)
+        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
+        self.ln_f = nn.LayerNorm(n_embd)
+        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
+        self.lm_head.weight = self.wte.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            m.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, idx, targets=None, return_logits=True):
+        x = self.wte(idx)
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+            if not return_logits:
+                logits = None
+        return logits, loss
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters()) - self.wte.weight.numel()
+
+
+# ─── LitData Dataset (same as train_300m.py) ─────────────────────────────────
+
+class LitDataDataset(torch.utils.data.Dataset):
+    def __init__(self, data_path, block_size=1024):
+        self.block_size = block_size
+        self.data_path = Path(data_path)
+        with open(self.data_path / "index.json") as f:
+            idx = json.load(f)
+        self.chunks_meta = idx["chunks"]
+        self._cum_blocks = []
+        total = 0
+        for c in self.chunks_meta:
+            n = c["dim"] // (block_size + 1)
+            total += n
+            self._cum_blocks.append(total)
+        self.total_blocks = total
+        self._chunk_cache = {}
+
+    def _load_chunk(self, chunk_idx):
+        if chunk_idx in self._chunk_cache:
+            return self._chunk_cache[chunk_idx]
+        meta = self.chunks_meta[chunk_idx]
+        with open(self.data_path / meta["filename"], "rb") as f:
+            raw = f.read()
+        num_items = struct.unpack_from("<I", raw, 0)[0]
+        header_bytes = (num_items + 2) * 4
+        tokens = torch.from_numpy(np.frombuffer(raw[header_bytes:], dtype=np.int32).copy())
+        if len(self._chunk_cache) >= 4:
+            del self._chunk_cache[next(iter(self._chunk_cache))]
+        self._chunk_cache[chunk_idx] = tokens
+        return tokens
+
+    def __len__(self):
+        return self.total_blocks
+
+    def __getitem__(self, idx):
+        chunk_idx = 0
+        for i, cum in enumerate(self._cum_blocks):
+            if idx < cum:
+                chunk_idx = i
+                break
+        prev = self._cum_blocks[chunk_idx - 1] if chunk_idx > 0 else 0
+        tokens = self._load_chunk(chunk_idx)
+        s = (idx - prev) * (self.block_size + 1)
+        e = s + self.block_size + 1
+        chunk = tokens[s:e]
+        if len(chunk) < self.block_size + 1:
+            pad = torch.zeros(self.block_size + 1, dtype=torch.int32)
+            pad[:len(chunk)] = chunk
+            chunk = pad
+        chunk = chunk.long()
+        return chunk[:self.block_size], chunk[1:self.block_size + 1]
+
+
+# ─── Smoke test logic ────────────────────────────────────────────────────────
+
+MODEL_CFG = {
+    "vocab_size": 50304,
+    "seq_len":    1024,
+    "n_layer":    20,
+    "n_embd":     1024,
+    "n_head":     16,
+}
+
+VALIDATION_PROMPTS = [
+    "The theory of general relativity describes gravity as",
+    "In machine learning, neural networks learn patterns by",
+    "The Amazon rainforest spans across several countries and",
+]
+
+
+def detect_hardware():
+    info = {
+        "cpu_cores": os.cpu_count() or 4,
+        "ram_gb": psutil.virtual_memory().total / 1024**3,
+    }
+    if torch.cuda.is_available():
+        props = torch.cuda.get_device_properties(0)
+        info.update({
+            "device": "cuda",
+            "gpu_name": props.name,
+            "vram_gb": props.total_memory / 1024**3,
+        })
+        if props.major >= 8:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+            info["dtype"] = torch.bfloat16
+        else:
+            info["dtype"] = torch.float16
+    else:
+        info.update({
+            "device": "cpu",
+            "gpu_name": "CPU",
+            "vram_gb": 0,
+            "dtype": torch.float32,
+        })
+    return info
+
+
+def fetch_dataset(data_path, hf_repo):
+    """Fetch dataset — either use local path or download from HF."""
+    dp = Path(data_path)
+    if (dp / "index.json").exists():
+        print(f"  Dataset found at {dp}")
+        return str(dp)
+
+    # Try auto-discovery in common locations
+    for candidate in [dp, Path("Base/data/litdata_pretrain_final"),
+                      Path("/workspace/data/litdata_pretrain_final")]:
+        if (candidate / "index.json").exists():
+            print(f"  Dataset found at {candidate}")
+            return str(candidate)
+
+    # Download from HF
+    print(f"  Dataset not found locally. Downloading from HF: {hf_repo}")
+    out_dir = Path("/workspace/data/litdata_pretrain_final")
+    if not out_dir.parent.exists():
+        out_dir = Path(".smoke_test_data")
+
+    import subprocess
+    result = subprocess.run(
+        [sys.executable, "fetch_data.py",
+         "--source", "huggingface",
+         "--hf_repo", hf_repo,
+         "--out_dir", str(out_dir),
+         "--hf_token", os.environ.get("HF_TOKEN", "")],
+        capture_output=False
+    )
+    if result.returncode != 0:
+        raise RuntimeError("Dataset download failed!")
+    return str(out_dir)
+
+
+@torch.no_grad()
+def run_quick_validation(model, tokenizer, device, dtype, seq_len=1024, max_new=32):
+    """Generate short text from a few prompts to verify generation works."""
+    model.eval()
+    for i, prompt in enumerate(VALIDATION_PROMPTS):
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+        if input_ids.shape[1] > seq_len:
+            input_ids = input_ids[:, -seq_len:]
+        generated = input_ids.clone()
+        for _ in range(max_new):
+            ctx = generated[:, -seq_len:] if generated.shape[1] > seq_len else generated
+            with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
+                logits, _ = model(ctx)
+            next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
+            generated = torch.cat([generated, next_token], dim=1)
+        output = tokenizer.decode(generated[0], skip_special_tokens=True)
+        continuation = output[len(prompt):].strip()[:120]
+        print(f"    [{i+1}] {prompt}")
+        print(f"        -> {continuation}")
+    model.train()
+
+
+def run_smoke_test(args):
+    tracker = StageTracker()
+
+    # ── Stage 1: Hardware ─────────────────────────────────────────────────────
+    tracker.start("Hardware Detection")
+    try:
+        hw = detect_hardware()
+        device = torch.device(hw["device"])
+        dtype = hw["dtype"]
+        tracker.ok(f"{hw['gpu_name']}, {hw['ram_gb']:.0f}GB RAM, dtype={dtype}")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 2: Dataset ──────────────────────────────────────────────────────
+    tracker.start("Dataset Fetch & Verify")
+    try:
+        data_path = fetch_dataset(args.data_path, args.hf_repo)
+        with open(Path(data_path) / "index.json") as f:
+            idx = json.load(f)
+        chunks = idx.get("chunks", [])
+        total_tokens = sum(c.get("dim", 0) for c in chunks)
+        present = sum(1 for c in chunks if (Path(data_path) / c["filename"]).exists())
+        missing = len(chunks) - present
+        if missing > 0:
+            tracker.fail(f"{missing} chunks missing out of {len(chunks)}")
+            return tracker.summary()
+        tracker.ok(f"{len(chunks)} chunks, {total_tokens:,} tokens, 0 missing")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 3: Model Init ───────────────────────────────────────────────────
+    tracker.start("Model Initialization (300M)")
+    try:
+        model = LUNAModel(
+            vocab_size=MODEL_CFG["vocab_size"],
+            block_size=MODEL_CFG["seq_len"],
+            n_layer=MODEL_CFG["n_layer"],
+            n_embd=MODEL_CFG["n_embd"],
+            n_head=MODEL_CFG["n_head"],
+        ).to(device)
+        n_params = sum(p.numel() for p in model.parameters())
+        tracker.ok(f"{n_params:,} parameters on {device}")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 4: Dataset Loading ──────────────────────────────────────────────
+    tracker.start("Dataset Loading (LitData)")
+    try:
+        dataset = LitDataDataset(data_path, block_size=MODEL_CFG["seq_len"])
+        loader = torch.utils.data.DataLoader(
+            dataset, batch_size=1, shuffle=True,
+            num_workers=0, pin_memory=False, drop_last=True,
+        )
+        x_sample, t_sample = next(iter(loader))
+        tracker.ok(f"{len(dataset):,} blocks, sample shape={list(x_sample.shape)}")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 5: Tokenizer ────────────────────────────────────────────────────
+    tracker.start("Tokenizer Loading")
+    try:
+        from transformers import AutoTokenizer
+        # Try common tokenizer locations
+        tok_dir = None
+        for candidate in [
+            Path("Base/checkpoints/EleutherAI/pythia-160m"),
+            Path("/workspace/Base/checkpoints/EleutherAI/pythia-160m"),
+            Path("/workspace/LUNA/Base/checkpoints/EleutherAI/pythia-160m"),
+        ]:
+            if candidate.exists():
+                tok_dir = candidate
+                break
+
+        if tok_dir is None:
+            # Fall back to downloading from HF
+            tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-160m")
+            tracker.ok(f"Downloaded from HF (EleutherAI/pythia-160m)")
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(str(tok_dir))
+            tracker.ok(f"Loaded from {tok_dir}")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 6: Forward + Backward Pass ──────────────────────────────────────
+    tracker.start(f"Training Loop ({args.steps} steps)")
+    try:
+        model.train()
+        optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
+        use_scaler = dtype == torch.float16
+        scaler = GradScaler(enabled=use_scaler)
+        data_iter = iter(loader)
+
+        losses = []
+        for step in range(1, args.steps + 1):
+            t0 = time.perf_counter()
+            try:
+                x, t = next(data_iter)
+            except StopIteration:
+                data_iter = iter(loader)
+                x, t = next(data_iter)
+
+            x = x.to(device, non_blocking=True)
+            t = t.to(device, non_blocking=True)
+
+            optimizer.zero_grad(set_to_none=True)
+            with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
+                _, loss = model(x, t, return_logits=False)
+            scaler.scale(loss).backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+
+            if device.type == "cuda":
+                torch.cuda.synchronize()
+
+            dt = time.perf_counter() - t0
+            loss_val = loss.item()
+            losses.append(loss_val)
+            tps = MODEL_CFG["seq_len"] / dt
+            vram = torch.cuda.max_memory_allocated() / 1024**3 if device.type == "cuda" else 0
+            print(f"    step {step}/{args.steps} | loss={loss_val:.4f} | "
+                  f"{tps:,.0f} tok/s | {dt:.2f}s | VRAM={vram:.1f}GB")
+
+        avg_loss = sum(losses) / len(losses)
+        tracker.ok(f"avg_loss={avg_loss:.4f}, all {args.steps} steps completed")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 7: Checkpoint Save ──────────────────────────────────────────────
+    tracker.start("Checkpoint Save")
+    try:
+        out_dir = Path(".smoke_test_output")
+        out_dir.mkdir(parents=True, exist_ok=True)
+        torch.save(model.state_dict(), out_dir / "lit_model.pth")
+        with open(out_dir / "model_config.json", "w") as f:
+            json.dump(MODEL_CFG, f, indent=2)
+        size_mb = (out_dir / "lit_model.pth").stat().st_size / 1024**2
+        tracker.ok(f"Saved to {out_dir} ({size_mb:.0f} MB)")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 8: Validation Generation ────────────────────────────────────────
+    tracker.start("Validation Generation (3 prompts)")
+    try:
+        run_quick_validation(model, tokenizer, device, dtype,
+                             seq_len=MODEL_CFG["seq_len"], max_new=32)
+        tracker.ok("Generated text from all prompts")
+    except Exception as e:
+        tracker.fail(str(e))
+        traceback.print_exc()
+        return tracker.summary()
+
+    # ── Stage 9: Cleanup ──────────────────────────────────────────────────────
+    tracker.start("Cleanup")
+    try:
+        if Path(".smoke_test_output").exists():
+            shutil.rmtree(".smoke_test_output")
+        if Path(".smoke_test_data").exists():
+            shutil.rmtree(".smoke_test_data")
+        del model, optimizer, scaler, dataset, loader
+        gc.collect()
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+        tracker.ok("Cleaned up temp files and freed memory")
+    except Exception as e:
+        tracker.fail(str(e))
+
+    # ── Summary ───────────────────────────────────────────────────────────────
+    return tracker.summary()
+
+
+# ─── Entry ────────────────────────────────────────────────────────────────────
+
+def parse_args():
+    p = argparse.ArgumentParser(description="LUNA 300M Smoke Test")
+    p.add_argument("--data_path", type=str,
+                   default="Base/data/litdata_pretrain_final",
+                   help="Local dataset path (auto-downloads from HF if not found)")
+    p.add_argument("--hf_repo", type=str,
+                   default="ASTERIZER/Luna_Dataset",
+                   help="HF dataset repo to download from if local not found")
+    p.add_argument("--steps", type=int, default=2,
+                   help="Number of training steps to run (default: 2)")
+    return p.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    print("=" * 60)
+    print("  LUNA 300M — SMOKE TEST")
+    print("  This runs 2 training steps to verify the full pipeline")
+    print("=" * 60)
+    passed = run_smoke_test(args)
+    sys.exit(0 if passed else 1)