Auto upload 2026-08-06T20:45:13.807790

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+wandb/run-20260806_202910-50qq810b/run-50qq810b.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20260806_203335-6n8s22c0/run-6n8s22c0.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20260806_203759-m4g42ash/run-m4g42ash.wandb filter=lfs diff=lfs merge=lfs -text
+wandb/run-20260806_204221-wp360ttn/run-wp360ttn.wandb filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1 @@
+# Activation

exp.py

@@ -0,0 +1,1060 @@
+"""
+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,
+)
+# CHANGE 1: LlamaSdpaAttention/LlamaFlashAttention2 no longer exist as separate
+# classes in modern transformers — there is a single LlamaAttention that
+# dispatches to eager/sdpa/flash_attention_2 based on config._attn_implementation.
+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))
+
+
+# =============================================================================
+# 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 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
+
+        # Handle special activations first
+        if self.mlp_type == "glu" and self.activation_name == "situglu":
+            # SiTUGLU is a complete GLU block (has its own gate & up)
+            self.situglu = SiTUGLU(self.hidden_size, effective_intermediate)
+            self.gate_proj = None
+            self.up_proj = None
+            self.act_fn = None
+        else:
+            self.situglu = None
+            # For both GLU and MLP we need an activation function (except for situglu)
+            if self.activation_name == "situglu":
+                raise ValueError(
+                    f"Activation 'situglu' is only allowed with mlp_type='glu'. "
+                    f"Received mlp_type='{self.mlp_type}'."
+                )
+            self.act_fn = GLUActivationRegistry.get(self.activation_name)
+
+            if self.mlp_type == "glu":
+                # GLU: two parallel projections
+                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
+                # Standard MLP: only one projection before activation
+                self.gate_proj = None
+                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:
+            # SiTUGLU handles its own gate and up
+            return self.down_proj(self.situglu(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
+        # CHANGE 2: LlamaAttention picks its backend (eager/sdpa/flash_attention_2)
+        # from config._attn_implementation internally.
+        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,)
+
+
+# CHANGE 3: Replace float mask with Boolean mask (required for Flash Attention)
+def _build_causal_mask(
+    attention_mask: Optional[torch.Tensor],
+    seq_len: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
+    """
+    Build a 4D Boolean causal + padding mask for SDPA.
+    - True  -> masked out (do NOT attend)
+    - False -> allowed to attend
+    This is the format expected by torch.nn.functional.scaled_dot_product_attention
+    when using Flash Attention kernels.
+    """
+    # Causal mask: upper triangular (future tokens) = True
+    causal = torch.triu(
+        torch.ones((seq_len, seq_len), dtype=torch.bool, device=device),
+        diagonal=1
+    )
+    causal = causal[None, None, :, :]  # (1, 1, seq_len, seq_len)
+
+    if attention_mask is None:
+        return causal
+
+    batch_size = attention_mask.shape[0]
+    causal = causal.expand(batch_size, 1, seq_len, seq_len).clone()
+
+    # Padding: where attention_mask == 0 -> masked out (True)
+    padding_mask = attention_mask[:, None, None, :].to(device) == 0
+    causal = causal | padding_mask  # OR combines both masks
+
+    return causal
+
+
+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)
+        # RoPE: computed once per forward pass and threaded into every
+        # decoder layer's attention. Without this, LlamaAttention receives
+        # position_embeddings=None and crashes on `cos, sin = position_embeddings`.
+        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,
+    ):
+        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
+        # (cos, sin) tuple, shared across all layers this forward pass
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # CHANGE 3: flash_attention_2's kernel does not accept a dense 4D
+        # bool/float mask — it wants None (pure causal) or the raw 2D padding
+        # mask, and handles causality internally. Only build the dense 4D
+        # mask for eager/sdpa backends.
+        if self.config._attn_implementation == "flash_attention_2":
+            causal_mask = attention_mask
+        else:
+            seq_len = hidden_states.shape[1]
+            causal_mask = _build_causal_mask(
+                attention_mask, seq_len, hidden_states.dtype, hidden_states.device
+            )
+
+        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
+    # Tells save_pretrained() that lm_head.weight and model.embed_tokens.weight
+    # are the SAME tensor on purpose (see tie_word_embeddings below), so it
+    # doesn't error out when it finds two names pointing at shared memory.
+    # NOTE: must be a dict (maps tied key -> source key), not a list —
+    # transformers' _get_tied_weight_keys() calls .keys() on this.
+    _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",
+):
+    """Concatenate and chunk TinyStories for causal LM. Fast path with multiprocessing."""
+    ds = load_dataset(dataset_name, split=split)
+
+    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

h.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# Re-occupy GPU3: relaunch the JSL-Med-9B-v3.2 vLLM server (port 8010).
+set -uo pipefail
+SERVE_SCRIPT=/mnt/data/zainulabideen/aws/serve_jsl_9b_v3_2.sh
+LOG=/mnt/data/zainulabideen/Data/RV3/fine_tune/logs/serve_jsl_9b_v3_2_resume.log
+echo ">> checking GPU3 is free before starting..."
+USED=$(nvidia-smi -i 3 --query-gpu=memory.used --format=csv,noheader,nounits)
+if [ "$USED" -gt 5000 ]; then
+  echo ">> GPU3 has ${USED} MiB in use — someone else may be using it. Run free-gpu3.sh first if that's stale."
+  exit 1
+fi
+echo ">> launching JSL-Med-9B-v3.2 on GPU3, port 8010..."
+nohup bash "$SERVE_SCRIPT" > "$LOG" 2>&1 &
+disown
+echo ">> launched, PID $!"
+echo ">> waiting for health check..."
+for i in $(seq 1 40); do
+  r=$(curl -s -o /dev/null -w "%{http_code}" http://localhost:8010/v1/models -H "Authorization: Bearer jsl-1" 2>/dev/null)
+  if [ "$r" == "200" ]; then
+    echo ">> HEALTHY after ${i}0s"
+    exit 0
+  fi
+  sleep 10
+done
+echo ">> WARNING: not healthy after 400s, check $LOG"

llm_analyzer_wandb.py

@@ -0,0 +1,570 @@
+"""
+Multi-LLM Activation & Loss Analyzer with WandB Logging
+
+Evaluates multiple language models on WikiText dataset,
+computing per-tensor and global activation statistics (mean, max_abs, std, norm)
+and logging to Weights & Biases with separate runs per model.
+
+Requirements:
+    pip install transformers torch datasets wandb tqdm
+
+Usage:
+    export WANDB_PROJECT="llm-activation-analysis"
+    export WANDB_API_KEY="your-key"
+    python llm_analyzer_wandb.py
+"""
+
+import math
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+from datasets import load_dataset
+from typing import List, Dict, Optional, Union, Tuple
+from dataclasses import dataclass, asdict
+from collections import defaultdict
+import json
+import warnings
+import os
+from tqdm import tqdm
+
+import wandb
+
+warnings.filterwarnings("ignore")
+
+
+@dataclass
+class TensorStats:
+    mean: float
+    max_abs: float
+    std: float
+    norm: float
+    numel: int
+
+
+@dataclass
+class ModelResult:
+    model_name: str
+    loss: float
+    perplexity: float
+    global_act: TensorStats
+    layer_acts: Dict[str, TensorStats]
+    num_tokens: int
+    num_layers: int
+    hidden_size: int
+    num_params: int
+
+
+class ActivationHookManager:
+    """Manages forward hooks to capture activations from every tensor."""
+
+    def __init__(self):
+        self.activations = {}
+        self.hooks = []
+        # Set once per batch via set_attention_mask(); used to exclude
+        # padding-token positions from activation statistics.
+        self._attention_mask: Optional[torch.Tensor] = None
+
+    def set_attention_mask(self, attention_mask: Optional[torch.Tensor]):
+        """Call once per batch before the forward pass so hooks can mask
+        out padding positions when computing stats."""
+        self._attention_mask = (
+            attention_mask.detach().cpu() if attention_mask is not None else None
+        )
+
+    def _make_hook(self, name: str):
+        def hook(module, input, output):
+            # Handle different output types
+            if isinstance(output, torch.Tensor):
+                tensor = output
+            elif isinstance(output, tuple) and isinstance(output[0], torch.Tensor):
+                tensor = output[0]
+            else:
+                return
+
+            # Detach and move to CPU to avoid GPU memory blowup
+            self.activations[name] = tensor.detach().cpu().float()
+        return hook
+
+    def register_hooks(self, model: torch.nn.Module):
+        """Register hooks on all modules that produce activations."""
+        for name, module in model.named_modules():
+            # Skip trivial containers
+            if len(list(module.children())) == 0 and hasattr(module, 'forward'):
+                hook = module.register_forward_hook(self._make_hook(name))
+                self.hooks.append(hook)
+
+    def clear(self):
+        self.activations.clear()
+
+    def remove_hooks(self):
+        for hook in self.hooks:
+            hook.remove()
+        self.hooks.clear()
+
+    def _select_real_tokens(self, tensor: torch.Tensor) -> torch.Tensor:
+        """
+        BUG FIX: previously every captured tensor (including activations at
+        padding-token positions) was flattened and used as-is. With
+        padding_side="left" and a small batch_size, the padding fraction
+        varies a lot batch-to-batch, so padding-token activations (which
+        are real, non-zero values — not zeros) were silently mixed into
+        mean/std/max_abs/norm, biasing exactly the saturation signal this
+        script exists to measure.
+
+        Here we mask out padding positions whenever a tensor's shape is
+        consistent with (batch, seq_len, ...) against the stored
+        attention_mask (batch, seq_len). Tensors that don't match that
+        shape (e.g. a module operating on the pooled/final dimension only)
+        are left as-is rather than guessing.
+        """
+        mask = self._attention_mask
+        if mask is None or tensor.dim() < 2:
+            return tensor.reshape(-1)
+        if tensor.shape[0] != mask.shape[0] or tensor.shape[1] != mask.shape[1]:
+            return tensor.reshape(-1)
+
+        bool_mask = mask.bool()
+        # Expand mask across any trailing dims (e.g. hidden_size) and select.
+        expand_shape = bool_mask.shape + (1,) * (tensor.dim() - 2)
+        bool_mask = bool_mask.view(expand_shape).expand_as(tensor)
+        return tensor[bool_mask].reshape(-1)
+
+    def compute_stats(self) -> Dict[str, TensorStats]:
+        """Compute statistics for all captured activations, excluding
+        padding-token positions where identifiable."""
+        stats = {}
+        for name, tensor in self.activations.items():
+            if tensor.numel() == 0:
+                continue
+            flat = self._select_real_tokens(tensor)
+            if flat.numel() == 0:
+                continue
+            stats[name] = TensorStats(
+                mean=flat.mean().item(),
+                max_abs=flat.abs().max().item(),
+                std=flat.std().item(),
+                norm=flat.norm().item(),
+                numel=flat.numel()
+            )
+        return stats
+
+
+class LLMAnalyzer:
+    def __init__(
+        self,
+        device: Optional[str] = None,
+        max_length: int = 512,
+        max_samples: int = 1000,  # number of wikitext samples to eval
+        batch_size: int = 4,
+        dtype: torch.dtype = torch.float16,
+        wandb_project: Optional[str] = None,
+    ):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.max_length = max_length
+        self.max_samples = max_samples
+        self.batch_size = batch_size
+        self.dtype = dtype if self.device == "cuda" else torch.float32
+        self.wandb_project = wandb_project or os.environ.get("WANDB_PROJECT", "llm-activation-analysis")
+        self._cache = {}
+
+    def load_dataset(self, split: str = "test"):
+        """Load Salesforce/wikitext dataset."""
+        print(f"[Dataset] Loading Salesforce/wikitext ({split}) ...")
+        ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split=split)
+        # Filter out empty lines
+        texts = [t for t in ds["text"] if len(t.strip()) > 50]
+        print(f"[Dataset] Loaded {len(texts)} non-empty samples")
+        return texts[:self.max_samples]
+
+    def load_model(self, model_name: str):
+        """Load model and tokenizer with caching."""
+        if model_name in self._cache:
+            return self._cache[model_name]
+
+        print(f"[Loading] {model_name} ...")
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            padding_side="left"
+        )
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            config=config,
+            torch_dtype=self.dtype,
+            device_map="auto" if self.device == "cuda" else None,
+            trust_remote_code=True,
+        )
+
+        if self.device == "cpu":
+            model = model.to(self.device)
+
+        model.eval()
+
+        num_params = sum(p.numel() for p in model.parameters())
+
+        self._cache[model_name] = (tokenizer, model, config, num_params)
+        print(f"[Loaded]  {model_name} | Params: {num_params/1e6:.1f}M | Layers: {config.num_hidden_layers} | Hidden: {config.hidden_size}")
+        return tokenizer, model, config, num_params
+
+    def compute(
+        self,
+        model_names: List[str],
+    ) -> List[ModelResult]:
+        """
+        Compute loss and per-tensor activation statistics for multiple models.
+        Logs each model as a separate WandB run.
+        """
+        texts = self.load_dataset()
+        results = []
+
+        for model_name in model_names:
+            try:
+                result = self._evaluate_model(model_name, texts)
+                results.append(result)
+            except Exception as e:
+                print(f"[Error] {model_name}: {e}")
+                import traceback
+                traceback.print_exc()
+                continue
+
+        return results
+
+    def _evaluate_model(
+        self,
+        model_name: str,
+        texts: List[str],
+    ) -> ModelResult:
+        tokenizer, model, config, num_params = self.load_model(model_name)
+
+        # Initialize WandB run for this model
+        run_name = model_name.replace("/", "-")
+        wandb.init(
+            project=self.wandb_project,
+            name=run_name,
+            config={
+                "model": model_name,
+                "max_length": self.max_length,
+                "max_samples": self.max_samples,
+                "batch_size": self.batch_size,
+                "dtype": str(self.dtype),
+                "num_params": num_params,
+                "num_layers": config.num_hidden_layers,
+                "hidden_size": config.hidden_size,
+            },
+            reinit=True
+        )
+
+        hook_mgr = ActivationHookManager()
+        hook_mgr.register_hooks(model)
+
+        total_loss = 0.0
+        total_tokens = 0
+
+        # Global activation accumulator
+        global_acts = []
+
+        # Per-layer activation accumulators
+        # We'll aggregate stats across batches, then compute final stats
+        layer_act_values = defaultdict(list)
+
+        num_batches = (len(texts) + self.batch_size - 1) // self.batch_size
+
+        for i in tqdm(range(0, len(texts), self.batch_size), desc=f"Eval {run_name}", total=num_batches):
+            batch_texts = texts[i:i + self.batch_size]
+
+            # Tokenize
+            inputs = tokenizer(
+                batch_texts,
+                return_tensors="pt",
+                truncation=True,
+                max_length=self.max_length,
+                padding=True
+            )
+
+            # Move to device
+            if self.device == "cuda" and hasattr(model, "device"):
+                # model is on auto device map
+                input_ids = inputs["input_ids"]
+                if hasattr(model, "device") and model.device != torch.device("meta"):
+                    input_ids = input_ids.to(model.device)
+                attention_mask = inputs.get("attention_mask")
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(input_ids.device)
+            else:
+                input_ids = inputs["input_ids"].to(self.device)
+                attention_mask = inputs.get("attention_mask")
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(self.device)
+
+            labels = input_ids.clone()
+
+            # BUG FIX: with padding_side="left" and no explicit position_ids,
+            # HF's default `position_ids = arange(seq_len)` is applied
+            # identically to every row in the batch, regardless of how much
+            # left-padding precedes the real tokens in that row (verified
+            # against transformers' LlamaModel.forward / GPT2Model.forward
+            # source — neither adjusts for padding when position_ids=None).
+            # That means a real token's absolute position (and therefore its
+            # RoPE rotation / absolute position embedding) depends on how
+            # much padding happened to precede it in this particular batch,
+            # not on its logical position within its own sequence. This
+            # silently corrupts logits -> loss -> perplexity, with the
+            # amount of corruption varying batch-to-batch. Fix: derive
+            # position_ids from attention_mask so they restart at 0 for the
+            # first real token of every row, and are stable (0) on padding.
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 0)
+            else:
+                position_ids = None
+
+            hook_mgr.set_attention_mask(attention_mask)
+
+            with torch.no_grad():
+                outputs = model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    labels=labels,
+                )
+
+            # --- Loss Computation ---
+            logits = outputs.logits
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            shift_mask = attention_mask[..., 1:].contiguous() if attention_mask is not None else None
+
+            loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+            token_losses = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+            if shift_mask is not None:
+                token_losses = token_losses * shift_mask.view(-1)
+                num_valid_tokens = shift_mask.sum().item()
+            else:
+                num_valid_tokens = token_losses.numel()
+
+            batch_loss = token_losses.sum().item()
+            total_loss += batch_loss
+            total_tokens += num_valid_tokens
+
+            # --- Activation Statistics ---
+            # Get activations captured by hooks
+            act_stats = hook_mgr.compute_stats()
+
+            for name, stats in act_stats.items():
+                # Skip dtype-limit fraction metrics entirely
+                # We only log mean, max_abs, std, norm
+
+                # For global: aggregate raw values
+                # We can't store all raw values due to memory, so we store running sums
+                # But for accurate std across all batches, we need a streaming algorithm
+                # For simplicity and correctness, we'll store per-batch stats and weight them
+                layer_act_values[name].append(asdict(stats))
+
+            hook_mgr.clear()
+
+            # Log per-batch metrics to wandb
+            if num_valid_tokens > 0:
+                batch_avg_loss = batch_loss / num_valid_tokens
+                wandb.log({
+                    "batch_loss": batch_avg_loss,
+                    "batch_perplexity": torch.exp(torch.tensor(batch_avg_loss)).item(),
+                    "batch_tokens": num_valid_tokens,
+                    "progress": i / len(texts)
+                }, step=i)
+
+        hook_mgr.remove_hooks()
+
+        # --- Final Aggregation ---
+        avg_loss = total_loss / max(total_tokens, 1)
+        perplexity = torch.exp(torch.tensor(avg_loss)).item()
+
+        # Aggregate per-tensor stats across all batches
+        # Weighted by numel for mean, max for max_abs, pooled std, pooled norm
+        final_layer_stats = {}
+
+        for name, batch_stats_list in layer_act_values.items():
+            total_numel = sum(s["numel"] for s in batch_stats_list)
+            if total_numel == 0:
+                continue
+
+            # Weighted mean
+            weighted_mean = sum(s["mean"] * s["numel"] for s in batch_stats_list) / total_numel
+
+            # Max abs across all batches
+            max_abs = max(s["max_abs"] for s in batch_stats_list)
+
+            # BUG FIX: the previous formula (weighted average of per-batch
+            # variances only) drops the between-batch term that accounts for
+            # per-batch means differing from the global mean. Whenever batch
+            # means differ (they will — different texts, different lengths),
+            # this systematically UNDERESTIMATES the true global std — in a
+            # quick numeric test with two batches of different means this was
+            # off by ~2.8x. Correct pooled-variance formula (population form,
+            # matches exp.py's StepAccumulator._merge_entry):
+            #   E[X^2] = weighted_avg(var_i + mean_i^2)
+            #   Var(X) = E[X^2] - mean_global^2
+            ex2 = sum(
+                s["numel"] * (s["std"] ** 2 + s["mean"] ** 2) for s in batch_stats_list
+            ) / total_numel
+            pooled_std = math.sqrt(max(0.0, ex2 - weighted_mean ** 2))
+
+            # Norm: sqrt(sum of squared norms / total_numel) * sqrt(total_numel)
+            # Actually norm^2 = sum(x_i^2), so pooled_norm = sqrt(sum(norm_i^2))
+            pooled_norm = (sum(s["norm"] ** 2 for s in batch_stats_list)) ** 0.5
+
+            final_layer_stats[name] = TensorStats(
+                mean=weighted_mean,
+                max_abs=max_abs,
+                std=pooled_std,
+                norm=pooled_norm,
+                numel=total_numel
+            )
+
+        # Compute global stats across all layers
+        if final_layer_stats:
+            all_numel = sum(s.numel for s in final_layer_stats.values())
+            global_mean = sum(s.mean * s.numel for s in final_layer_stats.values()) / all_numel
+            global_max_abs = max(s.max_abs for s in final_layer_stats.values())
+            # Same pooled-std correction as above — same bug was present here.
+            global_ex2 = sum(
+                s.numel * (s.std ** 2 + s.mean ** 2) for s in final_layer_stats.values()
+            ) / all_numel
+            global_std = math.sqrt(max(0.0, global_ex2 - global_mean ** 2))
+            global_norm = (sum(s.norm ** 2 for s in final_layer_stats.values())) ** 0.5
+
+            global_stats = TensorStats(
+                mean=global_mean,
+                max_abs=global_max_abs,
+                std=global_std,
+                norm=global_norm,
+                numel=all_numel
+            )
+        else:
+            global_stats = TensorStats(0.0, 0.0, 0.0, 0.0, 0)
+
+        result = ModelResult(
+            model_name=model_name,
+            loss=avg_loss,
+            perplexity=perplexity,
+            global_act=global_stats,
+            layer_acts=final_layer_stats,
+            num_tokens=total_tokens,
+            num_layers=config.num_hidden_layers,
+            hidden_size=config.hidden_size,
+            num_params=num_params
+        )
+
+        # --- WandB Logging ---
+        self._log_to_wandb(result)
+        wandb.finish()
+
+        return result
+
+    def _log_to_wandb(self, result: ModelResult):
+        """Log final metrics to WandB. No frac_near_dtype_limit."""
+
+        # Global metrics
+        wandb.log({
+            "final/loss": result.loss,
+            "final/perplexity": result.perplexity,
+            "final/num_tokens": result.num_tokens,
+
+            "train/global/act/mean": result.global_act.mean,
+            "train/global/act/max_abs": result.global_act.max_abs,
+            "train/global/act/std": result.global_act.std,
+            "train/global/act/norm": result.global_act.norm,
+            # Intentionally NOT logging frac_near_dtype_limit or frac_near_user_limit
+        })
+
+        # Per-tensor (per-layer) metrics
+        # Organize by layer for cleaner WandB UI
+        for tensor_name, stats in result.layer_acts.items():
+            # Clean name for wandb: replace dots with slashes
+            clean_name = tensor_name.replace(".", "/")
+
+            wandb.log({
+                f"train/{clean_name}/act/mean": stats.mean,
+                f"train/{clean_name}/act/max_abs": stats.max_abs,
+                f"train/{clean_name}/act/std": stats.std,
+                f"train/{clean_name}/act/norm": stats.norm,
+                # No frac_near_dtype_limit
+            })
+
+        # Also log as a wandb.Table for easy comparison
+        table_data = []
+        for tensor_name, stats in sorted(result.layer_acts.items()):
+            table_data.append([
+                tensor_name,
+                stats.mean,
+                stats.max_abs,
+                stats.std,
+                stats.norm,
+                stats.numel
+            ])
+
+        if table_data:
+            table = wandb.Table(
+                columns=["tensor_name", "mean", "max_abs", "std", "norm", "numel"],
+                data=table_data
+            )
+            wandb.log({"activation_table": table})
+
+    def print_report(self, results: List[ModelResult]):
+        """Pretty-print comparison report."""
+        print("\n" + "=" * 110)
+        print(f"{'Model':<35} {'Loss':>10} {'PPL':>10} {'ActMean':>12} {'ActMaxAbs':>12} {'ActStd':>12} {'Tokens':>8}")
+        print("-" * 110)
+
+        for r in results:
+            name = r.model_name.split("/")[-1][:33]
+            print(
+                f"{name:<35} "
+                f"{r.loss:>10.4f} "
+                f"{r.perplexity:>10.2f} "
+                f"{r.global_act.mean:>12.6f} "
+                f"{r.global_act.max_abs:>12.6f} "
+                f"{r.global_act.std:>12.6f} "
+                f"{r.num_tokens:>8}"
+            )
+
+        print("=" * 110)
+
+        # Print top 5 layers by max_abs for each model
+        print("\n[Per-Tensor Max Abs Top 5]")
+        for r in results:
+            name = r.model_name.split("/")[-1]
+            sorted_layers = sorted(r.layer_acts.items(), key=lambda x: x[1].max_abs, reverse=True)[:5]
+            print(f"\n  {name}:")
+            for tensor_name, stats in sorted_layers:
+                print(f"    {tensor_name:<50} max_abs={stats.max_abs:>10.4f}  mean={stats.mean:>10.6f}  std={stats.std:>10.4f}")
+
+    def export_json(self, results: List[ModelResult], path: str):
+        """Export results to JSON."""
+        data = []
+        for r in results:
+            entry = {
+                "model": r.model_name,
+                "loss": r.loss,
+                "perplexity": r.perplexity,
+                "num_tokens": r.num_tokens,
+                "num_layers": r.num_layers,
+                "hidden_size": r.hidden_size,
+                "num_params": r.num_params,
+                "global_act": asdict(r.global_act),
+                "layer_acts": {k: asdict(v) for k, v in r.layer_acts.items()}
+            }
+            data.append(entry)
+
+        with open(path, "w") as f:
+            json.dump(data, f, indent=2)
+        print(f"[Exported] Results saved to {path}")

out/glu-gelu-9L_run/README.md

@@ -0,0 +1,65 @@
+---
+library_name: transformers
+tags:
+- generated_from_trainer
+model-index:
+- name: ACTio-glu-gelu-9L
+  results: []
+---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+# ACTio-glu-gelu-9L
+
+This model is a fine-tuned version of [](https://huggingface.co/) on an unknown dataset.
+It achieves the following results on the evaluation set:
+- Loss: 0.2151
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 0.0003
+- train_batch_size: 64
+- eval_batch_size: 64
+- seed: 42
+- optimizer: Use OptimizerNames.ADAMW_TORCH with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: constant
+- training_steps: 2000
+
+### Training results
+
+| Training Loss | Epoch  | Step | Validation Loss |
+|:-------------:|:------:|:----:|:---------------:|
+| 3.3152        | 0.0135 | 200  | 3.1407          |
+| 0.8976        | 0.0270 | 400  | 0.8561          |
+| 0.4081        | 0.0404 | 600  | 0.4029          |
+| 0.3257        | 0.0539 | 800  | 0.3262          |
+| 0.2851        | 0.0674 | 1000 | 0.2846          |
+| 0.2629        | 0.0809 | 1200 | 0.2618          |
+| 0.2410        | 0.0944 | 1400 | 0.2429          |
+| 0.2452        | 0.1079 | 1600 | 0.2447          |
+| 0.2224        | 0.1213 | 1800 | 0.2213          |
+| 0.2135        | 0.1348 | 2000 | 0.2151          |
+
+
+### Framework versions
+
+- Transformers 5.15.0.dev0
+- Pytorch 2.6.0+cu124
+- Datasets 5.0.1
+- Tokenizers 0.22.2

out/glu-gelu-9L_run/checkpoint-1000/config.json

@@ -0,0 +1,35 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "tokenizer_name": "w-ahmad/tiny-stories-tokenizer",
+  "transformers_version": "5.15.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096
+}

out/glu-gelu-9L_run/checkpoint-1000/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e35b0e11a46d87fbedb0dbd610a7d8df1fe3ebbc9c047856d41f4791f8532b2
+size 4011496

out/glu-gelu-9L_run/checkpoint-1000/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10ec2f99857c3ef4d84564c7d553e90d3577d6f590e656cb57f71830f27ecc07
+size 8074746

out/glu-gelu-9L_run/checkpoint-1000/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95b6047bd8cc6f4cdf7c46dea47edb8e542435510070c6cd1e0a7d9ccf5fd7da
+size 14244

out/glu-gelu-9L_run/checkpoint-1000/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5f71c5fa69d56807b14256b78700be6088fcb1bda0ff8e0ea6f8709a7df2fde
+size 1064

out/glu-gelu-9L_run/checkpoint-1000/tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

out/glu-gelu-9L_run/checkpoint-1000/trainer_state.json

@@ -0,0 +1,424 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.06740815638692282,
+  "eval_steps": 200,
+  "global_step": 1000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.2578125,
+      "learning_rate": 0.0003,
+      "loss": 7.793476104736328,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 7.021186828613281,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.0,
+      "learning_rate": 0.0003,
+      "loss": 6.446919250488281,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.4609375,
+      "learning_rate": 0.0003,
+      "loss": 5.989651489257812,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 2.328125,
+      "learning_rate": 0.0003,
+      "loss": 5.436912536621094,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 2.671875,
+      "learning_rate": 0.0003,
+      "loss": 4.910789871215821,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 3.28125,
+      "learning_rate": 0.0003,
+      "loss": 4.441200637817383,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 2.625,
+      "learning_rate": 0.0003,
+      "loss": 4.035322952270508,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 3.765625,
+      "learning_rate": 0.0003,
+      "loss": 3.6691635131835936,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 2.296875,
+      "learning_rate": 0.0003,
+      "loss": 3.315158462524414,
+      "step": 200
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.140667200088501,
+      "eval_runtime": 8.3121,
+      "eval_samples_per_second": 1146.166,
+      "eval_steps_per_second": 17.926,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 2.9744741439819338,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 2.359375,
+      "learning_rate": 0.0003,
+      "loss": 2.6262161254882814,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 2.234375,
+      "learning_rate": 0.0003,
+      "loss": 2.3146018981933594,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 2.015625,
+      "learning_rate": 0.0003,
+      "loss": 1.9927532196044921,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.609375,
+      "learning_rate": 0.0003,
+      "loss": 1.6939159393310548,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.734375,
+      "learning_rate": 0.0003,
+      "loss": 1.4503639221191407,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 1.828125,
+      "learning_rate": 0.0003,
+      "loss": 1.2639431953430176,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 1.984375,
+      "learning_rate": 0.0003,
+      "loss": 1.1161455154418944,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 1.546875,
+      "learning_rate": 0.0003,
+      "loss": 0.9919822692871094,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 1.1484375,
+      "learning_rate": 0.0003,
+      "loss": 0.8975520133972168,
+      "step": 400
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 0.8561499714851379,
+      "eval_runtime": 8.3662,
+      "eval_samples_per_second": 1138.75,
+      "eval_steps_per_second": 17.81,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0003,
+      "loss": 0.7963034629821777,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.51171875,
+      "learning_rate": 0.0003,
+      "loss": 0.7076507568359375,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.439453125,
+      "learning_rate": 0.0003,
+      "loss": 0.6221071720123291,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.423828125,
+      "learning_rate": 0.0003,
+      "loss": 0.5685368537902832,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.5213922977447509,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.380859375,
+      "learning_rate": 0.0003,
+      "loss": 0.4871851444244385,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.36328125,
+      "learning_rate": 0.0003,
+      "loss": 0.4552904605865479,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0003,
+      "loss": 0.4377324104309082,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.400390625,
+      "learning_rate": 0.0003,
+      "loss": 0.4190248966217041,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.392578125,
+      "learning_rate": 0.0003,
+      "loss": 0.40812134742736816,
+      "step": 600
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "eval_loss": 0.40287962555885315,
+      "eval_runtime": 8.3848,
+      "eval_samples_per_second": 1136.228,
+      "eval_steps_per_second": 17.77,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0003,
+      "loss": 0.39397249221801756,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0003,
+      "loss": 0.38487560749053956,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.376953125,
+      "learning_rate": 0.0003,
+      "loss": 0.37073094844818116,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 0.36380269527435305,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 0.3727259159088135,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.44921875,
+      "learning_rate": 0.0003,
+      "loss": 0.36068120002746584,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.3359375,
+      "learning_rate": 0.0003,
+      "loss": 0.345484185218811,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.271484375,
+      "learning_rate": 0.0003,
+      "loss": 0.3413747549057007,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.279296875,
+      "learning_rate": 0.0003,
+      "loss": 0.3339890956878662,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.27734375,
+      "learning_rate": 0.0003,
+      "loss": 0.32567362785339354,
+      "step": 800
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "eval_loss": 0.32622790336608887,
+      "eval_runtime": 8.4003,
+      "eval_samples_per_second": 1134.122,
+      "eval_steps_per_second": 17.737,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.275390625,
+      "learning_rate": 0.0003,
+      "loss": 0.3216629266738892,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.291015625,
+      "learning_rate": 0.0003,
+      "loss": 0.3161738157272339,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.26171875,
+      "learning_rate": 0.0003,
+      "loss": 0.30926101207733153,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.27734375,
+      "learning_rate": 0.0003,
+      "loss": 0.30281963348388674,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.287109375,
+      "learning_rate": 0.0003,
+      "loss": 0.301728343963623,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.29296875,
+      "learning_rate": 0.0003,
+      "loss": 0.2986367464065552,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.306640625,
+      "learning_rate": 0.0003,
+      "loss": 0.294226336479187,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.33984375,
+      "learning_rate": 0.0003,
+      "loss": 0.28919248580932616,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.296875,
+      "learning_rate": 0.0003,
+      "loss": 0.2873764276504517,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.322265625,
+      "learning_rate": 0.0003,
+      "loss": 0.2850848436355591,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "eval_loss": 0.284612774848938,
+      "eval_runtime": 8.55,
+      "eval_samples_per_second": 1114.271,
+      "eval_steps_per_second": 17.427,
+      "step": 1000
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 290388443136000.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

out/glu-gelu-9L_run/checkpoint-1000/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5bc9b69db12445b4a7919b29e76d261b3c2acbea92402190b4ca5d146c69daa
+size 4920

out/glu-gelu-9L_run/checkpoint-1500/config.json

@@ -0,0 +1,35 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "tokenizer_name": "w-ahmad/tiny-stories-tokenizer",
+  "transformers_version": "5.15.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096
+}

out/glu-gelu-9L_run/checkpoint-1500/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02945cf70fed9c6e8088853f25a2118cdd716fe4cd30efc4b3fbc33b4554a983
+size 4011496

out/glu-gelu-9L_run/checkpoint-1500/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62950e32ae8361a13306e433bbe68d89378073f7f6089d0063c7023e0f942526
+size 8074746

out/glu-gelu-9L_run/checkpoint-1500/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2768285b45b2a0c05f6f50bbb8c0287fca6f62a8cde6d1b1f02151ac72ee8dc
+size 14244

out/glu-gelu-9L_run/checkpoint-1500/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b26a759cc207e28dea06eae7ad531c2dd78a09978f6972edbfc1566adccd7142
+size 1064

out/glu-gelu-9L_run/checkpoint-1500/tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

out/glu-gelu-9L_run/checkpoint-1500/trainer_state.json

@@ -0,0 +1,615 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.10111223458038422,
+  "eval_steps": 200,
+  "global_step": 1500,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.2578125,
+      "learning_rate": 0.0003,
+      "loss": 7.793476104736328,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 7.021186828613281,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.0,
+      "learning_rate": 0.0003,
+      "loss": 6.446919250488281,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.4609375,
+      "learning_rate": 0.0003,
+      "loss": 5.989651489257812,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 2.328125,
+      "learning_rate": 0.0003,
+      "loss": 5.436912536621094,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 2.671875,
+      "learning_rate": 0.0003,
+      "loss": 4.910789871215821,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 3.28125,
+      "learning_rate": 0.0003,
+      "loss": 4.441200637817383,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 2.625,
+      "learning_rate": 0.0003,
+      "loss": 4.035322952270508,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 3.765625,
+      "learning_rate": 0.0003,
+      "loss": 3.6691635131835936,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 2.296875,
+      "learning_rate": 0.0003,
+      "loss": 3.315158462524414,
+      "step": 200
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.140667200088501,
+      "eval_runtime": 8.3121,
+      "eval_samples_per_second": 1146.166,
+      "eval_steps_per_second": 17.926,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 2.9744741439819338,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 2.359375,
+      "learning_rate": 0.0003,
+      "loss": 2.6262161254882814,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 2.234375,
+      "learning_rate": 0.0003,
+      "loss": 2.3146018981933594,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 2.015625,
+      "learning_rate": 0.0003,
+      "loss": 1.9927532196044921,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.609375,
+      "learning_rate": 0.0003,
+      "loss": 1.6939159393310548,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.734375,
+      "learning_rate": 0.0003,
+      "loss": 1.4503639221191407,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 1.828125,
+      "learning_rate": 0.0003,
+      "loss": 1.2639431953430176,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 1.984375,
+      "learning_rate": 0.0003,
+      "loss": 1.1161455154418944,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 1.546875,
+      "learning_rate": 0.0003,
+      "loss": 0.9919822692871094,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 1.1484375,
+      "learning_rate": 0.0003,
+      "loss": 0.8975520133972168,
+      "step": 400
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 0.8561499714851379,
+      "eval_runtime": 8.3662,
+      "eval_samples_per_second": 1138.75,
+      "eval_steps_per_second": 17.81,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0003,
+      "loss": 0.7963034629821777,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.51171875,
+      "learning_rate": 0.0003,
+      "loss": 0.7076507568359375,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.439453125,
+      "learning_rate": 0.0003,
+      "loss": 0.6221071720123291,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.423828125,
+      "learning_rate": 0.0003,
+      "loss": 0.5685368537902832,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.5213922977447509,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.380859375,
+      "learning_rate": 0.0003,
+      "loss": 0.4871851444244385,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.36328125,
+      "learning_rate": 0.0003,
+      "loss": 0.4552904605865479,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0003,
+      "loss": 0.4377324104309082,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.400390625,
+      "learning_rate": 0.0003,
+      "loss": 0.4190248966217041,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.392578125,
+      "learning_rate": 0.0003,
+      "loss": 0.40812134742736816,
+      "step": 600
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "eval_loss": 0.40287962555885315,
+      "eval_runtime": 8.3848,
+      "eval_samples_per_second": 1136.228,
+      "eval_steps_per_second": 17.77,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0003,
+      "loss": 0.39397249221801756,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0003,
+      "loss": 0.38487560749053956,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.376953125,
+      "learning_rate": 0.0003,
+      "loss": 0.37073094844818116,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 0.36380269527435305,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 0.3727259159088135,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.44921875,
+      "learning_rate": 0.0003,
+      "loss": 0.36068120002746584,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.3359375,
+      "learning_rate": 0.0003,
+      "loss": 0.345484185218811,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.271484375,
+      "learning_rate": 0.0003,
+      "loss": 0.3413747549057007,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.279296875,
+      "learning_rate": 0.0003,
+      "loss": 0.3339890956878662,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.27734375,
+      "learning_rate": 0.0003,
+      "loss": 0.32567362785339354,
+      "step": 800
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "eval_loss": 0.32622790336608887,
+      "eval_runtime": 8.4003,
+      "eval_samples_per_second": 1134.122,
+      "eval_steps_per_second": 17.737,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.275390625,
+      "learning_rate": 0.0003,
+      "loss": 0.3216629266738892,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.291015625,
+      "learning_rate": 0.0003,
+      "loss": 0.3161738157272339,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.26171875,
+      "learning_rate": 0.0003,
+      "loss": 0.30926101207733153,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.27734375,
+      "learning_rate": 0.0003,
+      "loss": 0.30281963348388674,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.287109375,
+      "learning_rate": 0.0003,
+      "loss": 0.301728343963623,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.29296875,
+      "learning_rate": 0.0003,
+      "loss": 0.2986367464065552,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.306640625,
+      "learning_rate": 0.0003,
+      "loss": 0.294226336479187,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.33984375,
+      "learning_rate": 0.0003,
+      "loss": 0.28919248580932616,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.296875,
+      "learning_rate": 0.0003,
+      "loss": 0.2873764276504517,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.322265625,
+      "learning_rate": 0.0003,
+      "loss": 0.2850848436355591,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "eval_loss": 0.284612774848938,
+      "eval_runtime": 8.55,
+      "eval_samples_per_second": 1114.271,
+      "eval_steps_per_second": 17.427,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06875631951466127,
+      "grad_norm": 0.357421875,
+      "learning_rate": 0.0003,
+      "loss": 0.2803910255432129,
+      "step": 1020
+    },
+    {
+      "epoch": 0.07010448264239973,
+      "grad_norm": 0.37890625,
+      "learning_rate": 0.0003,
+      "loss": 0.27601516246795654,
+      "step": 1040
+    },
+    {
+      "epoch": 0.07145264577013818,
+      "grad_norm": 0.3203125,
+      "learning_rate": 0.0003,
+      "loss": 0.27540276050567625,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07280080889787664,
+      "grad_norm": 0.318359375,
+      "learning_rate": 0.0003,
+      "loss": 0.2708603858947754,
+      "step": 1080
+    },
+    {
+      "epoch": 0.0741489720256151,
+      "grad_norm": 0.328125,
+      "learning_rate": 0.0003,
+      "loss": 0.2696371555328369,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07549713515335356,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.26611151695251467,
+      "step": 1120
+    },
+    {
+      "epoch": 0.07684529828109202,
+      "grad_norm": 1.828125,
+      "learning_rate": 0.0003,
+      "loss": 0.27770705223083497,
+      "step": 1140
+    },
+    {
+      "epoch": 0.07819346140883048,
+      "grad_norm": 0.5390625,
+      "learning_rate": 0.0003,
+      "loss": 0.28163862228393555,
+      "step": 1160
+    },
+    {
+      "epoch": 0.07954162453656892,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.2673502445220947,
+      "step": 1180
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "grad_norm": 0.275390625,
+      "learning_rate": 0.0003,
+      "loss": 0.262909460067749,
+      "step": 1200
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "eval_loss": 0.26177507638931274,
+      "eval_runtime": 8.3534,
+      "eval_samples_per_second": 1140.5,
+      "eval_steps_per_second": 17.837,
+      "step": 1200
+    },
+    {
+      "epoch": 0.08223795079204584,
+      "grad_norm": 0.26953125,
+      "learning_rate": 0.0003,
+      "loss": 0.2581258535385132,
+      "step": 1220
+    },
+    {
+      "epoch": 0.0835861139197843,
+      "grad_norm": 0.26171875,
+      "learning_rate": 0.0003,
+      "loss": 0.25597529411315917,
+      "step": 1240
+    },
+    {
+      "epoch": 0.08493427704752275,
+      "grad_norm": 0.2578125,
+      "learning_rate": 0.0003,
+      "loss": 0.25467369556427,
+      "step": 1260
+    },
+    {
+      "epoch": 0.08628244017526121,
+      "grad_norm": 0.25,
+      "learning_rate": 0.0003,
+      "loss": 0.2502788543701172,
+      "step": 1280
+    },
+    {
+      "epoch": 0.08763060330299967,
+      "grad_norm": 0.259765625,
+      "learning_rate": 0.0003,
+      "loss": 0.25013630390167235,
+      "step": 1300
+    },
+    {
+      "epoch": 0.08897876643073811,
+      "grad_norm": 0.2470703125,
+      "learning_rate": 0.0003,
+      "loss": 0.24711339473724364,
+      "step": 1320
+    },
+    {
+      "epoch": 0.09032692955847657,
+      "grad_norm": 0.25,
+      "learning_rate": 0.0003,
+      "loss": 0.2463146924972534,
+      "step": 1340
+    },
+    {
+      "epoch": 0.09167509268621503,
+      "grad_norm": 0.32421875,
+      "learning_rate": 0.0003,
+      "loss": 0.2437490701675415,
+      "step": 1360
+    },
+    {
+      "epoch": 0.09302325581395349,
+      "grad_norm": 0.259765625,
+      "learning_rate": 0.0003,
+      "loss": 0.2432482957839966,
+      "step": 1380
+    },
+    {
+      "epoch": 0.09437141894169195,
+      "grad_norm": 0.265625,
+      "learning_rate": 0.0003,
+      "loss": 0.24097964763641358,
+      "step": 1400
+    },
+    {
+      "epoch": 0.09437141894169195,
+      "eval_loss": 0.24285024404525757,
+      "eval_runtime": 8.4015,
+      "eval_samples_per_second": 1133.969,
+      "eval_steps_per_second": 17.735,
+      "step": 1400
+    },
+    {
+      "epoch": 0.0957195820694304,
+      "grad_norm": 0.265625,
+      "learning_rate": 0.0003,
+      "loss": 0.2406073570251465,
+      "step": 1420
+    },
+    {
+      "epoch": 0.09706774519716886,
+      "grad_norm": 0.255859375,
+      "learning_rate": 0.0003,
+      "loss": 0.23866963386535645,
+      "step": 1440
+    },
+    {
+      "epoch": 0.09841590832490732,
+      "grad_norm": 0.267578125,
+      "learning_rate": 0.0003,
+      "loss": 0.2364635467529297,
+      "step": 1460
+    },
+    {
+      "epoch": 0.09976407145264576,
+      "grad_norm": 0.2578125,
+      "learning_rate": 0.0003,
+      "loss": 0.23464763164520264,
+      "step": 1480
+    },
+    {
+      "epoch": 0.10111223458038422,
+      "grad_norm": 0.24609375,
+      "learning_rate": 0.0003,
+      "loss": 0.23385767936706542,
+      "step": 1500
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 435582664704000.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

out/glu-gelu-9L_run/checkpoint-1500/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5bc9b69db12445b4a7919b29e76d261b3c2acbea92402190b4ca5d146c69daa
+size 4920

out/glu-gelu-9L_run/checkpoint-2000/config.json

@@ -0,0 +1,35 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "tokenizer_name": "w-ahmad/tiny-stories-tokenizer",
+  "transformers_version": "5.15.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096
+}

out/glu-gelu-9L_run/checkpoint-2000/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efa7bcb247c2e6fbd4d32ab423b71b076c710576b4833e9258447d136870faaf
+size 4011496

out/glu-gelu-9L_run/checkpoint-2000/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1d8ce91bbbfc275aa72005f5fd5506f144e86ade8681b0ecd59995238f0408b
+size 8074746

out/glu-gelu-9L_run/checkpoint-2000/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1a97db8e41139aa1239ba7fb79ddeb0af5998c6305a440c1fe182e6ad02f2f5
+size 14244

out/glu-gelu-9L_run/checkpoint-2000/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e91519ad666af6acc4beafd44aa617aefb90f632388a0f75cbf88f957400c190
+size 1064

out/glu-gelu-9L_run/checkpoint-2000/tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

out/glu-gelu-9L_run/checkpoint-2000/trainer_state.json

@@ -0,0 +1,814 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.13481631277384565,
+  "eval_steps": 200,
+  "global_step": 2000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.2578125,
+      "learning_rate": 0.0003,
+      "loss": 7.793476104736328,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 7.021186828613281,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.0,
+      "learning_rate": 0.0003,
+      "loss": 6.446919250488281,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.4609375,
+      "learning_rate": 0.0003,
+      "loss": 5.989651489257812,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 2.328125,
+      "learning_rate": 0.0003,
+      "loss": 5.436912536621094,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 2.671875,
+      "learning_rate": 0.0003,
+      "loss": 4.910789871215821,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 3.28125,
+      "learning_rate": 0.0003,
+      "loss": 4.441200637817383,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 2.625,
+      "learning_rate": 0.0003,
+      "loss": 4.035322952270508,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 3.765625,
+      "learning_rate": 0.0003,
+      "loss": 3.6691635131835936,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 2.296875,
+      "learning_rate": 0.0003,
+      "loss": 3.315158462524414,
+      "step": 200
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.140667200088501,
+      "eval_runtime": 8.3121,
+      "eval_samples_per_second": 1146.166,
+      "eval_steps_per_second": 17.926,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 2.9744741439819338,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 2.359375,
+      "learning_rate": 0.0003,
+      "loss": 2.6262161254882814,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 2.234375,
+      "learning_rate": 0.0003,
+      "loss": 2.3146018981933594,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 2.015625,
+      "learning_rate": 0.0003,
+      "loss": 1.9927532196044921,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.609375,
+      "learning_rate": 0.0003,
+      "loss": 1.6939159393310548,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.734375,
+      "learning_rate": 0.0003,
+      "loss": 1.4503639221191407,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 1.828125,
+      "learning_rate": 0.0003,
+      "loss": 1.2639431953430176,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 1.984375,
+      "learning_rate": 0.0003,
+      "loss": 1.1161455154418944,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 1.546875,
+      "learning_rate": 0.0003,
+      "loss": 0.9919822692871094,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 1.1484375,
+      "learning_rate": 0.0003,
+      "loss": 0.8975520133972168,
+      "step": 400
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 0.8561499714851379,
+      "eval_runtime": 8.3662,
+      "eval_samples_per_second": 1138.75,
+      "eval_steps_per_second": 17.81,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0003,
+      "loss": 0.7963034629821777,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.51171875,
+      "learning_rate": 0.0003,
+      "loss": 0.7076507568359375,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.439453125,
+      "learning_rate": 0.0003,
+      "loss": 0.6221071720123291,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.423828125,
+      "learning_rate": 0.0003,
+      "loss": 0.5685368537902832,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.5213922977447509,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.380859375,
+      "learning_rate": 0.0003,
+      "loss": 0.4871851444244385,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.36328125,
+      "learning_rate": 0.0003,
+      "loss": 0.4552904605865479,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0003,
+      "loss": 0.4377324104309082,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.400390625,
+      "learning_rate": 0.0003,
+      "loss": 0.4190248966217041,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.392578125,
+      "learning_rate": 0.0003,
+      "loss": 0.40812134742736816,
+      "step": 600
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "eval_loss": 0.40287962555885315,
+      "eval_runtime": 8.3848,
+      "eval_samples_per_second": 1136.228,
+      "eval_steps_per_second": 17.77,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0003,
+      "loss": 0.39397249221801756,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0003,
+      "loss": 0.38487560749053956,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.376953125,
+      "learning_rate": 0.0003,
+      "loss": 0.37073094844818116,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 0.36380269527435305,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 0.3727259159088135,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.44921875,
+      "learning_rate": 0.0003,
+      "loss": 0.36068120002746584,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.3359375,
+      "learning_rate": 0.0003,
+      "loss": 0.345484185218811,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.271484375,
+      "learning_rate": 0.0003,
+      "loss": 0.3413747549057007,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.279296875,
+      "learning_rate": 0.0003,
+      "loss": 0.3339890956878662,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.27734375,
+      "learning_rate": 0.0003,
+      "loss": 0.32567362785339354,
+      "step": 800
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "eval_loss": 0.32622790336608887,
+      "eval_runtime": 8.4003,
+      "eval_samples_per_second": 1134.122,
+      "eval_steps_per_second": 17.737,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.275390625,
+      "learning_rate": 0.0003,
+      "loss": 0.3216629266738892,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.291015625,
+      "learning_rate": 0.0003,
+      "loss": 0.3161738157272339,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.26171875,
+      "learning_rate": 0.0003,
+      "loss": 0.30926101207733153,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.27734375,
+      "learning_rate": 0.0003,
+      "loss": 0.30281963348388674,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.287109375,
+      "learning_rate": 0.0003,
+      "loss": 0.301728343963623,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.29296875,
+      "learning_rate": 0.0003,
+      "loss": 0.2986367464065552,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.306640625,
+      "learning_rate": 0.0003,
+      "loss": 0.294226336479187,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.33984375,
+      "learning_rate": 0.0003,
+      "loss": 0.28919248580932616,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.296875,
+      "learning_rate": 0.0003,
+      "loss": 0.2873764276504517,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.322265625,
+      "learning_rate": 0.0003,
+      "loss": 0.2850848436355591,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "eval_loss": 0.284612774848938,
+      "eval_runtime": 8.55,
+      "eval_samples_per_second": 1114.271,
+      "eval_steps_per_second": 17.427,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06875631951466127,
+      "grad_norm": 0.357421875,
+      "learning_rate": 0.0003,
+      "loss": 0.2803910255432129,
+      "step": 1020
+    },
+    {
+      "epoch": 0.07010448264239973,
+      "grad_norm": 0.37890625,
+      "learning_rate": 0.0003,
+      "loss": 0.27601516246795654,
+      "step": 1040
+    },
+    {
+      "epoch": 0.07145264577013818,
+      "grad_norm": 0.3203125,
+      "learning_rate": 0.0003,
+      "loss": 0.27540276050567625,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07280080889787664,
+      "grad_norm": 0.318359375,
+      "learning_rate": 0.0003,
+      "loss": 0.2708603858947754,
+      "step": 1080
+    },
+    {
+      "epoch": 0.0741489720256151,
+      "grad_norm": 0.328125,
+      "learning_rate": 0.0003,
+      "loss": 0.2696371555328369,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07549713515335356,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.26611151695251467,
+      "step": 1120
+    },
+    {
+      "epoch": 0.07684529828109202,
+      "grad_norm": 1.828125,
+      "learning_rate": 0.0003,
+      "loss": 0.27770705223083497,
+      "step": 1140
+    },
+    {
+      "epoch": 0.07819346140883048,
+      "grad_norm": 0.5390625,
+      "learning_rate": 0.0003,
+      "loss": 0.28163862228393555,
+      "step": 1160
+    },
+    {
+      "epoch": 0.07954162453656892,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.2673502445220947,
+      "step": 1180
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "grad_norm": 0.275390625,
+      "learning_rate": 0.0003,
+      "loss": 0.262909460067749,
+      "step": 1200
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "eval_loss": 0.26177507638931274,
+      "eval_runtime": 8.3534,
+      "eval_samples_per_second": 1140.5,
+      "eval_steps_per_second": 17.837,
+      "step": 1200
+    },
+    {
+      "epoch": 0.08223795079204584,
+      "grad_norm": 0.26953125,
+      "learning_rate": 0.0003,
+      "loss": 0.2581258535385132,
+      "step": 1220
+    },
+    {
+      "epoch": 0.0835861139197843,
+      "grad_norm": 0.26171875,
+      "learning_rate": 0.0003,
+      "loss": 0.25597529411315917,
+      "step": 1240
+    },
+    {
+      "epoch": 0.08493427704752275,
+      "grad_norm": 0.2578125,
+      "learning_rate": 0.0003,
+      "loss": 0.25467369556427,
+      "step": 1260
+    },
+    {
+      "epoch": 0.08628244017526121,
+      "grad_norm": 0.25,
+      "learning_rate": 0.0003,
+      "loss": 0.2502788543701172,
+      "step": 1280
+    },
+    {
+      "epoch": 0.08763060330299967,
+      "grad_norm": 0.259765625,
+      "learning_rate": 0.0003,
+      "loss": 0.25013630390167235,
+      "step": 1300
+    },
+    {
+      "epoch": 0.08897876643073811,
+      "grad_norm": 0.2470703125,
+      "learning_rate": 0.0003,
+      "loss": 0.24711339473724364,
+      "step": 1320
+    },
+    {
+      "epoch": 0.09032692955847657,
+      "grad_norm": 0.25,
+      "learning_rate": 0.0003,
+      "loss": 0.2463146924972534,
+      "step": 1340
+    },
+    {
+      "epoch": 0.09167509268621503,
+      "grad_norm": 0.32421875,
+      "learning_rate": 0.0003,
+      "loss": 0.2437490701675415,
+      "step": 1360
+    },
+    {
+      "epoch": 0.09302325581395349,
+      "grad_norm": 0.259765625,
+      "learning_rate": 0.0003,
+      "loss": 0.2432482957839966,
+      "step": 1380
+    },
+    {
+      "epoch": 0.09437141894169195,
+      "grad_norm": 0.265625,
+      "learning_rate": 0.0003,
+      "loss": 0.24097964763641358,
+      "step": 1400
+    },
+    {
+      "epoch": 0.09437141894169195,
+      "eval_loss": 0.24285024404525757,
+      "eval_runtime": 8.4015,
+      "eval_samples_per_second": 1133.969,
+      "eval_steps_per_second": 17.735,
+      "step": 1400
+    },
+    {
+      "epoch": 0.0957195820694304,
+      "grad_norm": 0.265625,
+      "learning_rate": 0.0003,
+      "loss": 0.2406073570251465,
+      "step": 1420
+    },
+    {
+      "epoch": 0.09706774519716886,
+      "grad_norm": 0.255859375,
+      "learning_rate": 0.0003,
+      "loss": 0.23866963386535645,
+      "step": 1440
+    },
+    {
+      "epoch": 0.09841590832490732,
+      "grad_norm": 0.267578125,
+      "learning_rate": 0.0003,
+      "loss": 0.2364635467529297,
+      "step": 1460
+    },
+    {
+      "epoch": 0.09976407145264576,
+      "grad_norm": 0.2578125,
+      "learning_rate": 0.0003,
+      "loss": 0.23464763164520264,
+      "step": 1480
+    },
+    {
+      "epoch": 0.10111223458038422,
+      "grad_norm": 0.24609375,
+      "learning_rate": 0.0003,
+      "loss": 0.23385767936706542,
+      "step": 1500
+    },
+    {
+      "epoch": 0.10246039770812268,
+      "grad_norm": 0.2451171875,
+      "learning_rate": 0.0003,
+      "loss": 0.23197500705718993,
+      "step": 1520
+    },
+    {
+      "epoch": 0.10380856083586114,
+      "grad_norm": 0.291015625,
+      "learning_rate": 0.0003,
+      "loss": 0.2319192886352539,
+      "step": 1540
+    },
+    {
+      "epoch": 0.1051567239635996,
+      "grad_norm": 0.30078125,
+      "learning_rate": 0.0003,
+      "loss": 0.2309124231338501,
+      "step": 1560
+    },
+    {
+      "epoch": 0.10650488709133805,
+      "grad_norm": 0.263671875,
+      "learning_rate": 0.0003,
+      "loss": 0.2315352201461792,
+      "step": 1580
+    },
+    {
+      "epoch": 0.10785305021907651,
+      "grad_norm": 1.0546875,
+      "learning_rate": 0.0003,
+      "loss": 0.2452239990234375,
+      "step": 1600
+    },
+    {
+      "epoch": 0.10785305021907651,
+      "eval_loss": 0.24468769133090973,
+      "eval_runtime": 8.3885,
+      "eval_samples_per_second": 1135.725,
+      "eval_steps_per_second": 17.762,
+      "step": 1600
+    },
+    {
+      "epoch": 0.10920121334681497,
+      "grad_norm": 0.400390625,
+      "learning_rate": 0.0003,
+      "loss": 0.23531816005706788,
+      "step": 1620
+    },
+    {
+      "epoch": 0.11054937647455342,
+      "grad_norm": 0.267578125,
+      "learning_rate": 0.0003,
+      "loss": 0.22987775802612304,
+      "step": 1640
+    },
+    {
+      "epoch": 0.11189753960229187,
+      "grad_norm": 0.234375,
+      "learning_rate": 0.0003,
+      "loss": 0.22664937973022461,
+      "step": 1660
+    },
+    {
+      "epoch": 0.11324570273003033,
+      "grad_norm": 0.2314453125,
+      "learning_rate": 0.0003,
+      "loss": 0.22559750080108643,
+      "step": 1680
+    },
+    {
+      "epoch": 0.11459386585776879,
+      "grad_norm": 0.2373046875,
+      "learning_rate": 0.0003,
+      "loss": 0.22569913864135743,
+      "step": 1700
+    },
+    {
+      "epoch": 0.11594202898550725,
+      "grad_norm": 0.228515625,
+      "learning_rate": 0.0003,
+      "loss": 0.22344863414764404,
+      "step": 1720
+    },
+    {
+      "epoch": 0.1172901921132457,
+      "grad_norm": 0.2314453125,
+      "learning_rate": 0.0003,
+      "loss": 0.22342846393585206,
+      "step": 1740
+    },
+    {
+      "epoch": 0.11863835524098416,
+      "grad_norm": 0.23046875,
+      "learning_rate": 0.0003,
+      "loss": 0.22171788215637206,
+      "step": 1760
+    },
+    {
+      "epoch": 0.11998651836872262,
+      "grad_norm": 0.22265625,
+      "learning_rate": 0.0003,
+      "loss": 0.2223306655883789,
+      "step": 1780
+    },
+    {
+      "epoch": 0.12133468149646107,
+      "grad_norm": 0.228515625,
+      "learning_rate": 0.0003,
+      "loss": 0.22238383293151856,
+      "step": 1800
+    },
+    {
+      "epoch": 0.12133468149646107,
+      "eval_loss": 0.2212776094675064,
+      "eval_runtime": 8.3797,
+      "eval_samples_per_second": 1136.917,
+      "eval_steps_per_second": 17.781,
+      "step": 1800
+    },
+    {
+      "epoch": 0.12268284462419952,
+      "grad_norm": 0.228515625,
+      "learning_rate": 0.0003,
+      "loss": 0.22141022682189943,
+      "step": 1820
+    },
+    {
+      "epoch": 0.12403100775193798,
+      "grad_norm": 0.232421875,
+      "learning_rate": 0.0003,
+      "loss": 0.21716527938842772,
+      "step": 1840
+    },
+    {
+      "epoch": 0.12537917087967643,
+      "grad_norm": 0.2353515625,
+      "learning_rate": 0.0003,
+      "loss": 0.21714138984680176,
+      "step": 1860
+    },
+    {
+      "epoch": 0.1267273340074149,
+      "grad_norm": 0.2294921875,
+      "learning_rate": 0.0003,
+      "loss": 0.2167724370956421,
+      "step": 1880
+    },
+    {
+      "epoch": 0.12807549713515334,
+      "grad_norm": 0.2255859375,
+      "learning_rate": 0.0003,
+      "loss": 0.2163745403289795,
+      "step": 1900
+    },
+    {
+      "epoch": 0.12942366026289182,
+      "grad_norm": 0.2333984375,
+      "learning_rate": 0.0003,
+      "loss": 0.21465139389038085,
+      "step": 1920
+    },
+    {
+      "epoch": 0.13077182339063026,
+      "grad_norm": 0.2236328125,
+      "learning_rate": 0.0003,
+      "loss": 0.21518900394439697,
+      "step": 1940
+    },
+    {
+      "epoch": 0.13211998651836873,
+      "grad_norm": 0.2255859375,
+      "learning_rate": 0.0003,
+      "loss": 0.21305620670318604,
+      "step": 1960
+    },
+    {
+      "epoch": 0.13346814964610718,
+      "grad_norm": 0.22265625,
+      "learning_rate": 0.0003,
+      "loss": 0.21399173736572266,
+      "step": 1980
+    },
+    {
+      "epoch": 0.13481631277384565,
+      "grad_norm": 0.220703125,
+      "learning_rate": 0.0003,
+      "loss": 0.21352832317352294,
+      "step": 2000
+    },
+    {
+      "epoch": 0.13481631277384565,
+      "eval_loss": 0.21510063111782074,
+      "eval_runtime": 8.4122,
+      "eval_samples_per_second": 1132.528,
+      "eval_steps_per_second": 17.712,
+      "step": 2000
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 580776886272000.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

out/glu-gelu-9L_run/checkpoint-2000/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5bc9b69db12445b4a7919b29e76d261b3c2acbea92402190b4ca5d146c69daa
+size 4920

out/glu-gelu-9L_run/checkpoint-500/config.json

@@ -0,0 +1,35 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "tokenizer_name": "w-ahmad/tiny-stories-tokenizer",
+  "transformers_version": "5.15.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096
+}

out/glu-gelu-9L_run/checkpoint-500/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87d2485366e666ffe60c69a356f612ec9466a7f4ca08911d72d1afb077fb92b8
+size 4011496

out/glu-gelu-9L_run/checkpoint-500/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9b4891f5d92679afc5c80fecf82636d6a4bc82ccaeb318adb74e33e82840daf
+size 8074746

out/glu-gelu-9L_run/checkpoint-500/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cf9097d4513154245c48236b6ec5137b7ee2a21c9f58f2cba798ea275c6026f
+size 14244

out/glu-gelu-9L_run/checkpoint-500/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e328dcb5672eefa48692126747c226aede5ae2e575011f8f5a2b20326a61f850
+size 1064

out/glu-gelu-9L_run/checkpoint-500/tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

out/glu-gelu-9L_run/checkpoint-500/trainer_state.json

@@ -0,0 +1,225 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.03370407819346141,
+  "eval_steps": 200,
+  "global_step": 500,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.2578125,
+      "learning_rate": 0.0003,
+      "loss": 7.793476104736328,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.0003,
+      "loss": 7.021186828613281,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.0,
+      "learning_rate": 0.0003,
+      "loss": 6.446919250488281,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.4609375,
+      "learning_rate": 0.0003,
+      "loss": 5.989651489257812,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 2.328125,
+      "learning_rate": 0.0003,
+      "loss": 5.436912536621094,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 2.671875,
+      "learning_rate": 0.0003,
+      "loss": 4.910789871215821,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 3.28125,
+      "learning_rate": 0.0003,
+      "loss": 4.441200637817383,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 2.625,
+      "learning_rate": 0.0003,
+      "loss": 4.035322952270508,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 3.765625,
+      "learning_rate": 0.0003,
+      "loss": 3.6691635131835936,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 2.296875,
+      "learning_rate": 0.0003,
+      "loss": 3.315158462524414,
+      "step": 200
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.140667200088501,
+      "eval_runtime": 8.3121,
+      "eval_samples_per_second": 1146.166,
+      "eval_steps_per_second": 17.926,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 2.453125,
+      "learning_rate": 0.0003,
+      "loss": 2.9744741439819338,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 2.359375,
+      "learning_rate": 0.0003,
+      "loss": 2.6262161254882814,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 2.234375,
+      "learning_rate": 0.0003,
+      "loss": 2.3146018981933594,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 2.015625,
+      "learning_rate": 0.0003,
+      "loss": 1.9927532196044921,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.609375,
+      "learning_rate": 0.0003,
+      "loss": 1.6939159393310548,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.734375,
+      "learning_rate": 0.0003,
+      "loss": 1.4503639221191407,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 1.828125,
+      "learning_rate": 0.0003,
+      "loss": 1.2639431953430176,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 1.984375,
+      "learning_rate": 0.0003,
+      "loss": 1.1161455154418944,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 1.546875,
+      "learning_rate": 0.0003,
+      "loss": 0.9919822692871094,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 1.1484375,
+      "learning_rate": 0.0003,
+      "loss": 0.8975520133972168,
+      "step": 400
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 0.8561499714851379,
+      "eval_runtime": 8.3662,
+      "eval_samples_per_second": 1138.75,
+      "eval_steps_per_second": 17.81,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0003,
+      "loss": 0.7963034629821777,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.51171875,
+      "learning_rate": 0.0003,
+      "loss": 0.7076507568359375,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.439453125,
+      "learning_rate": 0.0003,
+      "loss": 0.6221071720123291,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.423828125,
+      "learning_rate": 0.0003,
+      "loss": 0.5685368537902832,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.3671875,
+      "learning_rate": 0.0003,
+      "loss": 0.5213922977447509,
+      "step": 500
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 145194221568000.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

out/glu-gelu-9L_run/checkpoint-500/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5bc9b69db12445b4a7919b29e76d261b3c2acbea92402190b4ca5d146c69daa
+size 4920

out/glu-gelu-9L_run/config.json

@@ -0,0 +1,35 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "tokenizer_name": "w-ahmad/tiny-stories-tokenizer",
+  "transformers_version": "5.15.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096
+}

out/glu-gelu-9L_run/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efa7bcb247c2e6fbd4d32ab423b71b076c710576b4833e9258447d136870faaf
+size 4011496

out/glu-gelu-9L_run/tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

out/glu-gelu-9L_run/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5bc9b69db12445b4a7919b29e76d261b3c2acbea92402190b4ca5d146c69daa
+size 4920

out/glu-linear-9L_run/checkpoint-1000/config.json

@@ -0,0 +1,35 @@
+{
+  "activation": "linear",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "tokenizer_name": "w-ahmad/tiny-stories-tokenizer",
+  "transformers_version": "5.15.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096
+}