# ============================================================================
# ViuAI 500M — Supervised Fine-Tuning (SFT) Script
# ============================================================================
# Universal script: works on H200(140GB), 5090(32GB), L4(24GB), T4x2(16GBx2)
# Works on: Kaggle, Rental GPU, Any cloud provider
#
# Run commands:
#   Kaggle T4x2:      torchrun --nproc_per_node=2 train_sft.py
#   Single GPU:        torchrun --nproc_per_node=1 train_sft.py
#   Multi-GPU:         torchrun --nproc_per_node=<N> train_sft.py
# ============================================================================

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# Redirect HF cache to working directory (rental GPUs often have tiny OS partitions)
# This prevents "No space left on device" errors during checkpoint upload
if os.path.exists("/kaggle/working"):
    os.environ.setdefault("HF_HOME", "/kaggle/working/hf_cache")
else:
    os.environ.setdefault("HF_HOME", "./working/hf_cache")

import json, math, time, random, sys, contextlib, gc, threading

# Global list to track all background upload threads
_upload_threads = []
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 huggingface_hub import HfApi, hf_hub_download, create_repo
from datetime import timedelta

# ============================================================================
# 1. DDP Setup (works for single GPU too via torchrun --nproc_per_node=1)
# ============================================================================
def setup_ddp():
    if not torch.cuda.is_available():
        raise RuntimeError("SFT training requires at least one CUDA GPU.")
    dist.init_process_group(backend="nccl", timeout=timedelta(minutes=30))
    rank = dist.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = dist.get_world_size()
    torch.cuda.set_device(local_rank)
    return rank, local_rank, world_size

rank, local_rank, world_size = setup_ddp()
DEVICE = f"cuda:{local_rank}"
IS_MAIN = rank == 0
IS_LOCAL_MAIN = local_rank == 0

def log(msg):
    if IS_MAIN:
        print(msg, flush=True)

# ============================================================================
# 2. Authentication (Kaggle + Rental GPU dono support)
# ============================================================================
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    try:
        from kaggle_secrets import UserSecretsClient
        HF_TOKEN = UserSecretsClient().get_secret("HF_TOKEN")
        os.environ["HF_TOKEN"] = HF_TOKEN
    except Exception:
        pass
if not HF_TOKEN:
    raise RuntimeError(
        "HF_TOKEN not found!\n"
        "  Kaggle: Settings → Secrets → Add HF_TOKEN\n"
        "  Rental GPU: export HF_TOKEN='hf_your_token'"
    )
api = HfApi(token=HF_TOKEN)

# ============================================================================
# 3. Repository Configuration
# ============================================================================
MAIN_REPO      = os.environ.get("VIUAI_MODEL_REPO", "ViuAI/ViuAI-500M")
SFT_DATA_REPO  = os.environ.get("VIUAI_SFT_DATA_REPO", "ViuAI/viuai-500m-sft-tokenized")

def env_flag(name, default=False):
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "y", "on"}

# Fresh SFT is the safe default. Set RESUME_SFT=1 only to continue an
# interrupted SFT run from the configured SFT checkpoint directory.
RESUME_SFT = env_flag("RESUME_SFT", default=True)
RESUME_CHECKPOINT_FILE = "sft_ckpt_latest.pt" if RESUME_SFT else "__fresh_sft_start__.pt"
SFT_CHECKPOINT_DIR = os.environ.get("SFT_CHECKPOINT_DIR", "sft_checkpoints/sft_v2").strip("/")
if not SFT_CHECKPOINT_DIR or SFT_CHECKPOINT_DIR.startswith(".") or ".." in SFT_CHECKPOINT_DIR.split("/"):
    raise ValueError("SFT_CHECKPOINT_DIR must be a safe relative Hugging Face repository path.")

if IS_MAIN:
    create_repo(MAIN_REPO, repo_type="model", token=HF_TOKEN, exist_ok=True, private=True)
dist.barrier()

# ============================================================================
# 4. Working Directory (auto-detect environment)
# ============================================================================
if os.path.exists("/kaggle/working"):
    WORK = "/kaggle/working"
elif os.path.exists("/root/working"):
    WORK = "/root/working"
else:
    WORK = "./working"
os.makedirs(WORK, exist_ok=True)
os.makedirs(f"{WORK}/code", exist_ok=True)
sys.path.insert(0, f"{WORK}/code")

# ============================================================================
# 5. Download Model Code from HuggingFace
# ============================================================================
def with_retry(fn, *args, retries=4, delay=5, **kwargs):
    for attempt in range(1, retries + 1):
        try:
            return fn(*args, **kwargs)
        except Exception as e:
            if "404" in str(e) or "EntryNotFoundError" in type(e).__name__:
                raise
            if attempt == retries:
                raise
            log(f"  [retry {attempt}/{retries}] {type(e).__name__}: {e}")
            time.sleep(delay)
            delay *= 2

if IS_MAIN:
    log("📥 Downloading model code...")
try:
    download_on_main(hf_hub_download, MAIN_REPO, "code/config.py", local_dir=WORK, token=HF_TOKEN)
    download_on_main(hf_hub_download, MAIN_REPO, "code/model.py", local_dir=WORK, token=HF_TOKEN)
