epfl_final/SparseAtt/train.py

import argparse
import contextlib
import copy
import json
import os
import math
import time

# Persist torch.compile()'d kernels across runs so startup is fast after the first run.
os.environ.setdefault(
    "TORCHINDUCTOR_CACHE_DIR",
    os.path.expanduser("~/.cache/torch_inductor"),
)
import wandb
import numpy as np
from pathlib import Path
from datetime import datetime
from typing import List, Optional, Sequence, Tuple, Union
from dataclasses import dataclass, asdict
from datasets import load_from_disk

import torch
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim.lr_scheduler import LambdaLR
from sparseatt.config import Qwen3Config
from sparseatt.model import Qwen3ForCausalLM, load_tokenizer
import torch._inductor.config as inductor_config

try:
    from torchinfo import summary
except ImportError:
    summary = None

torch.set_float32_matmul_precision('high')


@dataclass
class TrainConfig:
    """Training configuration loaded from JSON."""
    # Model
    model_path: str = "../qwen3_06b_base"
    model_dtype: str = "bfloat16"
    init_from_scratch: bool = False
    
    # Data
    train_path: str = "./data/sample/sample.jsonl"
    val_path: Optional[str] = None
    max_seq_length: int = 512
    num_workers: int = 8
    # Optional data/sampler seed.  When omitted, the sampler follows the
    # global training seed; legacy-compatible configs may set this explicitly.
    data_seed: Optional[int] = None
    
    # Training
    output_dir: str = "./checkpoints"
    resume_from_checkpoint: Optional[str] = None  # Path to checkpoint dir to resume from
    force_new_run: bool = False  # Set True to start fresh even if a previous checkpoint exists
    num_epochs: int = 1
    max_steps: Optional[int] = None
    per_device_batch_size: int = 2
    gradient_accumulation_steps: int = 8
    learning_rate: float = 1e-5
    weight_decay: float = 0.01
    warmup_steps: int = 100
    lr_scheduler: str = "cosine"
    max_grad_norm: float = 1.0
    save_steps: int = 500
    eval_steps: int = 500
    eval_batches: int = 100
    logging_steps: int = 50
    seed: int = 42
    verbose: bool = False
    checkpoint_activations: bool = False
    # Explicit per-rank DataLoader cursor for a deliberate DDP world-size migration.
    resume_micro_step_override: Optional[int] = None
    
    # Distributed
    strategy: str = "ddp"
    mixed_precision: str = "bf16"
    
    # Wandb
    wandb_enabled: bool = True
    wandb_project: str = "qwen3-training"
    wandb_name: Optional[str] = None
    wandb_tags: Union[list, None] = None
    
    # Architecture overrides (applied on top of the base model config.json)
    hidden_size: Optional[int] = None
    intermediate_size: Optional[int] = None
    num_attention_heads: Optional[int] = None
    num_key_value_heads: Optional[int] = None
    num_hidden_layers: Optional[int] = None
    head_dim: Optional[int] = None
    attention_dropout: Optional[float] = None
    rope_theta: Optional[float] = None
    rope_train_position_scale_set: Optional[List[float]] = None
    positional: str = "rope"
    qk_norm: str = "on"
    qk_gain_fp32: bool = False

    # Attention mode
    attention_type: str = "dense"
    sparse_block_size: int = 128
    sparse_num_random_blocks: int = 1
    sparse_mask_pool_size: int = 1000
    sparse_val_dense_compare: bool = False
    sparse_layer_indices: Optional[List[int]] = None
    sparse_mask_cache_dir: Optional[str] = None
    sparse_mask_seed: int = 0
    sparse_mixed_nrb_set: Optional[List[int]] = None
    dense_mode: bool = False
    dense_step_intervals: Optional[List[List[int]]] = None
    
    @classmethod
    def from_json(cls, path: str) -> "TrainConfig":
        with open(path, "r") as f:
            data = json.load(f)
        
        config = cls()
        
        # Flatten nested config
        if "model" in data:
            config.model_path = data["model"].get("path", config.model_path)
            config.model_dtype = data["model"].get("dtype", config.model_dtype)
            config.init_from_scratch = data["model"].get("init_from_scratch", config.init_from_scratch)
            config.hidden_size = data["model"].get("hidden_size", config.hidden_size)
            config.intermediate_size = data["model"].get("intermediate_size", config.intermediate_size)
            config.num_attention_heads = data["model"].get("num_attention_heads", config.num_attention_heads)
            config.num_key_value_heads = data["model"].get("num_key_value_heads", config.num_key_value_heads)
            config.num_hidden_layers = data["model"].get("num_hidden_layers", config.num_hidden_layers)
            config.head_dim = data["model"].get("head_dim", config.head_dim)
            config.attention_dropout = data["model"].get("attention_dropout", config.attention_dropout)
            config.rope_theta = data["model"].get("rope_theta", config.rope_theta)
            config.rope_train_position_scale_set = data["model"].get(
                "rope_train_position_scale_set", config.rope_train_position_scale_set
            )
            config.positional = data["model"].get("positional", config.positional)
            config.qk_norm = data["model"].get("qk_norm", config.qk_norm)
            config.qk_gain_fp32 = data["model"].get("qk_gain_fp32", config.qk_gain_fp32)
        
        if "data" in data:
            config.train_path = data["data"].get("train_path", config.train_path)
            config.val_path = data["data"].get("val_path", config.val_path)
            config.max_seq_length = data["data"].get("max_seq_length", config.max_seq_length)
            config.num_workers = data["data"].get("num_workers", config.num_workers)
            config.data_seed = data["data"].get("seed", config.data_seed)
        
        if "training" in data:
            t = data["training"]
            config.output_dir = t.get("output_dir", config.output_dir)
            config.resume_from_checkpoint = t.get("resume_from_checkpoint", config.resume_from_checkpoint)
            config.force_new_run = t.get("force_new_run", config.force_new_run)
            config.num_epochs = t.get("num_epochs", config.num_epochs)
            config.max_steps = t.get("max_steps", config.max_steps)
            config.per_device_batch_size = t.get("per_device_batch_size", config.per_device_batch_size)
            config.gradient_accumulation_steps = t.get("gradient_accumulation_steps", config.gradient_accumulation_steps)
            config.learning_rate = t.get("learning_rate", config.learning_rate)
            config.weight_decay = t.get("weight_decay", config.weight_decay)
            config.warmup_steps = t.get("warmup_steps", config.warmup_steps)
            config.lr_scheduler = t.get("lr_scheduler", config.lr_scheduler)
            config.max_grad_norm = t.get("max_grad_norm", config.max_grad_norm)
            config.save_steps = t.get("save_steps", config.save_steps)
            config.eval_steps = t.get("eval_steps", config.eval_steps)
            config.eval_batches = t.get("eval_batches", config.eval_batches)
            config.logging_steps = t.get("logging_steps", config.logging_steps)
            config.seed = t.get("seed", config.seed)
            config.verbose = t.get("verbose", config.verbose)
            config.checkpoint_activations = t.get("checkpoint_activations", config.checkpoint_activations)
            config.resume_micro_step_override = t.get("resume_micro_step_override", config.resume_micro_step_override)
        if "distributed" in data:
            config.strategy = data["distributed"].get("strategy", config.strategy)
            config.mixed_precision = data["distributed"].get("mixed_precision", config.mixed_precision)
        
        if "wandb" in data:
            config.wandb_enabled = data["wandb"].get("enabled", config.wandb_enabled)
            config.wandb_project = data["wandb"].get("project", config.wandb_project)
            config.wandb_name = data["wandb"].get("name", config.wandb_name)
            config.wandb_tags = data["wandb"].get("tags", config.wandb_tags)
        
        attention_cfg = data.get("attention", {})
        config.attention_type = attention_cfg.get("type", data.get("attention_type", config.attention_type))
        config.sparse_block_size = attention_cfg.get("sparse_block_size", data.get("sparse_block_size", config.sparse_block_size))
        config.sparse_num_random_blocks = attention_cfg.get(
            "sparse_num_random_blocks",
            data.get("sparse_num_random_blocks", config.sparse_num_random_blocks),
        )
        config.sparse_mask_pool_size = attention_cfg.get(
            "sparse_mask_pool_size",
            data.get("sparse_mask_pool_size", config.sparse_mask_pool_size),
        )
        config.sparse_mask_seed = attention_cfg.get(
            "sparse_mask_seed", data.get("sparse_mask_seed", config.sparse_mask_seed)
        )
        config.sparse_mixed_nrb_set = attention_cfg.get(
            "sparse_mixed_nrb_set",
            data.get("sparse_mixed_nrb_set", config.sparse_mixed_nrb_set),
        )
        config.sparse_val_dense_compare = attention_cfg.get(
            "sparse_val_dense_compare",
            attention_cfg.get(
                "sparse_validate_dense_compare",
                data.get("sparse_val_dense_compare", config.sparse_val_dense_compare),
            ),
        )
        config.sparse_layer_indices = attention_cfg.get(
            "sparse_layer_indices",
            data.get("sparse_layer_indices", config.sparse_layer_indices),
        )
        config.sparse_mask_cache_dir = attention_cfg.get(
            "sparse_mask_cache_dir",
            data.get("sparse_mask_cache_dir", config.sparse_mask_cache_dir),
        )
        config.dense_mode = attention_cfg.get(
            "dense_mode",
            data.get("dense_mode", config.dense_mode),
        )
        config.dense_step_intervals = attention_cfg.get(
            "dense_step_intervals",
            attention_cfg.get(
                "dense_mode_step_intervals",
                data.get("dense_step_intervals", data.get("dense_mode_step_intervals", config.dense_step_intervals)),
            ),
        )
        
        return config


