code/train_sft.py

# ============================================================================
# ViuAI 500M — Supervised Fine-Tuning (SFT) Script (v7 Cognitive & RAG)
# ============================================================================
# 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 sys
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

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

# Redirect HF cache to working directory
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
import math
import time
import random
import contextlib
import gc
import threading
from datetime import timedelta
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 transformers import PreTrainedTokenizerFast

# Global list to track all background upload threads
_upload_threads = []

# Special tokens defined for ViuAI v7
SPECIAL_TOKENS = [
    "<|user|>", "<|assistant|>", "<|endofturn|>",
    "[THINK]", "[/THINK]",
    "<search>", "</search>", "<search_result>", "</search_result>"
]
SFT_VOCAB_SIZE = 64009

# ============================================================================
# 1. DDP Setup (works for multi-GPU & single GPU via torchrun)
# ============================================================================
def setup_ddp():
    if not torch.cuda.is_available():
        raise RuntimeError("SFT training requires at least one CUDA GPU.")
    
    # Standalone execution fallback (if not launched via torchrun)
    if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
        os.environ["RANK"] = "0"
        os.environ["LOCAL_RANK"] = "0"
        os.environ["WORLD_SIZE"] = "1"
        os.environ["MASTER_ADDR"] = "127.0.0.1"
        os.environ["MASTER_PORT"] = str(random.randint(29100, 29900))
    
    if not dist.is_initialized():
        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)
    device = torch.device(f"cuda:{local_rank}")
    return rank, local_rank, world_size, device

# ============================================================================
# 2. HF Token & Configuration Setup
# ============================================================================
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    # Check fallback environment variables or prompt notice
    _alt_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if _alt_token:
        HF_TOKEN = _alt_token
    else:
        print("⚠️ Warning: HF_TOKEN environment variable is not set. Private repo uploads will require authentication.", flush=True)

MAIN_REPO = "ViuAI/ViuAI-500M"
SFT_DATA_REPO = "ViuAI/viuai-500m-sft-tokenized"
api = HfApi(token=HF_TOKEN) if HF_TOKEN else HfApi()

# ============================================================================
# 3. Helper: Network Retry & DDP Download Logic
# ============================================================================
def with_retry(func, *args, max_retries=5, initial_backoff=5, **kwargs):
    backoff = initial_backoff
    for attempt in range(max_retries):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            if attempt == max_retries - 1:
                raise e
            time.sleep(backoff)
            backoff *= 2

def download_on_main(func, *args, fatal=True, is_main=True, device=None, **kwargs):
    err = None
    ok = 1
    if is_main:
        try:
            with_retry(func, *args, **kwargs)
        except Exception as e:
            err = e
            ok = 0
            if fatal:
                print(f"❌ Failed download on main rank: {e}", flush=True)
    if dist.is_initialized():
        ok_tensor = torch.tensor([ok], device=device or "cuda")
        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

# ============================================================================
# 4. Learning Rate Schedule (Cosine with Linear Warmup)
# ============================================================================
def get_lr(step, warmup_steps, total_steps, max_lr, min_lr):
    if step < warmup_steps:
        return max_lr * (step + 1) / max(1, warmup_steps)
    if step >= total_steps:
        return min_lr
    decay_ratio = (step - warmup_steps) / max(1, (total_steps - warmup_steps))
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (max_lr - min_lr)

# Memory allocator optimization & TF32 acceleration
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
if torch.cuda.is_available():
    torch.set_float32_matmul_precision('high')

# ============================================================================
# 5. SFT Loss Computation (Chunked Memory-Efficient CrossEntropy)
# ============================================================================
def compute_sft_loss(model, input_ids, labels, pad_id=None):
    out = model(input_ids, pad_id=pad_id)
    logits = out[0] if isinstance(out, tuple) else out
    shift_logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
    shift_labels = labels[:, 1:].reshape(-1)
    
    # Process cross-entropy in chunks to cap peak VRAM allocation to <1.5GB
    chunk_size = 16384
    total_tokens = shift_logits.size(0)
    if total_tokens > chunk_size:
        total_loss = torch.tensor(0.0, device=shift_logits.device, dtype=shift_logits.dtype)
        total_valid = 0
        for i in range(0, total_tokens, chunk_size):
            end = min(i + chunk_size, total_tokens)
            l_chunk = shift_logits[i:end]
            y_chunk = shift_labels[i:end]
            valid_mask = (y_chunk != -100)
            n_valid = int(valid_mask.sum().item())
            if n_valid > 0:
                loss_chunk = F.cross_entropy(l_chunk, y_chunk, ignore_index=-100, reduction='sum')
                total_loss = total_loss + loss_chunk
                total_valid += n_valid
        if total_valid > 0:
            return total_loss / total_valid
        else:
            return F.cross_entropy(shift_logits[:1], shift_labels[:1], ignore_index=-100)
    else:
        return F.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction='mean')