except NameError:
    # download_on_main is defined later, so we just use the original with_retry pattern for code download
    # Since we're downloading the code, we just do it manually with all_reduce
    ok = True
    err = ""
    if IS_MAIN:
        try:
            with_retry(hf_hub_download, MAIN_REPO, "code/config.py", local_dir=WORK, token=HF_TOKEN)
            with_retry(hf_hub_download, MAIN_REPO, "code/model.py", local_dir=WORK, token=HF_TOKEN)
        except Exception as e:
            ok = False
            err = str(e)
            print(f"Download failed: {err}")
    
    # Can't use DEVICE here since setup_ddp is done but DEVICE isn't assigned if this is before device assignment
    # Wait, DEVICE is defined right after setup_ddp.
    ok_tensor = torch.tensor([1 if ok else 0], device=DEVICE)
    dist.all_reduce(ok_tensor, op=dist.ReduceOp.MIN)
    if ok_tensor.item() == 0:
        raise RuntimeError(f"Model code download failed on main rank: {err}")
    dist.barrier()

from config import ViuAIConfig
from model import ViuAI

# ============================================================================
# 6. GPU Auto-Detection & Hyperparameters
# ============================================================================
mem_gb = torch.cuda.get_device_properties(local_rank).total_memory / 1e9
gpu_name = torch.cuda.get_device_name(local_rank)

# Check if GPU supports bfloat16 (Ampere+ architecture, compute capability >= 8.0)
compute_cap = torch.cuda.get_device_capability(local_rank)
supports_bf16 = compute_cap[0] >= 8  # Ampere (8.0), Hopper (9.0) etc.

if mem_gb >= 130:       # H200 (141GB) / H100 (140GB)
    MICRO_BATCH   = 64
    GRAD_ACCUM    = 1
    CONTEXT_LEN   = 2048
    AMP_DTYPE     = torch.bfloat16
    USE_COMPILE   = True
    EMPTY_CACHE_EVERY = 2000
elif mem_gb >= 70:     # A100 (80GB)
    MICRO_BATCH   = 32
    GRAD_ACCUM    = 2
    CONTEXT_LEN   = 2048
    AMP_DTYPE     = torch.bfloat16
    USE_COMPILE   = True
    EMPTY_CACHE_EVERY = 2000
elif mem_gb >= 28:     # RTX 5090 (32GB), A10G (24GB but reported higher)
    MICRO_BATCH   = 16       # increased from 2 -> 16, now feasible since 
                             # we're on a 32GB RTX 5090 instead of a 15.6GB 
                             # Kaggle T4. Matches the proven working 250M 
                             # config exactly.
                             # NOTE: If this causes a CUDA out-of-memory 
                             # (OOM) error (500M has ~2x the parameters of 
                             # 250M, so may need more headroom), fall back 
                             # to MICRO_BATCH=8, GRAD_ACCUM=8 instead 
                             # (keeps the same effective batch size of 64).
    GRAD_ACCUM    = 4        # reverted from 8 -> 4, matches 250M's proven 
                             # config. Effective batch size = 16 x 4 = 64 
                             # (same as the working 250M run).
    CONTEXT_LEN   = 1024     # unchanged
    AMP_DTYPE     = torch.bfloat16
    USE_COMPILE   = True
    EMPTY_CACHE_EVERY = 200
elif mem_gb >= 20:     # L4 (24GB), RTX 4090 (24GB)
    MICRO_BATCH   = 4
    GRAD_ACCUM    = 8
    CONTEXT_LEN   = 2048
    AMP_DTYPE     = torch.bfloat16 if supports_bf16 else torch.float16
    USE_COMPILE   = supports_bf16  # Only compile on Ampere+
    EMPTY_CACHE_EVERY = 100
else:                  # T4 (16GB), RTX 3060 (12GB)
    MICRO_BATCH   = 2        # unchanged — limited by T4 GPU VRAM (15.6GB)
    GRAD_ACCUM    = 16       # increased from 8 → 16, to raise effective 
                             # batch size from 32 to 64 (matches the 
                             # proven-stable 250M SFT run's effective batch,
                             # reduces gradient noise)
    CONTEXT_LEN   = 1024     # unchanged
    AMP_DTYPE     = torch.float32  # Changed from float16 to float32 for stability
    USE_COMPILE   = False  # T4 Turing has poor torch.compile support
    EMPTY_CACHE_EVERY = 50

# SFT-specific hyperparameters (same for all GPUs)
MAX_LR        = 5e-05    # Increased from 1e-05 to 5e-05 to quickly learn new chat tokens
                         # the PROVEN value from the working 250M run; 
                         # 3e-05 combined with our smaller batch size 
                         # caused training instability and garbage output
MIN_LR        = 1e-06    # unchanged
WARMUP_STEPS  = 50       # reverted from 150 → 50, matching the working 
                         # 250M run (no longer need extra-long warmup 
                         # since LR is back to a gentler, proven value)
EPOCHS        = 4        # reduced from 6 → 4 (250M used 3 successfully; 
                         # we keep slightly more since 500M has more 
                         # parameters to adapt, but avoid excessive 
                         # total training time)
WEIGHT_DECAY  = 0.01
GRAD_CLIP     = 1.0
EVAL_EVERY    = 1000
CKPT_EVERY    = 4000

# Vocab size: 64009 because new SFT data uses 9 special tokens (<|user|>, <|assistant|>, <|endofturn|>, <think>, </think>, <search>, </search>, <search_result>, </search_result>)
SFT_VOCAB_SIZE = 64009

