exp.py

"""
Tiny Llama GLU Lab — Consolidated training library.
One file: model definition, activation registry, stability monitoring,
time tracking, dataset builder, and Trainer factory.
"""

import math
import os
import time
import json
import re
from pathlib import Path
from itertools import chain
from typing import Dict, Callable, Optional, List, Any, Tuple

import torch
import torch.nn as nn
from transformers import (
    LlamaConfig,
    LlamaPreTrainedModel,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    DataCollatorForLanguageModeling,
    AutoTokenizer,
    set_seed,
)
from transformers.models.llama.modeling_llama import (
    LlamaAttention,
    LlamaRMSNorm,
    LlamaRotaryEmbedding,
)
from transformers.modeling_outputs import CausalLMOutputWithPast
from datasets import load_dataset


# =============================================================================
# 1. ACTIVATION REGISTRY
# =============================================================================

class GLUActivationRegistry:
    """Own every gating activation you test. Add new variants in one line."""
    _registry: Dict[str, Callable[[torch.Tensor], torch.Tensor]] = {}

    @classmethod
    def register(cls, name: str, fn: Callable[[torch.Tensor], torch.Tensor]) -> None:
        cls._registry[name] = fn

    @classmethod
    def get(cls, name: str) -> Callable[[torch.Tensor], torch.Tensor]:
        if name not in cls._registry:
            raise KeyError(
                f"Activation '{name}' not found. Available: {list(cls._registry.keys())}"
            )
        return cls._registry[name]


# Built-ins
GLUActivationRegistry.register("silu", nn.functional.silu)
GLUActivationRegistry.register("swish", nn.functional.silu)
GLUActivationRegistry.register("relu", nn.functional.relu)
GLUActivationRegistry.register("gelu", nn.functional.gelu)
GLUActivationRegistry.register("mish", nn.functional.mish)
GLUActivationRegistry.register("sigmoid", torch.sigmoid)
GLUActivationRegistry.register("tanh", torch.tanh)
GLUActivationRegistry.register("elu", nn.functional.elu)
GLUActivationRegistry.register("softplus", nn.functional.softplus)

# Identity activation – used for "linear" baseline
GLUActivationRegistry.register("linear", lambda x: x)
# For backward compatibility, keep "bilinear" as alias to "linear"
GLUActivationRegistry.register("bilinear", lambda x: x)

# NEW: s10 = x^2 * sigmoid(x)
GLUActivationRegistry.register("s10", lambda x: x * x * torch.sigmoid(x))

# NEW: w1a = x * tanh(x)  (SiLU with tanh instead of sigmoid)
GLUActivationRegistry.register("w1a", lambda x: x * torch.tanh(x))


# =============================================================================
# 2. CONFIG
# =============================================================================

class TinyLlamaConfig(LlamaConfig):
    """
    Exact Llama config plus two fields:
      - mlp_type: "glu" or "mlp" (standard Transformer MLP)
      - activation: name of the activation function to use inside the MLP block.
    Enforces pure MHA by requiring num_key_value_heads == num_attention_heads.
    """
    model_type = "tiny_llama"

    def __init__(
        self,
        mlp_type: str = "glu",          # default for backward compatibility, but config must override
        activation: str = "silu",       # default for backward compatibility
        **kwargs
    ):
        super().__init__(**kwargs)
        self.mlp_type = mlp_type
        self.activation = activation
        if self.num_key_value_heads != self.num_attention_heads:
            raise ValueError(
                f"Pure MHA required: num_key_value_heads ({self.num_key_value_heads}) "
                f"must equal num_attention_heads ({self.num_attention_heads})."
            )


# =============================================================================
# 3. MODEL
# =============================================================================