def _normalize_dense_step_intervals(
    raw_intervals: Optional[Sequence[Sequence[int]]],
) -> List[Tuple[int, int]]:
    if raw_intervals in (None, []):
        return []
    if not isinstance(raw_intervals, Sequence) or isinstance(raw_intervals, (str, bytes)):
        raise ValueError("dense_step_intervals must be a list of [start, end] pairs")

    normalized: List[Tuple[int, int]] = []
    for idx, interval in enumerate(raw_intervals):
        if not isinstance(interval, Sequence) or isinstance(interval, (str, bytes)) or len(interval) != 2:
            raise ValueError(
                f"dense_step_intervals[{idx}] must be a [start, end] pair"
            )
        start, end = interval
        if not isinstance(start, int) or not isinstance(end, int):
            raise ValueError(
                f"dense_step_intervals[{idx}] must contain integer step indices"
            )
        if start < 0 or end < 0:
            raise ValueError(
                f"dense_step_intervals[{idx}] must contain non-negative step indices"
            )
        if start > end:
            raise ValueError(
                f"dense_step_intervals[{idx}] start ({start}) must be <= end ({end})"
            )
        normalized.append((start, end))

    normalized.sort(key=lambda pair: pair[0])
    for idx in range(1, len(normalized)):
        prev_start, prev_end = normalized[idx - 1]
        start, end = normalized[idx]
        if start <= prev_end:
            raise ValueError(
                "dense_step_intervals must be non-overlapping when endpoints are inclusive; "
                f"got [{prev_start}, {prev_end}] and [{start}, {end}]"
            )
    return normalized


def _step_uses_dense_attention(step: int, dense_step_intervals: Sequence[Tuple[int, int]]) -> bool:
    if step <= 0:
        return False
    return any(start <= step <= end for start, end in dense_step_intervals)


def _active_dense_step_interval(
    step: int,
    dense_step_intervals: Sequence[Tuple[int, int]],
) -> Optional[Tuple[int, int]]:
    if step <= 0:
        return None
    for start, end in dense_step_intervals:
        if start <= step <= end:
            return start, end
    return None


def _set_model_dense_mode(model, enabled: bool) -> None:
    base_model = _unwrap_base_model(model)
    base_model.enable_dense_mode(enabled)


def _apply_dense_mode_for_step(
    model,
    config: TrainConfig,
    step: int,
    dense_step_intervals: Sequence[Tuple[int, int]],
) -> bool:
    if config.attention_type != "sparse_stochastic":
        return False
    enabled = config.dense_mode or _step_uses_dense_attention(step, dense_step_intervals)
    _set_model_dense_mode(model, enabled)
    return enabled


class TextDataset(Dataset):
    """Simple text dataset for language model training."""
    
    def __init__(
        self,
        data_path: str,
        tokenizer,
        max_seq_length: int = 512,
    ):
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        self.examples = []
        
        # Load data
        if data_path.endswith(".jsonl"):
            self._load_jsonl(data_path)
        else:
            # Assume it's a directory with HF dataset
            self._load_hf_dataset(data_path)
    
    def _load_jsonl(self, path: str):

        with open(path, "r") as f:
            for line in f:
                data = json.loads(line)
                text = data.get("text", data.get("content", ""))
                if text.strip():
                    self.examples.append(text)
    
    def _load_hf_dataset(self, path: str):
        
        dataset = load_from_disk(path)
        for example in dataset:
            text = example.get("text", example.get("content", ""))
            if text.strip():
                self.examples.append(text)
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, idx):
        text = self.examples[idx]
        
        # Tokenize
        encodings = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_seq_length,
            padding="max_length",
            return_tensors="pt",
        )
        
        input_ids = encodings["input_ids"].squeeze(0)
        attention_mask = encodings["attention_mask"].squeeze(0)
        
        # Labels are same as input_ids for causal LM
        labels = input_ids.clone()
        # Mask padding tokens in labels
        labels[attention_mask == 0] = -100
        
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }


