train_spider.py

#!/usr/bin/env python3
"""Spider-FLEXITOKENS training pipeline.

Byte-level pretraining on FineWeb-Edu with boundary predictor curriculum.
Architecture: RDT (2 prelude + 6 recurrent + 2 coda) with:
- SharedProjectionMoE (32 experts, top-2, shared_inter=6144, rank=256)
- MLA (Multi-Latent Attention) with compressed KV cache + sliding window
- Engram conditional memory at recurrent layers 1 and 4
- BoundaryPredictor + downsample/upsample for FlexiToken integration
- LTI Injection + ACT Halting + LoRA Adapter
- 256k context (YaRN factor=8.0), sliding_window=8192
- 272-token byte-level vocab (256 bytes + 16 specials)

Usage:
  Single GPU:
    python train_spider.py
  Multi-GPU:
    torchrun --nproc_per_node=$(python -c "import torch; print(torch.cuda.device_count())") train_spider.py
  Resume from checkpoint:
    python train_spider.py --resume checkpoints/spider-step5000.pt
  Quick smoke test:
    python train_spider.py --max_steps 50 --mock_data
"""

import os
import math
import re
import sys
import time
import argparse
from contextlib import nullcontext
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn import CrossEntropyLoss
from torch.utils.data import IterableDataset, DataLoader, get_worker_info

from datasets import load_dataset

try:
    import bitsandbytes as bnb
    AdamW8bit = bnb.optim.AdamW8bit
    Adam8bit = bnb.optim.Adam8bit
    _HAS_8BIT = True
except ImportError:
    _HAS_8BIT = False
    AdamW8bit = None
    Adam8bit = None

from spider import (
    SpiderConfig,
    SpiderForConditionalGeneration,
    SENTINEL_TOKENS,
)

try:
    from loguru import logger
    logger.remove()
    logger.add(sys.stderr, format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}")
    logger.add("train_spider.log", rotation="100 MB", retention="10 days",
               format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}")
except ImportError:
    import logging
    logging.basicConfig(level=logging.INFO)
    class _LoguruShim:
        def info(self, msg): logging.info(msg)
        def success(self, msg): logging.info(msg)
        def warning(self, msg): logging.warning(msg)
        def error(self, msg): logging.error(msg)
    logger = _LoguruShim()

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


# ============================================================================
# Byte-Level Dataset
# ============================================================================

BOS_ID = SENTINEL_TOKENS['BOS']  # 257
EOS_ID = SENTINEL_TOKENS['EOS']  # 258
PAD_ID = SENTINEL_TOKENS['PAD']  # 256


class ByteLevelDataset(IterableDataset):
    """Streaming byte-level dataset from FineWeb-Edu.

    Per D-23: FineWeb-Edu (English first), per-sample UTF-8 byte encoding.
    Per D-24: Curriculum ordering (English -> multilingual -> code -> math).
    Per D-34: Streaming only, no local download.

    Each sample is encoded as raw UTF-8 bytes with BOS/EOS sentinel tokens.
    Vocab: 272 tokens (256 bytes + 16 specials). Max 8192 bytes per sample.
    """

    def __init__(
        self,
        dataset_name: str = "HuggingFaceFW/fineweb-edu",
        subset: str = "sample-10BT",
        split: str = "train",
        seq_len: int = 8192,
        max_bytes: int = 8192,
        rank: int = 0,
        world_size: int = 1,
    ):
        self.seq_len = seq_len
        self.max_bytes = max_bytes
        self.dataset_name = dataset_name
        self.subset = subset
        self.split = split
        self.rank = rank
        self.world_size = world_size

    def _encode_sample(self, text: str) -> List[int]:
        """Encode text as UTF-8 bytes with BOS/EOS, truncated to max_bytes."""
        byte_ids = list(text.encode('utf-8'))[:self.max_bytes]
        return [BOS_ID] + byte_ids + [EOS_ID]

    def _pad_or_truncate(self, ids: List[int]) -> List[int]:
        """Pad or truncate to seq_len, mask padding with -100 for labels."""
        ids = ids[:self.seq_len]
        ids = ids + [PAD_ID] * (self.seq_len - len(ids))
        return ids

    def __iter__(self):
        worker = get_worker_info()
        num_workers = worker.num_workers if worker else 1
        worker_id = worker.id if worker else 0
        total_shards = self.world_size * num_workers
        shard_index = self.rank * num_workers + worker_id

        ds = load_dataset(
            self.dataset_name,
            name=self.subset,
            split=self.split,
            streaming=True,
        ).shard(num_shards=total_shards, index=shard_index)

        buf = []
        for sample in ds:
            text = sample.get("text", "")
            if not text:
                continue
            byte_ids = self._encode_sample(text)
            buf.extend(byte_ids)
            while len(buf) >= self.seq_len + 1:
                chunk = buf[:self.seq_len + 1]
                buf = buf[self.seq_len + 1:]
                x = torch.tensor(chunk[:-1], dtype=torch.long)
                y = torch.tensor(chunk[1:], dtype=torch.long)
                y[y == PAD_ID] = -100
                yield x, y


class MockByteLevelDataset(IterableDataset):
    """In-memory byte-level dataset for testing (no network required).

    Uses a fixed set of text samples in multiple languages to verify
    byte-level encoding, BOS/EOS placement, and multilingual handling.
    """

    SAMPLES = [
        "Hello world, this is a test of the byte-level encoding system.",
        "The quick brown fox jumps over the lazy dog.",
        "Spider is a recurrent latent reasoning architecture with engram memory.",
        "Boundary predictors learn to merge byte sequences into meaningful tokens.",
        "FineWeb-Edu contains high-quality educational content for pretraining.",
        "Это текст на русском языке для проверки многозычной поддержки.",
        "తెలుగు భాష యొక్క పరీక్ష కోసం నమూనా వచనం.",
        "中文文本用于测试多语言字节编码支持。",
        "def fibonacci(n): return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)",
        "The integral of x^2 from 0 to 1 equals 1/3.",
    ]

    def __init__(self, seq_len: int = 512, max_bytes: int = 512, num_samples: int = 1000):
        self.seq_len = seq_len
        self.max_bytes = max_bytes
        self.num_samples = num_samples

    def __iter__(self):
        buf = []
        count = 0
        while count < self.num_samples:
            for text in self.SAMPLES:
                byte_ids = list(text.encode('utf-8'))[:self.max_bytes]
                ids = [BOS_ID] + byte_ids + [EOS_ID]
                buf.extend(ids)
                while len(buf) >= self.seq_len + 1:
                    chunk = buf[:self.seq_len + 1]
                    buf = buf[self.seq_len + 1:]
                    x = torch.tensor(chunk[:-1], dtype=torch.long)
                    y = torch.tensor(chunk[1:], dtype=torch.long)
                    y[y == PAD_ID] = -100
                    yield x, y
                    count += 1
                    if count >= self.num_samples:
                        return