class SiTUGLU(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        beta1: float = 4.0,
        beta2: float = 25.0,
    ):
        super().__init__()
        self.beta1 = beta1
        self.beta2 = beta2
        self.W_g = nn.Linear(input_dim, hidden_dim, bias=False)
        self.W_u = nn.Linear(input_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g = self.W_g(x)
        gate = self.beta1 * torch.tanh(g / self.beta1) * torch.sigmoid(g)
        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
        return gate * up


class WaleedGLU(nn.Module):
    """
    Variant of SiTUGLU without the sigmoid in the gate.
    gate = beta1 * tanh(g / beta1)
    up   = beta2 * tanh(W_u(x) / beta2)
    output = gate * up
    """
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        beta1: float = 4.0,
        beta2: float = 25.0,
    ):
        super().__init__()
        self.beta1 = beta1
        self.beta2 = beta2
        self.W_g = nn.Linear(input_dim, hidden_dim, bias=False)
        self.W_u = nn.Linear(input_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g = self.W_g(x)
        gate = self.beta1 * torch.tanh(g / self.beta1)   # sigmoid removed
        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
        return gate * up


class TinyLlamaMLP(nn.Module):
    """
    MLP block supporting both:
      - GLU: down_proj( act(gate_proj(x)) * up_proj(x) )
      - Standard MLP: down_proj( act(up_proj(x)) )
    For MLP, intermediate_size is auto‑scaled by 1.5× to match GLU parameter count.
    """
    def __init__(self, config: TinyLlamaConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size  # base value
        self.mlp_type = config.mlp_type
        self.activation_name = config.activation

        # Determine effective dimensions
        if self.mlp_type == "glu":
            # GLU: keep original intermediate_size
            effective_intermediate = self.intermediate_size
        elif self.mlp_type == "mlp":
            # Standard MLP: scale by 1.5 to keep parameter count equal
            effective_intermediate = int(self.intermediate_size * 1.5)
            print(f"[MLP] Auto‑scaled intermediate_size from {self.intermediate_size} to {effective_intermediate} for parameter parity.")
        else:
            raise ValueError(f"Unknown mlp_type: {self.mlp_type}")

        # Store effective size for use in forward
        self.effective_intermediate = effective_intermediate

        # Initialize special GLU modules to None
        self.situglu = None
        self.waleed = None
        self.gate_proj = None
        self.up_proj = None
        self.act_fn = None

        if self.mlp_type == "glu":
            if self.activation_name == "situglu":
                self.situglu = SiTUGLU(self.hidden_size, effective_intermediate)
            elif self.activation_name == "waleed":
                self.waleed = WaleedGLU(self.hidden_size, effective_intermediate)
            else:
                # Normal GLU with an activation function
                self.act_fn = GLUActivationRegistry.get(self.activation_name)
                self.gate_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
                self.up_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
        else:  # mlp_type == "mlp"
            if self.activation_name in ("situglu", "waleed"):
                raise ValueError(
                    f"Activation '{self.activation_name}' requires a gated architecture (GLU). "
                    f"Please use mlp_type='glu'."
                )
            else:
                self.act_fn = GLUActivationRegistry.get(self.activation_name)
                self.up_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)

        # Down projection is always present, using the effective intermediate size
        self.down_proj = nn.Linear(effective_intermediate, self.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.situglu is not None:
            return self.down_proj(self.situglu(x))
        if self.waleed is not None:
            return self.down_proj(self.waleed(x))

        if self.mlp_type == "glu":
            # GLU: act(gate) * up
            gate = self.gate_proj(x)
            up = self.up_proj(x)
            hidden = self.act_fn(gate) * up
        else:  # mlp
            # Standard MLP: act(up)
            hidden = self.act_fn(self.up_proj(x))

        return self.down_proj(hidden)


class TinyLlamaDecoderLayer(nn.Module):
    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
        self.mlp = TinyLlamaMLP(config)
        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = LlamaRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_out = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            position_embeddings=position_embeddings,
        )[0]
        hidden_states = residual + attn_out

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return (hidden_states,)


# ----------------------------------------------------------------------------
# ATTENTION MASK – float mask with 0.0 / -inf (works with all backends)
# ----------------------------------------------------------------------------
def _build_causal_mask(
    attention_mask: Optional[torch.Tensor],
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
) -> torch.Tensor:
    """
    Build a 4D float attention mask for scaled_dot_product_attention.
    - 0.0 where attention is allowed
    - -inf where it is masked (causal future + padding)
    """
    min_value = torch.finfo(dtype).min

    # Causal mask: upper triangle (future) = -inf
    causal = torch.full((seq_len, seq_len), fill_value=min_value, dtype=dtype, device=device)
    causal = torch.triu(causal, diagonal=1)
    causal = causal[None, None, :, :]  # (1, 1, seq_len, seq_len)

    if attention_mask is None:
        batch_size = 1
        return causal.expand(batch_size, 1, seq_len, seq_len)

    batch_size = attention_mask.shape[0]
    causal = causal.expand(batch_size, 1, seq_len, seq_len).clone()

    # Padding: where attention_mask == 0, set to -inf
    padding = attention_mask[:, None, None, :].to(device) == 0  # (batch, 1, 1, seq_len)
    causal = causal.masked_fill(padding, min_value)

    return causal


# Global flag to print mask message only once
_MASK_PRINTED = False


class TinyLlamaModel(LlamaPreTrainedModel):
    config_class = TinyLlamaConfig

    def __init__(self, config: TinyLlamaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [TinyLlamaDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = LlamaRotaryEmbedding(config=config)
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        global _MASK_PRINTED

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if position_ids is None:
            seq_len = inputs_embeds.shape[1]
            position_ids = torch.arange(
                seq_len, device=inputs_embeds.device
            ).unsqueeze(0).expand(inputs_embeds.shape[0], -1)

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # Build float causal + padding mask (print only once)
        seq_len = hidden_states.shape[1]
        causal_mask = _build_causal_mask(
            attention_mask, seq_len, hidden_states.dtype, hidden_states.device
        )

        if not _MASK_PRINTED:
            print("[INFO] Causal mask (float with -inf) applied to all attention layers.")
            _MASK_PRINTED = True

        for decoder_layer in self.layers:
            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=position_ids,
                position_embeddings=position_embeddings,
            )
            hidden_states = layer_outputs[0]

        hidden_states = self.norm(hidden_states)
        if not return_dict:
            return (hidden_states,)
        return {"last_hidden_state": hidden_states, "hidden_states": None, "attentions": None}


class TinyLlamaForCausalLM(LlamaPreTrainedModel):
    config_class = TinyLlamaConfig
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: TinyLlamaConfig):
        super().__init__(config)
        self.model = TinyLlamaModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            return_dict=return_dict,
        )
        hidden_states = outputs["last_hidden_state"] if return_dict else outputs[0]
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)
            )

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, **kwargs
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]
        position_ids = kwargs.get("position_ids")
        if attention_mask is not None and position_ids is None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)
        return {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
        }


