achbj
feat: implement project scaffolding, training pipelines, and utility scripts for NovaMind model development
3dfccac | """ | |
| NovaMind-256M Training Script | |
| ============================== | |
| 3-phase training pipeline with FULL logging and tracking: | |
| LOGGING (every optimizer step): | |
| • Rich console table: step, loss (total/think/response), LR, grad_norm, | |
| tokens/sec, GPU mem, ETA, elapsed time | |
| • Local JSON-lines log (append-only, survives crashes) | |
| • CSV log (easy to open in Excel / pandas) | |
| • W&B integration (graceful fallback to local if unavailable) | |
| • Red-flag warnings (high grad_norm, loss stuck, OOM risk) | |
| CHECKPOINTING (aggressive — for Lightning.ai 1-hr sessions): | |
| • Rolling checkpoint every `save_every` steps (default 500, ~10 min) | |
| • Milestone checkpoint every `milestone_every` steps (kept forever) | |
| • Phase-final checkpoint at end of each phase | |
| • Every checkpoint is 100% resumable: model + optimizer + scaler + | |
| scheduler_state + step + tokens_seen + full loss histories | |
| TRAINING: | |
| • Tag-Aware Loss Curriculum (Novel Contribution #2) | |
| • BF16 + Flash Attention 2 on H100 | |
| • Gradient checkpointing (saves ~40% VRAM) | |
| • Separate weight-decay / no-decay param groups | |
| • Mid-epoch resume after Lightning.ai session expiry | |
| Usage: | |
| python train.py --phase pretrain --out_dir ./checkpoints | |
| python train.py --phase sft --resume ./checkpoints/novamind_phase1_final.pt | |
| python train.py --phase cot_sft --resume ./checkpoints/novamind_phase2_final.pt | |
| python train.py --phase pretrain --smoke_test # quick 50-step pipeline test | |
| """ | |
| import os, sys, json, csv, time, math, argparse, datetime, shutil | |
| from pathlib import Path | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torch.utils.data import Dataset, DataLoader | |
| from tqdm import tqdm | |
| # ── Optional W&B ────────────────────────────────────────────────────────────── | |
| try: | |
| import wandb | |
| WANDB_AVAILABLE = True | |
| except ImportError: | |
| WANDB_AVAILABLE = False | |
| from model import NovaMind256M, NovaMindConfig, tag_aware_loss | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # ARGUMENT PARSING | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| def parse_args(): | |
| p = argparse.ArgumentParser(description="NovaMind-256M Training") | |
| # Phase | |
| p.add_argument("--phase", required=True, choices=["pretrain", "sft", "cot_sft"]) | |
| # Checkpointing | |
| p.add_argument("--resume", default=None, help="Path to checkpoint to resume from") | |
| p.add_argument("--out_dir", default="./checkpoints", help="Output dir for checkpoints + logs") | |
| # Model (defaults match NovaMindConfig) | |
| p.add_argument("--d_model", type=int, default=1024) | |
| p.add_argument("--n_heads", type=int, default=16) | |
| p.add_argument("--n_kv_heads", type=int, default=4) | |
| p.add_argument("--n_layers", type=int, default=24) | |
| p.add_argument("--ff_dim", type=int, default=2304) # ← 2304, not 2816 | |
| p.add_argument("--vocab_size", type=int, default=32000) | |
| p.add_argument("--max_seq_len",type=int, default=2048) | |
| # Training hypers (None = use phase defaults) | |
| p.add_argument("--batch_size", type=int, default=None) | |
| p.add_argument("--grad_accum_steps", type=int, default=None) | |
| p.add_argument("--lr", type=float, default=None) | |
| p.add_argument("--min_lr_ratio", type=float, default=0.1) | |
| p.add_argument("--warmup_steps", type=int, default=None) | |
| p.add_argument("--epochs", type=int, default=None) | |
| p.add_argument("--max_steps", type=int, default=None) | |
| p.add_argument("--max_grad_norm", type=float, default=1.0) | |
| p.add_argument("--weight_decay", type=float, default=0.1) | |
| # Data | |
| p.add_argument("--num_workers", type=int, default=4) | |
| p.add_argument("--data_cap", type=int, default=None, | |
| help="Cap dataset samples (useful for smoke tests)") | |
| p.add_argument("--data_file", default=None, | |
| help="Path to pre-tokenized .pt file (from prepare_data_phase1b.py). " | |
| "Skips all HuggingFace downloading — GPU starts training immediately.") | |
| # Checkpoint schedule | |
| p.add_argument("--save_every", type=int, default=500, | |
| help="Rolling checkpoint every N optimizer steps (~10 min on H100)") | |
| p.add_argument("--milestone_every",type=int, default=2000, | |
| help="Permanent milestone checkpoint every N steps") | |
| p.add_argument("--val_every", type=int, default=500, | |
| help="Validation + sample generation every N steps") | |
| p.add_argument("--val_batches", type=int, default=50) | |
| # Logging | |
| p.add_argument("--log_every", type=int, default=10, | |
| help="Console + JSON/CSV log every N optimizer steps") | |
| p.add_argument("--wandb_project", default="novamind-256m") | |
| p.add_argument("--wandb_run", default=None) | |
| p.add_argument("--no_wandb", action="store_true") | |
| # Misc | |
| p.add_argument("--seed", type=int, default=42) | |
| p.add_argument("--compile", action="store_true", help="torch.compile (PyTorch 2.0+)") | |
| p.add_argument("--smoke_test", action="store_true", help="50-step end-to-end test") | |
| return p.parse_args() | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # PHASE DEFAULTS | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| PHASE_DEFAULTS = { | |
| "pretrain": { | |
| "lr": 3e-4, "warmup_steps": 2000, "epochs": 2, | |
| "batch_size": 4, "grad_accum": 8, | |
| "think_weight": 0.0, | |
| "desc": "Phase 1 — Language Pretraining (~5.5B tokens: FineWeb-Edu + TinyStories)", | |
| }, | |
| "sft": { | |
| "lr": 1e-4, "warmup_steps": 500, "epochs": 2, | |
| "batch_size": 4, "grad_accum": 4, | |
| "think_weight": 0.5, | |
| "desc": "Phase 2 — Instruction SFT (~300M tokens: Alpaca + DailyDialog)", | |
| }, | |
| "cot_sft": { | |
| "lr": 5e-5, "warmup_steps": 200, "epochs": 2, | |
| "batch_size": 2, "grad_accum": 8, | |
| "think_weight": 1.5, | |
| "desc": "Phase 3 — CoT SFT (~150M tokens: OpenHermes + GSM8K)", | |
| }, | |
| } | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # DATASETS | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| class PretrainDataset(Dataset): | |
| """Phase 1: raw text chunks, uniform loss weights.""" | |
| def __init__(self, chunks: torch.Tensor): | |
| self.chunks = chunks # [N, seq_len+1] | |
| def __len__(self): return self.chunks.shape[0] | |
| def __getitem__(self, idx): | |
| c = self.chunks[idx].clone() | |
| return c[:-1], c[1:].clone(), torch.ones(c.shape[0]-1, dtype=torch.float32) | |
| class TagAwareDataset(Dataset): | |
| """ | |
| Phase 2/3: returns (input_ids, targets, token_weights). | |
| Weights encode Tag-Aware Loss Curriculum — see model.py for full explanation. | |
| """ | |
| def __init__(self, chunks: torch.Tensor, think_weight: float, | |
| think_start_id=-1, think_end_id=-1, | |
| asst_start_id=-1, user_start_id=-1, eos_id=2): | |
| self.chunks = chunks | |
| self.tw = think_weight | |
| self.think_start = think_start_id | |
| self.think_end = think_end_id | |
| self.asst_start = asst_start_id | |
| self.user_start = user_start_id | |
| self.eos = eos_id | |
| def __len__(self): return self.chunks.shape[0] | |
| def __getitem__(self, idx): | |
| c = self.chunks[idx].clone() | |
| x, y = c[:-1], c[1:].clone() | |
| return x, y, self._weights(y) | |
| def _weights(self, tokens: torch.Tensor) -> torch.Tensor: | |
| T = len(tokens) | |
| w = torch.ones(T, dtype=torch.float32) | |
| if self.think_start < 0 and self.asst_start < 0: | |
| return w | |
| state = "default" | |
| for i, tok in enumerate(tokens.tolist()): | |
| if tok == self.user_start: state = "user" | |
| elif tok == self.think_start: state = "think" | |
| elif tok == self.think_end: state = "default" | |
| elif tok == self.asst_start: state = "assistant" | |
| elif tok == self.eos: state = "default" | |
| base_w = 0.0 if state == "user" else (self.tw if state == "think" else 1.0) | |
| # FIX: give EOS and </assistant> tokens 2x weight so the model learns | |
| # to stop reliably. This prevents runaway generation & </s> spam. | |
| stop_boost = 2.0 if tok in (self.eos, self.think_end) else 1.0 | |
| w[i] = base_w * stop_boost if base_w > 0.0 else 0.0 | |
| return w | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # DATA LOADING HELPERS | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| def pack_texts(texts, tokenizer, seq_len: int, desc="packing") -> torch.Tensor: | |
| """Document-boundary-aware token packing → [N, seq_len+1] stacked tensor.""" | |
| eos = tokenizer.eos_token_id | |
| chunks, current = [], [] | |
| for text in tqdm(texts, desc=desc, leave=False): | |
| ids = tokenizer.encode(text, add_special_tokens=False) + [eos] | |
| if len(ids) >= seq_len + 1: | |
| if current: | |
| chunks.append(current + [eos] * ((seq_len+1) - len(current))) | |
| current = [] | |
| for i in range(0, len(ids) - seq_len, seq_len): | |
| chunks.append(ids[i: i+seq_len+1]) | |
| else: | |
| if len(current) + len(ids) > seq_len + 1: | |
| chunks.append(current + [eos] * ((seq_len+1) - len(current))) | |
| current = [] | |
| current.extend(ids) | |
| if len(current) >= 2: | |
| chunks.append(current + [eos] * ((seq_len+1) - len(current))) | |
| if not chunks: | |
| return torch.empty(0, seq_len+1, dtype=torch.long) | |
| return torch.tensor(chunks, dtype=torch.long) | |
| def _tag_id(tokenizer, tag: str) -> int: | |
| ids = tokenizer.encode(tag, add_special_tokens=False) | |
| return ids[0] if ids else -1 | |
| def load_pretrain_data(tokenizer, seq_len, cap=None) -> torch.Tensor: | |
| from datasets import load_dataset | |
| print("📚 Loading FineWeb-Edu...") | |
| fw_cap = cap if cap else 2_500_000 | |
| fw_ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True) | |
| fw_texts = [row["text"] for i, row in enumerate(fw_ds) if i < fw_cap] | |
| fw = pack_texts(fw_texts, tokenizer, seq_len, "FineWeb-Edu") | |
| print(f" ✅ {len(fw):,} chunks") | |
| print("📚 Loading TinyStories...") | |
| ts_ds = load_dataset("roneneldan/TinyStories", split="train") | |
| if cap: ts_ds = ts_ds.select(range(min(cap//2, len(ts_ds)))) | |
| ts = pack_texts([r["text"] for r in ts_ds], tokenizer, seq_len, "TinyStories") | |
| print(f" ✅ {len(ts):,} chunks") | |
| all_c = torch.cat([fw, ts], dim=0) | |
| print(f" ✅ Phase 1 total: {len(all_c):,} chunks ({len(all_c)*seq_len/1e9:.2f}B tokens)") | |
| return all_c | |
| def load_sft_data(tokenizer, seq_len, think_weight, cap=None) -> TagAwareDataset: | |
| from datasets import load_dataset | |
| texts = [] | |
| SYS = ( | |
| "You are NovaMind, a helpful and concise AI assistant. " | |
| "Answer questions directly and clearly." | |
| ) | |
| def wrap(human: str, assistant: str, system: str = SYS) -> str: | |
| return ( | |
| f"<system>\n{system}\n</system>\n" | |
| f"<human>\n{human.strip()}\n</human>\n" | |
| f"<assistant>\n{assistant.strip()}\n</assistant>" | |
| ) | |
| # ── 0. Identity seeds (repeated to anchor model's self-knowledge) ───────── | |
| # Without these, the model will hallucinate a human identity ("I'm from NYC...") | |
| print("🪪 Injecting NovaMind identity seeds...") | |
| IDENTITY_SEEDS = [ | |
| ("Who are you?", | |
| "I am NovaMind, a helpful AI assistant. I'm here to answer questions and help with tasks."), | |
| ("What is your name?", | |
| "My name is NovaMind. I'm an AI assistant designed to be helpful and informative."), | |
| ("Are you a human?", | |
| "No, I am not a human. I am NovaMind, an AI language model. How can I help you today?"), | |
| ("Are you an AI?", | |
| "Yes, I am NovaMind, an AI assistant. I can answer questions, explain concepts, and help with many tasks."), | |
| ("Who created you?", | |
| "I was created as NovaMind, an AI assistant. I'm here to help you with questions and tasks."), | |
| ("What can you do?", | |
| "I can answer questions, explain concepts, help with writing, discuss topics, and assist with many tasks. What would you like help with?"), | |
| ("Hello!", "Hello! I'm NovaMind, your AI assistant. How can I help you today?"), | |
| ("Hi there", "Hi! I'm NovaMind. What can I help you with?"), | |
| ("Hey", "Hey! I'm NovaMind, an AI assistant. How can I help you?"), | |
| ("Good morning", "Good morning! I'm NovaMind. What can I help you with today?"), | |
| ("How are you?", | |
| "I'm doing well, thank you for asking! I'm NovaMind, your AI assistant. How can I help you?"), | |
| ("hey buddy", | |
| "Hey! I'm NovaMind, your AI assistant. What can I help you with today?"), | |
| ("What are you?", | |
| "I am NovaMind, an AI language model. I'm trained to answer questions and assist with a wide range of tasks."), | |
| ("Tell me about yourself", | |
| "I am NovaMind, an AI assistant. I can help answer questions, explain topics, assist with writing, and have conversations on many subjects."), | |
| ] | |
| # Repeat identity seeds MANY times so they completely override pre-training human identity. | |
| # 20 was not enough (drowned by 5B tokens of human text). 500x anchors it strongly. | |
| REPEAT = 500 | |
| for q, a in IDENTITY_SEEDS: | |
| for _ in range(REPEAT): | |
| texts.append(wrap(q, a)) | |
| print(f" ✅ {len(IDENTITY_SEEDS) * REPEAT:,} identity seed examples (x{REPEAT} repeats)") | |
| # ── 1. Alpaca (52K instruction-following examples, downloadable) ────────── | |
| print("💬 Loading Alpaca...") | |
| alpaca = load_dataset("tatsu-lab/alpaca", split="train") | |
| if cap: alpaca = alpaca.select(range(min(cap, len(alpaca)))) | |
| n_before = len(texts) | |
| for r in alpaca: | |
| u = r["instruction"].strip() | |
| if r["input"].strip(): | |
| u += f"\n\n{r['input'].strip()}" | |
| ans = r["output"].strip() | |
| if not ans or len(ans) < 5: # skip empty/junk answers | |
| continue | |
| texts.append(wrap(u, ans)) | |
| print(f" ✅ {len(texts) - n_before:,} Alpaca examples") | |
| # ── 2. OpenAssistant OASST1 (88K human conversations, downloadable) ─────── | |
| # Replaces ShareGPT which was streaming-only and caused the GPU credit waste. | |
| # OASST1 is high quality: human annotators rating conversations. | |
| print("💬 Loading OpenAssistant OASST1 (replaces ShareGPT)...") | |
| try: | |
| oasst = load_dataset("OpenAssistant/oasst1", split="train") | |
| # Build message map: id → row | |
| msg_map = {row["message_id"]: row for row in oasst} | |
| n_before = len(texts) | |
| oasst_cap = min(cap, 60_000) if cap else 60_000 | |
| count = 0 | |
| for row in oasst: | |
| if count >= oasst_cap: | |
| break | |
| # Only take assistant replies that have a prompter parent | |
| if row["role"] != "assistant": | |
| continue | |
| parent = msg_map.get(row["parent_id"]) | |
| if parent is None or parent["role"] != "prompter": | |
| continue | |
| q = parent["text"].strip() | |
| ans = row["text"].strip() | |
| if not q or not ans or len(ans) < 10: | |
| continue | |
| # Quality filter: skip very low-rated responses | |
| if row.get("rank", 0) is not None and row.get("rank", 0) > 3: | |
| continue | |
| texts.append(wrap(q, ans)) | |
| count += 1 | |
| print(f" ✅ {len(texts) - n_before:,} OASST1 examples") | |
| except Exception as e: | |
| print(f" ⚠️ OASST1 unavailable ({e}), skipping.") | |
| # ── 3. Dolly 15K (15K human-written instructions, high quality) ─────────── | |
| print("💬 Loading Databricks Dolly 15K...") | |
| try: | |
| dolly = load_dataset("databricks/databricks-dolly-15k", split="train") | |
| n_before = len(texts) | |
| dolly_cap = min(cap, 15_000) if cap else 15_000 | |
| for i, row in enumerate(dolly): | |
| if i >= dolly_cap: | |
| break | |
| q = row["instruction"].strip() | |
| if row.get("context", "").strip(): | |
| q += f"\n\nContext: {row['context'].strip()}" | |
| ans = row["response"].strip() | |
| if not ans or len(ans) < 10: | |
| continue | |
| texts.append(wrap(q, ans)) | |
| print(f" ✅ {len(texts) - n_before:,} Dolly examples") | |
| except Exception as e: | |
| print(f" ⚠️ Dolly unavailable ({e}), skipping.") | |
| # ── 4. DailyDialog (conversational back-and-forth) ───────────────────────── | |
| print("💬 Loading DailyDialog...") | |
| try: | |
| try: | |
| dd = load_dataset("agentlans/li2017dailydialog", split="train") | |
| except Exception: | |
| dd = load_dataset("DeepPavlov/daily_dialog", split="train") | |
| if cap: dd = dd.select(range(min(cap // 4, len(dd)))) | |
| n_before = len(texts) | |
| for item in dd: | |
| convs = item.get("conversations") or item.get("dialog") or [] | |
| if not convs: continue | |
| pairs = ( | |
| [(c.get("from", ""), c.get("value", "")) for c in convs] | |
| if isinstance(convs[0], dict) | |
| else [("human" if i % 2 == 0 else "gpt", c) for i, c in enumerate(convs)] | |
| ) | |
| t = f"<system>\n{SYS}\n</system>\n" | |
| t += "".join( | |
| f"<{'human' if 'human' in r.lower() else 'assistant'}>\n{v.strip()}\n" | |
| f"</{'human' if 'human' in r.lower() else 'assistant'}>\n" | |
| for r, v in pairs if v.strip() | |
| ) | |
| if t: texts.append(t) | |
| print(f" ✅ {len(texts) - n_before:,} DailyDialog examples") | |
| except Exception as e: | |
| print(f" ⚠️ DailyDialog unavailable ({e}), skipping.") | |
| print(f"\n 📊 Total SFT examples: {len(texts):,}") | |
| chunks = pack_texts(texts, tokenizer, seq_len, "SFT packing") | |
| print(f" ✅ {len(chunks):,} chunks packed") | |
| return TagAwareDataset(chunks, think_weight, | |
| _tag_id(tokenizer, "<think>"), _tag_id(tokenizer, "</think>"), | |
| _tag_id(tokenizer, "<assistant>"), _tag_id(tokenizer, "<human>"), | |
| tokenizer.eos_token_id) | |
| def load_cot_data(tokenizer, seq_len, think_weight, cap=None) -> TagAwareDataset: | |
| from datasets import load_dataset | |
| texts = [] | |
| print("🧠 Loading OpenHermes-2.5...") | |
| try: | |
| oh = load_dataset("teknium/OpenHermes-2.5", split="train") | |
| oh_cap = min(cap, 100_000) if cap else 100_000 | |
| for i, row in enumerate(oh): | |
| if i >= oh_cap: break | |
| t = "" | |
| for c in row.get("conversations", []): | |
| r, v = c.get("from",""), c.get("value","").strip() | |
| if r == "system": t += f"<system>\n{v}\n</system>\n" | |
| elif r == "human": t += f"<human>\n{v}\n</human>\n" | |
| elif r == "gpt": | |
| parts = v.split("\n\n") | |
| if len(parts) > 1 and any(k in v for k in ["Step ","##","Let me"]): | |
| t += f"<think>\n{chr(10).join(parts[:-1]).strip()}\n</think>\n" | |
| t += f"<assistant>\n{parts[-1].strip()}\n</assistant>\n" | |
| else: | |
| t += f"<assistant>\n{v}\n</assistant>\n" | |
| if t: texts.append(t) | |
| print(f" ✅ {len(texts):,} OpenHermes examples") | |
| except Exception as e: | |
| print(f" ⚠️ OpenHermes unavailable ({e}), skipping.") | |
| print("🧠 Loading GSM8K math reasoning...") | |
| gsm = load_dataset("openai/gsm8k", "main", split="train") | |
| if cap: gsm = gsm.select(range(min(cap//5, len(gsm)))) | |
| for row in gsm: | |
| q, a = row["question"].strip(), row["answer"].strip() | |
| if "####" in a: | |
| reasoning, final = a.split("####", 1) | |
| texts.append( | |
| f"<human>\n{q}\n</human>\n" | |
| f"<think>\n{reasoning.strip()}\n</think>\n" | |
| f"<assistant>\n{final.strip()}\n</assistant>" | |
| ) | |
| else: | |
| texts.append(f"<human>\n{q}\n</human>\n<assistant>\n{a}\n</assistant>") | |
| print(f" ✅ {len(texts):,} total CoT examples") | |
| chunks = pack_texts(texts, tokenizer, seq_len, "CoT packing") | |
| print(f" ✅ {len(chunks):,} chunks packed") | |
| return TagAwareDataset(chunks, think_weight, | |
| _tag_id(tokenizer,"<think>"), _tag_id(tokenizer,"</think>"), | |
| _tag_id(tokenizer,"<assistant>"), _tag_id(tokenizer,"<human>"), | |
| tokenizer.eos_token_id) | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # LEARNING RATE SCHEDULE | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| def get_lr(step, warmup, max_steps, peak, min_lr): | |
| """Linear warmup → cosine decay. Standard for all modern LLMs.""" | |
| if step < warmup: | |
| return peak * step / max(warmup, 1) | |
| if step >= max_steps: | |
| return min_lr | |
| progress = (step - warmup) / max(max_steps - warmup, 1) | |
| return min_lr + 0.5 * (peak - min_lr) * (1.0 + math.cos(math.pi * progress)) | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # COMPREHENSIVE LOGGER | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| class TrainingLogger: | |
| """ | |
| Full logging to three sinks simultaneously: | |
| 1. Console — rich formatted table printed every log_every steps | |
| 2. JSON-lines — one JSON object per log event, append-only (survives crashes) | |
| 3. CSV — spreadsheet-friendly, one row per log event | |
| 4. W&B — cloud tracking (optional, graceful fallback) | |
| Also maintains running averages so the console always shows the smoothed | |
| loss over the last `log_every` steps, not just the current noisy batch. | |
| """ | |
| # CSV columns — every metric we track | |
| CSV_FIELDS = [ | |
| "step", "epoch_frac", "phase", | |
| "loss_total", "loss_think", "loss_response", | |
| "lr", "grad_norm", | |
| "tokens_seen_B", "tokens_per_sec", | |
| "gpu_mem_gb", "gpu_mem_reserved_gb", | |
| "loss_weight_think", | |
| "val_loss", | |
| "time_elapsed_min", "eta_min", | |
| "timestamp", | |
| ] | |
| def __init__(self, out_dir: str, phase: str, config: dict, | |
| wandb_project="novamind-256m", wandb_run=None, | |
| use_wandb=True, log_every=10): | |
| self.out_dir = Path(out_dir) | |
| self.out_dir.mkdir(parents=True, exist_ok=True) | |
| self.phase = phase | |
| self.log_every = log_every | |
| self.use_wandb = use_wandb and WANDB_AVAILABLE | |
| # ── File paths ──────────────────────────────────────────────────── | |
| ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") | |
| self.jsonl_path = self.out_dir / f"log_{phase}_{ts}.jsonl" | |
| self.csv_path = self.out_dir / f"log_{phase}_{ts}.csv" | |
| # Keep a "latest" symlink so scripts can always find the current log | |
| self._symlink(self.jsonl_path, self.out_dir / f"log_{phase}_latest.jsonl") | |
| self._symlink(self.csv_path, self.out_dir / f"log_{phase}_latest.csv") | |
| # ── CSV: write header immediately ──────────────────────────────── | |
| with open(self.csv_path, "w", newline="") as f: | |
| csv.DictWriter(f, fieldnames=self.CSV_FIELDS, extrasaction="ignore").writeheader() | |
| # ── W&B ────────────────────────────────────────────────────────── | |
| if self.use_wandb: | |
| try: | |
| wandb.init( | |
| project=wandb_project, | |
| name=wandb_run or f"novamind-{phase}-{ts}", | |
| config=config, | |
| resume="allow", | |
| ) | |
| print("📡 W&B initialized:", wandb.run.url) | |
| except Exception as e: | |
| print(f"⚠️ W&B failed ({e}) — falling back to local logs only.") | |
| self.use_wandb = False | |
| # ── Running accumulators (reset every log_every steps) ──────────── | |
| self._acc_loss_total = 0.0 | |
| self._acc_loss_think = 0.0 | |
| self._acc_loss_response = 0.0 | |
| self._acc_grad_norm = 0.0 | |
| self._acc_tok_per_sec = 0.0 | |
| self._acc_count = 0 | |
| print(f"\n📊 Logger ready:") | |
| print(f" JSON-lines → {self.jsonl_path}") | |
| print(f" CSV → {self.csv_path}") | |
| if not self.use_wandb: | |
| print(f" W&B → disabled (install wandb for cloud tracking)") | |
| def _symlink(self, target: Path, link: Path): | |
| try: | |
| if link.exists() or link.is_symlink(): | |
| link.unlink() | |
| link.symlink_to(target.name) | |
| except Exception: | |
| pass # symlinks may fail on some systems, not critical | |
| def accumulate(self, loss_total: float, loss_think: float, loss_response: float, | |
| grad_norm: float, tok_per_sec: float): | |
| """Accumulate per-step values for smoothed logging.""" | |
| self._acc_loss_total += loss_total | |
| self._acc_loss_think += loss_think | |
| self._acc_loss_response += loss_response | |
| self._acc_grad_norm += grad_norm | |
| self._acc_tok_per_sec += tok_per_sec | |
| self._acc_count += 1 | |
| def log_step(self, step: int, epoch_frac: float, lr: float, | |
| tokens_seen: int, think_weight: float, | |
| t_start: float, max_steps: int, | |
| val_loss: Optional[float] = None): | |
| """ | |
| Called every log_every steps. Drains the accumulators, prints to | |
| console, and writes to JSON + CSV + W&B. | |
| """ | |
| n = max(self._acc_count, 1) | |
| avg_loss = self._acc_loss_total / n | |
| avg_think = self._acc_loss_think / n | |
| avg_response = self._acc_loss_response / n | |
| avg_gnorm = self._acc_grad_norm / n | |
| avg_tps = self._acc_tok_per_sec / n | |
| elapsed_min = (time.time() - t_start) / 60.0 | |
| steps_done = step | |
| steps_left = max(max_steps - steps_done, 0) | |
| secs_per_step = (time.time() - t_start) / max(steps_done, 1) | |
| eta_min = steps_left * secs_per_step / 60.0 | |
| gpu_mem = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0 | |
| gpu_res = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0.0 | |
| ts_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| # ── Build metrics dict ──────────────────────────────────────────── | |
| metrics = { | |
| "step": step, | |
| "epoch_frac": round(epoch_frac, 4), | |
| "phase": self.phase, | |
| "loss_total": round(avg_loss, 6), | |
| "loss_think": round(avg_think, 6), | |
| "loss_response": round(avg_response, 6), | |
| "lr": lr, | |
| "grad_norm": round(avg_gnorm, 4), | |
| "tokens_seen_B": round(tokens_seen / 1e9, 6), | |
| "tokens_per_sec": round(avg_tps, 1), | |
| "gpu_mem_gb": round(gpu_mem, 3), | |
| "gpu_mem_reserved_gb":round(gpu_res, 3), | |
| "loss_weight_think": think_weight, | |
| "val_loss": round(val_loss, 6) if val_loss is not None else None, | |
| "time_elapsed_min": round(elapsed_min, 2), | |
| "eta_min": round(eta_min, 1), | |
| "timestamp": ts_str, | |
| } | |
| # ── Console output ────────────────────────────────────────────── | |
| self._print_step(metrics) | |
| # ── Red-flag warnings ─────────────────────────────────────────── | |
| if avg_gnorm > 10.0: | |
| print(f" 🚨 HIGH GRAD NORM {avg_gnorm:.2f} — training may be unstable! Consider reducing LR.") | |
| if step > 200 and avg_loss > 4.0: | |
| print(f" 🚨 LOSS STILL HIGH ({avg_loss:.4f}) after {step} steps — check dataset & LR!") | |
| if gpu_res > 38.0: | |
| print(f" ⚠️ GPU MEM {gpu_res:.1f}GB / 40GB — close to OOM! Reduce batch_size or seq_len.") | |
| if avg_think > 0 and avg_response > 0 and avg_think > avg_response * 3: | |
| print(f" ℹ️ Think loss ({avg_think:.3f}) >> Response loss ({avg_response:.3f}) — normal in Phase 3.") | |
| # ── JSON-lines ────────────────────────────────────────────────── | |
| with open(self.jsonl_path, "a") as f: | |
| f.write(json.dumps(metrics) + "\n") | |
| # ── CSV ───────────────────────────────────────────────────────── | |
| with open(self.csv_path, "a", newline="") as f: | |
| csv.DictWriter(f, fieldnames=self.CSV_FIELDS, extrasaction="ignore").writerow(metrics) | |
| # ── W&B ───────────────────────────────────────────────────────── | |
| if self.use_wandb: | |
| try: | |
| wandb.log({ | |
| "train/loss": avg_loss, | |
| "train/loss_think": avg_think, | |
| "train/loss_response": avg_response, | |
| "train/grad_norm": avg_gnorm, | |
| "train/lr": lr, | |
| "train/tokens_seen_B": tokens_seen / 1e9, | |
| "train/tokens_per_sec":avg_tps, | |
| "train/gpu_mem_gb": gpu_mem, | |
| "train/eta_min": eta_min, | |
| **({"val/loss": val_loss} if val_loss is not None else {}), | |
| }, step=step) | |
| except Exception: | |
| pass | |
| # ── Reset accumulators ────────────────────────────────────────── | |
| self._acc_loss_total = self._acc_loss_think = self._acc_loss_response = 0.0 | |
| self._acc_grad_norm = self._acc_tok_per_sec = 0.0 | |
| self._acc_count = 0 | |
| return metrics | |
| def log_val(self, step: int, val_loss: float, samples: list[tuple[str,str]]): | |
| """Log validation loss and sample generations.""" | |
| metrics = {"step": step, "val_loss": round(val_loss, 6), | |
| "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} | |
| # Write to JSON log | |
| with open(self.jsonl_path, "a") as f: | |
| f.write(json.dumps({"type": "validation", **metrics}) + "\n") | |
| # Console | |
| print(f"\n{'─'*65}") | |
| print(f" 📊 VALIDATION step={step:,} val_loss={val_loss:.6f}") | |
| print(f"{'─'*65}") | |
| for q, a in samples: | |
| print(f" Q: {q}") | |
| resp = a[:200] + ("..." if len(a) > 200 else "") | |
| print(f" A: {resp}") | |
| print() | |
| print(f"{'─'*65}\n") | |
| if self.use_wandb: | |
| try: | |
| wandb.log({"val/loss": val_loss}, step=step) | |
| wandb.log({"samples": wandb.Table( | |
| columns=["prompt","response"], | |
| data=[[q, a] for q,a in samples], | |
| )}, step=step) | |
| except Exception: | |
| pass | |
| def _print_step(self, m: dict): | |
| """Rich console table row.""" | |
| bar_width = 20 | |
| filled = int(bar_width * m["step"] / max(m.get("_max_steps", m["step"]), 1)) | |
| # Progress indicator: ████████░░░░░░░░░░░░ | |
| progress_bar = "█" * filled + "░" * (bar_width - filled) | |
| print( | |
| f" step {m['step']:>7,} [{progress_bar}] " | |
| f"loss={m['loss_total']:.4f} " | |
| f"(think={m['loss_think']:.4f} resp={m['loss_response']:.4f}) | " | |
| f"lr={m['lr']:.2e} gnorm={m['grad_norm']:.3f} | " | |
| f"{m['tokens_per_sec']:>9,.0f} tok/s | " | |
| f"mem={m['gpu_mem_reserved_gb']:.1f}GB | " | |
| f"elapsed={m['time_elapsed_min']:.1f}m ETA={m['eta_min']:.0f}m" | |
| ) | |
| def finish(self): | |
| print(f"\n📊 Training complete. Logs saved:") | |
| print(f" {self.jsonl_path}") | |
| print(f" {self.csv_path}") | |
| if self.use_wandb: | |
| wandb.finish() | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # TAG-SPECIFIC LOSS COMPUTATION | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| def compute_tag_losses( | |
| logits: torch.Tensor, | |
| targets: torch.Tensor, | |
| token_weights: torch.Tensor, | |
| ) -> tuple[torch.Tensor, float, float]: | |
| """ | |
| Compute the weighted total loss PLUS separate think and response sub-losses | |
| for monitoring. The total loss is what we .backward() on; the sub-losses | |
| are purely for logging to understand training progress by token type. | |
| Returns: | |
| (total_loss, think_loss_item, response_loss_item) | |
| """ | |
| total_loss = tag_aware_loss(logits, targets, token_weights) | |
| with torch.no_grad(): | |
| per_tok = torch.nn.functional.cross_entropy( | |
| logits.view(-1, logits.size(-1)), | |
| targets.view(-1), | |
| ignore_index=-1, reduction="none", | |
| ) # [B*T] | |
| flat_w = token_weights.view(-1) | |
| valid = (targets.view(-1) != -1).float() | |
| # Think tokens: 0 < weight < 1 (phase 2) or weight > 1 (phase 3) | |
| # Response tokens: weight == 1.0 | |
| # User/system tokens: weight == 0.0 | |
| think_mask = ((flat_w > 0.0) & (flat_w != 1.0) & valid.bool()).float() | |
| response_mask = ((flat_w == 1.0) & valid.bool()).float() | |
| think_loss = (per_tok * think_mask).sum() / think_mask.sum().clamp(min=1e-8) | |
| response_loss = (per_tok * response_mask).sum() / response_mask.sum().clamp(min=1e-8) | |
| return total_loss, think_loss.item(), response_loss.item() | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # CHECKPOINTING | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| def save_checkpoint( | |
| out_dir: str, model: nn.Module, optimizer, scaler, | |
| step: int, phase: str, tokens_seen: int, epoch_frac: float, | |
| train_loss: float, think_loss: float, response_loss: float, | |
| lr: float, grad_norm: float, | |
| config: NovaMindConfig, | |
| train_loss_history: list, val_loss_history: list, | |
| milestone=False, phase_final=False, | |
| ): | |
| """ | |
| Save a 100% resumable checkpoint containing: | |
| • model weights (DataParallel-unwrapped) | |
| • optimizer state | |
| • scaler state (for BF16/FP16) | |
| • scheduler state (current LR, step) | |
| • step counter + tokens seen (for resuming data iteration) | |
| • full loss histories (for plotting) | |
| • all hyperparameters | |
| Checkpoint types: | |
| • Rolling : novamind_roll_latest.pt — always overwritten | |
| • Milestone: novamind_milestone_NNNNNNN.pt — kept forever | |
| • Final : novamind_{phase}_final.pt — kept forever | |
| """ | |
| out = Path(out_dir) | |
| out.mkdir(parents=True, exist_ok=True) | |
| model_state = (model.module.state_dict() if hasattr(model, "module") | |
| else model.state_dict()) | |
| ckpt = { | |
| # Model | |
| "model_state_dict": model_state, | |
| "config": config.to_dict(), | |
| # Optimizer / precision | |
| "optimizer_state_dict": optimizer.state_dict(), | |
| "scaler_state_dict": scaler.state_dict() if scaler is not None else None, | |
| # Training state (all needed to resume mid-epoch) | |
| "step": step, | |
| "phase": phase, | |
| "epoch_frac": epoch_frac, | |
| "tokens_seen": tokens_seen, | |
| "lr": lr, | |
| # Last-step metrics (for human inspection) | |
| "train_loss": train_loss, | |
| "think_loss": think_loss, | |
| "response_loss": response_loss, | |
| "grad_norm": grad_norm, | |
| # Full histories (for loss curves) | |
| "train_loss_history": train_loss_history, | |
| "val_loss_history": val_loss_history, | |
| # Meta | |
| "timestamp": datetime.datetime.now().isoformat(), | |
| "torch_version": torch.__version__, | |
| } | |
| # 1. Always save rolling (overwrites previous) | |
| roll = out / "novamind_roll_latest.pt" | |
| torch.save(ckpt, roll) | |
| mb = roll.stat().st_size / 1e6 | |
| print(f"\n{'─'*65}") | |
| print(f" 💾 CHECKPOINT step={step:,} | loss={train_loss:.4f} | " | |
| f"think={think_loss:.4f} | resp={response_loss:.4f}") | |
| print(f" tokens={tokens_seen/1e9:.3f}B | lr={lr:.2e} | size={mb:.0f} MB") | |
| print(f" → {roll}") | |
| # 2. Milestone — always copy (never overwrite) for permanent history | |
| if milestone: | |
| ms = out / f"novamind_milestone_{step:07d}.pt" | |
| shutil.copy2(roll, ms) | |
| print(f" → MILESTONE {ms.name}") | |
| # 3. Phase final | |
| if phase_final: | |
| fin = out / f"novamind_{phase}_final.pt" | |
| shutil.copy2(roll, fin) | |
| print(f" → PHASE FINAL {fin.name}") | |
| print(f"{'─'*65}\n") | |
| return roll | |
| def write_checkpoint_index(out_dir: str, step: int, loss: float, tokens_seen: int, ckpt_path: str): | |
| """ | |
| Maintain a human-readable checkpoint index file so you can quickly see | |
| what checkpoints exist and at what training state. | |
| """ | |
| index_path = Path(out_dir) / "checkpoint_index.jsonl" | |
| with open(index_path, "a") as f: | |
| f.write(json.dumps({ | |
| "step": step, "loss": round(loss, 6), | |
| "tokens_seen_B": round(tokens_seen/1e9, 4), | |
| "path": str(ckpt_path), | |
| "saved_at": datetime.datetime.now().isoformat(), | |
| }) + "\n") | |
| def load_checkpoint(path: str, model: nn.Module, optimizer, scaler, device): | |
| """Load checkpoint and return the full ckpt dict for state restoration.""" | |
| print(f"\n📂 Loading checkpoint: {path}") | |
| ckpt = torch.load(path, map_location=device, weights_only=False) | |
| # Handle DataParallel prefix mismatch | |
| state = ckpt["model_state_dict"] | |
| if all(k.startswith("module.") for k in state): | |
| state = {k[7:]: v for k, v in state.items()} | |
| (model.module if hasattr(model, "module") else model).load_state_dict(state) | |
| optimizer.load_state_dict(ckpt["optimizer_state_dict"]) | |
| if scaler is not None and ckpt.get("scaler_state_dict"): | |
| scaler.load_state_dict(ckpt["scaler_state_dict"]) | |
| print(f" ✅ Resumed: phase={ckpt['phase']} | step={ckpt['step']:,} | " | |
| f"tokens={ckpt['tokens_seen']/1e9:.3f}B | " | |
| f"last_loss={ckpt.get('train_loss', '?')}") | |
| return ckpt | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # VALIDATION | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| def estimate_val_loss(model, loader, device, num_batches=50, dtype=torch.bfloat16) -> float: | |
| model.eval() | |
| losses = [] | |
| for i, (x, y, w) in enumerate(loader): | |
| if i >= num_batches: break | |
| x, y, w = x.to(device), y.to(device), w.to(device) | |
| with torch.autocast(device_type=device.type, dtype=dtype, enabled=torch.cuda.is_available()): | |
| _, loss = model(x, y, w) | |
| losses.append(loss.item()) | |
| model.train() | |
| return sum(losses) / len(losses) if losses else float("nan") | |
| def generate_samples(model, tokenizer, device, prompts, max_new_tokens=100): | |
| model.eval() | |
| inner = model.module if hasattr(model, "module") else model | |
| results = [] | |
| for prompt in prompts: | |
| fmt = f"<human>\n{prompt}\n</human>\n<assistant>\n" | |
| ids = tokenizer.encode(fmt, return_tensors="pt").to(device) | |
| try: | |
| stop_ids = [tokenizer.encode("</assistant>", add_special_tokens=False)[0]] | |
| except Exception: | |
| stop_ids = None | |
| out = inner.generate(ids, max_new_tokens=max_new_tokens, | |
| temperature=0.7, top_k=40, top_p=0.9, | |
| stop_token_ids=stop_ids) | |
| full = tokenizer.decode(out[0], skip_special_tokens=False) | |
| resp = full.split("<assistant>\n")[-1].split("</assistant>")[0].strip() | |
| results.append((prompt, resp)) | |
| model.train() | |
| return results | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| # MAIN TRAINING LOOP | |
| # ══════════════════════════════════════════════════════════════════════════════ | |
| TEST_PROMPTS = [ | |
| "What is the capital of France?", | |
| "What is 15 multiplied by 17?", | |
| "Tell me a short joke.", | |
| "Explain what a neural network is in simple terms.", | |
| ] | |
| def main(): | |
| args = parse_args() | |
| torch.manual_seed(args.seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(args.seed) | |
| # ── Phase defaults (CLI args override if provided) ──────────────────────── | |
| pd = PHASE_DEFAULTS[args.phase] | |
| peak_lr = args.lr or pd["lr"] | |
| warmup_steps = args.warmup_steps or pd["warmup_steps"] | |
| batch_size = args.batch_size or pd["batch_size"] | |
| grad_accum = args.grad_accum_steps or pd["grad_accum"] | |
| n_epochs = args.epochs or pd["epochs"] | |
| think_weight = pd["think_weight"] | |
| min_lr = peak_lr * args.min_lr_ratio | |
| # ── Device & Precision ──────────────────────────────────────────────────── | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| use_cuda = device.type == "cuda" | |
| # H100 supports BF16 natively — more numerically stable than FP16 | |
| dtype = torch.bfloat16 if (use_cuda and torch.cuda.is_bf16_supported()) else torch.float16 | |
| # ── Banner ──────────────────────────────────────────────────────────────── | |
| print(f"\n{'═'*65}") | |
| print(f" 🚀 NovaMind-256M Training") | |
| print(f" Phase : {args.phase.upper()} — {pd['desc']}") | |
| print(f" Device: {device} Precision: {dtype}") | |
| if use_cuda: | |
| for i in range(torch.cuda.device_count()): | |
| props = torch.cuda.get_device_properties(i) | |
| print(f" GPU {i} : {props.name} ({props.total_memory/1e9:.1f} GB)") | |
| print(f" Peak LR: {peak_lr:.2e} → min {min_lr:.2e} | Warmup: {warmup_steps} steps") | |
| print(f" Batch: {batch_size} seqs × {args.max_seq_len} tokens × {grad_accum} accum = " | |
| f"{batch_size*args.max_seq_len*grad_accum:,} tok/update") | |
| print(f" Think-token loss weight: {think_weight}") | |
| print(f" Checkpoint every: {args.save_every} steps | Milestone every: {args.milestone_every} steps") | |
| print(f"{'═'*65}\n") | |
| Path(args.out_dir).mkdir(parents=True, exist_ok=True) | |
| # ── Tokenizer ───────────────────────────────────────────────────────────── | |
| from transformers import AutoTokenizer | |
| print("🔤 Loading tokenizer...") | |
| for tok_name in ["NousResearch/Llama-2-7b-hf", "meta-llama/Llama-2-7b-hf", "gpt2"]: | |
| try: | |
| tokenizer = AutoTokenizer.from_pretrained(tok_name) | |
| print(f" ✅ {tok_name} — vocab: {tokenizer.vocab_size:,}") | |
| break | |
| except Exception as e: | |
| print(f" ⚠️ {tok_name} failed: {e}") | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| # ── Dataset ─────────────────────────────────────────────────────────────── | |
| cap = 500 if args.smoke_test else args.data_cap | |
| print(f"\n📂 Loading {args.phase} dataset{' [SMOKE TEST: cap=500]' if args.smoke_test else ''}...") | |
| if args.phase == "pretrain": | |
| if getattr(args, 'data_file', None) and Path(args.data_file).exists(): | |
| print(f" ⚡ Loading pre-tokenized data from disk: {args.data_file}") | |
| chunks = torch.load(args.data_file, map_location="cpu", weights_only=True).long() | |
| print(f" ✅ {len(chunks):,} chunks ({len(chunks)*args.max_seq_len/1e9:.2f}B tokens) — no download needed!") | |
| else: | |
| if getattr(args, 'data_file', None): | |
| print(f" ⚠️ data_file not found: {args.data_file} — falling back to download") | |
| chunks = load_pretrain_data(tokenizer, args.max_seq_len, cap) | |
| dataset = PretrainDataset(chunks) | |
| elif args.phase == "sft": | |
| if getattr(args, 'data_file', None) and Path(args.data_file).exists(): | |
| # FIX: SFT phase now supports --data_file (pre-tokenized on CPU, saves H100 time) | |
| print(f" ⚡ Loading pre-tokenized SFT data from disk: {args.data_file}") | |
| chunks = torch.load(args.data_file, map_location="cpu", weights_only=True).long() | |
| print(f" ✅ {len(chunks):,} chunks ({len(chunks)*args.max_seq_len/1e9:.2f}B tokens)") | |
| dataset = TagAwareDataset( | |
| chunks, think_weight, | |
| _tag_id(tokenizer, "<think>"), _tag_id(tokenizer, "</think>"), | |
| _tag_id(tokenizer, "<assistant>"), _tag_id(tokenizer, "<human>"), | |
| tokenizer.eos_token_id, | |
| ) | |
| else: | |
| if getattr(args, 'data_file', None): | |
| print(f" ⚠️ data_file not found: {args.data_file} — falling back to download") | |
| dataset = load_sft_data(tokenizer, args.max_seq_len, think_weight, cap) | |
| else: | |
| dataset = load_cot_data(tokenizer, args.max_seq_len, think_weight, cap) | |
| n_val = max(1, int(0.05 * len(dataset))) # 5% held out for validation | |
| n_train = len(dataset) - n_val | |
| train_ds, val_ds = torch.utils.data.random_split( | |
| dataset, [n_train, n_val], | |
| generator=torch.Generator().manual_seed(args.seed), | |
| ) | |
| print(f" Train: {n_train:,} chunks | Val: {n_val:,} chunks") | |
| train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, | |
| num_workers=args.num_workers, pin_memory=use_cuda, | |
| persistent_workers=(args.num_workers > 0)) | |
| val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, | |
| num_workers=min(2, args.num_workers), pin_memory=use_cuda) | |
| print(f" DataLoader: {len(train_loader):,} batches × {batch_size}") | |
| # ── Model ───────────────────────────────────────────────────────────────── | |
| model_config = NovaMindConfig( | |
| vocab_size=args.vocab_size, d_model=args.d_model, | |
| n_heads=args.n_heads, n_kv_heads=args.n_kv_heads, | |
| n_layers=args.n_layers, ff_dim=args.ff_dim, | |
| max_seq_len=args.max_seq_len, | |
| ) | |
| print(f"\n🏗️ Building model...") | |
| model = NovaMind256M(model_config).to(device) | |
| model.gradient_checkpointing_enable() # saves ~40% VRAM at cost of ~20% speed | |
| if args.compile and hasattr(torch, "compile"): | |
| print("⚡ torch.compile enabled...") | |
| model = torch.compile(model) | |
| if use_cuda and torch.cuda.device_count() > 1: | |
| print(f"🖥️ DataParallel across {torch.cuda.device_count()} GPUs") | |
| model = nn.DataParallel(model) | |
| inner = model.module if hasattr(model, "module") else model | |
| inner.count_parameters(print_table=True) | |
| # ── Optimizer ───────────────────────────────────────────────────────────── | |
| # Separate param groups: weight decay only on >= 2D params (matrices), | |
| # NOT on biases or norms (1D params). Standard best practice. | |
| decay_p = [p for n,p in model.named_parameters() if p.requires_grad and p.dim()>=2] | |
| no_decay_p = [p for n,p in model.named_parameters() if p.requires_grad and p.dim()<2] | |
| optimizer = optim.AdamW( | |
| [{"params": decay_p, "weight_decay": args.weight_decay}, | |
| {"params": no_decay_p, "weight_decay": 0.0}], | |
| lr=peak_lr, betas=(0.9, 0.95), eps=1e-8, | |
| ) | |
| # GradScaler: needed for FP16, no-op for BF16 (H100) | |
| scaler = torch.amp.GradScaler(device.type, enabled=(dtype == torch.float16)) | |
| # ── Max steps ───────────────────────────────────────────────────────────── | |
| steps_per_epoch = math.ceil(len(train_loader) / grad_accum) | |
| max_steps = args.max_steps or (n_epochs * steps_per_epoch) | |
| if args.smoke_test: max_steps = 50 | |
| print(f"\n Steps/epoch: {steps_per_epoch:,} | Total max steps: {max_steps:,}") | |
| print(f" Estimated training time on H100: " | |
| f"~{max_steps * batch_size * args.max_seq_len * grad_accum / 220_000 / 3600:.1f} hrs\n") | |
| # ── Resume ──────────────────────────────────────────────────────────────── | |
| global_step = 0 | |
| tokens_seen = 0 | |
| train_loss_history = [] | |
| val_loss_history = [] | |
| if args.resume: | |
| ckpt = load_checkpoint(args.resume, model, optimizer, scaler, device) | |
| if ckpt.get("phase") == args.phase: | |
| global_step = ckpt["step"] | |
| tokens_seen = ckpt["tokens_seen"] | |
| train_loss_history = ckpt.get("train_loss_history", []) | |
| val_loss_history = ckpt.get("val_loss_history", []) | |
| # Restore LR to where we left off | |
| resumed_lr = get_lr(global_step, warmup_steps, max_steps, peak_lr, min_lr) | |
| for pg in optimizer.param_groups: | |
| pg["lr"] = resumed_lr | |
| else: | |
| print(f"\n ⏩ Transitioning from phase '{ckpt.get('phase')}' to '{args.phase}'. Resetting step counter.") | |
| tokens_seen = ckpt.get("tokens_seen", 0) | |
| # ── Logger ──────────────────────────────────────────────────────────────── | |
| cfg_for_log = {**vars(args), **model_config.to_dict(), | |
| "peak_lr": peak_lr, "min_lr": min_lr, | |
| "think_weight": think_weight, "max_steps": max_steps} | |
| logger = TrainingLogger( | |
| out_dir=args.out_dir, phase=args.phase, config=cfg_for_log, | |
| wandb_project=args.wandb_project, wandb_run=args.wandb_run, | |
| use_wandb=(not args.no_wandb), | |
| log_every=args.log_every, | |
| ) | |
| # ══════════════════════════════════════════════════════════════════════════ | |
| # TRAINING LOOP | |
| # ══════════════════════════════════════════════════════════════════════════ | |
| print(f"\n{'═'*65}") | |
| print(f" 🚦 TRAINING STARTED — Phase: {args.phase.upper()}") | |
| print(f" Resume from step: {global_step:,} | Target: {max_steps:,} steps") | |
| print(f"{'═'*65}\n") | |
| model.train() | |
| optimizer.zero_grad() | |
| t_start = time.time() | |
| t_batch = time.time() | |
| last_loss = float("nan") # track for checkpoint metadata | |
| last_think_loss = 0.0 | |
| last_response_loss = 0.0 | |
| last_grad_norm = 0.0 | |
| start_epoch = 1 | |
| skip_micro_steps = 0 | |
| if args.resume and ckpt.get("phase") == args.phase: | |
| total_micro_steps_done = global_step * grad_accum | |
| start_epoch = (total_micro_steps_done // len(train_loader)) + 1 | |
| skip_micro_steps = total_micro_steps_done % len(train_loader) | |
| for epoch in range(start_epoch, n_epochs + 1): | |
| if global_step >= max_steps: | |
| break | |
| print(f" 📅 Epoch {epoch}/{n_epochs}") | |
| for micro_step, (x, y, w) in enumerate(train_loader): | |
| if global_step >= max_steps: | |
| break | |
| # Skip already processed batches if resuming mid-epoch | |
| if epoch == start_epoch and micro_step < skip_micro_steps: | |
| continue | |
| x, y, w = x.to(device), y.to(device), w.to(device) | |
| batch_tokens = x.numel() | |
| t_b0 = time.time() | |
| # ── Forward + loss ───────────────────────────────────────────── | |
| with torch.autocast(device_type=device.type, dtype=dtype, | |
| enabled=use_cuda): | |
| logits, total_loss = model(x, y, w) | |
| # During training, model returns empty logits — recompute | |
| # logits only if needed for sub-loss breakdown (no_grad below) | |
| scaled_loss = total_loss.mean() / grad_accum | |
| # ── Backward ────────────────────────────────────────────────── | |
| if dtype == torch.float16: | |
| scaler.scale(scaled_loss).backward() | |
| else: | |
| scaled_loss.backward() | |
| # ── Compute per-tag sub-losses for logging (cheap, no_grad) ── | |
| # We need logits for this — call model again in eval mode for | |
| # a micro-subset to get the breakdown without full batch cost. | |
| with torch.no_grad(): | |
| with torch.autocast(device_type=device.type, dtype=dtype, enabled=use_cuda): | |
| logits_for_log, _ = (model.module if hasattr(model,"module") else model)( | |
| x[:1], None # just first sample, no loss | |
| ) | |
| # Compute sub-losses on that single sample | |
| _, think_l, response_l = compute_tag_losses( | |
| logits_for_log, y[:1], w[:1] | |
| ) | |
| # ── Optimizer step (every grad_accum micro-steps) ────────────── | |
| if (micro_step + 1) % grad_accum == 0: | |
| # Clip gradient norm — critical: must unscale first for FP16 | |
| if dtype == torch.float16: | |
| scaler.unscale_(optimizer) | |
| grad_norm = torch.nn.utils.clip_grad_norm_( | |
| model.parameters(), args.max_grad_norm | |
| ).item() | |
| if dtype == torch.float16: | |
| scaler.step(optimizer) | |
| scaler.update() | |
| else: | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| global_step += 1 | |
| tokens_seen += batch_tokens * grad_accum | |
| # Update LR | |
| lr = get_lr(global_step, warmup_steps, max_steps, peak_lr, min_lr) | |
| for pg in optimizer.param_groups: | |
| pg["lr"] = lr | |
| # Track last values for checkpointing | |
| last_loss = total_loss.mean().item() | |
| last_think_loss = think_l | |
| last_response_loss = response_l | |
| last_grad_norm = grad_norm | |
| tok_per_sec = batch_tokens * grad_accum / max(time.time() - t_b0, 1e-6) | |
| epoch_frac = (epoch - 1) + (micro_step / len(train_loader)) | |
| # ── Accumulate for smoothed logging ─────────────────────── | |
| logger.accumulate(last_loss, think_l, response_l, grad_norm, tok_per_sec) | |
| # ── Log to console + files every log_every steps ────────── | |
| if global_step % args.log_every == 0: | |
| m = logger.log_step( | |
| step=global_step, | |
| epoch_frac=epoch_frac, | |
| lr=lr, | |
| tokens_seen=tokens_seen, | |
| think_weight=think_weight, | |
| t_start=t_start, | |
| max_steps=max_steps, | |
| ) | |
| m["_max_steps"] = max_steps # used by progress bar | |
| train_loss_history.append({ | |
| "step": global_step, | |
| "loss": last_loss, | |
| "think_loss": think_l, | |
| "response_loss": response_l, | |
| }) | |
| # ── Validation every val_every steps ────────────────────── | |
| if global_step % args.val_every == 0: | |
| val_loss = estimate_val_loss(model, val_loader, device, | |
| args.val_batches, dtype) | |
| samples = generate_samples(inner, tokenizer, device, TEST_PROMPTS) | |
| logger.log_val(global_step, val_loss, samples) | |
| val_loss_history.append({"step": global_step, "val_loss": val_loss}) | |
| # ── Checkpoint every save_every steps ───────────────────── | |
| if global_step % args.save_every == 0: | |
| milestone = (global_step % args.milestone_every == 0) | |
| ckpt_path = save_checkpoint( | |
| args.out_dir, model, optimizer, scaler, | |
| step=global_step, phase=args.phase, | |
| tokens_seen=tokens_seen, epoch_frac=epoch_frac, | |
| train_loss=last_loss, | |
| think_loss=last_think_loss, | |
| response_loss=last_response_loss, | |
| lr=lr, grad_norm=last_grad_norm, | |
| config=model_config, | |
| train_loss_history=train_loss_history, | |
| val_loss_history=val_loss_history, | |
| milestone=milestone, | |
| ) | |
| write_checkpoint_index(args.out_dir, global_step, last_loss, | |
| tokens_seen, str(ckpt_path)) | |
| t_b0 = time.time() | |
| print(f"\n ✅ Epoch {epoch}/{n_epochs} done | " | |
| f"steps={global_step:,} | tokens={tokens_seen/1e9:.3f}B\n") | |
| # ── Phase Final Checkpoint ──────────────────────────────────────────────── | |
| save_checkpoint( | |
| args.out_dir, model, optimizer, scaler, | |
| step=global_step, phase=args.phase, | |
| tokens_seen=tokens_seen, epoch_frac=float(n_epochs), | |
| train_loss=last_loss, think_loss=last_think_loss, | |
| response_loss=last_response_loss, lr=min_lr, grad_norm=0.0, | |
| config=model_config, | |
| train_loss_history=train_loss_history, val_loss_history=val_loss_history, | |
| milestone=True, phase_final=True, | |
| ) | |
| write_checkpoint_index(args.out_dir, global_step, last_loss, | |
| tokens_seen, f"novamind_{args.phase}_final.pt") | |
| total_min = (time.time() - t_start) / 60 | |
| print(f"\n{'═'*65}") | |
| print(f" ✅ Phase {args.phase.upper()} COMPLETE") | |
| print(f" Steps : {global_step:,} / {max_steps:,}") | |
| print(f" Tokens seen : {tokens_seen/1e9:.3f}B") | |
| print(f" Total time : {total_min:.1f} min") | |
| print(f" Final loss : {last_loss:.6f}") | |
| print(f" Checkpoints : {args.out_dir}/") | |
| print(f"{'═'*65}\n") | |
| logger.finish() | |
| # Print what to run next | |
| next_phase = {"pretrain":"sft","sft":"cot_sft","cot_sft":None}[args.phase] | |
| if next_phase: | |
| print(f" ▶ Next: python train.py --phase {next_phase} " | |
| f"--resume {args.out_dir}/novamind_{args.phase}_final.pt\n") | |
| if __name__ == "__main__": | |
| main() | |