# Target effective batch logic removed to respect GPU-specific GRAD_ACCUM settings

effective_batch = MICRO_BATCH * GRAD_ACCUM * world_size
log(f"{'='*60}")
log(f"🖥️  GPU: {gpu_name} | VRAM: {mem_gb:.1f}GB | Compute: {compute_cap[0]}.{compute_cap[1]}")
log(f"⚙️  MICRO_BATCH={MICRO_BATCH} | GRAD_ACCUM={GRAD_ACCUM} | CTX={CONTEXT_LEN}")
log(f"⚙️  AMP={AMP_DTYPE} | torch.compile={'ON' if USE_COMPILE else 'OFF'}")
log(f"⚙️  Effective Batch Size: {effective_batch} | GPUs: {world_size}")
log(f"⚙️  LR: {MAX_LR} → {MIN_LR} | Epochs: {EPOCHS}")
log(f"{'='*60}")

torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision('high')
# ============================================================================
# 7. SFT Dataset (Pre-tokenized .npy with -100 label masking)
# ============================================================================
# Labels have -100 on user-turn tokens → loss is ONLY on assistant responses.
# input_ids and labels are at SAME alignment (no shift in the data files).
# The shift is handled in compute_sft_loss() below.
class NpySFTDataset:
    def __init__(self, ids_path, labels_path, offsets_path, ctx_len):
        self.ids = np.load(ids_path, mmap_mode='r')
        self.labels = np.load(labels_path, mmap_mode='r')
        self.offsets = np.load(offsets_path)
        self.ctx_len = ctx_len

        if self.ids.ndim != 1 or self.labels.ndim != 1:
            raise ValueError("SFT ids and labels must be one-dimensional arrays.")
        if len(self.ids) != len(self.labels):
            raise ValueError(
                f"SFT ids/labels length mismatch: {len(self.ids)} != {len(self.labels)}"
            )
        if self.offsets.ndim != 1 or len(self.offsets) == 0:
            raise ValueError("SFT offsets must be a non-empty one-dimensional array.")
        if int(self.offsets[0]) != 0 or np.any(np.diff(self.offsets) < 0):
            raise ValueError("SFT offsets must start at 0 and be monotonically increasing.")
        if int(self.offsets[-1]) > len(self.ids):
            raise ValueError("SFT offsets point beyond the ids array.")

        # Support both common formats: N start offsets, or N+1 offsets where
        # the final entry is the len(ids) sentinel.
        self.num_examples = len(self.offsets) - 1 if int(self.offsets[-1]) == len(self.ids) else len(self.offsets)
        if self.num_examples <= 0:
            raise ValueError("SFT dataset contains no examples.")

    def __len__(self):
        return self.num_examples

    def get(self, idx):
        start = int(self.offsets[idx])
        end = int(self.offsets[idx + 1]) if idx + 1 < len(self.offsets) else len(self.ids)
        if end <= start:
            raise ValueError(f"SFT example {idx} is empty or has invalid offsets.")

        # Keep the end of an overlong conversation. Assistant targets are
        # normally at the end, while prefix truncation can remove every label.
        start = max(start, end - self.ctx_len)
        length = end - start

        input_ids = np.zeros(self.ctx_len, dtype=np.int64)
        labels = np.full(self.ctx_len, -100, dtype=np.int64)

        input_ids[:length] = self.ids[start:start + length].astype(np.int64)
        labels[:length] = self.labels[start:start + length].astype(np.int64)
        return input_ids, labels

def get_sft_batch(dataset, batch_size, rng):
    batch_ids, batch_labels = [], []
    empty_examples = 0
    while len(batch_ids) < batch_size:
        idx = rng.randrange(len(dataset))
        try:
            ids, labels = dataset.get(idx)
        except ValueError:
            empty_examples += 1
            if empty_examples >= batch_size * 10:
                raise RuntimeError("Too many invalid/empty SFT examples encountered.")
            continue
            
        if not np.any(labels != -100):
            empty_examples += 1
            if empty_examples >= batch_size * 10:
                raise RuntimeError(
                    "Too many SFT examples have no unmasked assistant labels. "
                    "Check the data preprocessing and conversation template."
                )
            continue
        batch_ids.append(ids)
        batch_labels.append(labels)
    x = torch.from_numpy(np.stack(batch_ids)).pin_memory().to(DEVICE, non_blocking=True)
    y = torch.from_numpy(np.stack(batch_labels)).pin_memory().to(DEVICE, non_blocking=True)
    return x, y