# =============================================================================
# 4. MONITORING ENGINE
# =============================================================================

class StatsEngine:
    """Compute the unified 6-scalar signature for any tensor."""

    @staticmethod
    def compute(
        tensor: torch.Tensor, user_limit: float, dtype_ratio: float
    ) -> Dict[str, float]:
        with torch.no_grad():
            abs_t = tensor.abs()
            dtype_info = torch.finfo(tensor.dtype)
            dtype_limit = (
                dtype_ratio * dtype_info.max
                if not torch.isinf(torch.tensor(dtype_info.max))
                else float("inf")
            )

            return {
                "norm": tensor.norm(2).item(),
                "mean": tensor.mean().item(),
                "std": tensor.std().item(),
                "max_abs": abs_t.max().item(),
                "frac_near_dtype_limit": (
                    (abs_t > dtype_limit).float().mean().item()
                    if not math.isinf(dtype_limit)
                    else 0.0
                ),
                "frac_near_user_limit": (abs_t > user_limit).float().mean().item(),
            }


class StepAccumulator:
    """
    Stores per-tensor entries, then aggregates to layer-scope or global-scope
    using exact population formulas (no tensor retention).
    """

    def __init__(self):
        # name -> {numel, norm, mean, std, max_abs, frac_near_dtype_limit, frac_near_user_limit}
        self.tensors: Dict[str, Dict[str, float]] = {}

    def add(self, name: str, numel: int, stats: Dict[str, float]):
        # With gradient_accumulation_steps > 1, a single optimizer "step"
        # runs several forward/backward micro-batches, so the same hook
        # (e.g. a given layer's activation, or a given param's grad) fires
        # more than once before clear() is next called. Previously this
        # method did a plain overwrite, silently discarding every
        # micro-batch but the last. Merge instead, using the same exact
        # population formulas _aggregate() uses to combine tensors.
        new_entry = {"numel": numel, **stats}
        existing = self.tensors.get(name)
        self.tensors[name] = (
            new_entry if existing is None else self._merge_entry(existing, new_entry)
        )

    @staticmethod
    def _merge_entry(a: Dict[str, float], b: Dict[str, float]) -> Dict[str, float]:
        total_n = a["numel"] + b["numel"]
        if total_n == 0:
            return a
        norm = math.sqrt(a["norm"] ** 2 + b["norm"] ** 2)
        max_abs = max(a["max_abs"], b["max_abs"])
        mean = (a["mean"] * a["numel"] + b["mean"] * b["numel"]) / total_n
        ex2 = (
            a["numel"] * (a["std"] ** 2 + a["mean"] ** 2)
            + b["numel"] * (b["std"] ** 2 + b["mean"] ** 2)
        ) / total_n
        std = math.sqrt(max(0.0, ex2 - mean ** 2))
        frac_dtype = (
            a["frac_near_dtype_limit"] * a["numel"] + b["frac_near_dtype_limit"] * b["numel"]
        ) / total_n
        frac_user = (
            a["frac_near_user_limit"] * a["numel"] + b["frac_near_user_limit"] * b["numel"]
        ) / total_n
        return {
            "numel": total_n,
            "norm": norm,
            "mean": mean,
            "std": std,
            "max_abs": max_abs,
            "frac_near_dtype_limit": frac_dtype,
            "frac_near_user_limit": frac_user,
        }

    def clear(self):
        self.tensors.clear()

    def _aggregate(self, entries: Dict[str, Dict[str, float]]) -> Dict[str, float]:
        if not entries:
            return {}
        numels = [e["numel"] for e in entries.values()]
        total_n = sum(numels)

        # L2 norm
        norm = math.sqrt(sum(e["norm"] ** 2 for e in entries.values()))
        # Max abs
        max_abs = max(e["max_abs"] for e in entries.values())
        # Weighted mean
        mean = sum(e["mean"] * e["numel"] for e in entries.values()) / total_n
        # Pooled std: sqrt( E[σ² + μ²] - μ_global² )
        ex2 = (
            sum(e["numel"] * (e["std"] ** 2 + e["mean"] ** 2) for e in entries.values())
            / total_n
        )
        std = math.sqrt(max(0.0, ex2 - mean ** 2))
        # Weighted fractions
        frac_dtype = (
            sum(e["frac_near_dtype_limit"] * e["numel"] for e in entries.values())
            / total_n
        )
        frac_user = (
            sum(e["frac_near_user_limit"] * e["numel"] for e in entries.values())
            / total_n
        )

        return {
            "norm": norm,
            "mean": mean,
            "std": std,
            "max_abs": max_abs,
            "frac_near_dtype_limit": frac_dtype,
            "frac_near_user_limit": frac_user,
        }

    def get_global_stats(self) -> Dict[str, float]:
        return self._aggregate(self.tensors)

    def get_layer_stats(self, layer_prefix: str) -> Dict[str, float]:
        entries = {
            k: v for k, v in self.tensors.items() if k.startswith(layer_prefix + ".")
        }
        return self._aggregate(entries)