class PreTokenizedDataset(Dataset):
    """Dataset for pre-tokenized data (no tokenization at runtime)."""
    
    def __init__(self, data_path: str, max_seq_length: int = 2048):
        """
        Load pre-tokenized data.
        
        Args:
            data_path: Path to directory containing tokenized_data.pt or
                      path to HuggingFace dataset directory with tokenized data
        """
        self.data_path = data_path
        self.max_seq_length = max_seq_length
        # Check if it's a .pt file or directory
        if os.path.isfile(data_path) and data_path.endswith(".pt"):
            # Load from .pt file
            print(f"Loading pre-tokenized data from {data_path}...")
            self.data = torch.load(data_path)
            self.is_hf_dataset = False
        elif os.path.isdir(data_path):
            # Check for tokenized_data.pt in directory
            pt_path = os.path.join(data_path, "tokenized_data.pt")
            if os.path.exists(pt_path):
                print(f"Loading pre-tokenized data from {pt_path}...")
                self.data = torch.load(pt_path)
                self.is_hf_dataset = False
            else:
                # Try loading as HuggingFace dataset
                print(f"Loading pre-tokenized HF dataset from {data_path}...")
                self.data = load_from_disk(data_path)
                # Set format to torch to avoid repeated tensor conversions
                self.data.set_format(type="torch", columns=["input_ids", "attention_mask"])
                self.is_hf_dataset = True
        else:
            raise ValueError(f"Invalid data path: {data_path}")
        
        # Compute pack factor: how many stored sequences to concatenate per example.
        # Uses ceil so any max_seq_length >= stored_len works, not just exact multiples.
        self.pack_factor = 1
        self.stored_seq_length = max_seq_length
        if self.is_hf_dataset and len(self.data) > 0:
            stored_len = len(self.data[0]["input_ids"])
            if stored_len > 0 and stored_len < max_seq_length:
                self.pack_factor = math.ceil(max_seq_length / stored_len)
                self.stored_seq_length = stored_len

        n_packed = len(self.data) // self.pack_factor
        print(f"Loaded {len(self.data)} pre-tokenized examples "
              f"(stored_len={self.stored_seq_length}, pack_factor={self.pack_factor}, "
              f"packed_examples={n_packed})")

    def __len__(self):
        return len(self.data) // self.pack_factor

    def __getitem__(self, idx):
        if self.is_hf_dataset:
            if self.pack_factor > 1:
                # Concatenate pack_factor consecutive stored sequences into one
                start = idx * self.pack_factor
                ids_list, mask_list = [], []
                for i in range(self.pack_factor):
                    ex = self.data[start + i]
                    ids_list.append(ex["input_ids"])
                    mask_list.append(ex["attention_mask"])
                input_ids = torch.cat(ids_list, dim=0)[:self.max_seq_length]
                attention_mask = torch.cat(mask_list, dim=0)[:self.max_seq_length]
            else:
                example = self.data[idx]
                # Data is already in tensor format due to set_format()
                input_ids = example["input_ids"][:self.max_seq_length]
                attention_mask = example["attention_mask"][:self.max_seq_length]

            labels = input_ids.clone()
            # Mask padding tokens in labels
            labels[attention_mask == 0] = -100

            return {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "labels": labels,
            }
        else:
            # Data is already in tensor format
            return self.data[idx][:self.max_seq_length]


def get_lr_scheduler(optimizer, config: TrainConfig, num_training_steps: int):
    """Create learning rate scheduler."""
    
    num_warmup_steps = config.warmup_steps
    
    if config.lr_scheduler == "cosine":
        def lr_lambda(current_step):
            if current_step < num_warmup_steps:
                return float(current_step) / float(max(1, num_warmup_steps))
            progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
    elif config.lr_scheduler == "linear":
        def lr_lambda(current_step):
            if current_step < num_warmup_steps:
                return float(current_step) / float(max(1, num_warmup_steps))
            return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))
    else:  # constant
        def lr_lambda(current_step):
            if current_step < num_warmup_steps:
                return float(current_step) / float(max(1, num_warmup_steps))
            return 1.0
    
    return LambdaLR(optimizer, lr_lambda)