# ============================================================================
# 8. Checkpoint Key Remapping (handles torch.compile prefix + vocab resize)
# ============================================================================
def remap_checkpoint_keys(ckpt_state_dict, model_state_dict):
    """Remap checkpoint keys to match current model.
    
    Handles:
      1. torch.compile prefix: _orig_mod.X → X
      2. Vocab size mismatch: resize tok_emb.weight and head.weight
    """
    model_keys = set(model_state_dict.keys())

    # Normalize common DistributedDataParallel and torch.compile prefixes.
    stripped = {}
    n_stripped = 0
    for key, value in ckpt_state_dict.items():
        normalized_key = key
        for prefix in ("module._orig_mod.", "_orig_mod.", "module."):
            if normalized_key.startswith(prefix):
                normalized_key = normalized_key[len(prefix):]
                break
        if normalized_key != key:
            stripped[normalized_key] = value
            n_stripped += 1
        else:
            stripped[key] = value
    if n_stripped:
        log(f"  🔧 Stripped '_orig_mod.' from {n_stripped} keys (torch.compile artifact)")

    for emb_key in ["tok_emb.weight", "head.weight"]:
        if emb_key in stripped and emb_key in model_state_dict:
            ckpt_shape = stripped[emb_key].shape
            model_shape = model_state_dict[emb_key].shape
            if ckpt_shape != model_shape:
                log(f"  📐 Resizing '{emb_key}': {list(ckpt_shape)} → {list(model_shape)}")
                if len(ckpt_shape) != 2 or len(model_shape) != 2 or ckpt_shape[1] != model_shape[1]:
                    raise ValueError(
                        f"Cannot resize '{emb_key}' from {list(ckpt_shape)} to {list(model_shape)}"
                    )
                # Keep checkpoint surgery on CPU so loading does not create a
                # second large embedding tensor on the GPU.
                ckpt_weight = stripped[emb_key].detach().cpu()
                new_weight = model_state_dict[emb_key].detach().cpu().clone()
                min_vocab = min(ckpt_shape[0], model_shape[0])
                new_weight[:min_vocab] = ckpt_weight[:min_vocab]
                
                # Smart initialization: copy the mean embedding to new special tokens
                if model_shape[0] > ckpt_shape[0]:
                    log(f"  🧠 Smart init for new {model_shape[0] - ckpt_shape[0]} tokens (using mean embedding)")
                    mean_emb = ckpt_weight[:min_vocab].mean(dim=0)
                    for i in range(min_vocab, model_shape[0]):
                        new_weight[i] = mean_emb
                        
                stripped[emb_key] = new_weight

    return stripped

# ============================================================================
# 9. Checkpoint Save / Load
# ============================================================================
def save_checkpoint(model, optimizer, step, val_loss, filename, blocking=False, best_val_loss=None):
    """Save checkpoint locally + upload to HF in background.
    
    Args:
        blocking: If True, upload synchronously (used for final save).
                  If False, upload in background thread (used during training).
    """
    if not IS_MAIN:
        return
    ckpt_dir = f"{WORK}/{SFT_CHECKPOINT_DIR}"
    os.makedirs(ckpt_dir, exist_ok=True)
    path = f"{ckpt_dir}/{filename}"
    raw_model = model.module if hasattr(model, "module") else model
    if hasattr(raw_model, '_orig_mod'):
        raw_model = raw_model._orig_mod
    state_dict = raw_model.state_dict()
    clean_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}

    torch.save({"model": clean_dict, "optimizer": optimizer.state_dict(),
                "step": step, "val_loss": val_loss, "best_val_loss": best_val_loss}, path)
    log(f"  💾 Checkpoint saved locally @ step {step}, val_loss={val_loss:.4f}")

    if blocking:
        # Synchronous upload (for final save — must complete before exit)
        log(f"  ☁️ Uploading {filename} to HF (blocking, please wait)...")
        try:
            with_retry(api.upload_file, path_or_fileobj=path,
                       path_in_repo=f"{SFT_CHECKPOINT_DIR}/{filename}",
                       repo_id=MAIN_REPO, repo_type="model", token=HF_TOKEN)
            log(f"  ✅ Upload complete for {filename} @ step {step}")
        except Exception as e:
            log(f"  ❌ Upload FAILED for {filename}: {e}")
            log(f"  ⚠️ File is still saved locally at: {path}")
            log(f"  ⚠️ Please upload manually before shutting down GPU!")
    else:
        # Background upload (for mid-training saves - doesn't block GPU)
        if filename == "sft_ckpt_best.pt":
            log(f"  📌 Skipping background upload for {filename} as requested.")
            return

        def _bg_upload():
            try:
                with_retry(api.upload_file, path_or_fileobj=path,
                           path_in_repo=f"{SFT_CHECKPOINT_DIR}/{filename}",
                           repo_id=MAIN_REPO, repo_type="model", token=HF_TOKEN)
                log(f"  ☁️ [Background upload complete for {filename} @ step {step}]")
            except Exception as e:
                log(f"  ❌ [Background upload failed for {filename}: {e}]")

        t = threading.Thread(target=_bg_upload, daemon=False)  # daemon=False: won't be killed on exit
        t.start()
        _upload_threads.append(t)
        log(f"  ☁️ Uploading {filename} in background...")

def save_best_checkpoint(model, step, val_loss):
    """Save best SFT model locally only (no HF upload).
    
    Like pretraining: just saves locally and logs val_loss.
    Final upload happens at the end or manually.
    This avoids wasting GPU time on repeated uploads.
    """
    if not IS_MAIN:
        return
    ckpt_dir = f"{WORK}/{SFT_CHECKPOINT_DIR}"
    os.makedirs(ckpt_dir, exist_ok=True)
    path = f"{ckpt_dir}/sft_ckpt_best.pt"
    raw_model = model.module if hasattr(model, "module") else model
    if hasattr(raw_model, '_orig_mod'):
        raw_model = raw_model._orig_mod
    state_dict = raw_model.state_dict()
    clean_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}

    torch.save({"model": clean_dict, "step": step, "val_loss": val_loss}, path)
    log(f"  ✨ New best! val_loss={val_loss:.4f} @ step {step} (saved locally)")