class HookRegistry:
    """Attach and throttle forward/backward hooks."""

    def __init__(self, model: nn.Module):
        self.model = model
        self.handles: List[torch.utils.hooks.RemovableHandle] = []
        self.active = False

    def attach_forward(
        self,
        module_patterns: List[str],
        accumulator: StepAccumulator,
        user_limit: float,
        dtype_ratio: float,
    ):
        for name, module in self.model.named_modules():
            if any(re.search(p, name) for p in module_patterns):
                h = module.register_forward_hook(
                    self._make_forward_hook(name, accumulator, user_limit, dtype_ratio)
                )
                self.handles.append(h)

    def attach_backward(
        self, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
    ):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                h = param.register_hook(
                    self._make_backward_hook(
                        f"grad.{name}", accumulator, user_limit, dtype_ratio
                    )
                )
                self.handles.append(h)

    def _make_forward_hook(
        self, module_name: str, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
    ):
        def hook(module, inp, out):
            if not self.active:
                return

            # Modules can return a tensor, a tuple (take first item), or a
            # dict (e.g. TinyLlamaModel returns {"last_hidden_state": ...}).
            # Pull out the first real tensor we find; skip cleanly if none.
            if isinstance(out, dict):
                out_dict = out
                out = out_dict.get("last_hidden_state")
                if out is None:
                    out = next(
                        (v for v in out_dict.values() if torch.is_tensor(v)), None
                    )
            elif isinstance(out, (tuple, list)):
                out = out[0] if len(out) > 0 else None

            if not torch.is_tensor(out):
                return

            stats = StatsEngine.compute(out.detach(), user_limit, dtype_ratio)
            accumulator.add(f"act.{module_name}", out.numel(), stats)

        return hook

    def _make_backward_hook(
        self, param_name: str, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
    ):
        def hook(grad):
            if not self.active:
                return
            stats = StatsEngine.compute(grad.detach(), user_limit, dtype_ratio)
            accumulator.add(param_name, grad.numel(), stats)

        return hook

    def set_active(self, active: bool):
        self.active = active

    def clear(self):
        for h in self.handles:
            h.remove()
        self.handles.clear()