def setup_distributed():
    """Setup distributed training."""
    if "RANK" in os.environ:
        rank = int(os.environ["RANK"])
        local_rank = int(os.environ["LOCAL_RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(local_rank)
        
        return rank, local_rank, world_size
    else:
        return 0, 0, 1


def is_main_process(rank: int) -> bool:
    return rank == 0


def _unwrap_model(model):
    return model.module if hasattr(model, "module") else model


def _unwrap_base_model(model):
    unwrapped = _unwrap_model(model)
    return unwrapped._orig_mod if hasattr(unwrapped, "_orig_mod") else unwrapped


def _model_attention_mask(attention_mask):
    """Map a batch attention mask to what the model should receive.

    Returns ``None`` for fully-unpadded (packed) batches so dense / dense_mode
    attention layers take SDPA's ``is_causal`` fast path (keeps the Flash backend
    eligible and applies the causal upper-triangle skip). Returns the mask
    unchanged when padding is present so the explicit additive mask is built.

    The padding check runs here in eager mode — deliberately outside the compiled
    model — so it never introduces a graph break / recompile in the hot forward.
    Loss masking is unaffected: labels already carry ``-100`` for padded positions,
    so dropping an all-ones attention mask does not change the loss.
    """
    if attention_mask is None:
        return None
    return attention_mask if bool((attention_mask == 0).any()) else None


def _sample_sparse_indices(model, config: TrainConfig, mask_rng: Optional[np.random.RandomState]):
    if config.attention_type != "sparse_stochastic":
        return None
    if mask_rng is None:
        raise ValueError("mask_rng must be provided for sparse_stochastic attention")

    base_model = _unwrap_base_model(model)
    layers = base_model.model.layers
    pool_size = int(base_model.config.sparse_mask_pool_size)
    indices = mask_rng.randint(0, pool_size, size=len(layers))
    # Pre-select BlockMask objects outside the compiled graph so Dynamo never
    # sees list-indexing with a dynamic integer.  flex_attention's HOP
    # registration guards only on BlockMask structural metadata (Q_LEN,
    # KV_LEN, BLOCK_SIZE) — not on which random blocks are chosen — so
    # different masks across steps cause zero recompilations.
    return [
        layer.self_attn.pick_mask(int(idx)) if layer.is_sparse else None
        for layer, idx in zip(layers, indices)
    ]


def _build_dense_eval_model(train_model, device: torch.device, dtype: torch.dtype):
    train_base = _unwrap_base_model(train_model)
    dense_cfg = copy.deepcopy(train_base.config)
    dense_cfg.attention_type = "dense"
    return Qwen3ForCausalLM(dense_cfg).to(device=device, dtype=dtype)


def _sync_dense_eval_model_weights(train_model, dense_eval_model):
    train_base = _unwrap_base_model(train_model)
    missing, unexpected = dense_eval_model.load_state_dict(train_base.state_dict(), strict=False)
    if missing or unexpected:
        print(f"[sync_dense_eval] missing={len(missing)} unexpected={len(unexpected)}")
        if missing: print("    missing[:10]:", missing[:10])
        if unexpected: print("    unexpected[:10]:", unexpected[:10])


def evaluate(
    model,
    eval_loader,
    device,
    config: TrainConfig,
    rank: int,
    mask_rng: Optional[np.random.RandomState] = None,
    use_sparse_indices: bool = True,
    dense_step_intervals: Sequence[Tuple[int, int]] = (),
    step_for_mode: Optional[int] = None,
):
    """Evaluate model on validation set."""
    model.eval()
    total_loss = 0.0
    total_tokens = 0
    num_batches = 0
    
    # Divide the target by world_size so the work is distributed evenly
    world_size = dist.get_world_size() if dist.is_initialized() else 1
    
    # Use math.ceil to ensure you don't round down to 0 on small batch settings
    max_eval_batches = math.ceil(
        (config.eval_batches * config.gradient_accumulation_steps) / world_size
    )
    
    autocast_dtype = torch.bfloat16 if config.mixed_precision == "bf16" else torch.float16
    dense_active = False
    if step_for_mode is not None:
        dense_active = _apply_dense_mode_for_step(model, config, step_for_mode, dense_step_intervals)

    with torch.no_grad():
        for batch in eval_loader:
            if num_batches >= max_eval_batches:
                break
                
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)
            sparse_indices = (
                _sample_sparse_indices(model, config, mask_rng)
                if use_sparse_indices and not dense_active
                else None
            )
            
            with torch.amp.autocast("cuda", dtype=autocast_dtype):
                outputs = model(
                    input_ids=input_ids,
                    attention_mask=_model_attention_mask(attention_mask),
                    labels=labels,
                    sparse_indices=sparse_indices,
                )
                loss = outputs["loss"]
            
            total_loss += loss.item()
            total_tokens += attention_mask.sum().item()
            num_batches += 1
    
    avg_loss = total_loss / num_batches if num_batches > 0 else 0.0

    # evaluate() is called only on rank 0 (gated by is_main_process in the train loop).
    # All-reduce across ranks is therefore unreachable and has been removed to avoid
    # potential deadlocks if the gating is ever adjusted.
    
    # Calculate perplexity
    perplexity = math.exp(avg_loss) if avg_loss < 100 else float('inf')
    
    model.train()
    return avg_loss, perplexity


def run_validation_step(
    model,
    val_loader,
    device,
    config: TrainConfig,
    rank: int,
    mask_rng: Optional[np.random.RandomState],
    global_step: int,
    dense_eval_model=None,
    dense_step_intervals: Sequence[Tuple[int, int]] = (),
):
    """Run validation across all ranks to prevent DDP deadlocks, logging only on rank 0."""
    active_dense_interval = _active_dense_step_interval(global_step, dense_step_intervals)
    validation_uses_dense_attention = config.dense_mode or active_dense_interval is not None
    
    # CRITICAL: Every GPU must run this so they stay synchronized in the timeline
    eval_loss, eval_perplexity = evaluate(
        model,
        val_loader,
        device,
        config,
        rank,
        mask_rng,
        dense_step_intervals=dense_step_intervals,
        step_for_mode=global_step,
    )

    dense_eval_loss = None
    dense_eval_perplexity = None
    if dense_eval_model is not None:
        _sync_dense_eval_model_weights(model, dense_eval_model)
        dense_eval_loss, dense_eval_perplexity = evaluate(
            dense_eval_model,
            val_loader,
            device,
            config,
            rank,
            mask_rng=None,
            use_sparse_indices=False,
        )

    # GATING ONLY THE OUTPUTS: Only rank 0 prints and logs to WandB
    if rank == 0:
        print(f"Validation at step {global_step}: Loss = {eval_loss:.4f}, Perplexity = {eval_perplexity:.2f}")
        
        log_dict = {
            "eval/loss": eval_loss,
            "eval/perplexity": eval_perplexity,
            "eval/uses_dense_attention": int(validation_uses_dense_attention),
            "eval/uses_sparse_attention": int(not validation_uses_dense_attention),
        }

        if active_dense_interval is not None:
            interval_start, interval_end = active_dense_interval
            log_dict.update({
                "eval/dense_interval_active": 1,
                "eval/dense_interval_start": interval_start,
                "eval/dense_interval_end": interval_end,
            })
        else:
            log_dict.update({
                "eval/dense_interval_active": 0,
                "eval/dense_interval_start": -1,
                "eval/dense_interval_end": -1,
            })
        
        if dense_eval_model is not None:
            print(
                f"Dense validation compare at step {global_step}: "
                f"Loss = {dense_eval_loss:.4f}, Perplexity = {dense_eval_perplexity:.2f}"
            )
            log_dict.update({
                "eval_dense_compare/loss": dense_eval_loss,
                "eval_dense_compare/perplexity": dense_eval_perplexity,
            })

        if config.wandb_enabled:
            wandb.log(log_dict, step=global_step)
            
        return log_dict
        
    return None


def load_checkpoint(checkpoint_path: str, model, optimizer, scheduler, device, rank: int, config: TrainConfig | None = None):
    """Load model checkpoint and training state.

    Returns
    -------
    (global_step, resume_epoch, resume_micro_step)
        resume_epoch      – epoch index to start the training loop from
        resume_micro_step – number of micro-batches already consumed in that epoch;
                            the loop should skip this many batches before training.
    """
    if is_main_process(rank):
        print(f"Loading checkpoint from {checkpoint_path}...")

    checkpoint_dir = Path(checkpoint_path)

    # Load model weights
    model_path = checkpoint_dir / "model.pt"
    if not model_path.exists():
        raise FileNotFoundError(f"Model checkpoint not found: {model_path}")

    model_state = torch.load(model_path, map_location=device, weights_only=True)
    # Normalise checkpoint keys: strip "_orig_mod." prefix if present (checkpoint
    # was saved from a torch.compile()-wrapped model).
    if any(k.startswith("_orig_mod.") for k in model_state):
        model_state = {k.removeprefix("_orig_mod."): v for k, v in model_state.items()}
    # Drop _rt_pool entries — they are non-learned buffers rebuilt each run and
    # their shape depends on sparse_mask_pool_size which can differ across configs.
    model_state = {k: v for k, v in model_state.items() if not k.endswith("._rt_pool")}
    # Unwrap DDP, then torch.compile's OptimizedModule, so plain keys always match.
    target = model.module if hasattr(model, "module") else model
    if hasattr(target, "_orig_mod"):
        target = target._orig_mod
    target.load_state_dict(model_state)

    # Load training state (optimizer, scheduler, step, epoch, micro_step)
    training_state_path = checkpoint_dir / "training_state.pt"
    if not training_state_path.exists():
        raise FileNotFoundError(f"Training state not found: {training_state_path}")

    training_state = torch.load(training_state_path, map_location=device)
    optimizer.load_state_dict(training_state["optimizer"])
    scheduler.load_state_dict(training_state["scheduler"])
    # Override both optimizer and scheduler LR state only when the configured
    # schedule base really changed.  For an exact resume, the config contains
    # the original peak LR while the checkpoint optimizer contains the current
    # (scheduled) LR.  Unconditionally assigning the peak LR here would create
    # a one-step spike before scheduler.step() restores the schedule.
    if config is not None:
        new_lr = config.learning_rate
        checkpoint_base_lrs = list(getattr(scheduler, "base_lrs", []))
        same_schedule_base = bool(checkpoint_base_lrs) and all(
            math.isclose(base_lr, new_lr, rel_tol=1e-12, abs_tol=0.0)
            for base_lr in checkpoint_base_lrs
        )
        if same_schedule_base:
            if is_main_process(rank):
                current_lrs = [pg["lr"] for pg in optimizer.param_groups]
                print(
                    "Preserving checkpoint LR state for exact resume: "
                    f"current={current_lrs}, schedule_base={checkpoint_base_lrs}"
                )
        else:
            optimizer.defaults["lr"] = new_lr
            for pg in optimizer.param_groups:
                pg["lr"] = new_lr
                pg["initial_lr"] = new_lr
            if hasattr(scheduler, "base_lrs"):
                scheduler.base_lrs = [new_lr for _ in scheduler.base_lrs]
            if hasattr(scheduler, "_last_lr"):
                scheduler._last_lr = [new_lr for _ in scheduler._last_lr]
            if is_main_process(rank):
                print(f"Overriding checkpoint LR with config LR: {new_lr}")
    step = training_state["step"]
    resume_epoch = training_state.get("epoch", 0)
    resume_micro_step = training_state.get("micro_step", 0)

    # Restore per-rank RNG states so stochastic ops are exactly reproducible.
    rng_path = checkpoint_dir / f"rng_state_rank{rank}.pt"
    if rng_path.exists():
        # weights_only=False is required because the file contains numpy arrays
        # (torch CPU RNG state and numpy RandomState). The file is always written
        # by this codebase so the source is trusted.
        rng_state = torch.load(rng_path, map_location="cpu", weights_only=False)
        torch.set_rng_state(rng_state["torch_cpu"])
        if torch.cuda.is_available() and rng_state.get("torch_cuda") is not None:
            torch.cuda.set_rng_state(rng_state["torch_cuda"])
        np.random.set_state(rng_state["numpy"])
        # mask_rng is a local RandomState; its state is returned so the caller
        # can apply it after construction.
    else:
        if is_main_process(rank):
            print(f"Warning: no RNG state file found for rank {rank}; RNG will not be restored.")
        rng_state = None

    if is_main_process(rank):
        print(f"Resumed from checkpoint at step {step} (epoch {resume_epoch}, micro_step {resume_micro_step})")

    return step, resume_epoch, resume_micro_step, rng_state


def save_checkpoint(
    model,
    optimizer,
    scheduler,
    step: int,
    config: TrainConfig,
    rank: int,
    epoch: int = 0,
    micro_step: int = 0,
    mask_rng: Optional[np.random.RandomState] = None,
    loader_generator: Optional[torch.Generator] = None,
    sampler_seed: Optional[int] = None,
):
    """Save model checkpoint, dataloader position, and per-rank RNG states."""
    checkpoint_dir = Path(config.output_dir) / f"checkpoint-{step}"
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    # Every rank saves its own RNG state and DataLoader generator state for exact resume.
    rng_state = {
        "torch_cpu": torch.get_rng_state(),
        "torch_cuda": torch.cuda.get_rng_state() if torch.cuda.is_available() else None,
        "numpy": np.random.get_state(),
        "mask_rng": mask_rng.get_state() if mask_rng is not None else None,
    }
    torch.save(rng_state, checkpoint_dir / f"rng_state_rank{rank}.pt")
    if loader_generator is not None:
        torch.save(loader_generator.get_state(), checkpoint_dir / f"loader_gen_state_rank{rank}.pt")

    if not is_main_process(rank):
        return

    # Get model state dict (handle DDP)
    model_to_save = model.module if hasattr(model, "module") else model

    # Save model
    torch.save(model_to_save.state_dict(), checkpoint_dir / "model.pt")

    # Save optimizer, scheduler, and dataloader position
    torch.save({
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "step": step,
        "epoch": epoch,
        "micro_step": micro_step,
        "sampler_seed": sampler_seed,
    }, checkpoint_dir / "training_state.pt")

    # Save config
    with open(checkpoint_dir / "config.json", "w") as f:
        json.dump(vars(model_to_save.config), f, indent=2, default=str)

    print(f"Saved checkpoint to {checkpoint_dir}")


def find_latest_checkpoint(output_dir: str) -> Optional[str]:
    """Return the path of the most recent checkpoint inside output_dir/run_*/, or None."""
    output_path = Path(output_dir)
    if not output_path.exists():
        return None
    required_files = ("model.pt", "training_state.pt", "config.json")
    run_dirs = sorted(output_path.glob("run_*"), reverse=True)  # newest name first
    for run_dir in run_dirs:
        checkpoints = sorted(
            (
                checkpoint
                for checkpoint in run_dir.glob("checkpoint-*")
                if all(
                    (checkpoint / filename).is_file()
                    and (checkpoint / filename).stat().st_size > 0
                    for filename in required_files
                )
            ),
            key=lambda p: int(p.name.split("-")[-1]),
            reverse=True,
        )
        if checkpoints:
            return str(checkpoints[0])
    return None


def resolve_sampler_seed(config: TrainConfig) -> int:
    """Choose the sampler seed without changing legacy resume order.

    Checkpoints written before the sampler-seed refactor have no marker and
    used DistributedSampler's default seed (0). New runs use config.seed and
    persist that choice in training_state.pt for subsequent exact resumes.
    """
    if not config.resume_from_checkpoint:
        return int(config.data_seed if config.data_seed is not None else config.seed)
    state_path = Path(config.resume_from_checkpoint) / "training_state.pt"
    try:
        state = torch.load(state_path, map_location="cpu", weights_only=True)
        value = state.get("sampler_seed")
        if value is not None:
            return int(value)
    except Exception:
        pass
    return int(config.data_seed if config.data_seed is not None else 0)


def train(config: TrainConfig):
    """Main training function."""
    # Setup distributed
    rank, local_rank, world_size = setup_distributed()
    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
    
    # Set seed
    torch.manual_seed(config.seed + rank)
    mask_rng = np.random.RandomState(config.seed + 1 + rank)
    # mask_rng state is restored after load_checkpoint (below) once resume_mask_rng_state is known.

    if config.attention_type not in {"dense", "sparse_stochastic"}:
        raise ValueError(
            f"Unsupported attention_type={config.attention_type}. Expected 'dense' or 'sparse_stochastic'."
        )
    if config.dense_mode and config.attention_type != "sparse_stochastic":
        raise ValueError(
            "dense_mode=True requires attention_type='sparse_stochastic' so the sparse layers exist to switch over."
        )
    dense_step_intervals = _normalize_dense_step_intervals(config.dense_step_intervals)
    config.dense_step_intervals = [list(interval) for interval in dense_step_intervals] or None
    if dense_step_intervals and config.attention_type != "sparse_stochastic":
        raise ValueError(
            "dense_step_intervals requires attention_type='sparse_stochastic'."
        )
    if dense_step_intervals and config.dense_mode:
        raise ValueError(
            "dense_step_intervals cannot be combined with dense_mode=True; use one dense fallback mechanism."
        )
    
    # Auto-detect latest checkpoint to resume from if not explicitly specified
    if not config.resume_from_checkpoint and not config.force_new_run:
        latest_ckpt = find_latest_checkpoint(config.output_dir)
        if latest_ckpt is not None:
            config.resume_from_checkpoint = latest_ckpt
            if is_main_process(rank):
                print(f"Auto-resuming from latest checkpoint: {latest_ckpt}")

    # Generate unique run ID and update output directory
    if config.resume_from_checkpoint and not config.force_new_run:
        # Truly resuming the same run — reuse its output directory so wandb
        # and optimizer state continue from where they left off.
        config.output_dir = str(Path(config.resume_from_checkpoint).parent)
        if is_main_process(rank):
            print(f"Resuming: Using existing output directory {config.output_dir}")
    else:
        # Either a completely fresh run or a fork (force_new_run=True with a
        # resume_from_checkpoint for weights only) — always create a new dir
        # so a new wandb run is created instead of hijacking the old one.
        run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        config.output_dir = os.path.join(config.output_dir, f"run_{run_id}")
        if is_main_process(rank):
            print(f"New run: Checkpoints will be saved to {config.output_dir}")
            Path(config.output_dir).mkdir(parents=True, exist_ok=True)

    # Persist wandb run ID early (before heavy setup) so if the process crashes
    # during model/dataset loading, recovery will continue in the same wandb run.
    # Actual wandb.init() is deferred to just before training loop to ensure
    # wall-clock timing only measures training time, not setup time.
    wandb_run_id_path = None
    prior_wandb_id = None
    if config.wandb_enabled and is_main_process(rank):
        wandb_run_id_path = Path(config.output_dir) / "wandb_run_id.txt"
        if wandb_run_id_path.exists():
            prior_wandb_id = wandb_run_id_path.read_text().strip()

    if is_main_process(rank):
        print("Running config:")
        print(json.dumps(asdict(config), indent=2, default=str))
    
    # Load tokenizer
    if is_main_process(rank):
        print(f"Loading tokenizer from {config.model_path}...")
    tokenizer = load_tokenizer(config.model_path)
    
    # Set padding token if not set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Load model
    dtype = torch.bfloat16 if config.model_dtype == "bfloat16" else torch.float16
    attention_overrides = {
        "attention_type": config.attention_type,
        "positional": config.positional,
        "qk_norm": config.qk_norm,
        "qk_gain_fp32": config.qk_gain_fp32,
        "sparse_block_size": int(config.sparse_block_size),
        "sparse_num_random_blocks": int(config.sparse_num_random_blocks),
        "sparse_mask_pool_size": int(config.sparse_mask_pool_size),
        "sparse_mask_seed": int(config.sparse_mask_seed),
        "sparse_mixed_nrb_set": config.sparse_mixed_nrb_set,
        "sparse_max_seq_length": int(config.max_seq_length),
        "sparse_layer_indices": config.sparse_layer_indices,
        "sparse_mask_cache_dir": config.sparse_mask_cache_dir,
    }
    # Add any architecture overrides specified in the train config
    arch_override_keys = [
        "hidden_size", "intermediate_size", "num_attention_heads",
        "num_key_value_heads", "num_hidden_layers", "head_dim",
        "attention_dropout", "rope_theta", "qk_norm", "qk_gain_fp32",
        "rope_train_position_scale_set",
    ]
    for key in arch_override_keys:
        val = getattr(config, key, None)
        if val is not None:
            attention_overrides[key] = val
    
    # Determine model loading strategy
    if config.resume_from_checkpoint:
        # When resuming, load config from checkpoint
        if is_main_process(rank):
            print(f"Resuming from checkpoint: {config.resume_from_checkpoint}")
        checkpoint_config_path = os.path.join(config.resume_from_checkpoint, "config.json")
        model_config = Qwen3Config.from_json(checkpoint_config_path)
        model_config.checkpoint_activations = config.checkpoint_activations
        for key, value in attention_overrides.items():
            if hasattr(model_config, key):
                setattr(model_config, key, value)
        model = Qwen3ForCausalLM(model_config)
        model = model.to(device=device, dtype=dtype)
        
        # Model weights will be loaded later along with optimizer/scheduler
    elif config.init_from_scratch:
        if is_main_process(rank):
            print(f"Initializing model from scratch (random weights)...")
        # Load config from pretrained path but initialize with random weights
        config_path = os.path.join(config.model_path, "config.json")
        model_config = Qwen3Config.from_json(config_path)
        model_config.checkpoint_activations = config.checkpoint_activations
        for key, value in attention_overrides.items():
            if hasattr(model_config, key):
                setattr(model_config, key, value)
        model = Qwen3ForCausalLM(model_config)
        model = model.to(device=device, dtype=dtype)
    else:
        if is_main_process(rank):
            print(f"Loading pretrained model from {config.model_path}...")
        model = Qwen3ForCausalLM.from_pretrained(
            config.model_path,
            device=device,
            dtype=dtype,
            config_overrides=attention_overrides,
        )
        model.config.checkpoint_activations = config.checkpoint_activations
        model.model.checkpoint_activations = config.checkpoint_activations

    if config.dense_mode:
        model.enable_dense_mode()
        if is_main_process(rank):
            print("dense_mode=True: sparse attention layers will use standard causal SDPA.")

    # ``model.to(..., dtype=bf16)`` above intentionally converts the whole
    # model. Restore only Q/K RMSNorm gains to FP32 for this experiment; their
    # forward method casts the value back to the activation dtype.
    if config.qk_gain_fp32:
        qk_gain_params = []
        for module in model.modules():
            if getattr(module, "weight_fp32", False):
                module.weight.data = module.weight.data.float()
                qk_gain_params.append(module.weight)
        if len(qk_gain_params) != 56:
            raise RuntimeError(f"expected 56 Q/K gain tensors, found {len(qk_gain_params)}")
    else:
        qk_gain_params = []
    
    # Print model architecture BEFORE torch.compile so torchinfo hooks don't
    # survive into the compiled region (which would cause recompile storms).
    if is_main_process(rank):
        print()
        print("\n" + "=" * 80)
        print("MODEL ARCHITECTURE")
        print("=" * 80)
        print(f"Model config:")
        print(f"  vocab_size: {model.config.vocab_size}")
        print(f"  hidden_size: {model.config.hidden_size}")
        print(f"  num_hidden_layers: {model.config.num_hidden_layers}")
        print(f"  num_attention_heads: {model.config.num_attention_heads}")
        print(f"  num_key_value_heads: {model.config.num_key_value_heads}")
        print(f"  head_dim: {model.config.head_dim}")
        print(f"  intermediate_size: {model.config.intermediate_size}")
        print(f"  max_position_embeddings: {model.config.max_position_embeddings}")
        print(f"  rope_theta: {model.config.rope_theta}")
        print(f"  rope_train_position_scale_set: {model.config.rope_train_position_scale_set}")
        print(f"  attention_type: {model.config.attention_type}")
        if model.config.attention_type == "sparse_stochastic":
            print(f"  sparse_block_size: {model.config.sparse_block_size}")
            print(f"  sparse_num_random_blocks: {model.config.sparse_num_random_blocks}")
            print(f"  sparse_mask_pool_size: {model.config.sparse_mask_pool_size}")
            print(f"  sparse_max_seq_length: {model.config.sparse_max_seq_length}")
        
        # Count parameters
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"\nParameter count:")
        print(f"  Total parameters: {total_params:,}")
        print(f"  Trainable parameters: {trainable_params:,}")
        print(f"  Non-trainable parameters: {total_params - trainable_params:,}")
        print(f"  Model size (BF16): {total_params * 2 / 1e9:.2f} GB")
        
        # Print detailed layer summary if torchinfo is available
        if summary is not None and config.verbose:
            print("\n" + "-" * 80)
            print("LAYER-BY-LAYER SUMMARY")
            print("-" * 80)
            try:
                summary(
                    model,
                    input_size=(config.per_device_batch_size, config.max_seq_length),
                    dtypes=[torch.long],
                    col_names=["input_size", "output_size", "num_params", "trainable"],
                    depth=3,
                    verbose=config.verbose,
                )
            except Exception as e:
                print(f"Could not generate detailed summary: {e}")
            if getattr(model.config, "tie_word_embeddings", False):
                print(f"  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight "
                      f"({model.config.vocab_size * model.config.hidden_size:,} params). "
                      f"True unique parameter count: {total_params:,} ({total_params/1e6:.1f}M)")
        
        print("=" * 80 + "\n")

    if config.attention_type == "sparse_stochastic":
        try:
            backends = getattr(inductor_config, "max_autotune_gemm_backends", "")
            if isinstance(backends, str):
                backend_set = {token.strip() for token in backends.split(",") if token.strip()}
                if "ATEN" not in backend_set:
                    backend_set.add("ATEN")
                    inductor_config.max_autotune_gemm_backends = ",".join(sorted(backend_set))
        except Exception as exc:
            if is_main_process(rank):
                print(f"Warning: could not set sparse compile backend fallback: {exc}")

    if os.environ.get("SPARSEATT_DISABLE_COMPILE", "0") != "1":
        model = torch.compile(model)
    elif is_main_process(rank):
        print("torch.compile disabled for this run")
    
    dense_eval_model = None
    if is_main_process(rank) and config.attention_type == "sparse_stochastic" and config.sparse_val_dense_compare:
        dense_eval_model = _build_dense_eval_model(model, device, dtype)
        print("Sparse validation compare enabled: running both sparse and dense validation passes")
    
    # Wrap with DDP
    if world_size > 1:
        # broadcast_buffers=False: _rt_pool is int16 (NCCL unsupported) and is
        # rebuilt independently per rank each run, so broadcasting it is wrong anyway.
        model = DDP(model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False)

    # Create dataset and dataloader
    if is_main_process(rank):
        print(f"Loading dataset from {config.train_path}...")
    
    # Check if data is pre-tokenized
    is_pretokenized = False
    if os.path.isdir(config.train_path):
        # Check for tokenized_data.pt or HF dataset with tokenized columns
        if os.path.exists(os.path.join(config.train_path, "tokenized_data.pt")):
            is_pretokenized = True
        elif os.path.exists(os.path.join(config.train_path, "dataset_info.json")):
            # HF dataset - check if it has tokenized columns
            try:
                dataset_test = load_from_disk(config.train_path)
                if "input_ids" in dataset_test.column_names:
                    is_pretokenized = True
            except:
                pass
    elif config.train_path.endswith(".pt"):
        is_pretokenized = True
    
    if is_pretokenized:
        if is_main_process(rank):
            print("Using pre-tokenized dataset (no runtime tokenization)")
        full_dataset = PreTokenizedDataset(config.train_path, max_seq_length=config.max_seq_length)
    else:
        if is_main_process(rank):
            print("Using raw text dataset (tokenizing at runtime)")
        full_dataset = TextDataset(
            config.train_path,
            tokenizer,
            max_seq_length=config.max_seq_length,
        )
    
    # Do 95/5 train/val split
    train_size = int(0.95 * len(full_dataset))
    val_size = len(full_dataset) - train_size
    
    # Set seed for reproducible split
    generator = torch.Generator().manual_seed(config.seed)
    train_dataset, val_dataset = torch.utils.data.random_split(
        full_dataset, [train_size, val_size], generator=generator
    )
    
    if is_main_process(rank):
        print(f"Dataset split: {train_size} train, {val_size} validation")
    
    # Create train dataloader
    # Use a dedicated generator (seeded deterministically) for the DataLoader shuffle
    # so that its RNG is independent from the global training RNG and can be
    # saved/restored cleanly across checkpoints.
    loader_generator = torch.Generator()
    loader_generator.manual_seed(config.seed + rank)
    sampler_seed = resolve_sampler_seed(config)
    if is_main_process(rank):
        source = "checkpoint marker" if config.resume_from_checkpoint and sampler_seed != 0 else ("legacy checkpoint fallback" if config.resume_from_checkpoint else "run seed")
        print(f"DistributedSampler seed: {sampler_seed} ({source})")

    if world_size > 1:
        # Tie the sampler permutation to the same run seed used for model
        # initialization and the train/validation split. This makes fixed-seed
        # Dense/Mixed/ALiBi runs consume the same data order when their DDP
        # world size and loader geometry match.
        train_sampler = DistributedSampler(
            train_dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=True,
            seed=sampler_seed,
        )
    else:
        train_sampler = None

    train_loader = DataLoader(
        train_dataset,
        batch_size=config.per_device_batch_size,
        sampler=train_sampler,
        shuffle=(train_sampler is None),
        generator=loader_generator if train_sampler is None else None,
        num_workers=config.num_workers,
        pin_memory=True,
        drop_last=True,
        persistent_workers=True if config.num_workers > 0 else False,
        prefetch_factor=2 if config.num_workers > 0 else None,
    )
    
    # Create validation dataloader
    if world_size > 1:
        val_sampler = DistributedSampler(
            val_dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=False,
            seed=sampler_seed,
        )
    else:
        val_sampler = None
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config.per_device_batch_size,
        sampler=val_sampler,
        shuffle=False,
        num_workers=config.num_workers,
        pin_memory=True,
        drop_last=False,
        persistent_workers=True if config.num_workers > 0 else False,
        prefetch_factor=2 if config.num_workers > 0 else None,
    )
    
    # Calculate training steps
    # len(train_loader) is per-rank micro-batches/epoch (DistributedSampler divides data;
    # drop_last=True in the loader already floors to full batches).
    micro_batches_per_epoch = len(train_loader)
    num_update_steps_per_epoch = micro_batches_per_epoch // config.gradient_accumulation_steps
    num_training_steps = config.max_steps or (num_update_steps_per_epoch * config.num_epochs)
    effective_batch_size = config.per_device_batch_size * world_size * config.gradient_accumulation_steps
    tokens_per_step = effective_batch_size * config.max_seq_length

    if is_main_process(rank):
        print(f"Training examples (packed, train split):  {len(train_dataset)}")
        print(f"Validation examples (packed, val split):  {len(val_dataset)}")
        print(f"Per-device batch size:                    {config.per_device_batch_size}")
        print(f"World size (GPUs):                        {world_size}")
        print(f"Gradient accumulation steps:              {config.gradient_accumulation_steps}")
        print(f"Effective batch size (seqs/step):         {effective_batch_size}")
        print(f"Sequence length:                          {config.max_seq_length}")
        print(f"Tokens per optimizer step:                {tokens_per_step:,}")
        print(f"Micro-batches per epoch (per rank):       {micro_batches_per_epoch}")
        print(f"Optimizer steps per epoch:                {num_update_steps_per_epoch}")
        print(f"Total optimizer steps:                    {num_training_steps}"
              + (f"  (capped by max_steps; {num_training_steps/num_update_steps_per_epoch:.2f} epochs)" if config.max_steps else ""))
    
    # Create optimizer
    normal_params = [p for p in model.parameters() if all(p is not q for q in qk_gain_params)]
    optimizer = torch.optim.AdamW(
        ([{"params": normal_params}, {"params": qk_gain_params}] if qk_gain_params else normal_params),
        lr=config.learning_rate,
        weight_decay=config.weight_decay,
        betas=(0.9, 0.95),
        fused=True
    )
    
    # Create scheduler
    scheduler = get_lr_scheduler(optimizer, config, num_training_steps)
    
    # Load checkpoint if resuming
    global_step = 0
    resume_epoch = 0
    resume_micro_step = 0
    resume_mask_rng_state = None
    if config.resume_from_checkpoint:
        global_step, resume_epoch, resume_micro_step, _rng_state = load_checkpoint(
            config.resume_from_checkpoint,
            model,
            optimizer,
            scheduler,
            device,
            rank,
            config=config,
        )
        if _rng_state is not None:
            resume_mask_rng_state = _rng_state.get("mask_rng")
        if resume_mask_rng_state is not None:
            mask_rng.set_state(resume_mask_rng_state)
        if config.resume_micro_step_override is not None:
            if config.resume_micro_step_override < 0:
                raise ValueError("resume_micro_step_override must be non-negative")
            if is_main_process(rank):
                print(f"Overriding resume micro-step {resume_micro_step} -> {config.resume_micro_step_override} for DDP migration")
            resume_micro_step = config.resume_micro_step_override
        if is_main_process(rank):
            print(f"Resuming training from step {global_step}")
    
    # Mixed precision
    scaler = torch.amp.GradScaler("cuda") if config.mixed_precision == "fp16" else None
    autocast_dtype = torch.bfloat16 if config.mixed_precision == "bf16" else torch.float16
    
    # Initialize wandb just before training to ensure wall-clock timing
    # only measures actual training time (not mask creation, model loading, etc.)
    if config.wandb_enabled and is_main_process(rank):
        if prior_wandb_id:
            print(f"Resuming wandb run {prior_wandb_id}")
            wandb.init(
                project=config.wandb_project,
                name=config.wandb_name,
                tags=config.wandb_tags,
                config=vars(config),
                id=prior_wandb_id,
                resume="must",
            )
        else:
            wandb.init(
                project=config.wandb_project,
                name=config.wandb_name,
                tags=config.wandb_tags,
                config=vars(config),
            )
            wandb_run_id_path.write_text(wandb.run.id)
    
    # Training loop
    model.train()
    total_loss = 0.0
    start_time = time.time()
    last_log_time = start_time
    last_log_step = global_step
    _resume_initial_global_step = global_step
    
    if is_main_process(rank):
        print("Starting training...")
    
    # Restore loader_generator state from checkpoint so the DataLoader shuffle
    # on resume is identical to what the original run would have produced.
    resume_loader_gen_state = None
    if config.resume_from_checkpoint:
        _gen_path = Path(config.resume_from_checkpoint) / f"loader_gen_state_rank{rank}.pt"
        if _gen_path.exists():
            resume_loader_gen_state = torch.load(_gen_path, map_location="cpu", weights_only=True)
            loader_generator.set_state(resume_loader_gen_state)
        else:
            if is_main_process(rank):
                print("Warning: no loader_gen_state file found; DataLoader shuffle order on resume may differ from original run.")

    # On resume, run one validation pass at the restored global_step before any
    # new training work, so the baseline metric is logged at the correct step.
    if config.resume_from_checkpoint and global_step > 0:
        run_validation_step(
            model=model,
            val_loader=val_loader,
            device=device,
            config=config,
            rank=rank,
            mask_rng=mask_rng,
            global_step=global_step,
            dense_eval_model=dense_eval_model,
            dense_step_intervals=dense_step_intervals,
        )

    last_dense_interval_mode = None

    for epoch in range(config.num_epochs):
        # Skip epochs that were fully completed before the checkpoint.
        if epoch < resume_epoch:
            continue

        if train_sampler is not None:
            train_sampler.set_epoch(epoch)

        # How many micro-batches to skip at the start of the resume epoch.
        skip_micro_batches = resume_micro_step if epoch == resume_epoch else 0

        for step, batch in enumerate(train_loader):
            if step < skip_micro_batches:
                # Fast-forward: advance the DataLoader without training.
                continue
            active_step = global_step + 1
            dense_interval_active = _apply_dense_mode_for_step(
                model,
                config,
                active_step,
                dense_step_intervals,
            )
            if dense_interval_active != last_dense_interval_mode:
                if is_main_process(rank):
                    mode_name = "dense" if dense_interval_active else "sparse_stochastic"
                    source = "dense step interval" if dense_interval_active else "configured sparse attention"
                    print(f"Switching attention mode for step {active_step}: {mode_name} ({source})")
                last_dense_interval_mode = dense_interval_active
            # Move batch to device with non-blocking transfer
            input_ids = batch["input_ids"].to(device, non_blocking=True)
            attention_mask = batch["attention_mask"].to(device, non_blocking=True)
            labels = batch["labels"].to(device, non_blocking=True)
            sparse_indices = None if dense_interval_active else _sample_sparse_indices(model, config, mask_rng)
            
            # Forward pass with autocast
            with torch.amp.autocast("cuda", dtype=autocast_dtype):
                outputs = model(
                    input_ids=input_ids,
                    attention_mask=_model_attention_mask(attention_mask),
                    labels=labels,
                    sparse_indices=sparse_indices,
                )
                loss = outputs["loss"] / config.gradient_accumulation_steps
            
            # Backward pass — suppress DDP gradient sync on non-final accumulation steps
            # (avoids 7/8 wasted all-reduces when gradient_accumulation_steps=8)
            is_last_accum_step = (step + 1) % config.gradient_accumulation_steps == 0
            sync_ctx = contextlib.nullcontext() if is_last_accum_step else (
                model.no_sync() if hasattr(model, "no_sync") else contextlib.nullcontext()
            )
            with sync_ctx:
                if scaler is not None:
                    scaler.scale(loss).backward()
                else:
                    loss.backward()
            
            # Accumulate UNSCALED loss for logging
            total_loss += outputs["loss"].item()
            
            # Gradient accumulation
            if (step + 1) % config.gradient_accumulation_steps == 0:
                # Gradient clipping
                if scaler is not None:
                    scaler.unscale_(optimizer)
                
                torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
                
                # Optimizer step
                if scaler is not None:
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    optimizer.step()
                
                scheduler.step()
                optimizer.zero_grad(set_to_none=True)
                
                global_step += 1
                
                # Logging
                if global_step % config.logging_steps == 0 and is_main_process(rank):
                    avg_loss = total_loss / (config.logging_steps * config.gradient_accumulation_steps)
                    lr = scheduler.get_last_lr()[0]
                    
                    # Calculate speed metrics
                    current_time = time.time()
                    
                    # Instantaneous throughput (since last log)
                    time_since_last_log = current_time - last_log_time
                    steps_since_last_log = global_step - last_log_step
                    steps_per_sec = steps_since_last_log / time_since_last_log if time_since_last_log > 0 else 0
                    samples_since_last_log = steps_since_last_log * config.gradient_accumulation_steps * config.per_device_batch_size * world_size
                    samples_per_sec = samples_since_last_log / time_since_last_log if time_since_last_log > 0 else 0
                    ms_per_step = (time_since_last_log / steps_since_last_log * 1000) if steps_since_last_log > 0 else 0
                    
                    # Update last log tracking
                    last_log_time = current_time
                    last_log_step = global_step
                    
                    print(f"Step {global_step}/{num_training_steps} | Loss: {avg_loss:.4f} | LR: {lr:.2e} | "
                          f"{ms_per_step:.1f}ms/step | {samples_per_sec:.1f} samples/s")
                    
                    if config.wandb_enabled:
                        wandb.log({
                            "train/loss": avg_loss,
                            "train/learning_rate": lr,
                            "train/epoch": global_step / num_update_steps_per_epoch,
                            "train/global_step": global_step,
                            "speed/steps_per_sec": steps_per_sec,
                            "speed/samples_per_sec": samples_per_sec,
                            "speed/ms_per_step": ms_per_step,
                        }, step=global_step)
                    
                    total_loss = 0.0
                
                # Evaluation
                if global_step % config.eval_steps == 0:
                    run_validation_step(
                        model=model,
                        val_loader=val_loader,
                        device=device,
                        config=config,
                        rank=rank,
                        mask_rng=mask_rng,
                        global_step=global_step,
                        dense_eval_model=dense_eval_model,
                        dense_step_intervals=dense_step_intervals,
                    )
                
                # Save checkpoint
                if global_step % config.save_steps == 0:
                    save_checkpoint(
                        model, optimizer, scheduler, global_step, config, rank,
                        epoch=epoch, micro_step=step + 1, mask_rng=mask_rng,
                        loader_generator=loader_generator,
                        sampler_seed=sampler_seed,
                    )
                
                # Check max steps
                if config.max_steps and global_step >= config.max_steps:
                    break
        
        if config.max_steps and global_step >= config.max_steps:
            break
    
    # Do not overwrite the source checkpoint if a resume exits without a new optimizer step.
    if not config.resume_from_checkpoint or global_step > _resume_initial_global_step:
        save_checkpoint(
            model, optimizer, scheduler, global_step, config, rank,
            epoch=epoch, micro_step=step + 1, mask_rng=mask_rng,
            loader_generator=loader_generator,
            sampler_seed=sampler_seed,
        )
    elif is_main_process(rank):
        print(f"Skipping final save: no new optimizer step after resume at {global_step}")
    
    if config.wandb_enabled and is_main_process(rank):
        wandb.finish()
    
    if world_size > 1:
        dist.destroy_process_group()
    
    if is_main_process(rank):
        print("Training complete!")


def main():
    parser = argparse.ArgumentParser(description="Train Qwen3 model")
    parser.add_argument(
        "--config",
        type=str,
        default="/mloscratch/homes/navasard/sparse/SparseAtt/configs/train_config_dense_resume.json",
        help="Path to training config JSON",
    )
    args = parser.parse_args()

    config = TrainConfig.from_json(args.config)
    train(config)

if __name__ == "__main__":
    main()