def wait_for_uploads():
    """Wait for ALL background uploads to finish before exiting."""
    if not _upload_threads:
        return
    log(f"\n⏳ Waiting for {len(_upload_threads)} background upload(s) to complete...")
    for i, t in enumerate(_upload_threads):
        t.join(timeout=600)  # Wait up to 10 minutes per upload
        if t.is_alive():
            log(f"  ⚠️ Upload thread {i+1} timed out after 10 minutes!")
        else:
            log(f"  ✅ Upload {i+1}/{len(_upload_threads)} complete.")
    log(f"✅ All background uploads finished.")

def load_sft_checkpoint(model, optimizer):
    """Load checkpoint: tries SFT first, then falls back to base pretrained."""

    if not RESUME_SFT:
        log("Fresh SFT requested: loading the base pretrained checkpoint.")
    
    # Try 1: Resume from existing SFT checkpoint
    try:
        if not RESUME_SFT:
            raise FileNotFoundError("Fresh SFT run requested.")
        log("🔍 Looking for existing SFT checkpoint...")
        if IS_LOCAL_MAIN:
            try:
                with_retry(hf_hub_download, MAIN_REPO,
                                  f"{SFT_CHECKPOINT_DIR}/{RESUME_CHECKPOINT_FILE}",
                                  repo_type="model", local_dir=WORK, token=HF_TOKEN)
            except Exception:
                pass
        dist.barrier()
        path = f"{WORK}/{SFT_CHECKPOINT_DIR}/{RESUME_CHECKPOINT_FILE}"
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        
        # Get the actual model (unwrap DDP + compile)
        raw_model = model.module if hasattr(model, "module") else model
        if hasattr(raw_model, '_orig_mod'):
            raw_model = raw_model._orig_mod
        
        remapped = remap_checkpoint_keys(ckpt["model"], raw_model.state_dict())
        missing, unexpected = raw_model.load_state_dict(remapped, strict=False)
        allowed_missing = {"head.weight"}
        bad_missing = [m for m in missing if m not in allowed_missing]
        if bad_missing:
            raise RuntimeError(f"Missing critical checkpoint keys: {bad_missing}")
        if unexpected:
            log(f"Warning: unexpected checkpoint keys ignored: {unexpected}")
        raw_model.head.weight = raw_model.tok_emb.weight
        log("  🔗 Re-tied head.weight with tok_emb.weight")
        
        if "optimizer" in ckpt:
            try:
                optimizer.load_state_dict(ckpt["optimizer"])
            except Exception:
                log("  ⚠️ Could not load optimizer state (architecture change?), using fresh optimizer")
        
        step = ckpt.get("step", 0)
        val_loss = ckpt.get("val_loss", -1)
        best_val_loss = ckpt.get("best_val_loss", None)
        log(f"✅ Resumed SFT from step {step} (val_loss={val_loss:.4f})")
        
        del ckpt
        gc.collect()
        torch.cuda.empty_cache()
        return step, val_loss, best_val_loss
    except Exception as e:
        is_missing = isinstance(e, FileNotFoundError) or "No such file" in str(e)
        if RESUME_SFT and not is_missing:
            raise RuntimeError(
                "RESUME_SFT=1 was requested, but the SFT checkpoint could not be loaded due to an error. "
                "Fix the checkpoint or unset RESUME_SFT to start a fresh SFT run."
            ) from e
        log(f"  (No SFT checkpoint found. Falling back to base pretrained model.)")
        log(f"  (No SFT checkpoint found: {type(e).__name__})")

    # Try 2: Load BASE pretrained weights (30k step checkpoint)
    log("📥 Loading BASE 30k pretrained weights...")
    try:
        if IS_LOCAL_MAIN:
            with_retry(hf_hub_download, MAIN_REPO,
                                  "checkpoints/ckpt_latest.pt",
                              repo_type="model", local_dir=WORK, token=HF_TOKEN)
        dist.barrier()
        path = f"{WORK}/checkpoints/ckpt_latest.pt"
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        
        raw_model = model.module if hasattr(model, "module") else model
        if hasattr(raw_model, '_orig_mod'):
            raw_model = raw_model._orig_mod
        
        remapped = remap_checkpoint_keys(ckpt["model"], raw_model.state_dict())
        missing, unexpected = raw_model.load_state_dict(remapped, strict=False)
        allowed_missing = {"head.weight"}
        bad_missing = [m for m in missing if m not in allowed_missing]
        if bad_missing:
            raise RuntimeError(f"Missing critical checkpoint keys: {bad_missing}")
        if unexpected:
            log(f"Warning: unexpected checkpoint keys ignored: {unexpected}")
        raw_model.head.weight = raw_model.tok_emb.weight
        log("  🔗 Re-tied head.weight with tok_emb.weight")
        
        pretrain_step = ckpt.get("step", "?")
        pretrain_loss = ckpt.get("val_loss", -1)
        log(f"✅ Loaded BASE pretrained weights (pretrain step={pretrain_step}, "
            f"val_loss={pretrain_loss:.4f})")
        log("   Starting SFT from step 0 (optimizer is fresh for SFT)")
        
        del ckpt
        gc.collect()
        torch.cuda.empty_cache()
        return 0, -1, None
    except Exception as e:
        log(f"❌ FATAL: Cannot load base pretrained weights!")
        log(f"   Error: {type(e).__name__}: {e}")
        raise RuntimeError(
            "SFT requires a pretrained base model but none could be loaded.\n"
            f"Checked: {MAIN_REPO}/checkpoints/ckpt_latest.pt\n"
            f"         {MAIN_REPO}/{SFT_CHECKPOINT_DIR}/sft_ckpt_latest.pt"
        )