class StabilityMonitorCallback(TrainerCallback):
    """
    Full stability instrumentation: grad / param / act statistics
    at global, per-layer, and per-tensor scope.
    """

    def __init__(
        self,
        model: nn.Module,
        monitor_every_n_steps: int = 10,
        module_patterns: Optional[List[str]] = None,
        user_limits: Optional[Dict[str, float]] = None,
        dtype_proximity_ratio: float = 0.9,
        log_scope: Optional[Dict[str, bool]] = None,
        monitor_during_eval: bool = False,
    ):
        self.model = model
        self.monitor_every_n_steps = monitor_every_n_steps
        self.module_patterns = module_patterns or [".*mlp.*", ".*self_attn.*"]
        self.user_limits = user_limits or {"grad": 1.0, "param": 100.0, "act": 50.0}
        self.dtype_ratio = dtype_proximity_ratio
        self.log_scope = log_scope or {
            "global": True,
            "per_layer": True,
            "per_tensor": False,
        }
        self.monitor_during_eval = monitor_during_eval

        self.accumulator = StepAccumulator()
        self.hooks = HookRegistry(model)
        self.hooks.attach_forward(
            self.module_patterns,
            self.accumulator,
            self.user_limits["act"],
            self.dtype_ratio,
        )
        self.hooks.attach_backward(
            self.accumulator, self.user_limits["grad"], self.dtype_ratio
        )

        self.pending_metrics: Optional[Dict[str, float]] = None

    def _should_monitor(self, state) -> bool:
        return state.global_step % self.monitor_every_n_steps == 0

    def on_step_begin(self, args, state, control, **kwargs):
        if self._should_monitor(state):
            self.accumulator.clear()
            self.hooks.set_active(True)

    def on_step_end(self, args, state, control, **kwargs):
        if not self.hooks.active:
            return

        # Parameter stats (post-optimizer step)
        for name, param in self.model.named_parameters():
            stats = StatsEngine.compute(
                param.data, self.user_limits["param"], self.dtype_ratio
            )
            self.accumulator.add(f"param.{name}", param.numel(), stats)

        self.hooks.set_active(False)
        self.pending_metrics = self._build_metrics()

    @staticmethod
    def _kind_of(name: str) -> str:
        """Classify a tensor key by its source: activation, gradient, or parameter."""
        if name.startswith("act."):
            return "act"
        if name.startswith("grad."):
            return "grad"
        if name.startswith("param."):
            return "param"
        return "other"

    @staticmethod
    def _strip_kind(name: str) -> str:
        if name.startswith("act."):
            return name[4:]
        if name.startswith(("grad.", "param.")):
            return name[5:]
        return name

    def _build_metrics(self, scope: str = "train") -> Dict[str, float]:
        metrics: Dict[str, float] = {}

        # --- Global (split by kind: act / grad / param — never pooled together) ---
        if self.log_scope.get("global", True):
            by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
            for k, v in self.accumulator.tensors.items():
                by_kind.setdefault(self._kind_of(k), {})[k] = v

            for kind, entries in by_kind.items():
                stats = self.accumulator._aggregate(entries)
                for kk, vv in stats.items():
                    metrics[f"{scope}/global/{kind}/{kk}"] = vv

        # --- Per-layer (group by model.layers.{i}, split by kind) ---
        if self.log_scope.get("per_layer", True):
            layer_prefixes = set()
            for name in self.accumulator.tensors:
                clean = self._strip_kind(name)
                parts = clean.split(".")
                for i, p in enumerate(parts):
                    if p == "layers" and i + 1 < len(parts):
                        prefix = ".".join(parts[: i + 2])
                        layer_prefixes.add(prefix)

            for prefix in layer_prefixes:
                by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
                for k, v in self.accumulator.tensors.items():
                    clean = self._strip_kind(k)
                    if clean.startswith(prefix + ".") or clean == prefix:
                        by_kind.setdefault(self._kind_of(k), {})[k] = v

                safe = prefix.replace(".", "_")
                for kind, entries in by_kind.items():
                    if not entries:
                        continue
                    stats = self.accumulator._aggregate(entries)
                    for kk, vv in stats.items():
                        metrics[f"{scope}/layer_{safe}/{kind}/{kk}"] = vv

        # --- Per-tensor ---
        if self.log_scope.get("per_tensor", False):
            for name, stats in self.accumulator.tensors.items():
                safe = name.replace(".", "_")
                for kk, vv in stats.items():
                    if kk == "numel":
                        continue
                    metrics[f"{scope}/tensor_{safe}/{kk}"] = vv

        return metrics

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None and self.pending_metrics is not None:
            logs.update(self.pending_metrics)
            self.pending_metrics = None

    def on_prediction_step(self, args, state, control, **kwargs):
        """Fires once per eval/predict batch. Trainer.evaluate() calls this
        for every batch in the eval loop, then calls self.log(output.metrics)
        (which dispatches on_log to every callback, including the wandb/
        tensorboard reporting callbacks) BEFORE on_evaluate() runs. So to get
        eval-time stats into that same on_log dispatch, we have to build
        pending_metrics here, not in on_evaluate — by the time on_evaluate
        fires, self.log() has already happened and it's too late.
        """
        if not self.monitor_during_eval:
            return
        if not self.hooks.active:
            # First batch of this eval pass: start a fresh accumulation and
            # snapshot parameter stats once (they don't change during eval).
            self.accumulator.clear()
            self.hooks.set_active(True)
            for name, param in self.model.named_parameters():
                stats = StatsEngine.compute(
                    param.data, self.user_limits["param"], self.dtype_ratio
                )
                self.accumulator.add(f"param.{name}", param.numel(), stats)
        self.pending_metrics = self._build_metrics(scope="eval")

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        self.hooks.set_active(False)
        self.accumulator.clear()


