#!/usr/bin/env python3 """ Multilingual 3.14B GPT Training (Arabic-Rebalanced Data) Scaled from 1B v2/v3 validated recipe: - Architecture: dim=3072, depth=26, heads=24, ~3.14B params - Data: training-data-v2 (4.48B tokens, multi-epoch) - Schedule: WSD-LINEAR (validated at 1B) - LR: 3e-4 (scaled from 1B's 5e-4 via sqrt rule) """ import os, sys, json, math, time, copy import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.checkpoint import checkpoint as torch_checkpoint # ============ MODEL CONFIG ============ VOCAB_SIZE = 32000 DIM = 3072 DEPTH = 26 N_HEADS = 24 MAX_SEQ_LEN = 2048 ROPE_THETA = 10000 DROPOUT = 0.05 # Slightly less than 1B (0.1) — larger models need less regularization # ============ TRAINING CONFIG ============ TOTAL_STEPS = 20000 # ~10.5B tokens (2.3 epochs of 4.48B) WARMUP_STEPS = 600 # 3% warmup STABLE_END = 14000 # 70% at peak LR MIN_LR_RATIO = 0.03 BATCH_PER_GPU = 2 # Per-GPU batch size GRAD_ACCUM = 16 # With 8 GPUs: 8*4*8 = 256 seqs = 524K tokens/step # Total: 20000 * 524288 = 10.49B tokens ADAMW_LR = 3e-4 # Scaled: 5e-4 * sqrt(502M/3142M) ≈ 2e-4, being slightly aggressive ADAMW_BETAS = (0.9, 0.98) ADAMW_WD = 0.02 ADAMW_EPS = 1e-8 LABEL_SMOOTHING = 0.06 GRAD_CLIP = 1.0 SWA_START_FRAC = 0.40 # Start at step 8000 SWA_FREQ = 40 # Every 40 steps (scaled from 1B's 20) EVAL_EVERY = 500 # Eval every 500 steps SAVE_EVERY = 2000 # Checkpoint every 2000 steps LOG_EVERY = 50 DATA_DIR = "/tmp/training-data" CKPT_DIR = "/tmp/checkpoints" LOG_FILE = "/tmp/training.log" EVAL_FILE = "/tmp/eval_results.json" S3_BUCKET = "autoresearch-dashboard-196766918360" S3_PREFIX = "multilingual-7b" VERSION = "3b-v1" # ============ MODEL ============ class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).type_as(x) * self.weight class SwiGLU(nn.Module): def __init__(self, dim, hidden_dim): super().__init__() self.gate = nn.Linear(dim, hidden_dim, bias=False) self.up = nn.Linear(dim, hidden_dim, bias=False) self.down = nn.Linear(hidden_dim, dim, bias=False) def forward(self, x): return self.down(F.silu(self.gate(x)) * self.up(x)) def apply_rope(x, cos, sin): x1, x2 = x[..., ::2], x[..., 1::2] return torch.stack((x1*cos - x2*sin, x1*sin + x2*cos), dim=-1).flatten(-2) class Attention(nn.Module): def __init__(self, dim, n_heads, dropout=0.0): super().__init__() self.n_heads = n_heads self.head_dim = dim // n_heads self.qkv = nn.Linear(dim, 3*dim, bias=False) self.proj = nn.Linear(dim, dim, bias=False) self.attn_dropout = dropout def forward(self, x, cos, sin): B, T, C = x.shape qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=self.attn_dropout if self.training else 0.0) return self.proj(y.transpose(1, 2).contiguous().view(B, T, C)) class Block(nn.Module): def __init__(self, dim, n_heads, mlp_dim, dropout=0.0): super().__init__() self.ln1 = RMSNorm(dim) self.attn = Attention(dim, n_heads, dropout) self.ln2 = RMSNorm(dim) self.mlp = SwiGLU(dim, mlp_dim) self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity() def forward(self, x, cos, sin): x = x + self.drop(self.attn(self.ln1(x), cos, sin)) x = x + self.drop(self.mlp(self.ln2(x))) return x class GPT(nn.Module): def __init__(self, vocab_size=VOCAB_SIZE, dim=DIM, depth=DEPTH, n_heads=N_HEADS, max_seq_len=MAX_SEQ_LEN, rope_theta=ROPE_THETA, dropout=DROPOUT): super().__init__() self.tok_emb = nn.Embedding(vocab_size, dim) mlp_dim = ((int(2 * dim * 4 / 3) + 63) // 64) * 64 # = 8192 for dim=3072 self.blocks = nn.ModuleList([Block(dim, n_heads, mlp_dim, dropout) for _ in range(depth)]) self.ln_f = RMSNorm(dim) self.head = nn.Linear(dim, vocab_size, bias=False) self.head.weight = self.tok_emb.weight # weight tying hd = dim // n_heads freqs = 1.0 / (rope_theta ** (torch.arange(0, hd, 2).float() / hd)) angles = torch.outer(torch.arange(max_seq_len).float(), freqs) self.register_buffer('rope_cos', angles.cos()) self.register_buffer('rope_sin', angles.sin()) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx): B, T = idx.shape x = self.tok_emb(idx) cos = self.rope_cos[:T][None, None] sin = self.rope_sin[:T][None, None] for block in self.blocks: if self.training: x = torch_checkpoint(block, x, cos, sin, use_reentrant=False) else: x = block(x, cos, sin) return self.head(self.ln_f(x)) # ============ WSD LINEAR SCHEDULE ============ def wsd_lr_linear(step, total_steps, warmup_steps, stable_end, min_lr_ratio, base_lr): if step < warmup_steps: return base_lr * (step + 1) / max(warmup_steps, 1) elif step < stable_end: return base_lr else: progress = (step - stable_end) / max(total_steps - stable_end, 1) return base_lr * (1.0 - progress * (1.0 - min_lr_ratio)) # ============ DATA LOADING ============ class BinaryDataset: def __init__(self, path, seq_len): self.data = np.memmap(path, dtype=np.uint16, mode='r') self.seq_len = seq_len self.n_tokens = len(self.data) def get_batch(self, batch_size, device, rng): ix = torch.from_numpy(rng.integers(0, self.n_tokens - self.seq_len - 1, size=(batch_size,))) x = torch.stack([torch.from_numpy(self.data[i:i+self.seq_len].astype(np.int64)) for i in ix]) y = torch.stack([torch.from_numpy(self.data[i+1:i+1+self.seq_len].astype(np.int64)) for i in ix]) return x.to(device), y.to(device) def load_val_data(path, seq_len, max_batches=20, batch_size=8): data = np.memmap(path, dtype=np.uint16, mode='r') n_tokens = len(data) batches = [] stride = seq_len + 1 all_starts = list(range(0, n_tokens - stride, stride)) max_samples = max_batches * batch_size if len(all_starts) > max_samples: step_size = len(all_starts) // max_samples all_starts = all_starts[::step_size][:max_samples] for i in range(0, len(all_starts), batch_size): batch_starts = all_starts[i:i+batch_size] if len(batch_starts) < batch_size: break x = torch.stack([torch.from_numpy(data[s:s+seq_len].astype(np.int64)) for s in batch_starts]) y = torch.stack([torch.from_numpy(data[s+1:s+1+seq_len].astype(np.int64)) for s in batch_starts]) batches.append((x, y)) return batches @torch.no_grad() def evaluate(model, val_batches, device): model.eval() total_loss = 0.0 total_tokens = 0 for x, y in val_batches: x, y = x.to(device), y.to(device) with torch.amp.autocast('cuda', dtype=torch.bfloat16): logits = model(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1), reduction='sum') total_loss += loss.item() total_tokens += y.numel() model.train() return (total_loss / total_tokens) / math.log(2) if total_tokens > 0 else float('inf') class Logger: def __init__(self, log_file, rank): self.rank = rank if rank == 0: self.f = open(log_file, 'w') def log(self, msg): if self.rank == 0: ts = time.strftime('%Y-%m-%d %H:%M:%S') line = f"[{ts}] {msg}" print(line, flush=True) self.f.write(line + '\n') self.f.flush() def close(self): if self.rank == 0: self.f.close() class SWAState: def __init__(self): self.avg_state = None self.n_averaged = 0 def update(self, model): state = {k: v.cpu().float().clone() for k, v in model.module.state_dict().items()} if self.avg_state is None: self.avg_state = state self.n_averaged = 1 else: n = self.n_averaged for k in self.avg_state: self.avg_state[k] = (self.avg_state[k] * n + state[k]) / (n + 1) self.n_averaged += 1 # ============ MAIN ============ def main(): dist.init_process_group('nccl') rank = dist.get_rank() world_size = dist.get_world_size() local_rank = int(os.environ.get('LOCAL_RANK', 0)) device = torch.device(f'cuda:{local_rank}') torch.cuda.set_device(device) effective_batch = BATCH_PER_GPU * GRAD_ACCUM * world_size tokens_per_step = effective_batch * MAX_SEQ_LEN logger = Logger(LOG_FILE, rank) logger.log(f"=== Multilingual 3.14B Training (Arabic-Rebalanced) ===") logger.log(f"World size: {world_size}, Batch/GPU: {BATCH_PER_GPU}, Grad accum: {GRAD_ACCUM}") logger.log(f"Effective batch: {effective_batch} seqs = {tokens_per_step:,} tokens/step") logger.log(f"Total steps: {TOTAL_STEPS} = {TOTAL_STEPS * tokens_per_step:,} tokens") logger.log(f"Schedule: WSD-LINEAR | warmup={WARMUP_STEPS} | stable_end={STABLE_END} | total={TOTAL_STEPS}") logger.log(f"AdamW LR={ADAMW_LR}, betas={ADAMW_BETAS}, WD={ADAMW_WD}") logger.log(f"Label smoothing={LABEL_SMOOTHING}, min_lr={MIN_LR_RATIO}, grad_clip={GRAD_CLIP}") logger.log(f"SWA: start={int(TOTAL_STEPS*SWA_START_FRAC)}, freq={SWA_FREQ}") logger.log(f"Model: dim={DIM}, depth={DEPTH}, heads={N_HEADS}, dropout={DROPOUT}") os.makedirs(CKPT_DIR, exist_ok=True) # Data logger.log("Loading training data...") train_ds = BinaryDataset(f"{DATA_DIR}/train.bin", MAX_SEQ_LEN) logger.log(f"Train tokens: {train_ds.n_tokens:,}") logger.log("Loading validation data...") val_batches = load_val_data(f"{DATA_DIR}/val.bin", MAX_SEQ_LEN) val_lang_batches = {} for lang in ['en', 'ar', 'he', 'fa']: vpath = f"{DATA_DIR}/val_{lang}.bin" if os.path.exists(vpath): val_lang_batches[lang] = load_val_data(vpath, MAX_SEQ_LEN) logger.log(f" val_{lang}: {len(val_lang_batches[lang])} batches") # Model logger.log("Creating model...") torch.manual_seed(42) model = GPT().to(device) n_params = sum(p.numel() for p in model.parameters()) n_params_no_emb = n_params - model.tok_emb.weight.numel() logger.log(f"Model params: {n_params:,} (non-embedding: {n_params_no_emb:,})") logger.log(f"GPU memory after model: {torch.cuda.memory_allocated(device)/1e9:.1f} GB") model = DDP(model, device_ids=[local_rank]) try: import bitsandbytes as bnb optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=ADAMW_LR, weight_decay=ADAMW_WD, betas=ADAMW_BETAS, eps=ADAMW_EPS) logger.log("Using 8-bit AdamW (bitsandbytes)") except ImportError: optimizer = torch.optim.AdamW(model.parameters(), lr=ADAMW_LR, weight_decay=ADAMW_WD, betas=ADAMW_BETAS, eps=ADAMW_EPS) logger.log("Using standard AdamW (bitsandbytes not available)") swa = SWAState() swa_start_step = int(TOTAL_STEPS * SWA_START_FRAC) rng = np.random.default_rng(42 + rank) scaler = torch.amp.GradScaler('cuda') best_val_bpb = float('inf') eval_results = [] tokens_processed = 0 start_time = time.time() logger.log("Starting training...") for step in range(1, TOTAL_STEPS + 1): model.train() lr = wsd_lr_linear(step, TOTAL_STEPS, WARMUP_STEPS, STABLE_END, MIN_LR_RATIO, ADAMW_LR) for g in optimizer.param_groups: g['lr'] = lr optimizer.zero_grad() accum_loss = 0.0 for micro in range(GRAD_ACCUM): x, y = train_ds.get_batch(BATCH_PER_GPU, device, rng) with torch.amp.autocast('cuda', dtype=torch.bfloat16): logits = model(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1), label_smoothing=LABEL_SMOOTHING) / GRAD_ACCUM scaler.scale(loss).backward() accum_loss += loss.item() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) scaler.step(optimizer) scaler.update() tokens_processed += tokens_per_step if step >= swa_start_step and step % SWA_FREQ == 0 and rank == 0: swa.update(model) if step % LOG_EVERY == 0 and rank == 0: elapsed = time.time() - start_time tps = tokens_processed / elapsed bpb = accum_loss / math.log(2) phase = "warmup" if step < WARMUP_STEPS else ("stable" if step < STABLE_END else "decay") mem = torch.cuda.max_memory_allocated(device) / 1e9 logger.log(f"Step {step}/{TOTAL_STEPS} [{phase}] | Loss: {accum_loss:.4f} | " f"BPB: {bpb:.4f} | LR: {lr:.6f} | Tokens: {tokens_processed:,} | " f"TPS: {tps:,.0f} | SWA: {swa.n_averaged} | Mem: {mem:.1f}GB | {elapsed/60:.1f}min") if step % EVAL_EVERY == 0 or step == TOTAL_STEPS: if rank == 0: logger.log(f"--- Evaluation at step {step} ---") combined_bpb = evaluate(model.module, val_batches, device) logger.log(f" Combined val BPB: {combined_bpb:.4f}") result = {"step": step, "tokens": tokens_processed, "combined_bpb": combined_bpb} for lang, batches in val_lang_batches.items(): lang_bpb = evaluate(model.module, batches, device) result[f"{lang}_bpb"] = lang_bpb logger.log(f" {lang} val BPB: {lang_bpb:.4f}") eval_results.append(result) with open(EVAL_FILE, 'w') as f: json.dump(eval_results, f, indent=2) if combined_bpb < best_val_bpb: best_val_bpb = combined_bpb torch.save(model.module.state_dict(), f"{CKPT_DIR}/best_model.pt") logger.log(f" New best! BPB: {combined_bpb:.4f}") dist.barrier() if step % SAVE_EVERY == 0 and rank == 0: ckpt = { 'step': step, 'model': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'scaler': scaler.state_dict(), 'best_val_bpb': best_val_bpb, 'tokens_processed': tokens_processed, 'eval_results': eval_results, 'swa_n': swa.n_averaged, } torch.save(ckpt, f"{CKPT_DIR}/ckpt_step_{step}.pt") logger.log(f"Saved checkpoint at step {step}") # Background upload os.system(f"aws s3 cp {CKPT_DIR}/best_model.pt s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}_best_model.pt --quiet &") os.system(f"aws s3 cp {EVAL_FILE} s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}_eval_results.json --quiet &") os.system(f"aws s3 cp {LOG_FILE} s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}_training.log --quiet &") # Finalize if rank == 0: torch.save(model.module.state_dict(), f"{CKPT_DIR}/final_model.pt") logger.log("Saved final model") if swa.avg_state is not None and swa.n_averaged > 0: logger.log(f"Evaluating SWA model ({swa.n_averaged} checkpoints)...") swa_model = GPT().to(device) swa_load = {k: v.float().to(device) for k, v in swa.avg_state.items()} swa_model.load_state_dict(swa_load) swa_bpb = evaluate(swa_model, val_batches, device) logger.log(f"SWA model combined BPB: {swa_bpb:.4f} (vs best raw: {best_val_bpb:.4f})") swa_result = {"step": "swa", "combined_bpb": swa_bpb, "n_averaged": swa.n_averaged} for lang, batches in val_lang_batches.items(): lang_bpb = evaluate(swa_model, batches, device) swa_result[f"{lang}_bpb"] = lang_bpb logger.log(f" SWA {lang} BPB: {lang_bpb:.4f}") eval_results.append(swa_result) torch.save(swa_load, f"{CKPT_DIR}/swa_model.pt") with open(EVAL_FILE, 'w') as f: json.dump(eval_results, f, indent=2) del swa_model # Final S3 upload logger.log("Uploading all artifacts to S3...") os.system(f"aws s3 sync {CKPT_DIR}/ s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}/") os.system(f"aws s3 cp {LOG_FILE} s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}_training.log") os.system(f"aws s3 cp {EVAL_FILE} s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}_eval_results.json") elapsed = time.time() - start_time logger.log(f"=== Training complete! Total time: {elapsed/3600:.2f}h ===") logger.log(f"Best combined BPB: {best_val_bpb:.4f}") logger.log(f"Total tokens: {tokens_processed:,}") logger.close() dist.destroy_process_group() if __name__ == '__main__': main()