# ============================================================================
# 10. Learning Rate Schedule (Cosine with Warmup)
# ============================================================================
def get_lr(step, warmup, total, max_lr, min_lr):
    if step < warmup:
        return max_lr * (step + 1) / warmup
    if step > total:
        return min_lr
    ratio = (step - warmup) / max(total - warmup, 1)
    return min_lr + 0.5 * (1 + math.cos(math.pi * ratio)) * (max_lr - min_lr)

# ============================================================================
# 11. SFT Loss Function (CRITICAL: handles shift mismatch)
# ============================================================================
# The 500M model.py does NOT shift logits/labels internally.
# (It was removed because pretraining dataloader already shifts: y = tokens[i+1:])
# But SFT data has labels at SAME position as input_ids.
# So we MUST shift here:
#   - Input to model: x[:, :-1]  (all tokens except the last)
#   - Labels:         y[:, 1:]   (all labels shifted left by 1)
#
# This makes: logits[t] (predicts next token after 0..t) ↔ labels[t+1] ✓
def compute_sft_loss(model, x, y, pad_id=0):
    """Forward pass with proper shift for SFT."""
    x_in = x[:, :-1].contiguous()       # [B, T-1] — drop last token
    y_shifted = y[:, 1:].contiguous()    # [B, T-1] — shift labels left by 1
    
    use_pad_id = None
    if pad_id is not None and (x_in == pad_id).any():
        use_pad_id = pad_id

    logits, loss = model(x_in, targets=y_shifted, pad_id=use_pad_id)
    
    if not (y_shifted != -100).any():
        return torch.nan_to_num(logits.float()).sum() * 0.0

    return loss

# ============================================================================
# 12. Model Initialization
# ============================================================================
cfg = ViuAIConfig(
    vocab_size=SFT_VOCAB_SIZE,   # 64005 (includes <|user|>, <|assistant|>, <|endofturn|>, <think>, </think>)
    context_length=CONTEXT_LEN,
    use_checkpoint=True,          # Gradient checkpointing (saves VRAM)
    z_loss_weight=0.0,            # Disable z-loss for SFT
    attn_dropout=0.1,             # Regularization for SFT
    resid_dropout=0.1,
)
raw_model = ViuAI(cfg).to(DEVICE)
if IS_MAIN:
    log(f"🧠 Model params: {raw_model.num_params()/1e6:.1f}M | vocab_size={SFT_VOCAB_SIZE}")

# Optionally compile (significant speedup on Ampere+ GPUs)
if USE_COMPILE:
    log("⚡ Compiling model with torch.compile (this takes 1-2 mins on first step)...")
    raw_model = torch.compile(raw_model)

model = DDP(raw_model, device_ids=[local_rank])
decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
    if not param.requires_grad:
        continue
    if any(k in name for k in ["norm", "tok_emb"]):
        no_decay_params.append(param)
    else:
        decay_params.append(param)

optimizer = torch.optim.AdamW(
    [
        {"params": decay_params, "weight_decay": WEIGHT_DECAY},
        {"params": no_decay_params, "weight_decay": 0.0},
    ],
    lr=MAX_LR,
    betas=(0.9, 0.95)
)

# Load checkpoint (SFT resume or base pretrained)
start_step, prev_val_loss, best_val_loss_loaded = load_sft_checkpoint(model, optimizer)

state_t = torch.tensor([
    start_step, 
    prev_val_loss if isinstance(prev_val_loss, (int, float)) else -1.0, 
    best_val_loss_loaded if isinstance(best_val_loss_loaded, (int, float)) else -1.0
], dtype=torch.float32, device=DEVICE)
dist.broadcast(state_t, src=0)

start_step = int(state_t[0].item())
prev_val_loss = state_t[1].item() if state_t[1].item() != -1.0 else None
best_val_loss = state_t[2].item() if state_t[2].item() != -1.0 else None

if best_val_loss is None and isinstance(prev_val_loss, (int, float)) and prev_val_loss >= 0:
    best_val_loss = prev_val_loss

# ============================================================================
# 13. Download & Load SFT Data
# ============================================================================
log("📥 Downloading SFT data (.npy files)...")
# v2 dataset uses sft_v2_* prefix; v1 used sft_* prefix.
# Set env var SFT_DATA_VERSION=v1 to use the old dataset.
_DATA_VER = os.environ.get("SFT_DATA_VERSION", "v2")
_PFX      = f"sft_{_DATA_VER}_" if _DATA_VER != "v1_legacy" else "sft_"
if _DATA_VER == "v1":   # legacy v1 filenames had no version prefix
    _PFX = "sft_"
