Upload training_scripts/train_multilingual_3b_fsdp.py with huggingface_hub

training_scripts/train_multilingual_3b_fsdp.py

@@ -0,0 +1,601 @@
+#!/usr/bin/env python3
+"""
+Multilingual 3.14B GPT Training — FSDP Version (Arabic-Rebalanced Data)
+
+Converted from DDP script. Key changes:
+- FullyShardedDataParallel (FSDP) with FULL_SHARD strategy
+- Mixed precision: bf16 compute, fp32 reduce
+- transformer_auto_wrap_policy wrapping each Block
+- Activation checkpointing via FSDP's apply_activation_checkpointing
+- FULL_STATE_DICT for checkpoint saving
+- SWA simplified (gather full state dict on rank 0)
+
+Architecture: dim=3072, depth=26, heads=24, ~3.14B params
+Data: training-data-v2 (4.48B tokens, multi-epoch)
+Schedule: WSD-LINEAR
+LR: 3e-4
+
+Run with:
+  /opt/pytorch/bin/torchrun --nproc_per_node=8 --master_port=29500 train_multilingual_3b_fsdp.py
+"""
+
+import os, sys, json, math, time, copy, functools, datetime
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
+
+# FSDP imports
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    ShardingStrategy,
+    MixedPrecision,
+    StateDictType,
+    FullStateDictConfig,
+    BackwardPrefetch,
+    CPUOffload,
+)
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+
+# Activation checkpointing for FSDP
+try:
+    from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+        apply_activation_checkpointing,
+        checkpoint_wrapper,
+        CheckpointImpl,
+    )
+    HAS_FSDP_AC = True
+except ImportError:
+    HAS_FSDP_AC = False
+
+# ============ MODEL CONFIG ============
+VOCAB_SIZE = 32000
+DIM = 3072
+DEPTH = 26
+N_HEADS = 24
+MAX_SEQ_LEN = 2048
+ROPE_THETA = 10000
+DROPOUT = 0.05
+
+# ============ TRAINING CONFIG ============
+TOTAL_STEPS = 20000
+WARMUP_STEPS = 600
+STABLE_END = 14000
+MIN_LR_RATIO = 0.03
+
+BATCH_PER_GPU = 4          # FSDP uses less memory per GPU → can increase from 2 to 4
+GRAD_ACCUM = 8             # With 8 GPUs: 8*4*8 = 256 seqs = 524K tokens/step
+# Total: 20000 * 524288 = 10.49B tokens
+
+ADAMW_LR = 3e-4
+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
+SWA_FREQ = 40
+
+EVAL_EVERY = 500
+SAVE_EVERY = 500
+LOG_EVERY = 50
+RESUME_STEP = int(os.environ.get('RESUME_STEP', '0'))  # Set via env to resume from checkpoint
+
+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-fsdp"
+
+# ============ 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:
+            # NOTE: With FSDP activation checkpointing applied externally via
+            # apply_activation_checkpointing, we do NOT need manual torch_checkpoint here.
+            # FSDP's checkpoint wrapper handles it.
+            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):
+    """Evaluate model. Works with both FSDP-wrapped and unwrapped models."""
+    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:
+    """SWA for FSDP: gathers full state dict on rank 0 before averaging."""
+    def __init__(self):
+        self.avg_state = None
+        self.n_averaged = 0
+
+    def update(self, model):
+        """Gather full state dict from FSDP model and update running average on rank 0."""
+        save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
+            state = model.state_dict()
+
+        # Only rank 0 gets the full state dict with rank0_only=True
+        if dist.get_rank() != 0:
+            return
+
+        if self.avg_state is None:
+            self.avg_state = {k: v.float().clone() for k, v in state.items()}
+            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].float()) / (n + 1)
+            self.n_averaged += 1
+
+# ============ FSDP HELPERS ============
+def get_fsdp_wrap_policy():
+    """Wrap each Block as a separate FSDP unit."""
+    return functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={Block},
+    )
+
+def get_mixed_precision():
+    """bf16 for compute, fp32 for reduce (gradient all-reduce in fp32 for stability)."""
+    return MixedPrecision(
+        param_dtype=torch.bfloat16,
+        reduce_dtype=torch.float32,
+        buffer_dtype=torch.bfloat16,
+    )
+
+def save_fsdp_checkpoint(model, optimizer, scaler, step, best_val_bpb,
+                         tokens_processed, eval_results, swa_n, ckpt_dir, logger):
+    """Save full state dict checkpoint from FSDP model (rank 0 only)."""
+    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
+        model_state = model.state_dict()
+
+    if dist.get_rank() == 0:
+        ckpt = {
+            'step': step,
+            'model': model_state,
+            'best_val_bpb': best_val_bpb,
+            'tokens_processed': tokens_processed,
+            'eval_results': eval_results,
+            'swa_n': swa_n,
+            # NOTE: We don't save optimizer/scaler state for simplicity with FSDP.
+            # For full resumability, use FSDP's ShardedStateDictConfig instead.
+        }
+        torch.save(ckpt, f"{ckpt_dir}/ckpt_step_{step}.pt")
+        logger.log(f"Saved checkpoint at step {step}")
+    dist.barrier()
+
+def save_fsdp_model(model, path, logger):
+    """Save just the model state dict (rank 0 only)."""
+    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
+        state = model.state_dict()
+    if dist.get_rank() == 0:
+        torch.save(state, path)
+        logger.log(f"Saved model to {path}")
+    dist.barrier()
+
+# ============ MAIN ============
+def main():
+    dist.init_process_group('nccl', timeout=datetime.timedelta(minutes=30))
+    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 — FSDP (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}")
+    logger.log(f"FSDP: FULL_SHARD, MixedPrecision(bf16/fp32), Block-level wrapping")
+
+    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 — create on CPU first, then FSDP wraps and shards to GPUs
+    logger.log("Creating model...")
+    torch.manual_seed(42)
+    model = GPT()
+    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:,})")
+
+    # Wrap with FSDP
+    logger.log("Wrapping model with FSDP...")
+    wrap_policy = get_fsdp_wrap_policy()
+    mixed_precision = get_mixed_precision()
+
+    model = FSDP(
+        model,
+        sharding_strategy=ShardingStrategy.FULL_SHARD,
+        mixed_precision=mixed_precision,
+        auto_wrap_policy=wrap_policy,
+        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+        device_id=local_rank,
+        limit_all_gathers=True,
+        use_orig_params=True,  # Required for weight decay on specific params
+    )
+
+    logger.log(f"FSDP wrapped. GPU memory: {torch.cuda.memory_allocated(device)/1e9:.1f} GB")
+
+    # Apply activation checkpointing to each Block within FSDP
+    if HAS_FSDP_AC:
+        check_fn = lambda submodule: isinstance(submodule, Block)
+        apply_activation_checkpointing(
+            model,
+            checkpoint_wrapper_fn=checkpoint_wrapper,
+            check_fn=check_fn,
+        )
+        logger.log("Applied FSDP activation checkpointing to Block layers")
+    else:
+        logger.log("WARNING: FSDP activation checkpointing not available, using manual checkpointing")
+
+    # Optimizer — use standard AdamW (bitsandbytes may not work well with FSDP sharded params)
+    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)
+
+    # Use ShardedGradScaler for FSDP (handles sharded gradients correctly)
+    scaler = ShardedGradScaler()
+
+    best_val_bpb = float('inf')
+    eval_results = []
+    tokens_processed = 0
+    start_step = 1
+
+    # Resume from checkpoint if requested
+    if RESUME_STEP > 0:
+        ckpt_path = f"{CKPT_DIR}/ckpt_step_{RESUME_STEP}.pt"
+        if not os.path.exists(ckpt_path):
+            # Try downloading from S3
+            if rank == 0:
+                logger.log(f"Downloading checkpoint from S3 for step {RESUME_STEP}...")
+                os.system(f"aws s3 cp s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}/ckpt_step_{RESUME_STEP}.pt {ckpt_path}")
+        dist.barrier()
+        if os.path.exists(ckpt_path):
+            if rank == 0:
+                logger.log(f"Loading checkpoint from step {RESUME_STEP}...")
+            ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+            # Load model weights into FSDP
+            with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
+                model.load_state_dict(ckpt['model'])
+            start_step = ckpt['step'] + 1
+            tokens_processed = ckpt.get('tokens_processed', 0)
+            best_val_bpb = ckpt.get('best_val_bpb', float('inf'))
+            eval_results = ckpt.get('eval_results', [])
+            swa.n_averaged = ckpt.get('swa_n', 0)
+            if rank == 0:
+                logger.log(f"Resumed from step {RESUME_STEP}. Tokens: {tokens_processed:,}, Best BPB: {best_val_bpb:.4f}")
+                logger.log(f"NOTE: Optimizer state reset (fresh AdamW). LR schedule continues from step {start_step}.")
+            del ckpt
+            # Advance RNG to approximately correct position
+            rng = np.random.default_rng(42 + rank + RESUME_STEP * 1000)
+            dist.barrier()
+        else:
+            if rank == 0:
+                logger.log(f"WARNING: No checkpoint found for step {RESUME_STEP}, starting from scratch")
+
+    start_time = time.time()
+    logger.log(f"Starting training from step {start_step}...")
+
+    for step in range(start_step, 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)
+
+            # For FSDP with gradient accumulation, we need to use no_sync() for all
+            # micro-steps except the last one to avoid unnecessary all-reduce
+            ctx = model.no_sync() if micro < GRAD_ACCUM - 1 else nullcontext()
+            with ctx:
+                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)
+        # FSDP clip_grad_norm_ works on the sharded params
+        model.clip_grad_norm_(GRAD_CLIP)
+        scaler.step(optimizer)
+        scaler.update()
+        tokens_processed += tokens_per_step
+
+        # SWA: gather full state dict and average on rank 0
+        if step >= swa_start_step and step % SWA_FREQ == 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")
+
+        # Save checkpoint BEFORE eval to avoid FSDP state dict deadlocks
+        if step % SAVE_EVERY == 0:
+            save_fsdp_checkpoint(
+                model, optimizer, scaler, step, best_val_bpb,
+                tokens_processed, eval_results, swa.n_averaged, CKPT_DIR, logger
+            )
+            if rank == 0:
+                os.system(f"aws s3 sync {CKPT_DIR}/ s3://{S3_BUCKET}/{S3_PREFIX}/checkpoints/{VERSION}/ --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 &")
+
+        if step % EVAL_EVERY == 0 or step == TOTAL_STEPS:
+            dist.barrier()  # Ensure all ranks are synced before eval
+            # All ranks participate in eval (FSDP needs all ranks for forward pass)
+            combined_bpb = evaluate(model, val_batches, device)
+
+            if rank == 0:
+                logger.log(f"--- Evaluation at step {step} ---")
+                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, batches, device)
+                if rank == 0:
+                    result[f"{lang}_bpb"] = lang_bpb
+                    logger.log(f"  {lang} val BPB: {lang_bpb:.4f}")
+
+            dist.barrier()  # All ranks must finish eval before proceeding
+
+            # Determine if new best (rank 0 decides, broadcast to all)
+            is_new_best = False
+            if rank == 0:
+                eval_results.append(result)
+                with open(EVAL_FILE, 'w') as f:
+                    json.dump(eval_results, f, indent=2)
+                is_new_best = combined_bpb < best_val_bpb
+                if is_new_best:
+                    best_val_bpb = combined_bpb
+                    logger.log(f"  New best! BPB: {combined_bpb:.4f}")
+
+            # All ranks participate in model save (FSDP state dict gather requires it)
+            save_flag = torch.tensor([1 if is_new_best else 0], device=device)
+            dist.broadcast(save_flag, src=0)
+            if save_flag.item() == 1:
+                save_fsdp_model(model, f"{CKPT_DIR}/best_model.pt", logger)
+            dist.barrier()
+
+    # Finalize
+    save_fsdp_model(model, f"{CKPT_DIR}/final_model.pt", logger)
+
+    if rank == 0:
+        # SWA evaluation
+        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()
+
+# Need nullcontext for no_sync
+from contextlib import nullcontext
+
+if __name__ == '__main__':
+    main()