# ============================================================================
# 6. Dataset Loader & Shuffled Epoch Sampler (.npy memmap)
# ============================================================================
class NpySFTDataset:
    def __init__(self, ids_path, labels_path, offsets_path, context_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.context_len = context_len
        self.n_examples = len(self.offsets) - 1

    def __len__(self):
        return self.n_examples

    def get(self, idx):
        start = int(self.offsets[idx])
        end = int(self.offsets[idx + 1])
        
        ex_ids = np.zeros(self.context_len, dtype=np.int64)
        ex_labels = np.full(self.context_len, -100, dtype=np.int64)

        length = min(end - start, self.context_len)
        ex_ids[:length] = self.ids[start:start + length]
        ex_labels[:length] = self.labels[start:start + length]

        return torch.from_numpy(ex_ids), torch.from_numpy(ex_labels)

class ShuffledEpochSampler:
    """
    True epoch-based shuffling without replacement across DDP ranks.
    Ensures every training example is seen exactly once per epoch, precisely honoring EPOCHS.
    Seamlessly advances across epochs and supports resuming from any step without duplicating data.
    """
    def __init__(self, dataset_len, rank=0, world_size=1, seed=1337, initial_step=0, effective_batch=1):
        self.dataset_len = dataset_len
        self.rank = rank
        self.world_size = world_size
        self.base_seed = seed
        self.effective_batch = effective_batch
        self.epoch = 0
        self.indices = []
        self.ptr = 0
        self._new_epoch()
        
        # Exact mathematical stride fast-forward if resuming from a specific step
        if initial_step > 0:
            calls_per_rank = initial_step * (effective_batch // world_size)
            rank_samples_per_epoch = max(1, dataset_len // world_size)
            epochs_to_skip = calls_per_rank // rank_samples_per_epoch
            
            for _ in range(epochs_to_skip):
                self._new_epoch()
            
            calls_remaining = calls_per_rank % rank_samples_per_epoch
            self.ptr = self.rank + (calls_remaining * self.world_size)

    def _new_epoch(self):
        rng = random.Random(self.base_seed + self.epoch * 1000)
        perm = list(range(self.dataset_len))
        rng.shuffle(perm)
        self.indices = perm
        self.ptr = self.rank
        self.epoch += 1

    def next_index(self):
        if self.ptr >= len(self.indices):
            self._new_epoch()
        idx = self.indices[self.ptr]
        self.ptr += self.world_size
        return idx

def get_sft_batch(dataset, batch_size, sampler_or_rng, device="cuda"):
    batch_ids, batch_labels = [], []
    attempts = 0
    max_attempts = batch_size * 25
    
    while len(batch_ids) < batch_size and attempts < max_attempts:
        attempts += 1
        if hasattr(sampler_or_rng, "next_index"):
            idx = sampler_or_rng.next_index()
        elif hasattr(sampler_or_rng, "randint"):
            idx = sampler_or_rng.randint(0, len(dataset) - 1)
        else:
            idx = random.randint(0, len(dataset) - 1)
            
        ids, labels = dataset.get(idx)
        if not (labels != -100).any():
            continue
        batch_ids.append(ids)
        batch_labels.append(labels)
        
    if len(batch_ids) == 0:
        ids, labels = dataset.get(0)
        batch_ids = [ids] * batch_size
        batch_labels = [labels] * batch_size
        
    return (
        torch.stack(batch_ids).to(device, non_blocking=True),
        torch.stack(batch_labels).to(device, non_blocking=True)
    )

# ============================================================================
# 7. Checkpoint Remapping & Serialization
# ============================================================================
def remap_checkpoint_keys(ckpt_state_dict, model_state_dict):
    clean_ckpt = {}
    for k, v in ckpt_state_dict.items():
        k = k.replace("_orig_mod.", "").replace("module.", "")
        clean_ckpt[k] = v

    remapped = {}
    for model_key in model_state_dict.keys():
        clean_key = model_key.replace("_orig_mod.", "").replace("module.", "")
        if model_key in clean_ckpt:
            remapped[model_key] = clean_ckpt[model_key]
        elif clean_key in clean_ckpt:
            remapped[model_key] = clean_ckpt[clean_key]
        else:
            print(f"⚠️ Unmapped key: {model_key}", flush=True)
    return remapped

def save_sft_checkpoint(model, optimizer, step, val_loss, best_val_loss, fname, scaler=None, is_main=True, work=".", blocking=False):
    if not is_main:
        return
    local_path = os.path.join(work, fname)
    raw_model = model.module if hasattr(model, "module") else model
    checkpoint = {
        "step": step,
        "model_state_dict": raw_model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scaler_state_dict": scaler.state_dict() if scaler is not None else None,
        "val_loss": val_loss,
        "best_val_loss": best_val_loss,
        "timestamp": time.time(),
    }
    torch.save(checkpoint, local_path)
    print(f"💾 Saved checkpoint to {local_path} (step {step}, val_loss: {val_loss:.4f})", flush=True)

    def _async_upload():
        if not HF_TOKEN:
            return
        try:
            with_retry(
                api.upload_file,
                path_or_fileobj=local_path,
                path_in_repo=f"sft_checkpoints/sft_v7/{fname}",
                repo_id=MAIN_REPO,
                repo_type="model"
            )
            print(f"☁️ Uploaded {fname} to Hugging Face successfully.", flush=True)
        except Exception as e:
            print(f"⚠️ Failed to upload checkpoint {fname}: {e}", flush=True)

    if blocking:
        _async_upload()
    else:
        thread = threading.Thread(target=_async_upload, daemon=True)
        thread.start()
        _upload_threads.append(thread)

def load_sft_checkpoint(model, optimizer, scaler=None, resume_sft=False, resume_checkpoint_file="sft_ckpt_latest.pt", is_main=True, work=".", device="cuda"):
    """
    Deadlock-free checkpoint loader across all DDP ranks:
    1. All ranks participate collectively in download_on_main() (Rank 0 downloads, all ranks wait at barrier).
    2. Every rank loads weights locally from shared disk, remaps prefixes (_orig_mod/module), and resizes vocab.
    """
    if resume_sft:
        if is_main:
            print(f"🔄 Resuming SFT training from {resume_checkpoint_file}...", flush=True)
        try:
            download_on_main(
                hf_hub_download, MAIN_REPO, f"sft_checkpoints/sft_v7/{resume_checkpoint_file}",
                local_dir=work, is_main=is_main, device=device, token=HF_TOKEN
            )
            ckpt_path = os.path.join(work, "sft_checkpoints", "sft_v7", resume_checkpoint_file)
            ckpt = torch.load(ckpt_path, map_location=device)
            raw_model = model.module if hasattr(model, "module") else model
            
            # Robust key remapping & strict=False for cross-GPU resume (e.g. torch.compile <-> standard)
            remapped_sd = remap_checkpoint_keys(ckpt["model_state_dict"], raw_model.state_dict())
            
            # Resize embedding and head if needed
            model_sd = raw_model.state_dict()
            for k in ["tok_emb.weight", "head.weight"]:
                if k in remapped_sd and k in model_sd and remapped_sd[k].shape != model_sd[k].shape:
                    old_w = remapped_sd[k]
                    new_w = model_sd[k].clone()
                    min_v = min(old_w.shape[0], new_w.shape[0])
                    new_w[:min_v] = old_w[:min_v]
                    remapped_sd[k] = new_w
                    if is_main:
                        print(f"  Resized parameter '{k}' from {old_w.shape} to {new_w.shape}", flush=True)

            raw_model.load_state_dict(remapped_sd, strict=False)
            
            if "optimizer_state_dict" in ckpt:
                try:
                    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
                except Exception as e_opt:
                    if is_main:
                        print(f"⚠️ Optimizer state load warning (optimizer re-initialized): {e_opt}", flush=True)

            if scaler is not None and "scaler_state_dict" in ckpt and ckpt["scaler_state_dict"] is not None:
                try:
                    scaler.load_state_dict(ckpt["scaler_state_dict"])
                except Exception as e_sc:
                    if is_main:
                        print(f"⚠️ Scaler state load notice: {e_sc}", flush=True)

            step = ckpt.get("step", 0)
            val_loss = ckpt.get("val_loss", None)
            best_val_loss = ckpt.get("best_val_loss", None)
            if is_main:
                print(f"✅ SFT checkpoint loaded! Resuming from step {step} (val_loss: {val_loss})", flush=True)
            return step, val_loss, best_val_loss
        except Exception as e:
            if is_main:
                print(f"⚠️ Failed to load SFT resume checkpoint: {e}", flush=True)
                print("  Falling back to loading base pre-trained weights...", flush=True)

    if is_main:
        print("📥 Loading base pre-trained checkpoint (ckpt_latest.pt)...", flush=True)
    try:
        download_on_main(
            hf_hub_download, MAIN_REPO, "checkpoints/ckpt_latest.pt",
            local_dir=work, is_main=is_main, device=device, token=HF_TOKEN
        )
        ckpt_path = os.path.join(work, "checkpoints", "ckpt_latest.pt")
        ckpt = torch.load(ckpt_path, map_location=device)

        if "model_state_dict" in ckpt:
            state_dict = ckpt["model_state_dict"]
        elif "model" in ckpt:
            state_dict = ckpt["model"]
        else:
            state_dict = ckpt
        raw_model = model.module if hasattr(model, "module") else model
        remapped_sd = remap_checkpoint_keys(state_dict, raw_model.state_dict())

        # Resize embedding and head if vocab grew (e.g. 64000 -> 64009)
        model_sd = raw_model.state_dict()
        for k in ["tok_emb.weight", "head.weight", "_orig_mod.tok_emb.weight", "_orig_mod.head.weight"]:
            if k in remapped_sd and k in model_sd and remapped_sd[k].shape != model_sd[k].shape:
                old_w = remapped_sd[k]
                new_w = model_sd[k].clone()
                min_v = min(old_w.shape[0], new_w.shape[0])
                new_w[:min_v] = old_w[:min_v]
                remapped_sd[k] = new_w
                if is_main:
                    print(f"  Resized parameter '{k}' from {old_w.shape} to {new_w.shape}", flush=True)

        raw_model.load_state_dict(remapped_sd, strict=False)
        if is_main:
            print(f"✅ Base pre-trained weights ({len(remapped_sd)} tensors) successfully loaded into SFT model!", flush=True)
        del ckpt
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return 0, -1, None
    except Exception as e:
        if is_main:
            print(f"⚠️ Could not load base weights: {e}", flush=True)
            print("  Starting SFT with randomly initialized weights.", flush=True)
        return 0, -1, None

# ============================================================================
# 8. Main Execution Function
# ============================================================================
def main():
    rank, local_rank, world_size, DEVICE = setup_ddp()
    IS_MAIN = (rank == 0)

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

    # Initialize HF repo on Rank 0 if token is present
    if IS_MAIN and HF_TOKEN:
        try:
            create_repo(repo_id=MAIN_REPO, repo_type="model", exist_ok=True, token=HF_TOKEN)
        except Exception as e:
            print(f"⚠️ Warning during repo creation: {e}", flush=True)
    dist.barrier()

    _resume_env = os.environ.get("RESUME_SFT")
    RESUME_SFT = _resume_env.lower() in ("1", "true", "yes") if _resume_env is not None else False
    RESUME_CHECKPOINT_FILE = os.environ.get("RESUME_CHECKPOINT_FILE", "sft_ckpt_latest.pt")
    SFT_CHECKPOINT_DIR = os.environ.get("SFT_CHECKPOINT_DIR", "sft_checkpoints/sft_v7")

    WORK = "/kaggle/working" if os.path.exists("/kaggle/working") else "."

    log("📥 Downloading tokenizer & model architecture code...")
    download_on_main(hf_hub_download, MAIN_REPO, "tokenizer/tokenizer.json", local_dir=WORK, is_main=IS_MAIN, device=DEVICE, token=HF_TOKEN)
    download_on_main(hf_hub_download, MAIN_REPO, "code/config.py", local_dir=WORK, is_main=IS_MAIN, device=DEVICE, token=HF_TOKEN)
    download_on_main(hf_hub_download, MAIN_REPO, "code/model.py", local_dir=WORK, is_main=IS_MAIN, device=DEVICE, token=HF_TOKEN)

    # Add WORK and code dirs to sys.path
    if WORK not in sys.path:
        sys.path.insert(0, WORK)

    from config import ViuAIConfig
    from model import ViuAI

    # 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)
    compute_cap = torch.cuda.get_device_capability(local_rank)
    supports_bf16 = compute_cap[0] >= 8

    if mem_gb >= 130:       # H200 (141GB) / H100 (140GB)
        MICRO_BATCH   = 32
        GRAD_ACCUM    = 4
        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 = 1000
    elif mem_gb >= 30:     # RTX 5090 (32GB) / RTX 4090 (24GB) / L4 (24GB)
        MICRO_BATCH   = 16
        GRAD_ACCUM    = 4
        CONTEXT_LEN   = 2048
        AMP_DTYPE     = torch.bfloat16 if supports_bf16 else torch.float16
        USE_COMPILE   = True
        EMPTY_CACHE_EVERY = 500
    elif mem_gb >= 20:     # RTX 3090 (24GB) / RTX 4090 (24GB) / L4 (24GB) / A10G (24GB)
        MICRO_BATCH   = 8
        GRAD_ACCUM    = 8
        CONTEXT_LEN   = 2048
        AMP_DTYPE     = torch.bfloat16 if supports_bf16 else torch.float16
        USE_COMPILE   = True
        EMPTY_CACHE_EVERY = 200
    else:                  # Kaggle T4x2 (16GB each) / T4 single
        MICRO_BATCH   = 4
        GRAD_ACCUM    = 16
        CONTEXT_LEN   = 2048
        AMP_DTYPE     = torch.float16
        USE_COMPILE   = False
        EMPTY_CACHE_EVERY = 50

    # Allow custom environment overrides
    if "MICRO_BATCH" in os.environ:
        MICRO_BATCH = int(os.environ["MICRO_BATCH"])
    if "GRAD_ACCUM" in os.environ:
        GRAD_ACCUM = int(os.environ["GRAD_ACCUM"])

    MAX_LR        = 2e-5
    MIN_LR        = 2e-6
    WEIGHT_DECAY  = 0.01
    GRAD_CLIP     = 1.0
    EPOCHS        = 3

    effective_batch = MICRO_BATCH * GRAD_ACCUM * world_size
    tokens_per_step = effective_batch * CONTEXT_LEN

    if IS_MAIN:
        log("=" * 65)
        log(f"🚀 ViuAI-500M SFT v7 Training Initialized")
        log(f"   GPU: {gpu_name} ({mem_gb:.1f} GB) x {world_size}")
        log(f"   Micro batch: {MICRO_BATCH} | Grad accum: {GRAD_ACCUM} | Effective batch: {effective_batch}")
        log(f"   Context length: {CONTEXT_LEN} | Tokens/step: {tokens_per_step:,}")
        log(f"   Precision: {AMP_DTYPE} | Torch compile: {USE_COMPILE}")
        log(f"   Max LR: {MAX_LR:.2e} | Min LR: {MIN_LR:.2e}")
        log(f"   Vocab size: {SFT_VOCAB_SIZE} | Checkpoint dir: {SFT_CHECKPOINT_DIR}")
        log(f"   Resume mode: {RESUME_SFT}")
        log("=" * 65)

    # Build Model & Optimizer
    cfg = ViuAIConfig(
        vocab_size=SFT_VOCAB_SIZE,
        context_length=CONTEXT_LEN,
        use_checkpoint=True,
        z_loss_weight=0.0,
        attn_dropout=0.0,
        resid_dropout=0.0,
    )
    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}")

    if USE_COMPILE:
        log("⚡ Compiling model with torch.compile...")
        raw_model = torch.compile(raw_model)

    model = DDP(raw_model, device_ids=[local_rank])
    
    # Bulletproof parameter grouping: handles weight tying without duplicate param error
    decay_params = []
    no_decay_params = []
    seen_param_ids = set()

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if id(param) in seen_param_ids:
            continue
        seen_param_ids.add(id(param))

        if any(k in name for k in ["norm", "tok_emb", "head"]):
            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)
    )

    scaler = torch.amp.GradScaler('cuda', enabled=(AMP_DTYPE == torch.float16))

    # 1. Collective Checkpoint Loading on all Ranks
    start_step, prev_val_loss, best_val_loss_loaded = load_sft_checkpoint(
        model, optimizer, scaler=scaler, resume_sft=RESUME_SFT, resume_checkpoint_file=RESUME_CHECKPOINT_FILE,
        is_main=IS_MAIN, work=WORK, device=DEVICE
    )

    # 2. Defensive Model & Buffer Broadcast to ensure 100% parameter equality
    for param in raw_model.parameters():
        dist.broadcast(param.data, src=0)
    for buf in raw_model.buffers():
        dist.broadcast(buf.data, src=0)

    # Broadcast training scalar states across all ranks
    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

    # Download & Load SFT Data
    log("📥 Downloading SFT data (.npy files)...")
    _DATA_VER = os.environ.get("SFT_DATA_VERSION", "v7")
    _PFX = f"sft_{_DATA_VER}_" if _DATA_VER != "v1_legacy" else "sft_"
    if _DATA_VER == "v1":
        _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}'")

    for fname in SFT_DATA_FILES:
        download_on_main(hf_hub_download, SFT_DATA_REPO, fname, repo_type="dataset", local_dir=WORK, is_main=IS_MAIN, device=DEVICE, 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)}")

    # Validate PAD_ID & Special Tokens against downloaded tokenizer
    PAD_ID = 0
    tok_json_path = os.path.join(WORK, "tokenizer", "tokenizer.json")
    if not os.path.exists(tok_json_path):
        tok_json_path = os.path.join(WORK, "tokenizer.json")
    if os.path.exists(tok_json_path):
        try:
            _loaded_tok = PreTrainedTokenizerFast(tokenizer_file=tok_json_path, pad_token="<pad>", bos_token="<bos>", eos_token="<eos>", unk_token="<unk>")
            _loaded_tok.add_tokens(SPECIAL_TOKENS, special_tokens=True)
            assert len(_loaded_tok) == SFT_VOCAB_SIZE, f"Tokenizer vocab size mismatch: {len(_loaded_tok)} != {SFT_VOCAB_SIZE}"
            assert _loaded_tok.pad_token_id == 0, f"Expected pad_token_id=0, got {_loaded_tok.pad_token_id}"
            PAD_ID = _loaded_tok.pad_token_id
            if IS_MAIN:
                log(f"✅ Verified tokenizer vocab size={len(_loaded_tok)} and PAD_ID={PAD_ID}")
        except Exception as e:
            if IS_MAIN:
                log(f"ℹ️ Tokenizer check notice: {e}. Using PAD_ID={PAD_ID}")

    if IS_MAIN:
        _s_ids, _s_labels = train_ds.get(0)
        _has_active = (_s_labels != -100).any()
        assert _has_active, "SFT FATAL: All tokens are masked (-100)! Check data pipeline."
        log(f"✅ Data sanity check passed. Using PAD_ID={PAD_ID}")

    # Training Schedule Math & True Shuffled Epoch Sampler
    train_sampler = ShuffledEpochSampler(
        len(train_ds), rank=rank, world_size=world_size,
        seed=1337, initial_step=start_step, effective_batch=effective_batch
    )
    val_rng = random.Random(999 + rank)

    # AUTO MAX STEPS: Exactly 3 full epochs automatically calculated from dataset size & effective batch
    TOTAL_STEPS = max(1, math.ceil(len(train_ds) / effective_batch) * EPOCHS)
    
    _max_steps_env = os.environ.get("MAX_TRAIN_STEPS")
    if _max_steps_env:
        MAX_TRAIN_STEPS = int(_max_steps_env)
        end_step = min(TOTAL_STEPS, start_step + MAX_TRAIN_STEPS)
    else:
        MAX_TRAIN_STEPS = TOTAL_STEPS
        end_step = TOTAL_STEPS

    WARMUP_STEPS = min(100, max(10, int(TOTAL_STEPS * 0.03)))

    log(f"📈 [AUTO-MAX] SFT Total Steps for {EPOCHS} epochs: {TOTAL_STEPS} | Effective Batch: {effective_batch}")
    log(f"   Warmup: {WARMUP_STEPS} steps | Training Schedule: step {start_step} → {end_step} (Single save at step {end_step})")

    def evaluate(val_batches=40):
        model.eval()
        total_val_loss = 0.0
        with torch.no_grad():
            for _ in range(val_batches):
                x, y = get_sft_batch(val_ds, MICRO_BATCH, val_rng, device=DEVICE)
                with torch.autocast(device_type="cuda", dtype=AMP_DTYPE):
                    loss = compute_sft_loss(model, x, y, pad_id=PAD_ID)
                total_val_loss += loss.item()
                del x, y, loss
        
        loss_tensor = torch.tensor([total_val_loss / val_batches], device=DEVICE)
        dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
        avg_loss = loss_tensor.item() / world_size
        model.train()
        return avg_loss

    # SFT Training Loop
    log("\n" + "=" * 65)
    log(f"🔥 SFT v7 TRAINING RUNNING: Step {start_step} → {end_step}")
    log("=" * 65)

    model.train()
    ema_tok_s = None
    VAL_EVERY   = 200
    SAVE_EVERY  = 500

    global_t0 = time.time()
    t0 = time.time()

    for step in range(start_step, end_step):
        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, train_sampler, device=DEVICE)
            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
                scaler.scale(loss).backward()
            accum_loss += loss.item()
            del x, y, loss

        # Synchronized NaN/Inf Loss Check
        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)
            scaler.update()
            continue

        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
        
        # Synchronized Finite Grad Norm Check
        is_finite_grad = torch.tensor([1 if torch.isfinite(grad_norm) else 0], device=DEVICE)
        dist.all_reduce(is_finite_grad, op=dist.ReduceOp.MIN)
        if is_finite_grad.item() == 0 or 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)
            scaler.update()
            continue

        scaler.step(optimizer)
        scaler.update()

        # Logging
        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"VRAM {mem_alloc:.1f}/{mem_reserved:.1f}GB")

        # Periodic Evaluation (Calculated only - No mid-training checkpoint save)
        if step > 0 and step % VAL_EVERY == 0:
            val_loss = evaluate()
            is_best = (best_val_loss is None) or (val_loss < best_val_loss)
            if is_best:
                best_val_loss = val_loss
                star = " ⭐️ BEST!"
            else:
                star = ""
            log(f"\n📊 [Validation @ step {step}] val_loss: {val_loss:.4f}{star}\n")

        # Periodic Memory Defragmentation
        if step > 0 and step % EMPTY_CACHE_EVERY == 0:
            torch.cuda.empty_cache()

    # Final Checkpoints & Training Wrap-up (Single Save at Max Step)
    log("\n" + "=" * 65)
    log("🎉 SFT v7 TRAINING COMPLETE!")
    log("=" * 65)

    final_val_loss = evaluate()
    best_str = f"{best_val_loss:.4f}" if best_val_loss is not None else f"{final_val_loss:.4f}"
    log(f"📊 Final validation loss: {final_val_loss:.4f} (Best achieved: {best_str})")

    log("💾 Saving final SFT model checkpoint at max step...")
    save_sft_checkpoint(model, optimizer, end_step, final_val_loss, best_val_loss, "sft_ckpt_final.pt", scaler=scaler, is_main=IS_MAIN, work=WORK, blocking=True)
    save_sft_checkpoint(model, optimizer, end_step, final_val_loss, best_val_loss, "sft_ckpt_latest.pt", scaler=scaler, is_main=IS_MAIN, work=WORK, blocking=True)

    log("⏳ Waiting for all background upload threads to complete...")
    for t in _upload_threads:
        t.join(timeout=300)
    log("✅ All checkpoints uploaded to Hugging Face successfully!")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()