SFT_DATA_FILES = [
    f"{_PFX}train_ids.npy", f"{_PFX}train_labels.npy", f"{_PFX}train_offsets.npy",
    f"{_PFX}val_ids.npy",   f"{_PFX}val_labels.npy",   f"{_PFX}val_offsets.npy",
]
log(f"📂 SFT data version: {_DATA_VER} | prefix: '{_PFX}'")
def download_on_main(fn, *args, fatal=True, **kwargs):
    ok = True
    err = ""
    if IS_LOCAL_MAIN:
        try:
            with_retry(fn, *args, **kwargs)
        except Exception as e:
            ok = False
            err = f"{type(e).__name__}: {e}"
            log(f"Download failed: {err}")
    
    ok_tensor = torch.tensor([1 if ok else 0], device=DEVICE)
    dist.all_reduce(ok_tensor, op=dist.ReduceOp.MIN)
    if ok_tensor.item() == 0:
        if fatal:
            raise RuntimeError(f"Download failed on main rank: {err}")
        return False
    dist.barrier()
    return True

for fname in SFT_DATA_FILES:
    download_on_main(hf_hub_download, SFT_DATA_REPO, fname, repo_type="dataset", local_dir=WORK, token=HF_TOKEN)
log("✅ SFT data ready.")

train_ds = NpySFTDataset(
    f"{WORK}/{_PFX}train_ids.npy", f"{WORK}/{_PFX}train_labels.npy",
    f"{WORK}/{_PFX}train_offsets.npy", CONTEXT_LEN
)
val_ds = NpySFTDataset(
    f"{WORK}/{_PFX}val_ids.npy", f"{WORK}/{_PFX}val_labels.npy",
    f"{WORK}/{_PFX}val_offsets.npy", CONTEXT_LEN
)
if IS_MAIN:
    log(f"📊 Train examples: {len(train_ds)}, Val examples: {len(val_ds)}")

# ============================================================================
# 14. Sanity Checks
# ============================================================================
if IS_MAIN:
    _s_ids, _s_labels = train_ds.get(0)
    _has_masked = (_s_labels == -100).any()
    _has_active = (_s_labels != -100).any()
    assert _has_active, (
        "SFT FATAL: ALL tokens are masked (-100)! "
        "No assistant tokens to train on. Your data pipeline is broken."
    )
    if not _has_masked:
        log("WARNING: sample 0 has no -100 labels; it will train on every token in that sample.")
    _max_id = int(_s_ids[_s_ids > 0].max()) if (_s_ids > 0).any() else 0
    _masked_count = (_s_labels == -100).sum()
    _active_count = ((_s_labels != -100) & (_s_labels != 0)).sum()
    log(f"✅ Data sanity check passed:")
    log(f"   Masked (user) tokens: {_masked_count} | Active (assistant) tokens: {_active_count}")
    log(f"   Max token ID in data: {_max_id} | Model vocab: {SFT_VOCAB_SIZE}")
    if _max_id >= SFT_VOCAB_SIZE:
        raise ValueError(
            f"SFT token ID ({_max_id}) is outside model vocab ({SFT_VOCAB_SIZE}). "
            "Use the tokenizer that created this dataset."
        )
        log(f"⚠️ WARNING: Max token ID ({_max_id}) >= vocab ({SFT_VOCAB_SIZE})! Will crash!")
    del _s_ids, _s_labels

# ============================================================================
# 15. Training Step Math
# ============================================================================
PAD_ID = 0
rng = random.Random(1337 + rank)
val_rng = random.Random(999)

# Optional cap on training steps, read from environment variable.
# Set MAX_TRAIN_STEPS in the shell/notebook before running to do a 
# limited test run (e.g. "2000"). Leave unset for full training 
# (runs all steps for EPOCHS as originally planned).
_max_steps_env = os.environ.get("MAX_TRAIN_STEPS")
MAX_TRAIN_STEPS = int(_max_steps_env) if _max_steps_env else None