# ============================================================================
# Curriculum Scheduler
# ============================================================================

class CurriculumScheduler:
    """Training curriculum scheduler per D-24 and D-25.

    Manages dataset switching across training phases and boundary predictor
    curriculum mode (fixed top-k vs adaptive threshold).

    Phases:
        0-30%:  English (FineWeb-Edu), fixed top-k BP (D-25)
        30-50%: English + multilingual, adaptive BP
        50-70%: English + multilingual + code, adaptive BP
        70-90%: English + multilingual + code + math, adaptive BP
        90-100%: Mixed + multimodal, adaptive BP
    """

    def __init__(
        self,
        total_steps: int,
        fixed_compression_k: float = 3.3,
        adaptive_threshold: float = 0.5,
    ):
        self.total_steps = total_steps
        self.fixed_compression_k = fixed_compression_k
        self.adaptive_threshold = adaptive_threshold
        self.curriculum_switch_step = int(0.3 * total_steps)

    def get_phase(self, step: int) -> int:
        if step < int(0.3 * self.total_steps):
            return 1
        elif step < int(0.5 * self.total_steps):
            return 2
        elif step < int(0.7 * self.total_steps):
            return 3
        elif step < int(0.9 * self.total_steps):
            return 4
        else:
            return 5

    def is_fixed_bp(self, step: int) -> bool:
        """Return True if BP should use fixed top-k boundaries (D-25)."""
        return step < self.curriculum_switch_step

    def get_fixed_k(self, seq_len: int) -> int:
        """Number of boundary positions for fixed top-k (3.3x compression)."""
        return max(1, int(seq_len / self.fixed_compression_k))

    def get_boundaries(
        self,
        soft_boundaries: torch.Tensor,
        step: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute hard boundaries based on curriculum phase.

        During fixed phase (first 30% of steps): top-k boundaries with
        straight-through estimator. During adaptive phase: threshold-based.

        Args:
            soft_boundaries: [B, L] boundary probabilities from BoundaryPredictor
            step: Current training step

        Returns:
            Tuple of (soft_boundaries, hard_boundaries), each [B, L]
        """
        if self.is_fixed_bp(step):
            k = self.get_fixed_k(soft_boundaries.shape[-1])
            topk_vals, topk_idx = soft_boundaries.topk(k, dim=-1)
            hard_boundaries = torch.zeros_like(soft_boundaries)
            hard_boundaries.scatter_(-1, topk_idx, 1.0)
            hard_boundaries = (
                hard_boundaries - soft_boundaries.detach() + soft_boundaries
            )
        else:
            hard_boundaries = (soft_boundaries > self.adaptive_threshold).float()
            hard_boundaries = (
                hard_boundaries - soft_boundaries.detach() + soft_boundaries
            )

        return soft_boundaries, hard_boundaries


# ============================================================================
# BP Loss (D-26)
# ============================================================================

def compute_bp_loss(
    soft_boundaries: torch.Tensor,
    hard_boundaries: torch.Tensor,
    seq_len: int,
    binomial_weight: float = 0.1,
    pred_prior: float = 0.303,
) -> torch.Tensor:
    """Compute boundary predictor loss per D-26: BCE + binomial prior.

    During fixed phase: BCE on boundary decisions vs uniform target.
    During adaptive phase: binomial prior loss only.

    Args:
        soft_boundaries: [B, L] boundary probabilities
        hard_boundaries: [B, L] binary boundary decisions
        seq_len: Sequence length
        binomial_weight: Weight for binomial prior term (0.1 per D-26)
        pred_prior: Expected fraction of boundary positions (1/3.3 ≈ 0.303)

    Returns:
        Scalar BP loss tensor
    """
    B = soft_boundaries.shape[0]

    # BCE loss: encourage boundary probability to match expected compression
    target_rate = 1.0 / 3.3
    target = torch.full_like(soft_boundaries, target_rate)
    bce_loss = F.binary_cross_entropy(soft_boundaries, target)

    # Binomial prior: regularize number of predicted boundaries
    sum_preds = hard_boundaries.sum(dim=-1)  # [B]
    binomial = torch.distributions.binomial.Binomial(
        total_count=float(seq_len),
        probs=pred_prior,
    )
    log_prob = binomial.log_prob(sum_preds)
    binomial_loss = -log_prob.mean() / seq_len

    return bce_loss + binomial_weight * binomial_loss


# ============================================================================
# Recurrent Monitor (drift/collapse detection)
# ============================================================================

class RecurrentMonitor:
    """Monitors recurrent dynamics across loops during training.

    Catches representation drift, expert collapse, and engram instability
    before they corrupt training. Per CONTEXT: representation drift across
    loops is the #1 failure mode for recurrent architectures.

    Logged metrics (every log_interval steps):
    - loop_norms: L2 norm of hidden states after each loop (drift detection)
    - routing_entropy: entropy of expert routing weights per loop (collapse detection)
    - engram_norms: L2 norm of engram residuals at layers 1 and 4 (memory stability)
    - halt_distribution: fraction of tokens halting at each loop (ACT health)
    - loop_grad_norms: gradient norms per recurrent layer (gradient health)
    """

    def __init__(
        self,
        drift_threshold: float = 10.0,
        collapse_threshold: float = 1.0,
    ):
        self.drift_threshold = drift_threshold
        self.collapse_threshold = collapse_threshold

    def compute_routing_entropy(self, router_logits: torch.Tensor) -> float:
        """Compute routing entropy from router logits.

        Args:
            router_logits: [B, L, num_experts] raw router logits

        Returns:
            Scalar entropy value (higher = more diverse routing)
        """
        p = F.softmax(router_logits, dim=-1).mean(dim=(0, 1))  # [num_experts]
        entropy = -(p * (p + 1e-10).log()).sum().item()
        return entropy

    def check_health(self, metrics: Dict, step: int) -> List[str]:
        """Check for drift, collapse, or instability.

        Args:
            metrics: Dict with keys: loop_norms, routing_entropy, engram_norms, halt_distribution
            step: Current training step

        Returns:
            List of warning strings (empty if healthy)
        """
        warnings = []

        # Drift detection: hidden norm ratio between first and last loop
        norms = metrics.get('loop_norms', [])
        if len(norms) >= 2 and norms[0] > 0:
            drift_ratio = norms[-1] / norms[0]
            if drift_ratio > self.drift_threshold:
                warnings.append(
                    f"DRIFT WARNING step {step}: loop norm ratio {drift_ratio:.1f}x "
                    f"(loop_1={norms[0]:.2f}, loop_{len(norms)}={norms[-1]:.2f})"
                )

        # Collapse detection: low routing entropy
        entropies = metrics.get('routing_entropy', [])
        if entropies and min(entropies) < self.collapse_threshold:
            warnings.append(
                f"COLLAPSE WARNING step {step}: routing entropy {min(entropies):.2f} "
                f"< threshold {self.collapse_threshold}"
            )

        return warnings


# ============================================================================
# BP Curriculum Trainer
# ============================================================================

class BPCurriculumTrainer:
    """Training wrapper for Spider-FLEXITOKENS with BP curriculum.

    Manages:
    - BP freeze/unfreeze during warmup (D-27)
    - Fixed -> adaptive boundary curriculum (D-25)
    - Dual loss: LM CE + MoE aux + BP (BCE + binomial prior) (D-26)
    - Per-loop gradient clipping for expert cores
    - RecurrentMonitor integration for drift/collapse detection
    """

    def __init__(
        self,
        model: SpiderForConditionalGeneration,
        optimizer: torch.optim.Optimizer,
        engram_optimizer: Optional[torch.optim.Optimizer],
        curriculum: CurriculumScheduler,
        monitor: RecurrentMonitor,
        warmup_steps: int,
        base_lr: float,
        bp_loss_weight: float = 0.1,
        grad_clip: float = 1.0,
        expert_core_grad_clip: float = 0.5,
    ):
        self.model = model
        self.optimizer = optimizer
        self.engram_optimizer = engram_optimizer
        self.curriculum = curriculum
        self.monitor = monitor
        self.warmup_steps = warmup_steps
        self.base_lr = base_lr
        self.bp_loss_weight = bp_loss_weight
        self.grad_clip = grad_clip
        self.expert_core_grad_clip = expert_core_grad_clip
        self._bp_frozen = False
        self.bp_optimizer = None

    def freeze_bp(self):
        """Freeze boundary predictor params during warmup (D-27)."""
        for name, param in self.model.named_parameters():
            if 'boundary_predictor' in name:
                param.requires_grad = False
        self._bp_frozen = True

    def unfreeze_bp(self):
        """Unfreeze BP at 0.1x base LR after warmup (D-27)."""
        bp_param_names = set()
        bp_params = []
        for name, param in self.model.named_parameters():
            if 'boundary_predictor' in name:
                param.requires_grad = True
                bp_params.append(param)
                bp_param_names.add(name)
        self._bp_frozen = False

        # Create separate optimizer for BP params with 0.1x base LR (D-27)
        bp_lr = self.base_lr * 0.1
        self.bp_optimizer = torch.optim.Adam(
            bp_params, lr=bp_lr, betas=(0.9, 0.95), eps=1e-8
        )

    def train_step(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor,
        step: int,
        n_loops: int = 6,
        amp_ctx: Optional[nullcontext] = None,
        sdpa_ctx: Optional[nullcontext] = None,
    ) -> Tuple[torch.Tensor, Dict]:
        """Single training step with dual loss and monitoring.

        Args:
            input_ids: [B, L] byte-level token IDs
            labels: [B, L] target token IDs (with -100 for padding)
            step: Current training step
            n_loops: Number of recurrent loops
            amp_ctx: Optional autocast context
            sdpa_ctx: Optional SDPA kernel context

        Returns:
            Tuple of (total_loss, metrics_dict)
        """
        amp_ctx = amp_ctx or nullcontext()
        sdpa_ctx = sdpa_ctx or nullcontext()

        # BP freeze/unfreeze logic (D-27)
        if step == 0 and self.warmup_steps > 0:
            self.freeze_bp()
        if self._bp_frozen and step >= self.warmup_steps:
            self.unfreeze_bp()

        with amp_ctx, sdpa_ctx:
            # Override boundaries based on curriculum phase
            output = self.model(input_ids, labels=labels, n_loops=n_loops)

            lm_loss = output['loss']
            aux_loss = output['aux_loss']
            soft_boundaries = output['soft_boundaries']
            hard_boundaries = output['hard_boundaries']

            # Apply curriculum override for hard_boundaries
            soft_boundaries, hard_boundaries = self.curriculum.get_boundaries(
                soft_boundaries, step
            )

            # BP dual loss (D-26)
            seq_len = input_ids.shape[-1]
            if not self._bp_frozen:
                bp_loss = compute_bp_loss(soft_boundaries, hard_boundaries, seq_len)
            else:
                bp_loss = torch.tensor(0.0, device=input_ids.device)

            # Total loss: LM + MoE aux + BP
            if isinstance(aux_loss, torch.Tensor):
                total_loss = lm_loss + self.model.config.router_aux_loss_coef * aux_loss
            else:
                total_loss = lm_loss + self.model.config.router_aux_loss_coef * aux_loss
            total_loss = total_loss + self.bp_loss_weight * bp_loss

        # Collect monitoring metrics
        metrics = {
            'lm_loss': lm_loss.item() if isinstance(lm_loss, torch.Tensor) else lm_loss,
            'aux_loss': aux_loss.item() if isinstance(aux_loss, torch.Tensor) else aux_loss,
            'bp_loss': bp_loss.item() if isinstance(bp_loss, torch.Tensor) else bp_loss,
            'bp_frozen': self._bp_frozen,
            'curriculum_phase': self.curriculum.get_phase(step),
            'is_fixed_bp': self.curriculum.is_fixed_bp(step),
        }

        return total_loss, metrics

    def clip_gradients(self) -> float:
        """Clip gradients: global + per-loop expert core clipping.

        Standard: clip_grad_norm_(all params, max_norm=1.0)
        Expert cores: tighter clip at 0.5 to prevent drift.
        """
        # Global gradient clipping
        grad_norm = nn.utils.clip_grad_norm_(
            self.model.parameters(), max_norm=self.grad_clip
        )

        # Per-loop expert core clipping (tighter)
        expert_core_params = []
        for name, param in self.model.named_parameters():
            if ('W_gate' in name or 'W_transform' in name) and param.grad is not None:
                expert_core_params.append(param)

        if expert_core_params:
            nn.utils.clip_grad_norm_(
                expert_core_params, max_norm=self.expert_core_grad_clip
            )

        return grad_norm.item() if isinstance(grad_norm, torch.Tensor) else float(grad_norm)


# ============================================================================
# LR Schedule
# ============================================================================

def get_lr(step: int, warmup: int, total: int, max_lr: float, min_lr: float) -> float:
    """Cosine learning rate with linear warmup."""
    if step < warmup:
        return max_lr * step / warmup
    if step >= total:
        return min_lr
    decay = (step - warmup) / (total - warmup)
    return min_lr + 0.5 * (max_lr - min_lr) * (1.0 + math.cos(math.pi * decay))


# ============================================================================
# Checkpointing
# ============================================================================

def save_step_checkpoint(model, optimizer, step, epoch, cfg, ckpt_dir, master, ddp=False, trainer=None, current_best_loss=float("inf")):
    """Save full checkpoint (model + optimizer) and keep only the last 2."""
    if ddp:
        from torch.distributed.fsdp import (
            FullyShardedDataParallel as FSDP,
            StateDictType,
            FullStateDictConfig,
        )
        with FSDP.state_dict_type(
            model,
            StateDictType.FULL_STATE_DICT,
            FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
        ):
            model_state = model.state_dict()
            optim_state = FSDP.optim_state_dict(model, optimizer)
    else:
        model_state = model.state_dict()
        optim_state = optimizer.state_dict()

    if not master:
        return None, 0

    os.makedirs(ckpt_dir, exist_ok=True)
    ckpt_path = os.path.join(ckpt_dir, f"spider-step{step}.pt")
    tmp_path = ckpt_path + ".tmp"
    torch.save(
        {
            "step": step,
            "epoch": epoch,
            "model_state_dict": model_state,
            "optimizer_state_dict": optim_state,
            "cfg": cfg,
            "bp_optimizer_state_dict": (
                trainer.bp_optimizer.state_dict() if trainer and trainer.bp_optimizer else None
            ),
            "best_loss": current_best_loss,
        },
        tmp_path,
    )
    os.replace(tmp_path, ckpt_path)
    size_mb = os.path.getsize(ckpt_path) / (1024 * 1024)

    # Keep only the last 2 step checkpoints
    step_pattern = re.compile(r"spider-step\d+\.pt$")
    step_ckpts = sorted(
        [os.path.join(ckpt_dir, f) for f in os.listdir(ckpt_dir) if step_pattern.search(f)],
        key=os.path.getmtime,
    )
    while len(step_ckpts) > 2:
        old = step_ckpts.pop(0)
        os.remove(old)

    return ckpt_path, size_mb


def save_full_checkpoint(model, optimizer, step, epoch, cfg, ckpt_dir, master, ddp=False, ckpt_name="full", trainer=None, current_best_loss=float("inf")):
    """Save full checkpoint with custom name."""
    if ddp:
        from torch.distributed.fsdp import (
            FullyShardedDataParallel as FSDP,
            StateDictType,
            FullStateDictConfig,
        )
        with FSDP.state_dict_type(
            model,
            StateDictType.FULL_STATE_DICT,
            FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
        ):
            model_state = model.state_dict()
            optim_state = FSDP.optim_state_dict(model, optimizer)
    else:
        model_state = model.state_dict()
        optim_state = optimizer.state_dict()

    if not master:
        return None, 0

    os.makedirs(ckpt_dir, exist_ok=True)
    final_path = os.path.join(ckpt_dir, f"spider-{ckpt_name}.pt")
    tmp_path = final_path + ".tmp"
    torch.save(
        {
            "step": step,
            "epoch": epoch,
            "model_state_dict": model_state,
            "optimizer_state_dict": optim_state,
            "cfg": cfg,
            "bp_optimizer_state_dict": (
                trainer.bp_optimizer.state_dict() if trainer and trainer.bp_optimizer else None
            ),
            "best_loss": current_best_loss,
        },
        tmp_path,
    )
    os.replace(tmp_path, final_path)
    size_mb = os.path.getsize(final_path) / (1024 * 1024)
    return final_path, size_mb


def load_checkpoint(model, optimizer, path, ddp=False):
    """Load model + optimizer state from checkpoint.

    Handles cross-optimizer resume (e.g. 8bit Adam on local → standard AdamW
    on remote): if optimizer state dict keys mismatch, we skip the optimizer
    state and log a warning. The model weights always load successfully.

    Returns: (step, epoch, bp_optim_state, saved_best_loss)
    """
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    model.load_state_dict(ckpt["model_state_dict"])
    try:
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    except (ValueError, KeyError, RuntimeError) as e:
        logger.warning(
            f"Optimizer state mismatch (likely 8bit→standard cross-resume): {e}. "
            f"Skipping optimizer state — optimizer will reinitialize."
        )
    bp_optim_state = ckpt.get("bp_optimizer_state_dict", None)
    saved_best_loss = ckpt.get("best_loss", float("inf"))
    return int(ckpt["step"]), int(ckpt.get("epoch", 0)), bp_optim_state, saved_best_loss


# ============================================================================
# DeepSpeed Config (fallback for RTX 4060 8GB)
# ============================================================================

DEEPSPEED_ZERO3_CONFIG = {
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True,
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True,
        },
        "overlap_comm": True,
        "contiguous_gradients": True,
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": 1.0,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}


# ============================================================================
# Precision Mode (MXFP8 / NVFP4 / FP8_DYNAMIC / BF16)
# ============================================================================

import enum

class PrecisionMode(enum.Enum):
    BF16 = "bf16"
    FP8_DYNAMIC = "fp8_dynamic"
    MXFP8 = "mxfp8"
    NVFP4 = "nvfp4"


def detect_precision_mode() -> PrecisionMode:
    """Auto-detect best available precision mode based on GPU + libraries.

    Fallback chain: MXFP8/NVFP4 → FP8_DYNAMIC → BF16

    - MXFP8: Requires Blackwell+ (sm120+), torchao with float8 training,
      block-wise scaling (128x128). Best accuracy among FP8 options.
    - NVFP4: Requires Blackwell+ (sm120+), fbgemm-gpu-genai with NVFP4
      kernels. Most aggressive compression (4-bit weights).
    - FP8_DYNAMIC: Requires Ada Lovelace+ (sm89+), torchao float8.
      Row-wise dynamic scaling. Good speed/accuracy tradeoff.
    - BF16: Fallback for all GPUs. Standard mixed precision.
    """
    if not torch.cuda.is_available():
        return PrecisionMode.BF16

    cc = torch.cuda.get_device_capability()
    major, minor = cc

    # Check for torchao float8 training support
    _has_torchao_fp8 = False
    try:
        from torchao.float8 import convert_to_float8_training
        _has_torchao_fp8 = True
    except ImportError:
        pass

    # Check for fbgemm NVFP4 support
    _has_nvfp4 = False
    try:
        from torchao.quantization import NVFP4Config  # type: ignore[attr-defined]
        _has_nvfp4 = True
    except (ImportError, AttributeError):
        try:
            import fbgemm_gpu.genai  # type: ignore[import-untyped]
            _has_nvfp4 = True
        except (ImportError, ModuleNotFoundError):
            pass

    # Blackwell+ (sm120+): MXFP8 or NVFP4
    if major >= 12:
        if _has_torchao_fp8:
            return PrecisionMode.MXFP8
        if _has_nvfp4:
            return PrecisionMode.NVFP4

    # Ada Lovelace+ (sm89+): FP8 dynamic
    if (major, minor) >= (8, 9) and _has_torchao_fp8:
        return PrecisionMode.FP8_DYNAMIC

    return PrecisionMode.BF16


def configure_fp8_training(model, mode: PrecisionMode):
    """Apply torchao float8 training conversion to model.

    FP8 training swaps nn.Linear layers with Float8Linear, which performs
    dynamic quantization of activations and weights to float8_e4m3fn during
    forward/backward, with high-precision accumulation.

    Two recipes:
    - MXFP8 (rowwise_with_gw_hp): Row-wise scaling + high-precision grad weight.
      Best accuracy. Requires sm120+ hardware.
    - FP8_DYNAMIC (rowwise): Row-wise dynamic scaling. Good tradeoff.
      Requires sm89+ hardware.

    Gradient computation stays in bf16/fp32 for stability.
    """
    from torchao.float8 import convert_to_float8_training, Float8LinearConfig

    if mode == PrecisionMode.MXFP8:
        recipe_name = "rowwise_with_gw_hp"
    elif mode == PrecisionMode.FP8_DYNAMIC:
        recipe_name = "rowwise"
    else:
        return model

    base = Float8LinearConfig.from_recipe_name(recipe_name)
    config = Float8LinearConfig(
        cast_config_input=base.cast_config_input,
        cast_config_weight=base.cast_config_weight,
        cast_config_grad_output=base.cast_config_grad_output,
        cast_config_input_for_grad_weight=base.cast_config_input_for_grad_weight,
        cast_config_weight_for_grad_input=base.cast_config_weight_for_grad_input,
        cast_config_grad_output_for_grad_weight=base.cast_config_grad_output_for_grad_weight,
        gemm_config_output=base.gemm_config_output,
        gemm_config_grad_input=base.gemm_config_grad_input,
        gemm_config_grad_weight=base.gemm_config_grad_weight,
        enable_fsdp_float8_all_gather=base.enable_fsdp_float8_all_gather,
        round_scales_to_power_of_2=base.round_scales_to_power_of_2,
        pad_inner_dim=True,
    )

    def module_filter_fn(mod, fqn):
        skip = any(s in fqn for s in (
            "boundary_predictor",
            "loop_embedding",
            "engram",
            "layernorm",
            "norm",
            "embed_tokens",
            "lm_head",
            "halt_predictor",
            "gate",
        ))
        return not skip

    model = convert_to_float8_training(
        model,
        module_filter_fn=module_filter_fn,
        config=config,
    )
    return model


def configure_nvfp4_training(model):
    """Apply NVFP4 weight-only quantization for training on Blackwell.

    NVFP4 uses 4-bit floating-point weights with 8-bit scaling factors.
    Activations stay in bf16/fp8. Requires fbgemm-gpu-genai kernels.

    Falls back to FP8_DYNAMIC if NVFP4 kernels unavailable.
    """
    try:
        from torchao.quantization import NVFP4Config, quantize_
        quantize_(model, NVFP4Config())
        return model
    except (ImportError, AttributeError, RuntimeError):
        logger.warning("NVFP4 not available, falling back to FP8_DYNAMIC")
        return configure_fp8_training(model, PrecisionMode.FP8_DYNAMIC)


def try_unsloth():
    """Attempt to apply Unsloth patches. Returns (available, FastLanguageModel)."""
    try:
        from unsloth import FastLanguageModel
        return True, FastLanguageModel
    except (ImportError, Exception):
        return False, None


# ============================================================================
# Main Training Loop
# ============================================================================

def parse_args():
    parser = argparse.ArgumentParser(description="Spider-FLEXITOKENS training")
    parser.add_argument("--resume", type=str, default="", help="Path to checkpoint to resume from")
    parser.add_argument("--max_steps", type=int, default=0, help="Override max training steps")
    parser.add_argument("--mock_data", action="store_true", help="Use mock data (no network)")
    parser.add_argument("--seq_len", type=int, default=0, help="Override sequence length")
    parser.add_argument("--micro_batch", type=int, default=0, help="Override micro batch size")
    parser.add_argument("--n_loops", type=int, default=0, help="Override number of loops")
    parser.add_argument("--lr", type=float, default=0, help="Override learning rate")
    parser.add_argument("--ckpt_dir", type=str, default="checkpoints-spider", help="Checkpoint directory")
    parser.add_argument("--no_unsloth", action="store_true", help="Skip Unsloth even if available")
    parser.add_argument(
        "--precision", type=str, default="auto",
        choices=["auto", "bf16", "fp8_dynamic", "mxfp8", "nvfp4"],
        help="Training precision: auto (detect), bf16, fp8_dynamic, mxfp8, nvfp4",
    )
    return parser.parse_args()


def main():
    global best_loss
    args = parse_args()

    # ------------------------------------------------------------------
    # Distributed init
    # ------------------------------------------------------------------
    ddp = int(os.environ.get("RANK", -1)) != -1
    if ddp:
        dist.init_process_group("nccl")
        rank = int(os.environ["RANK"])
        local_rank = int(os.environ["LOCAL_RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        device = f"cuda:{local_rank}"
        torch.cuda.set_device(device)
    else:
        rank = local_rank = 0
        world_size = 1
        device = "cuda" if torch.cuda.is_available() else "cpu"
    master = rank == 0

    # ------------------------------------------------------------------
    # Hyperparameters
    # ------------------------------------------------------------------
    seq_len = args.seq_len or int(os.environ.get("SEQ_LEN", "2048"))
    micro_batch = args.micro_batch or int(os.environ.get("MICRO_BATCH", "4"))
    target_tokens = int(os.environ.get("TARGET_TOKENS", "10_000_000_000"))
    grad_accum = int(os.environ.get("GRAD_ACCUM", "1"))
    n_loops = args.n_loops or int(os.environ.get("N_LOOPS", "6"))
    lr = args.lr or float(os.environ.get("LR", "3e-4"))
    wd = 0.1
    warmup_steps = 200
    log_every = 10
    ckpt_every = int(os.environ.get("CKPT_EVERY", "500"))
    ckpt_dir = args.ckpt_dir

    global_batch_tok = world_size * micro_batch * grad_accum * seq_len
    total_steps = target_tokens // global_batch_tok
    if args.max_steps > 0:
        total_steps = min(total_steps, args.max_steps)

    if master:
        logger.info(
            f"[Spider-FLEXITOKENS] hidden=2048 | 6 recurrent | 32 experts top-2 | "
            f"n_loops={n_loops} | seq_len={seq_len} | micro_batch={micro_batch} | "
            f"grad_accum={grad_accum} | global_batch_tokens={global_batch_tok:,} | "
            f"total_steps={total_steps:,}"
        )
        logger.info(
            f"Byte-level vocab: 272 | Context: 256k (YaRN-8) | "
            f"Sliding window: 8192 | BP curriculum: fixed 30% -> adaptive | "
            f"Gradient checkpointing: enabled | Precision: {prec_mode.value}"
        )

    # ------------------------------------------------------------------
    # Model + Precision Mode
    # ------------------------------------------------------------------
    cfg = SpiderConfig()
    bf16_ok = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    amp_dtype = torch.bfloat16 if bf16_ok else torch.float16

    # Resolve precision mode: CLI override or auto-detect
    if args.precision == "auto":
        prec_mode = detect_precision_mode()
    else:
        prec_mode = PrecisionMode(args.precision)

    if master:
        logger.info(f"Precision mode: {prec_mode.value}")

    model = SpiderForConditionalGeneration(cfg).to(amp_dtype)
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()

    # Apply FP8/MXFP8/NVFP4 quantization (before Unsloth, before FSDP)
    if prec_mode in (PrecisionMode.MXFP8, PrecisionMode.FP8_DYNAMIC):
        try:
            model = configure_fp8_training(model, prec_mode)
            if master:
                logger.info(f"torchao FP8 training enabled: {prec_mode.value}")
        except Exception as e:
            if master:
                logger.warning(f"FP8 training setup failed ({e}), falling back to BF16")
            prec_mode = PrecisionMode.BF16
    elif prec_mode == PrecisionMode.NVFP4:
        try:
            model = configure_nvfp4_training(model)
            if master:
                logger.info("NVFP4 training enabled")
        except Exception as e:
            if master:
                logger.warning(f"NVFP4 setup failed ({e}), falling back to FP8_DYNAMIC")
            try:
                model = configure_fp8_training(model, PrecisionMode.FP8_DYNAMIC)
                prec_mode = PrecisionMode.FP8_DYNAMIC
                if master:
                    logger.info("Fallback: FP8_DYNAMIC training enabled")
            except Exception as e2:
                if master:
                    logger.warning(f"FP8 fallback also failed ({e2}), using BF16")
                prec_mode = PrecisionMode.BF16

    # Unsloth (optional, per D-35): applies MoE kernel optimizations,
    # gradient checkpointing, and memory-efficient attention
    use_unsloth = False
    if not args.no_unsloth and not ddp:
        use_unsloth_available, FastLanguageModel_cls = try_unsloth()
        if use_unsloth_available:
            try:
                # Unsloth patches: SDPA optimization, memory-efficient GC
                # For MoE: set UNSLOTH_MOE_BACKEND=grouped_mm (default)
                os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
                use_unsloth = True
                if master:
                    logger.info("Unsloth MoE + training patches applied")
            except Exception as e:
                if master:
                    logger.warning(f"Unsloth patching failed: {e}")
    if not use_unsloth and master:
        logger.info("Unsloth not available, using standard PyTorch training")

    if ddp:
        from torch.distributed.fsdp import (
            FullyShardedDataParallel as FSDP,
            ShardingStrategy,
            MixedPrecision,
        )
        from torch.distributed.fsdp.wrap import ModuleWrapPolicy
        from spider import SpiderDenseLayer, SpiderRecurrentLayer

        mp_policy = MixedPrecision(
            param_dtype=amp_dtype,
            reduce_dtype=amp_dtype,
            buffer_dtype=amp_dtype,
        )
        wrap_policy = ModuleWrapPolicy({SpiderDenseLayer, SpiderRecurrentLayer})
        model = FSDP(
            model,
            sharding_strategy=ShardingStrategy.FULL_SHARD,
            mixed_precision=mp_policy,
            auto_wrap_policy=wrap_policy,
            device_id=local_rank,
        )
    else:
        model = model.to(device)

    if master:
        n_params = sum(p.numel() for p in model.parameters())
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logger.info(
            f"Parameters: {n_params:,} total | {trainable:,} trainable | "
            f"Precision: {prec_mode.value} | AMP dtype: {amp_dtype}"
        )

    # ------------------------------------------------------------------
    # Optimizer — 8-bit Adam for BF16 on small GPUs; standard AdamW for FP8+
    # When FP8/MXFP8/NVFP4 is active, weight memory is already halved,
    # so 8-bit Adam is less critical and can conflict with Float8Linear.
    # Dual optimizer for Engram embeddings (per mythos pattern).
    # ------------------------------------------------------------------
    engram_params_list = [
        p for n, p in model.named_parameters()
        if 'engram' in n and 'embed' in n and 'proj' not in n
    ]
    backbone_params = [
        p for n, p in model.named_parameters()
        if not ('engram' in n and 'embed' in n and 'proj' not in n)
    ]

    use_8bit_optimizer = _HAS_8BIT and prec_mode == PrecisionMode.BF16

    if use_8bit_optimizer:
        optimizer = AdamW8bit(
            backbone_params, lr=lr, weight_decay=wd,
            betas=(0.9, 0.95), eps=1e-8,
        )
        if engram_params_list:
            engram_optimizer = Adam8bit(
                engram_params_list, lr=lr * 5,
                betas=(0.9, 0.95), eps=1e-8,
            )
        else:
            engram_optimizer = None
        if master:
            logger.info("Optimizer: 8-bit AdamW (bf16 mode, saves ~50% optimizer VRAM)")
    else:
        optimizer = torch.optim.AdamW(
            backbone_params, lr=lr, weight_decay=wd,
            betas=(0.9, 0.95), foreach=True, eps=1e-8,
        )
        if engram_params_list:
            engram_optimizer = torch.optim.Adam(
                engram_params_list, lr=lr * 5,
                betas=(0.9, 0.95), eps=1e-8,
            )
        else:
            engram_optimizer = None
        if master:
            logger.info(f"Optimizer: standard AdamW ({prec_mode.value} mode)")

    # ------------------------------------------------------------------
    # Curriculum + Monitor + Trainer
    # ------------------------------------------------------------------
    curriculum = CurriculumScheduler(total_steps=total_steps)
    monitor = RecurrentMonitor()
    trainer = BPCurriculumTrainer(
        model=model,
        optimizer=optimizer,
        engram_optimizer=engram_optimizer,
        curriculum=curriculum,
        monitor=monitor,
        warmup_steps=warmup_steps,
        base_lr=lr,
    )

    # ------------------------------------------------------------------
    # Resume from checkpoint
    # ------------------------------------------------------------------
    start_step = 0
    start_epoch = 1
    bp_optim_state_to_load = None
    if args.resume and os.path.exists(args.resume):
        if master:
            logger.info(f"Resuming from checkpoint: {args.resume}")
        start_step, start_epoch, bp_optim_state_to_load, saved_best = load_checkpoint(
            model, optimizer, args.resume, ddp
        )
        best_loss = saved_best
        if master:
            logger.info(f"Resumed at step {start_step}, epoch {start_epoch}, best_loss={best_loss:.4f}")
    else:
        # Auto-resume from latest checkpoint in ckpt_dir
        existing_ckpts = sorted(
            [os.path.join(ckpt_dir, f) for f in os.listdir(ckpt_dir)
             if f.startswith("spider-") and f.endswith(".pt") and not f.endswith(".tmp")]
        ) if os.path.isdir(ckpt_dir) else []
        if existing_ckpts:
            latest = existing_ckpts[-1]
            if master:
                logger.info(f"Auto-resuming from: {latest}")
            start_step, start_epoch, bp_optim_state_to_load, saved_best = load_checkpoint(
            model, optimizer, latest, ddp
        )
        best_loss = saved_best
        if master:
                logger.info(f"Resumed at step {start_step}, epoch {start_epoch}, best_loss={best_loss:.4f}")

    # Restore BP optimizer state if available (after trainer is created,
    # BP optimizer is initialized during first unfreeze_bp() call)
    if bp_optim_state_to_load and trainer.bp_optimizer:
        try:
            trainer.bp_optimizer.load_state_dict(bp_optim_state_to_load)
            if master:
                logger.info("Restored BP optimizer state from checkpoint")
        except (ValueError, KeyError, RuntimeError) as e:
            if master:
                logger.warning(f"BP optimizer state mismatch, skipping: {e}")

    # ------------------------------------------------------------------
    # Dataset + DataLoader
    # ------------------------------------------------------------------
    if args.mock_data:
        dataset = MockByteLevelDataset(seq_len=seq_len)
    else:
        dataset = ByteLevelDataset(
            seq_len=seq_len,
            rank=rank,
            world_size=world_size,
        )

    loader = DataLoader(
        dataset,
        batch_size=micro_batch,
        num_workers=4 if not args.mock_data else 0,
        pin_memory=True,
        prefetch_factor=1 if not args.mock_data else None,
    )

    # ------------------------------------------------------------------
    # AMP + SDPA contexts
    # ------------------------------------------------------------------
    amp_ctx = (
        torch.amp.autocast(device_type="cuda", dtype=amp_dtype)
        if "cuda" in device
        else nullcontext()
    )
    amp_ctx = nullcontext() if ddp else amp_ctx

    try:
        from torch.nn.attention import sdpa_kernel
        sdpa_ctx = sdpa_kernel(enable_flash=True, enable_mem_efficient=True, enable_math=True)
    except Exception:
        sdpa_ctx = nullcontext()

    # ------------------------------------------------------------------
    # Training loop
    # ------------------------------------------------------------------
    if master:
        os.makedirs(ckpt_dir, exist_ok=True)

    model.train()
    data_iter = iter(loader)
    t0 = time.perf_counter()
    step = start_step
    epoch = start_epoch
    tokens_in_epoch = 0
    tokens_per_epoch = target_tokens

    while step < total_steps:
        cur_lr = get_lr(step, warmup_steps, total_steps, lr, lr * 0.1)
        for g in optimizer.param_groups:
            g["lr"] = cur_lr
        if engram_optimizer:
            for g in engram_optimizer.param_groups:
                g["lr"] = cur_lr * 5

        optimizer.zero_grad()
        if engram_optimizer:
            engram_optimizer.zero_grad()
        if trainer.bp_optimizer:
            trainer.bp_optimizer.zero_grad()
        loss_accum = 0.0
        metrics_accum = {}

        for micro_step in range(grad_accum):
            try:
                x, y = next(data_iter)
            except StopIteration:
                data_iter = iter(loader)
                x, y = next(data_iter)

            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            sync = (
                nullcontext()
                if (not ddp or micro_step == grad_accum - 1)
                else model.no_sync()
            )
            with sync:
                total_loss, metrics = trainer.train_step(
                    x, y, step, n_loops=n_loops,
                    amp_ctx=amp_ctx, sdpa_ctx=sdpa_ctx,
                )
                total_loss = total_loss / grad_accum
                total_loss.backward()

            if master and step == start_step and micro_step == 0:
                peak_vram = torch.cuda.max_memory_allocated() / 1024**3
                logger.info(f"First forward+backward | Peak VRAM: {peak_vram:.1f}GB")

            loss_accum += total_loss.item()
            for k, v in metrics.items():
                if k not in metrics_accum:
                    metrics_accum[k] = 0.0
                if isinstance(v, (int, float)):
                    metrics_accum[k] += v / grad_accum

        # Gradient clipping
        grad_norm = trainer.clip_gradients()
        optimizer.step()
        if engram_optimizer:
            engram_optimizer.step()
        if trainer.bp_optimizer:
            for g in trainer.bp_optimizer.param_groups:
                g["lr"] = cur_lr * 0.1
            trainer.bp_optimizer.step()
        step += 1
        tokens_in_epoch += global_batch_tok

        # Health checks
        if master and step % log_every == 0:
            health_warnings = monitor.check_health(metrics_accum, step)
            for w in health_warnings:
                logger.warning(w)

        # Logging
        if master and step % log_every == 0:
            dt = time.perf_counter() - t0
            tok_per_sec = global_batch_tok * log_every / dt
            tokens_seen = step * global_batch_tok
            bp_status = "FIXED" if metrics_accum.get('is_fixed_bp', True) else "ADAPTIVE"
            bp_frozen = "FROZEN" if metrics_accum.get('bp_frozen', False) else "ACTIVE"
            logger.info(
                f"Epoch {epoch} | step {step:6d}/{total_steps} | "
                f"loss {loss_accum:.4f} | lm {metrics_accum.get('lm_loss', 0):.4f} | "
                f"aux {metrics_accum.get('aux_loss', 0):.4f} | "
                f"bp {metrics_accum.get('bp_loss', 0):.4f} [{bp_status}/{bp_frozen}] | "
                f"gnorm {float(grad_norm):.2f} | lr {cur_lr:.2e} | "
                f"{tok_per_sec / 1e6:.2f}M tok/s | {tokens_seen / 1e9:.2f}B tokens"
            )
            t0 = time.perf_counter()

        # Checkpointing
        if step % ckpt_every == 0:
            ckpt_path, size_mb = save_step_checkpoint(
                model, optimizer, step, epoch, cfg, ckpt_dir, master, ddp, trainer,
                current_best_loss=best_loss,
            )
            if master and ckpt_path:
                logger.info(f"Saved step checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")

        # Epoch boundary
        if tokens_in_epoch >= tokens_per_epoch:
            epoch_loss = loss_accum
            if master:
                logger.info(f"Epoch {epoch} complete | loss={epoch_loss:.4f}")
            ckpt_path, size_mb = save_full_checkpoint(
                model, optimizer, step, epoch, cfg, ckpt_dir, master, ddp, f"ep{epoch}", trainer,
                current_best_loss=best_loss,
            )
            if master and ckpt_path:
                logger.info(f"Saved epoch checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")

            if epoch_loss < best_loss:
                best_loss = epoch_loss
                ckpt_path, size_mb = save_full_checkpoint(
                    model, optimizer, step, epoch, cfg, ckpt_dir, master, ddp, "best", trainer,
                    current_best_loss=best_loss,
                )
                if master and ckpt_path:
                    logger.info(f"Saved best checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")

            epoch += 1
            tokens_in_epoch = 0

    # Final checkpoint
    if step > start_step and master:
        ckpt_path, size_mb = save_full_checkpoint(
            model, optimizer, step, epoch, cfg, ckpt_dir, master, ddp, f"final-ep{epoch}", trainer,
            current_best_loss=best_loss,
        )
        if ckpt_path:
            logger.info(f"Saved final checkpoint: {os.path.basename(ckpt_path)} ({size_mb:.0f}MB)")

    if ddp:
        dist.barrier()
        dist.destroy_process_group()

    if master:
        logger.info("Training complete.")


if __name__ == "__main__":
    best_loss = float("inf")
    main()