class TimeTrackerCallback(TrainerCallback):
    """Precise training & eval timing with remaining-time estimates."""

    def __init__(self):
        self.step_start: Optional[float] = None
        self.epoch_start: Optional[float] = None
        self.total_train_time = 0.0
        self.step_times: List[float] = []

    def on_epoch_begin(self, args, state, control, **kwargs):
        self.epoch_start = time.perf_counter()

    def on_step_begin(self, args, state, control, **kwargs):
        self.step_start = time.perf_counter()

    def on_step_end(self, args, state, control, **kwargs):
        if self.step_start is not None:
            dt = time.perf_counter() - self.step_start
            self.step_times.append(dt)
            self.total_train_time += dt
            self.step_start = None

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return

        logs["train/total_time_seconds"] = self.total_train_time

        if self.step_times:
            recent = self.step_times[-100:]
            logs["train/time_per_step_avg"] = sum(recent) / len(recent)

        if self.epoch_start is not None:
            logs["train/epoch_time_elapsed"] = time.perf_counter() - self.epoch_start

        if state.max_steps and state.global_step > 0:
            avg = self.total_train_time / state.global_step
            remaining = (state.max_steps - state.global_step) * avg
            logs["train/estimated_remaining_minutes"] = remaining / 60.0

    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        pass