TOTAL_STEPS = max(1, (len(train_ds) // effective_batch) * EPOCHS)
end_step = TOTAL_STEPS
if MAX_TRAIN_STEPS is not None:
    end_step = min(end_step, start_step + MAX_TRAIN_STEPS)

log(f"📈 Total SFT steps for {EPOCHS} epochs (full plan): {TOTAL_STEPS}")
if MAX_TRAIN_STEPS is not None:
    log(f"   ⚠️ MAX_TRAIN_STEPS cap active: {MAX_TRAIN_STEPS}")
log(f"   Resumed from step: {start_step} | This run will go to: {end_step}")
log(f"   Training: step {start_step} → {end_step}")

# ============================================================================
# 16. Evaluation Function
# ============================================================================
@torch.no_grad()
def evaluate():
    # Call the wrapped module directly: only rank 0 evaluates, so DDP's
    # collective hooks must not be entered by a single process.
    eval_model = model.module if hasattr(model, "module") else model
    eval_model.eval()
    losses = []
    for _ in range(20):
        x, y = get_sft_batch(val_ds, MICRO_BATCH, val_rng)
        with torch.autocast(device_type="cuda", dtype=AMP_DTYPE):
            loss = compute_sft_loss(eval_model, x, y, pad_id=PAD_ID)
        if torch.isfinite(loss):
            losses.append(loss.item())
        del x, y, loss
    eval_model.train()
    torch.cuda.empty_cache()
    if not losses:
        return float("inf")
    return sum(losses) / len(losses)

# ============================================================================
# 17. Training Loop
# ============================================================================
log(f"\n{'='*60}")
log(f"🚀 Starting SFT Training...")
log(f"{'='*60}\n")

model.train()
val_loss = None
# best_val_loss is already loaded above
ema_tok_s = None
global_t0 = time.time()
t0 = time.time()

for step in range(start_step, end_step):
    # Keep the original cosine schedule when resuming; restarting warm-up on
    # every resume causes an unintended learning-rate jump.
    lr = get_lr(step, WARMUP_STEPS, TOTAL_STEPS, MAX_LR, MIN_LR)
    for g in optimizer.param_groups:
        g["lr"] = lr

    optimizer.zero_grad(set_to_none=True)
    accum_loss = 0.0

    for micro_i in range(GRAD_ACCUM):
        x, y = get_sft_batch(train_ds, MICRO_BATCH, rng)
        sync_ctx = model.no_sync() if micro_i < GRAD_ACCUM - 1 else contextlib.nullcontext()
        with sync_ctx:
            with torch.autocast(device_type="cuda", dtype=AMP_DTYPE):
                loss = compute_sft_loss(model, x, y, pad_id=PAD_ID)
                loss = loss / GRAD_ACCUM
            loss.backward()
        accum_loss += loss.item()
        del x, y, loss

    # NaN/Inf protection (synced across all ranks)
    is_nan = torch.tensor([0 if math.isfinite(accum_loss) else 1], device=DEVICE)
    dist.all_reduce(is_nan, op=dist.ReduceOp.MAX)
    if is_nan.item() > 0:
        if IS_MAIN:
            log(f"⚠️ NaN/Inf loss at step {step}! Skipping optimizer step.")
        optimizer.zero_grad(set_to_none=True)
        continue

    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
    
    if not torch.isfinite(grad_norm):
        if IS_MAIN:
            log(f"⚠️ Non-finite grad norm at step {step}. Skipping optimizer step.")
        optimizer.zero_grad(set_to_none=True)
        continue

    optimizer.step()

    # ---- Logging (every 10 steps) ----
    if IS_MAIN and step % 10 == 0:
        dt = time.time() - t0
        tok_s = (MICRO_BATCH * GRAD_ACCUM * CONTEXT_LEN * world_size * 10) / dt if dt > 0 else 0
        ema_tok_s = tok_s if ema_tok_s is None else 0.7 * ema_tok_s + 0.3 * tok_s
        t0 = time.time()
        mem_alloc = torch.cuda.memory_allocated(DEVICE) / 1e9
        mem_reserved = torch.cuda.memory_reserved(DEVICE) / 1e9
        
        steps_done = step - start_step
        if steps_done > 0:
            elapsed = time.time() - global_t0
            eta_s = elapsed / steps_done * (end_step - step)
            eta_str = f"{eta_s/60:.0f}min" if eta_s < 3600 else f"{eta_s/3600:.1f}hr"
        else:
            eta_str = "..."
        log(f"SFT step {step:5d}/{end_step} | loss {accum_loss:.4f} | lr {lr:.2e} | "
            f"{ema_tok_s:.0f} tok/s | ETA {eta_str} | "
            f"mem {mem_alloc:.1f}GB/{mem_reserved:.1f}GB")

    # ---- Memory cleanup (important for T4!) ----
    if step > start_step and step % EMPTY_CACHE_EVERY == 0:
        torch.cuda.empty_cache()

    # ---- Evaluation ----
    if step > start_step and step % EVAL_EVERY == 0:
        dist.barrier()
        val_loss = evaluate() if IS_MAIN else None
        if IS_MAIN and val_loss is not None:
            log(f"  >> val_loss {val_loss:.4f}")
            if best_val_loss is None or val_loss < best_val_loss:
                best_val_loss = val_loss
                # User requested to skip best checkpoint to save time
                # save_best_checkpoint(model, step, val_loss)
        dist.barrier()

    # ---- Save latest checkpoint (background upload) ----
    if step > start_step and step % CKPT_EVERY == 0:
        if IS_MAIN:
            wait_for_uploads()
            save_checkpoint(model, optimizer, step,
                            val_loss if val_loss is not None else -1,
                            "sft_ckpt_latest.pt", blocking=False, best_val_loss=best_val_loss)
        dist.barrier()

# ============================================================================
# 18. Final Save & Cleanup
# ============================================================================
dist.barrier()
final_val = evaluate() if IS_MAIN else None

# Final save: BLOCKING (synchronous) — must complete before program exits
if IS_MAIN:
    wait_for_uploads()  # Added to prevent race condition on final save
    final_val_for_save = final_val if final_val is not None and math.isfinite(final_val) else -1
    save_checkpoint(model, optimizer, end_step,
                    final_val_for_save,
                    "sft_ckpt_latest.pt", blocking=True, best_val_loss=best_val_loss)

if IS_MAIN and os.path.exists(f"{WORK}/{SFT_CHECKPOINT_DIR}/sft_ckpt_best.pt"):
    log(f"  📌 Skipping best checkpoint upload to HF as requested. File saved locally at: {WORK}/{SFT_CHECKPOINT_DIR}/sft_ckpt_best.pt")

# Wait for any remaining background uploads from mid-training saves
if IS_MAIN:
    wait_for_uploads()

if IS_MAIN:
    log(f"\n{'='*60}")
    log(f"✨ 500M SFT Training Complete!")
    log(f"   Final val_loss: {final_val if final_val is not None else -1:.4f}")
    log(f"   Best val_loss:  {best_val_loss if best_val_loss is not None else -1:.4f}")
    log(f"   Checkpoints saved to: {MAIN_REPO}/{SFT_CHECKPOINT_DIR}/")
    log(f"{'='*60}")

dist.barrier()
dist.destroy_process_group()