class MetricsLoggerCallback(TrainerCallback):
    """Persist every logged dict as JSONL in the output dir."""

    def __init__(self, output_dir: str):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.log_file = self.output_dir / "training_log.jsonl"

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return
        entry = {
            "step": state.global_step,
            "epoch": state.epoch,
            "timestamp": time.time(),
            **logs,
        }
        with open(self.log_file, "a") as f:
            f.write(json.dumps(entry, default=str) + "\n")


# =============================================================================
# 5. DATA & TRAINER FACTORY
# =============================================================================

def build_dataset(
    tokenizer,
    max_seq_len: int = 512,
    split: str = "train",
    dataset_name: str = "roneneldan/TinyStories",
    max_samples: Optional[int] = None,   # NEW: limit number of samples
):
    """Concatenate and chunk TinyStories for causal LM. Fast path with multiprocessing."""
    ds = load_dataset(dataset_name, split=split)
    
    # Optionally limit the number of samples
    if max_samples is not None and split == "train":
        ds = ds.select(range(min(max_samples, len(ds))))
        print(f"[Dataset] Using first {len(ds)} samples for training (max_samples={max_samples})")

    def tokenize(examples):
        # add_special_tokens=False so we control separators ourselves.
        # Without an explicit boundary token, group_texts() below would
        # concatenate unrelated stories back-to-back with nothing marking
        # where one ends and the next begins, teaching the model spurious
        # cross-document continuations. Append EOS to each example so every
        # packed chunk still carries a clear "new document" signal.
        out = tokenizer(examples["text"], add_special_tokens=False)
        eos_id = tokenizer.eos_token_id
        out["input_ids"] = [ids + [eos_id] for ids in out["input_ids"]]
        if "attention_mask" in out:
            out["attention_mask"] = [mask + [1] for mask in out["attention_mask"]]
        return out

    tokenized = ds.map(
        tokenize,
        batched=True,
        num_proc=4,
        remove_columns=ds.column_names,
        desc=f"Tokenizing {split}",
    )

    def group_texts(examples):
        # chain.from_iterable is O(total) instead of O(n²)
        concatenated = {
            k: list(chain.from_iterable(examples[k])) for k in examples.keys()
        }
        total_length = len(concatenated[list(examples.keys())[0]])
        total_length = (total_length // max_seq_len) * max_seq_len
        result = {
            k: [t[i : i + max_seq_len] for i in range(0, total_length, max_seq_len)]
            for k, t in concatenated.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result

    return tokenized.map(
        group_texts,
        batched=True,
        batch_size=10000,
        num_proc=4,
        desc=f"Chunking {split}",
    )


def create_trainer(
    model,
    tokenizer,
    config: Dict[str, Any],
    train_dataset,
    eval_dataset=None,
):
    """Assemble HF Trainer with all custom callbacks."""
    tc = config.get("training", {})
    mc = config.get("monitor", {})

    # Allow override of run_name via config (used by sweep)
    run_name = tc.get("run_name", None)

    args = TrainingArguments(
        output_dir=tc.get("output_dir", "./out"),
        run_name=run_name,                     # explicit run name for WandB
        num_train_epochs=tc.get("num_train_epochs", 3),
        per_device_train_batch_size=tc.get("per_device_train_batch_size", 16),
        per_device_eval_batch_size=tc.get("per_device_eval_batch_size", 16),
        gradient_accumulation_steps=tc.get("gradient_accumulation_steps", 4),
        learning_rate=tc.get("learning_rate", 3e-4),
        weight_decay=tc.get("weight_decay", 0.0),
        max_grad_norm=tc.get("max_grad_norm", 1.0),
        optim=tc.get("optim", "adamw_torch"),
        warmup_steps=tc.get("warmup_steps", 0),
        lr_scheduler_type=tc.get("lr_scheduler_type", "cosine"),
        bf16=tc.get("bf16", True),
        logging_steps=tc.get("logging_steps", 10),
        eval_strategy=tc.get("eval_strategy", "steps"),
        eval_steps=tc.get("eval_steps", 500),
        save_strategy=tc.get("save_strategy", "steps"),
        save_steps=tc.get("save_steps", 1000),
        load_best_model_at_end=tc.get("load_best_model_at_end", False),
        report_to=tc.get("report_to", "tensorboard"),
        push_to_hub=tc.get("push_to_hub", False),
        hub_model_id=tc.get("hub_model_id", None),
        # Never commit real tokens to config files. Prefer an explicit value
        # in the config only if someone deliberately put one there; normal
        # case is HF_TOKEN in the environment.
        hub_token=tc.get("hub_token") or os.environ.get("HF_TOKEN"),
        max_steps=tc.get("max_steps", -1),
        seed=tc.get("seed", 42),
        data_seed=tc.get("data_seed", 42),
        remove_unused_columns=False,
    )

    callbacks = [
        TimeTrackerCallback(),
    ]

    if mc.get("enabled", True):
        callbacks.append(
            StabilityMonitorCallback(
                model=model,
                monitor_every_n_steps=mc.get("monitor_every_n_steps", 10),
                module_patterns=mc.get("module_patterns", [".*mlp.*", ".*self_attn.*"]),
                user_limits=mc.get(
                    "user_limits", {"grad": 1.0, "param": 100.0, "act": 50.0}
                ),
                dtype_proximity_ratio=mc.get("dtype_proximity_ratio", 0.9),
                log_scope=mc.get(
                    "log_scope",
                    {"global": True, "per_layer": True, "per_tensor": False},
                ),
                monitor_during_eval=mc.get("monitor_during_eval", False),
            )
        )

    # Must be added LAST: on_log() writes the shared `logs` dict to disk, so
    # every callback that injects keys into that dict (e.g. StabilityMonitor's
    # max_abs/norm/std stats) needs to run BEFORE this one, not after.
    callbacks.append(MetricsLoggerCallback(args.output_dir))

    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=collator,
        callbacks=callbacks,
    )

    # Trainer.__init__ builds its callback list as:
    #   [DEFAULT_CALLBACKS..., <report_to integrations, e.g. WandbCallback>,
    #    *our callbacks (TimeTracker, StabilityMonitor, MetricsLogger)]
    # All callbacks share the *same* `logs` dict object on on_log(), and are
    # invoked in that list order. That means WandbCallback.on_log() was
    # reading `logs` and shipping it off BEFORE TimeTrackerCallback /
    # StabilityMonitorCallback ever mutated it with their train/global/*,
    # train/layer_*, eval/* keys — so those metrics only ever reached the
    # local training_log.jsonl (written by our MetricsLoggerCallback, which
    # happens to run after within our own sublist) and never wandb.
    # Move every reporting-integration callback to the very end so all of
    # our metric-producing callbacks mutate `logs` first.
    try:
        from transformers.integrations import get_reporting_integration_callbacks

        reporting_types = tuple(get_reporting_integration_callbacks(args.report_to))
    except Exception:
        reporting_types = ()

    if reporting_types:
        handler = trainer.callback_handler
        reporting_cbs = [cb for cb in handler.callbacks if isinstance(cb, reporting_types)]
        other_cbs = [cb for cb in handler.callbacks if not isinstance(cb, reporting_types)]
        handler.callbacks = other_cbs + reporting_cbs

    return trainer