Auto upload zain 2026-08-18T15:12:06.051058

zain/Activation/README.md

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

zain/Activation/exp.py

@@ -0,0 +1,953 @@
+# =====================================================================
+# exp.py 
+# =====================================================================
+
+import math
+import os
+import time
+import json
+import re
+from pathlib import Path
+from itertools import chain
+from typing import Dict, Callable, Optional, List, Any, Tuple
+
+import torch
+import torch.nn as nn
+from transformers import (
+    LlamaConfig,
+    LlamaPreTrainedModel,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    AutoTokenizer,
+    set_seed,
+)
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention,
+    LlamaRMSNorm,
+    LlamaRotaryEmbedding,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from datasets import load_dataset
+from huggingface_hub import snapshot_download
+
+
+# =============================================================================
+# 1. ACTIVATION REGISTRY
+# =============================================================================
+
+class GLUActivationRegistry:
+    _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("sigmoid", torch.sigmoid)
+GLUActivationRegistry.register("tanh", torch.tanh)
+GLUActivationRegistry.register("softplus", nn.functional.softplus)
+GLUActivationRegistry.register("linear", lambda x: x)
+GLUActivationRegistry.register("s10", lambda x: x * x * torch.sigmoid(x))
+GLUActivationRegistry.register("w1a", lambda x: x * torch.tanh(x))
+
+
+# =============================================================================
+# 2. CONFIG
+# =============================================================================
+
+class TinyLlamaConfig(LlamaConfig):
+    model_type = "tiny_llama"
+
+    def __init__(
+        self,
+        mlp_type: str = "glu",
+        activation: str = "silu",
+        waleed_beta: float = 10.0,
+        powlu_m: float = 3.0,          # <-- NEW: PowLU hyperparameter
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.mlp_type = mlp_type
+        self.activation = activation
+        self.waleed_beta = waleed_beta
+        self.powlu_m = powlu_m          # <-- store it
+        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 (with PowLU)
+# =============================================================================
+
+class TinyLlamaMLP(nn.Module):
+    override_active = False
+    override_value = -100.0
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.mlp_type = config.mlp_type
+        self.activation_name = config.activation
+        self.waleed_beta = getattr(config, "waleed_beta", 10.0)
+        self.powlu_m = getattr(config, "powlu_m", 3.0)   # <-- store m
+
+        if self.mlp_type == "glu":
+            effective_intermediate = self.intermediate_size
+        elif self.mlp_type == "mlp":
+            effective_intermediate = int(self.intermediate_size * 1.5)
+            print(f"[MLP] Auto‑scaled intermediate_size from {self.intermediate_size} to {effective_intermediate}")
+        else:
+            raise ValueError(f"Unknown mlp_type: {self.mlp_type}")
+
+        self.effective_intermediate = effective_intermediate
+
+        if self.mlp_type == "glu":
+            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:
+            self.up_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
+
+        self.down_proj = nn.Linear(effective_intermediate, self.hidden_size, bias=False)
+
+        # ---- Activation flags ----
+        self.is_powlu = (self.activation_name == "powlu")          # <-- NEW
+        self.is_situglu = self.activation_name in ("situglu", "situglu_low")
+        self.is_waleed = self.activation_name in ("waleed", "waleedglu_low")
+        self.is_waleed10 = self.activation_name in ("waleed10", "silu-waleed10")
+        self.has_sigmoid_gate = self.activation_name.startswith("situglu")
+
+        # ---- Activation function (for non-PowLU) ----
+        if self.mlp_type == "glu":
+            if self.is_powlu:
+                # PowLU handled in forward, no act_fn needed
+                self.act_fn = None
+            elif self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
+                self.act_fn = None  # custom handling
+            elif self.activation_name == "waleed10":
+                self.act_fn = GLUActivationRegistry.get("linear")
+            elif self.activation_name == "silu-waleed10":
+                self.act_fn = GLUActivationRegistry.get("silu")
+            else:
+                self.act_fn = GLUActivationRegistry.get(self.activation_name)
+        else:  # mlp
+            if self.is_powlu:
+                raise ValueError("PowLU is only supported with mlp_type='glu'")
+            if self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
+                raise ValueError(f"Activation '{self.activation_name}' requires GLU.")
+            elif self.activation_name == "waleed10":
+                self.act_fn = GLUActivationRegistry.get("linear")
+            elif self.activation_name == "silu-waleed10":
+                self.act_fn = GLUActivationRegistry.get("silu")
+            else:
+                self.act_fn = GLUActivationRegistry.get(self.activation_name)
+
+        # beta values for situglu/waleed
+        if self.activation_name in ("situglu_low", "waleedglu_low"):
+            self.beta1 = 2.5
+            self.beta2 = 4.0
+        else:
+            self.beta1 = 4.0
+            self.beta2 = 25.0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.mlp_type == "glu":
+            gate = self.gate_proj(x)
+            up = self.up_proj(x)
+
+            # ----- PowLU branch (FIXED: safe sqrt/pow to avoid NaN-gradient leak) -----
+            if self.is_powlu:
+                m = self.powlu_m
+                sig_gate = torch.sigmoid(gate)
+                pos_mask = gate > 0
+
+                # Clamp ONLY for the pow/sqrt computation so negative values never
+                # reach sqrt(); the real `gate` tensor is left untouched everywhere else.
+                safe_gate = torch.clamp(gate, min=1e-6)
+
+                # For positive: gate^(m/(sqrt(gate)+1)) * sigmoid(gate)
+                powlu_pos = safe_gate ** (m / (torch.sqrt(safe_gate) + 1)) * sig_gate
+                # For non-positive: gate * sigmoid(gate)  (= SiLU)
+                powlu_neg = gate * sig_gate
+                activated_gate = torch.where(pos_mask, powlu_pos, powlu_neg)
+                hidden = activated_gate * up
+
+            # ----- situglu / waleed branches (unchanged) -----
+            elif self.is_situglu or self.is_waleed:
+                if self.has_sigmoid_gate:
+                    gate = self.beta1 * torch.tanh(gate / self.beta1) * torch.sigmoid(gate)
+                else:
+                    gate = self.beta1 * torch.tanh(gate / self.beta1)
+                up = self.beta2 * torch.tanh(up / self.beta2)
+                hidden = gate * up
+
+            # ----- Standard GLU (ReLU, GELU, SiLU, etc.) -----
+            else:
+                hidden = self.act_fn(gate) * up
+
+            out = self.down_proj(hidden)
+
+        else:  # mlp
+            hidden = self.act_fn(self.up_proj(x))
+            out = self.down_proj(hidden)
+
+        # Optional Waleed clipping after down projection
+        if self.is_waleed10:
+            out = self.waleed_beta * torch.tanh(out / self.waleed_beta)
+
+        # Override hook (unchanged)
+        if TinyLlamaMLP.override_active and out.requires_grad:
+            out.register_hook(lambda grad: TinyLlamaMLP.override_value * torch.ones_like(grad))
+
+        return out
+
+
+# ----------------------------------------------------------------------------
+# DECODER LAYER, ATTENTION MASK, MODEL (unchanged)
+# ----------------------------------------------------------------------------
+
+class TinyLlamaDecoderLayer(nn.Module):
+    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
+        self.mlp = TinyLlamaMLP(config)
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.residual_pre_attn = nn.Identity()
+        self.residual_post_attn = nn.Identity()
+        self.residual_post_mlp = nn.Identity()
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, position_embeddings=None, **kwargs):
+        residual = hidden_states
+        hidden_states = self.residual_pre_attn(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
+        hidden_states = self.residual_post_attn(hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        hidden_states = self.residual_post_mlp(hidden_states)
+        return (hidden_states,)
+
+
+def _build_causal_mask(attention_mask, seq_len, dtype, device):
+    min_value = torch.finfo(dtype).min
+    causal = torch.full((seq_len, seq_len), fill_value=min_value, dtype=dtype, device=device)
+    causal = torch.triu(causal, diagonal=1)
+    causal = causal[None, None, :, :]
+    if attention_mask is None:
+        batch_size = 1
+        return causal.expand(batch_size, 1, seq_len, seq_len)
+    batch_size = attention_mask.shape[0]
+    causal = causal.expand(batch_size, 1, seq_len, seq_len).clone()
+    padding = attention_mask[:, None, None, :].to(device) == 0
+    causal = causal.masked_fill(padding, min_value)
+    return causal
+
+
+_MASK_PRINTED = False
+
+
+class TinyLlamaModel(LlamaPreTrainedModel):
+    config_class = TinyLlamaConfig
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([TinyLlamaDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LlamaRotaryEmbedding(config=config)
+        self.post_init()
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, return_dict=None, **kwargs):
+        global _MASK_PRINTED
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if position_ids is None:
+            seq_len = inputs_embeds.shape[1]
+            position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0).expand(inputs_embeds.shape[0], -1)
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        seq_len = hidden_states.shape[1]
+        causal_mask = _build_causal_mask(attention_mask, seq_len, hidden_states.dtype, hidden_states.device)
+        if not _MASK_PRINTED:
+            print("[INFO] Causal mask (float with -inf) applied to all attention layers.")
+            _MASK_PRINTED = True
+        for decoder_layer in self.layers:
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                position_embeddings=position_embeddings,
+            )
+            hidden_states = layer_outputs[0]
+        hidden_states = self.norm(hidden_states)
+        if not return_dict:
+            return (hidden_states,)
+        return {"last_hidden_state": hidden_states, "hidden_states": None, "attentions": None}
+
+
+class TinyLlamaForCausalLM(LlamaPreTrainedModel):
+    config_class = TinyLlamaConfig
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__(config)
+        self.model = TinyLlamaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
+                labels=None, return_dict=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. HF CHECKPOINT FETCHER (unchanged)
+# =============================================================================
+
+def fetch_latest_checkpoint_from_hub(
+    repo_id: str,
+    subpath: str,
+    variant: str,
+    checkpoint_step: Optional[int] = None
+) -> str:
+    remote_prefix = f"{subpath}/{variant}_run" if subpath else f"{variant}_run"
+    if checkpoint_step is not None:
+        full_remote_path = f"{remote_prefix}/checkpoint-{checkpoint_step}"
+        print(f"[Hub] Fetching specific checkpoint: {repo_id}/{full_remote_path}")
+        local_root = snapshot_download(
+            repo_id=repo_id,
+            allow_patterns=[f"{full_remote_path}/*"],
+            local_dir_use_symlinks=False,
+        )
+        checkpoint_local_path = os.path.join(local_root, full_remote_path)
+        if not os.path.exists(checkpoint_local_path):
+            raise RuntimeError(f"Downloaded checkpoint not found at {checkpoint_local_path}")
+        return checkpoint_local_path
+
+    print(f"[Hub] Downloading entire run folder: {repo_id}/{remote_prefix}")
+    local_root = snapshot_download(
+        repo_id=repo_id,
+        allow_patterns=[f"{remote_prefix}/*"],
+        local_dir_use_symlinks=False,
+    )
+    run_local_path = os.path.join(local_root, remote_prefix)
+    if not os.path.exists(run_local_path):
+        raise RuntimeError(f"Run folder not found at {run_local_path}")
+
+    checkpoints = []
+    for item in os.listdir(run_local_path):
+        if item.startswith("checkpoint-") and os.path.isdir(os.path.join(run_local_path, item)):
+            match = re.match(r"checkpoint-(\d+)", item)
+            if match:
+                step = int(match.group(1))
+                checkpoints.append((step, item))
+
+    if not checkpoints:
+        raise RuntimeError(f"No checkpoint folders found in {run_local_path}")
+
+    latest_step, latest_name = max(checkpoints, key=lambda x: x[0])
+    print(f"[Hub] Latest checkpoint found: step {latest_step}")
+    return os.path.join(run_local_path, latest_name)
+
+
+# =============================================================================
+# 5. RESUME + FREEZE + OVERRIDE CALLBACK (unchanged)
+# =============================================================================
+
+class ResumeFreezeOverrideCallback(TrainerCallback):
+    def __init__(self, override_value: Optional[float] = None):
+        self.override_value = override_value
+        self._trainer = None  # Will be set manually
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        # Try multiple ways to get the trainer
+        trainer = kwargs.get('trainer')
+        if trainer is None:
+            trainer = getattr(self, '_trainer', None)
+        if trainer is None:
+            trainer = getattr(self, 'trainer', None)
+        if trainer is None:
+            raise ValueError("Trainer not accessible in callback")
+
+        model = trainer.model
+
+        # ----- FREEZE ALL EXCEPT MLP PROJECTIONS -----
+        for name, param in model.named_parameters():
+            if any(x in name for x in ["gate_proj", "up_proj", "down_proj"]):
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+
+        trainable_params = [p for p in model.parameters() if p.requires_grad]
+
+        # ----- REBUILD OPTIMIZER -----
+        from torch.optim import AdamW
+        new_optimizer = AdamW(
+            trainable_params,
+            lr=args.learning_rate,
+            weight_decay=args.weight_decay,
+            betas=(args.adam_beta1, args.adam_beta2),
+            eps=args.adam_epsilon,
+        )
+        trainer.optimizer = new_optimizer
+
+        # ----- KEEP EXISTING SCHEDULER (re‑attach) -----
+        if trainer.lr_scheduler is not None:
+            trainer.lr_scheduler.optimizer = new_optimizer
+
+        # ----- ACTIVATE OVERRIDE -----
+        if self.override_value is not None:
+            TinyLlamaMLP.override_active = True
+            TinyLlamaMLP.override_value = self.override_value
+            print(f"[Override] Activated with value = {self.override_value}")
+        else:
+            print("[Freeze] MLP projections frozen; no gradient override.")
+
+
+# =============================================================================
+# 6. MONITORING (unchanged)
+# =============================================================================
+
+class StatsEngine:
+    @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(),
+                "min": tensor.min().item(),
+                "max": tensor.max().item(),
+                "range": tensor.max().item() - tensor.min().item(),
+            }
+
+
+class StepAccumulator:
+    def __init__(self):
+        self.tensors: Dict[str, Dict[str, float]] = {}
+
+    def add(self, name: str, numel: int, stats: Dict[str, float]):
+        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
+        t_min = min(a.get("min", float("inf")), b.get("min", float("inf")))
+        t_max = max(a.get("max", float("-inf")), b.get("max", float("-inf")))
+        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,
+            "min": t_min,
+            "max": t_max,
+            "range": t_max - t_min,
+        }
+
+    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)
+        norm = math.sqrt(sum(e["norm"] ** 2 for e in entries.values()))
+        max_abs = max(e["max_abs"] for e in entries.values())
+        mean = sum(e["mean"] * e["numel"] for e in entries.values()) / total_n
+        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))
+        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
+        t_min = min(e.get("min", float("inf")) for e in entries.values())
+        t_max = max(e.get("max", float("-inf")) for e in entries.values())
+        return {
+            "norm": norm,
+            "mean": mean,
+            "std": std,
+            "max_abs": max_abs,
+            "frac_near_dtype_limit": frac_dtype,
+            "frac_near_user_limit": frac_user,
+            "min": t_min,
+            "max": t_max,
+            "range": t_max - t_min,
+        }
+
+    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:
+    def __init__(self, model: nn.Module):
+        self.model = model
+        self.handles = []
+        self.active = False
+
+    def attach_forward(self, module_patterns, accumulator, user_limit, dtype_ratio):
+        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, param_patterns, accumulator, user_limit, dtype_ratio):
+        for name, param in self.model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if param_patterns and not any(re.search(p, name) for p in param_patterns):
+                continue
+            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, accumulator, user_limit, dtype_ratio):
+        def hook(module, inp, out):
+            if not self.active:
+                return
+            if isinstance(out, dict):
+                out = out.get("last_hidden_state")
+            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, accumulator, user_limit, dtype_ratio):
+        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):
+    def __init__(
+        self,
+        model: nn.Module,
+        monitor_every_n_steps: int = 10,
+        module_patterns: Optional[List[str]] = None,
+        param_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.*", ".*residual.*"]
+        self.param_patterns = param_patterns or self.module_patterns
+        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.param_patterns,
+            self.accumulator,
+            self.user_limits["grad"],
+            self.dtype_ratio,
+        )
+        self.pending_metrics = None
+
+    def _should_monitor(self, state):
+        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
+        for name, param in self.model.named_parameters():
+            if self.param_patterns and not any(re.search(p, name) for p in self.param_patterns):
+                continue
+            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()
+
+    def _kind_of(self, name: str) -> str:
+        if name.startswith("act."):
+            return "act"
+        if name.startswith("grad."):
+            return "grad"
+        if name.startswith("param."):
+            return "param"
+        return "other"
+
+    def _strip_kind(self, 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 = {}
+        if self.log_scope.get("global", True):
+            by_kind = {}
+            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
+
+        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 = {}
+                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
+
+        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)
+            try:
+                import wandb
+                if wandb.run is not None:
+                    wandb.log(self.pending_metrics, step=state.global_step)
+            except ImportError:
+                pass
+            self.pending_metrics = None
+
+    def on_prediction_step(self, args, state, control, **kwargs):
+        if not self.monitor_during_eval:
+            return
+        if not self.hooks.active:
+            self.accumulator.clear()
+            self.hooks.set_active(True)
+            for name, param in self.model.named_parameters():
+                if self.param_patterns and not any(re.search(p, name) for p in self.param_patterns):
+                    continue
+                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()
+
+
+# =============================================================================
+# 7. TIME TRACKER (unchanged)
+# =============================================================================
+
+class TimeTrackerCallback(TrainerCallback):
+    def __init__(self):
+        self.step_start = None
+        self.epoch_start = None
+        self.total_train_time = 0.0
+        self.step_times = []
+
+    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
+
+
+# =============================================================================
+# 8. METRICS LOGGER (unchanged)
+# =============================================================================
+
+class MetricsLoggerCallback(TrainerCallback):
+    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")
+
+
+# =============================================================================
+# 9. DATA & TRAINER FACTORY (unchanged)
+# =============================================================================
+
+def build_dataset(
+    tokenizer,
+    max_seq_len: int = 512,
+    split: str = "train",
+    dataset_name: str = "roneneldan/TinyStories",
+    max_samples: Optional[int] = None,
+):
+    ds = load_dataset(dataset_name, split=split)
+    if max_samples is not None and split == "train":
+        ds = ds.select(range(min(max_samples, len(ds))))
+
+    def tokenize(examples):
+        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)
+
+    def group_texts(examples):
+        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)
+
+
+def create_trainer(
+    model,
+    tokenizer,
+    config: Dict[str, Any],
+    train_dataset,
+    eval_dataset=None,
+):
+    tc = config.get("training", {})
+    mc = config.get("monitor", {})
+    go = config.get("gradient_override", {})
+
+    wandb_project = tc.get("wandb_project")
+    if wandb_project:
+        os.environ["WANDB_PROJECT"] = wandb_project
+
+    args = TrainingArguments(
+        output_dir=tc.get("output_dir", "./out"),
+        run_name=tc.get("run_name", None),
+        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),
+        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()]
+
+    # ---- FREEZE + OVERRIDE CALLBACK ----
+    freeze_mlp = tc.get("freeze_mlp", False)
+    override_enabled = go.get("enabled", False)
+    if freeze_mlp or override_enabled:
+        override_value = go.get("value", -100.0) if override_enabled else None
+        callbacks.append(ResumeFreezeOverrideCallback(override_value=override_value))
+
+    # ---- MONITORING ----
+    if mc.get("enabled", True):
+        module_patterns = mc.get("module_patterns", [".*mlp.*", ".*residual.*"])
+        callbacks.append(
+            StabilityMonitorCallback(
+                model=model,
+                monitor_every_n_steps=mc.get("monitor_every_n_steps", 10),
+                module_patterns=module_patterns,
+                param_patterns=module_patterns,
+                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),
+            )
+        )
+
+    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,
+    )
+
+    return trainer

zain/Activation/grad.py

@@ -0,0 +1,571 @@
+#!/usr/bin/env python3
+"""
+Standalone script to extract Adam RMS and SNR from HF Hub checkpoints.
+Uses exact model definitions from the activation analyzer (exp.py).
+Normalizes names to match the analyzer's naming (gate_proj, up_proj, etc.).
+Configured for: w-ahmad/STRESS-TEST2 / zain/Activation/out
+"""
+
+import math
+import os
+import json
+import tempfile
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Any
+
+import torch
+import torch.nn as nn
+import wandb
+from huggingface_hub import HfApi, hf_hub_download, list_repo_tree
+from transformers import LlamaConfig, LlamaPreTrainedModel
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention, LlamaRMSNorm, LlamaRotaryEmbedding,
+)
+
+# ============================================================================
+# EXACT MODEL DEFINITIONS (copied from the activation analyzer)
+# ============================================================================
+
+class GLUActivationRegistry:
+    _registry: Dict[str, callable] = {}
+
+    @classmethod
+    def register(cls, name: str, fn: callable) -> None:
+        cls._registry[name] = fn
+
+    @classmethod
+    def get(cls, name: str) -> callable:
+        if name not in cls._registry:
+            raise KeyError(f"Activation '{name}' not found")
+        return cls._registry[name]
+
+# Register all activations used in the analyzer
+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)
+GLUActivationRegistry.register("linear", lambda x: x)
+GLUActivationRegistry.register("s10", lambda x: x * x * torch.sigmoid(x))
+GLUActivationRegistry.register("w1a", lambda x: x * torch.tanh(x))
+GLUActivationRegistry.register("powlu", lambda x: x)   # dummy for param extraction
+
+
+class TinyLlamaConfig(LlamaConfig):
+    model_type = "tiny_llama"
+
+    def __init__(
+        self,
+        mlp_type: str = "glu",
+        activation: str = "silu",
+        waleed_beta: float = 10.0,
+        powlu_m: float = 3.0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.mlp_type = mlp_type
+        self.activation = activation
+        self.waleed_beta = waleed_beta
+        self.powlu_m = powlu_m
+        if self.num_key_value_heads != self.num_attention_heads:
+            raise ValueError("Pure MHA required.")
+
+
+class SiTUGLU(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, beta1: float = 4.0, beta2: float = 25.0):
+        super().__init__()
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.W_g = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.W_u = nn.Linear(input_dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        g = self.W_g(x)
+        gate = self.beta1 * torch.tanh(g / self.beta1) * torch.sigmoid(g)
+        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
+        return gate * up
+
+
+class WaleedGLU(nn.Module):
+    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)
+        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
+        return gate * up
+
+
+class TinyLlamaMLP(nn.Module):
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.mlp_type = config.mlp_type
+        self.activation_name = config.activation
+
+        if self.mlp_type == "glu":
+            effective_intermediate = self.intermediate_size
+        else:
+            effective_intermediate = int(self.intermediate_size * 1.5)
+
+        self.effective_intermediate = effective_intermediate
+        self.situglu = None
+        self.waleed = None
+        self.gate_proj = None
+        self.up_proj = None
+        self.act_fn = None
+
+        if self.mlp_type == "glu":
+            if self.activation_name == "situglu":
+                self.situglu = SiTUGLU(self.hidden_size, effective_intermediate)
+            elif self.activation_name == "waleed":
+                self.waleed = WaleedGLU(self.hidden_size, effective_intermediate)
+            else:
+                # Standard GLU (including waleed10, silu-waleed10, powlu, situglu_low, waleedglu_low)
+                self.gate_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
+                self.up_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
+                # Activation handling (not needed for param extraction, but avoid errors)
+                if self.activation_name == "waleed10":
+                    self.act_fn = GLUActivationRegistry.get("linear")
+                elif self.activation_name == "silu-waleed10":
+                    self.act_fn = GLUActivationRegistry.get("silu")
+                elif self.activation_name in ("situglu_low", "waleedglu_low"):
+                    self.act_fn = None
+                else:
+                    self.act_fn = GLUActivationRegistry.get(self.activation_name)
+        else:
+            # MLP (non-GLU) – not used in your runs, but kept for completeness
+            if self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
+                raise ValueError(f"Activation '{self.activation_name}' requires GLU.")
+            self.up_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
+            if self.activation_name == "waleed10":
+                self.act_fn = GLUActivationRegistry.get("linear")
+            elif self.activation_name == "silu-waleed10":
+                self.act_fn = GLUActivationRegistry.get("silu")
+            else:
+                self.act_fn = GLUActivationRegistry.get(self.activation_name)
+
+        self.down_proj = nn.Linear(effective_intermediate, self.hidden_size, bias=False)
+
+    def forward(self, x):
+        raise NotImplementedError
+
+
+class TinyLlamaDecoderLayer(nn.Module):
+    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
+        self.mlp = TinyLlamaMLP(config)
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class TinyLlamaModel(LlamaPreTrainedModel):
+    config_class = TinyLlamaConfig
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [TinyLlamaDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LlamaRotaryEmbedding(config=config)
+        self.post_init()
+
+
+class TinyLlamaForCausalLM(LlamaPreTrainedModel):
+    config_class = TinyLlamaConfig
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__(config)
+        self.model = TinyLlamaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(self, **kwargs):
+        raise NotImplementedError
+
+
+# ============================================================================
+# NAME NORMALIZATION – matches activation analyzer exactly
+# ============================================================================
+
+def canonical_param_name(name: str) -> str:
+    """
+    Normalize parameter names to match the activation analyzer's naming.
+    - W_g -> gate_proj, W_u -> up_proj (for situglu/waleed)
+    - Keep gate_proj and up_proj unchanged (they already are canonical)
+    - Leave all other names as-is.
+    This allows direct overlays between RMS/SNR and activation stats.
+    """
+    # WaleedGLU  → use gate_proj/up_proj
+    if "mlp.waleed.W_g" in name:
+        return name.replace("mlp.waleed.W_g", "mlp.waleed.gate_proj")
+    if "mlp.waleed.W_u" in name:
+        return name.replace("mlp.waleed.W_u", "mlp.waleed.up_proj")
+
+    # SiTUGLU → use gate_proj/up_proj
+    if "mlp.situglu.W_g" in name:
+        return name.replace("mlp.situglu.W_g", "mlp.situglu.gate_proj")
+    if "mlp.situglu.W_u" in name:
+        return name.replace("mlp.situglu.W_u", "mlp.situglu.up_proj")
+
+    # For standard GLU, names already have gate_proj/up_proj – keep them.
+    # Everything else (attention, norms, embeddings, lm_head) stays as-is.
+    return name
+
+
+# ============================================================================
+# Stats helpers
+# ============================================================================
+
+def tensor_stats(tensor: torch.Tensor) -> Dict[str, float]:
+    with torch.no_grad():
+        return {
+            "numel": tensor.numel(),
+            "mean": tensor.mean().item(),
+            "std": tensor.std().item(),
+            "max_abs": tensor.abs().max().item(),
+            "norm": tensor.norm().item(),
+        }
+
+
+def aggregate_stats(stats_list: List[Dict[str, float]]) -> Dict[str, float]:
+    if not stats_list:
+        return {"numel": 0, "mean": 0.0, "std": 0.0, "max_abs": 0.0, "norm": 0.0}
+
+    total_numel = sum(s["numel"] for s in stats_list)
+    if total_numel == 0:
+        return {"numel": 0, "mean": 0.0, "std": 0.0, "max_abs": 0.0, "norm": 0.0}
+
+    mean = sum(s["mean"] * s["numel"] for s in stats_list) / total_numel
+    ex2 = sum(s["numel"] * (s["std"]**2 + s["mean"]**2) for s in stats_list) / total_numel
+    std = math.sqrt(max(0.0, ex2 - mean**2))
+    norm = math.sqrt(sum(s["norm"]**2 for s in stats_list))
+    max_abs = max(s["max_abs"] for s in stats_list)
+
+    return {"numel": total_numel, "mean": mean, "std": std, "max_abs": max_abs, "norm": norm}
+
+
+def get_param_names_from_model(config: TinyLlamaConfig) -> List[str]:
+    model = TinyLlamaForCausalLM(config)
+    return [name for name, _ in model.named_parameters()]
+
+
+# ============================================================================
+# HF Hub helpers – configured for your repo
+# ============================================================================
+
+REPO_ID = "w-ahmad/art"
+BASE_PATH = "out"
+
+
+def list_run_directories() -> List[str]:
+    api = HfApi()
+    items = list_repo_tree(repo_id=REPO_ID, path_in_repo=BASE_PATH, recursive=False)
+    return [item.path for item in items if item.path.endswith("_run")]
+
+
+def list_checkpoints_for_run(run_path: str) -> List[str]:
+    api = HfApi()
+    items = list_repo_tree(repo_id=REPO_ID, path_in_repo=run_path, recursive=False)
+    ckpts = [item.path for item in items if item.path.split("/")[-1].startswith("checkpoint-")]
+    ckpts.sort(key=lambda x: int(x.split("-")[-1]))
+    return ckpts
+
+
+def download_file_from_hub(file_path: str, local_dir: str) -> str:
+    return hf_hub_download(
+        repo_id=REPO_ID,
+        filename=file_path,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+    )
+
+
+# ============================================================================
+# Core analysis
+# ============================================================================
+
+def analyze_checkpoint_from_hub(
+    ckpt_path: str,
+    config: TinyLlamaConfig,
+    param_names: List[str],
+    temp_dir: str,
+) -> Dict[str, Any]:
+    opt_file = download_file_from_hub(f"{ckpt_path}/optimizer.pt", temp_dir)
+    opt_state = torch.load(opt_file, map_location="cpu")
+    state = opt_state["state"]
+    param_groups = opt_state["param_groups"]
+
+    param_ids = []
+    for group in param_groups:
+        param_ids.extend(group["params"])
+
+    if len(param_ids) != len(param_names):
+        raise ValueError(
+            f"Param ID mismatch: {len(param_ids)} vs {len(param_names)}"
+        )
+
+    id_to_name = {pid: param_names[i] for i, pid in enumerate(param_ids)}
+
+    # Per-tensor RMS + SNR (canonical names)
+    per_tensor_rms: Dict[str, Dict[str, float]] = {}
+    per_tensor_snr: Dict[str, Dict[str, float]] = {}
+
+    for pid, raw_name in id_to_name.items():
+        if pid not in state:
+            continue
+
+        name = canonical_param_name(raw_name)
+
+        # RMS
+        v_t = state[pid]["exp_avg_sq"]
+        rms = torch.sqrt(v_t + 1e-8)
+        per_tensor_rms[name] = tensor_stats(rms)
+
+        # SNR = |m_t| / (rms + eps)
+        m_t = state[pid]["exp_avg"]
+        snr = m_t.abs() / (rms + 1e-8)
+        snr = snr.clamp(max=1e4)
+        per_tensor_snr[name] = tensor_stats(snr)
+
+    # Layer aggregation
+    def build_layer_stats(per_tensor_dict: Dict[str, Dict[str, float]]) -> Dict[str, Dict[str, float]]:
+        layers: Dict[str, List[Dict[str, float]]] = {}
+        for name, stats in per_tensor_dict.items():
+            parts = name.split(".")
+            layer_key = None
+            for i, p in enumerate(parts):
+                if p == "layers" and i + 1 < len(parts):
+                    layer_key = f"layers.{parts[i+1]}"
+                    break
+            if layer_key is None:
+                layer_key = "root"
+            layers.setdefault(layer_key, []).append(stats)
+        return {k: aggregate_stats(v) for k, v in layers.items()}
+
+    layer_rms = build_layer_stats(per_tensor_rms)
+    layer_snr = build_layer_stats(per_tensor_snr)
+
+    global_rms = aggregate_stats(list(per_tensor_rms.values()))
+    global_snr = aggregate_stats(list(per_tensor_snr.values()))
+
+    return {
+        "step": int(ckpt_path.split("-")[-1]),
+        "per_tensor_rms": per_tensor_rms,
+        "per_tensor_snr": per_tensor_snr,
+        "per_layer_rms": layer_rms,
+        "per_layer_snr": layer_snr,
+        "global_rms": global_rms,
+        "global_snr": global_snr,
+    }
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--wandb_project", type=str, default="llm-gradient-rms-analysis")
+    parser.add_argument("--wandb_entity", type=str, default=None)
+    parser.add_argument("--limit_runs", type=str, nargs="+", default=None)
+    parser.add_argument("--limit_checkpoints", type=str, nargs="+", default=None)
+    parser.add_argument("--cache_dir", type=str, default=None)
+    parser.add_argument("--log_all_tensors", action="store_true", default=False,
+                        help="Log every tensor's RMS/SNR (may hit WandB limits).")
+    args = parser.parse_args()
+
+    print(f"Listing run directories in {REPO_ID}/{BASE_PATH} ...")
+    run_paths = list_run_directories()
+    print(f"Found {len(run_paths)} run directories.")
+
+    if args.limit_runs:
+        run_paths = [r for r in run_paths if r.split("/")[-1] in args.limit_runs]
+        print(f"Filtered to {len(run_paths)} runs.")
+
+    if args.cache_dir:
+        cache_dir = Path(args.cache_dir)
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        use_temp = False
+    else:
+        cache_dir = Path(tempfile.mkdtemp(prefix="hf_cache_"))
+        use_temp = True
+
+    try:
+        for run_path in run_paths:
+            run_name = run_path.split("/")[-1]
+            print(f"\nProcessing run: {run_name}")
+
+            ckpt_paths = list_checkpoints_for_run(run_path)
+            if not ckpt_paths:
+                print(f"  No checkpoints found, skipping.")
+                continue
+
+            first_ckpt = ckpt_paths[0]
+            config_file = download_file_from_hub(f"{first_ckpt}/config.json", str(cache_dir))
+            config = TinyLlamaConfig.from_pretrained(config_file)
+            param_names = get_param_names_from_model(config)
+            print(f"  Model: {len(param_names)} params, {config.num_hidden_layers} layers.")
+            print(f"  waleed_beta: {config.waleed_beta}, powlu_m: {config.powlu_m}")
+
+            if args.log_all_tensors and len(param_names) > 200:
+                print(f"  ⚠️  {len(param_names)} tensors — may hit WandB limits.")
+
+            wandb.init(
+                project=args.wandb_project,
+                entity=args.wandb_entity,
+                name=run_name,
+                config={
+                    "run_name": run_name,
+                    "mlp_type": config.mlp_type,
+                    "activation": config.activation,
+                    "num_layers": config.num_hidden_layers,
+                    "hidden_size": config.hidden_size,
+                    "waleed_beta": config.waleed_beta,
+                    "powlu_m": config.powlu_m,
+                },
+                reinit=True,
+            )
+
+            for ckpt in ckpt_paths:
+                step = int(ckpt.split("-")[-1])
+                if args.limit_checkpoints and str(step) not in args.limit_checkpoints:
+                    continue
+                print(f"  Analyzing {ckpt} ...")
+                try:
+                    result = analyze_checkpoint_from_hub(ckpt, config, param_names, str(cache_dir))
+                except Exception as e:
+                    print(f"    Error: {e}, skipping.")
+                    continue
+
+                log_dict: Dict[str, Any] = {}
+
+                # GLOBAL  (RMS + SNR)
+                gr = result["global_rms"]
+                gs = result["global_snr"]
+
+                log_dict["global/rms_norm"] = gr["norm"]
+                log_dict["global/rms_mean"] = gr["mean"]
+                log_dict["global/rms_std"] = gr["std"]
+                log_dict["global/rms_max_abs"] = gr["max_abs"]
+
+                log_dict["global/snr_mean"] = gs["mean"]
+                log_dict["global/snr_std"] = gs["std"]
+                log_dict["global/snr_max_abs"] = gs["max_abs"]
+                log_dict["global/snr_norm"] = gs["norm"]
+
+                # PER-LAYER
+                for layer_name, stats in result["per_layer_rms"].items():
+                    safe = layer_name.replace(".", "_")
+                    log_dict[f"layers/{safe}/rms_norm"] = stats["norm"]
+
+                for layer_name, stats in result["per_layer_snr"].items():
+                    safe = layer_name.replace(".", "_")
+                    log_dict[f"layers/{safe}/snr_mean"] = stats["mean"]
+
+                # PER-TENSOR (canonical names → all variants overlay)
+                if args.log_all_tensors:
+                    # RMS
+                    for name, stats in result["per_tensor_rms"].items():
+                        safe = name.replace(".", "_")
+                        log_dict[f"tensors/{safe}/rms_norm"] = stats["norm"]
+                        log_dict[f"tensors/{safe}/rms_mean"] = stats["mean"]
+                        log_dict[f"tensors/{safe}/rms_std"] = stats["std"]
+                        log_dict[f"tensors/{safe}/rms_max_abs"] = stats["max_abs"]
+
+                    # SNR
+                    for name, stats in result["per_tensor_snr"].items():
+                        safe = name.replace(".", "_")
+                        log_dict[f"tensors/{safe}/snr_mean"] = stats["mean"]
+                        log_dict[f"tensors/{safe}/snr_std"] = stats["std"]
+                        log_dict[f"tensors/{safe}/snr_max_abs"] = stats["max_abs"]
+                        log_dict[f"tensors/{safe}/snr_norm"] = stats["norm"]
+                else:
+                    # Top/bottom 5 by RMS norm
+                    sorted_tensors = sorted(
+                        result["per_tensor_rms"].items(),
+                        key=lambda x: x[1]["norm"],
+                        reverse=True,
+                    )
+                    top5 = sorted_tensors[:5]
+                    bottom5 = sorted_tensors[-5:] if len(sorted_tensors) >= 5 else sorted_tensors
+
+                    for i, (name, stats) in enumerate(top5):
+                        log_dict[f"top_{i+1}/name"] = name
+                        log_dict[f"top_{i+1}/rms_norm"] = stats["norm"]
+                        log_dict[f"top_{i+1}/rms_mean"] = stats["mean"]
+                    for i, (name, stats) in enumerate(bottom5):
+                        log_dict[f"bottom_{i+1}/name"] = name
+                        log_dict[f"bottom_{i+1}/rms_norm"] = stats["norm"]
+
+                # TABLES
+                rms_table_data = []
+                for name, stats in result["per_tensor_rms"].items():
+                    rms_table_data.append([
+                        name, stats["mean"], stats["std"],
+                        stats["max_abs"], stats["norm"], stats["numel"],
+                    ])
+                if rms_table_data:
+                    log_dict["rms_table"] = wandb.Table(
+                        columns=["param_name", "mean", "std", "max_abs", "norm", "numel"],
+                        data=rms_table_data,
+                    )
+
+                snr_table_data = []
+                for name, stats in result["per_tensor_snr"].items():
+                    snr_table_data.append([
+                        name, stats["mean"], stats["std"],
+                        stats["max_abs"], stats["norm"], stats["numel"],
+                    ])
+                if snr_table_data:
+                    log_dict["snr_table"] = wandb.Table(
+                        columns=["param_name", "mean", "std", "max_abs", "norm", "numel"],
+                        data=snr_table_data,
+                    )
+
+                wandb.log(log_dict, step=step)
+
+            wandb.finish()
+
+    finally:
+        if use_temp and cache_dir.exists():
+            shutil.rmtree(cache_dir, ignore_errors=True)
+
+    print("\nAll done.")
+
+
+if __name__ == "__main__":
+    main()

zain/Activation/llm_analyzer_wandb.py

@@ -0,0 +1,1307 @@
+"""
+Multi-LLM Activation & Loss Analyzer with WandB Logging
+GPU-only, FP32, exact training‑style tokenization (concatenate + chunk).
+Supports PowLU, custom name normalization, full‑dataset evaluation,
+and now captures pre‑down_proj inputs & layer outputs.
+"""
+
+import math
+import os
+import tempfile
+import shutil
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import (
+    LlamaConfig,
+    LlamaPreTrainedModel,
+    AutoTokenizer,
+    AutoConfig,
+    AutoModelForCausalLM,
+)
+from transformers.models.llama.modeling_llama import (
+    LlamaAttention,
+    LlamaRMSNorm,
+    LlamaRotaryEmbedding,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from datasets import load_dataset as hf_load_dataset
+from typing import List, Dict, Optional, 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")
+
+# ============================================================================
+# REPO CONFIGURATION
+# ============================================================================
+
+REPO_ID = "w-ahmad/STRESS-TEST3"
+BASE_PATH = "zain/Activation/out"          # <-- NO trailing slash
+
+# ============================================================================
+# 1. ACTIVATION REGISTRY (exact copy from exp.py)
+# ============================================================================
+
+class GLUActivationRegistry:
+    _registry: Dict[str, callable] = {}
+
+    @classmethod
+    def register(cls, name: str, fn: callable) -> None:
+        cls._registry[name] = fn
+
+    @classmethod
+    def get(cls, name: str) -> callable:
+        if name not in cls._registry:
+            raise KeyError(f"Activation '{name}' not found.")
+        return cls._registry[name]
+
+# Register 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("sigmoid", torch.sigmoid)
+GLUActivationRegistry.register("tanh", torch.tanh)
+GLUActivationRegistry.register("softplus", nn.functional.softplus)
+GLUActivationRegistry.register("linear", lambda x: x)
+GLUActivationRegistry.register("s10", lambda x: x * x * torch.sigmoid(x))
+GLUActivationRegistry.register("w1a", lambda x: x * torch.tanh(x))
+
+# ============================================================================
+# 2. CUSTOM CONFIG & MODEL (with PowLU)
+# ============================================================================
+
+class TinyLlamaConfig(LlamaConfig):
+    model_type = "tiny_llama"
+
+    def __init__(
+        self,
+        mlp_type: str = "glu",
+        activation: str = "silu",
+        waleed_beta: float = 10.0,
+        powlu_m: float = 3.0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.mlp_type = mlp_type
+        self.activation = activation
+        self.waleed_beta = waleed_beta
+        self.powlu_m = powlu_m
+        if self.num_key_value_heads != self.num_attention_heads:
+            raise ValueError("Pure MHA required.")
+
+
+class SiTUGLU(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, beta1: float = 4.0, beta2: float = 25.0):
+        super().__init__()
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.W_g = nn.Linear(input_dim, hidden_dim, bias=False)
+        self.W_u = nn.Linear(input_dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        g = self.W_g(x)
+        gate = self.beta1 * torch.tanh(g / self.beta1) * torch.sigmoid(g)
+        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
+        return gate * up
+
+
+class WaleedGLU(nn.Module):
+    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)
+        up = self.beta2 * torch.tanh(self.W_u(x) / self.beta2)
+        return gate * up
+
+
+class TinyLlamaMLP(nn.Module):
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.mlp_type = config.mlp_type
+        self.activation_name = config.activation
+        self.waleed_beta = getattr(config, "waleed_beta", 10.0)
+        self.powlu_m = getattr(config, "powlu_m", 3.0)
+
+        if self.mlp_type == "glu":
+            effective_intermediate = self.intermediate_size
+        elif self.mlp_type == "mlp":
+            effective_intermediate = int(self.intermediate_size * 1.5)
+        else:
+            raise ValueError(f"Unknown mlp_type: {self.mlp_type}")
+
+        self.effective_intermediate = effective_intermediate
+
+        if self.mlp_type == "glu":
+            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:
+            self.up_proj = nn.Linear(self.hidden_size, effective_intermediate, bias=False)
+        self.down_proj = nn.Linear(effective_intermediate, self.hidden_size, bias=False)
+
+        # ---- Activation flags ----
+        self.is_powlu = (self.activation_name == "powlu")
+        self.is_situglu = self.activation_name in ("situglu", "situglu_low")
+        self.is_waleed = self.activation_name in ("waleed", "waleedglu_low")
+        self.is_waleed10 = self.activation_name in ("waleed10", "silu-waleed10")
+        self.has_sigmoid_gate = self.activation_name.startswith("situglu")
+
+        # ---- Activation function (for non-PowLU) ----
+        if self.mlp_type == "glu":
+            if self.is_powlu:
+                self.act_fn = None
+            elif self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
+                self.act_fn = None
+            elif self.activation_name == "waleed10":
+                self.act_fn = GLUActivationRegistry.get("linear")
+            elif self.activation_name == "silu-waleed10":
+                self.act_fn = GLUActivationRegistry.get("silu")
+            else:
+                self.act_fn = GLUActivationRegistry.get(self.activation_name)
+        else:  # mlp
+            if self.is_powlu:
+                raise ValueError("PowLU is only supported with mlp_type='glu'")
+            if self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
+                raise ValueError(f"Activation '{self.activation_name}' requires GLU.")
+            elif self.activation_name == "waleed10":
+                self.act_fn = GLUActivationRegistry.get("linear")
+            elif self.activation_name == "silu-waleed10":
+                self.act_fn = GLUActivationRegistry.get("silu")
+            else:
+                self.act_fn = GLUActivationRegistry.get(self.activation_name)
+
+        # beta values for situglu/waleed
+        if self.activation_name in ("situglu_low", "waleedglu_low"):
+            self.beta1 = 2.5
+            self.beta2 = 4.0
+        else:
+            self.beta1 = 4.0
+            self.beta2 = 25.0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.mlp_type == "glu":
+            gate = self.gate_proj(x)
+            up = self.up_proj(x)
+
+            # ----- PowLU branch -----
+            if self.is_powlu:
+                m = self.powlu_m
+                sig_gate = torch.sigmoid(gate)
+                pos_mask = gate > 0
+                powlu_pos = gate ** (m / (torch.sqrt(gate) + 1)) * sig_gate
+                powlu_neg = gate * sig_gate
+                activated_gate = torch.where(pos_mask, powlu_pos, powlu_neg)
+                hidden = activated_gate * up
+
+            # ----- situglu / waleed branches -----
+            elif self.is_situglu or self.is_waleed:
+                if self.has_sigmoid_gate:
+                    gate = self.beta1 * torch.tanh(gate / self.beta1) * torch.sigmoid(gate)
+                else:
+                    gate = self.beta1 * torch.tanh(gate / self.beta1)
+                up = self.beta2 * torch.tanh(up / self.beta2)
+                hidden = gate * up
+
+            # ----- Standard GLU -----
+            else:
+                hidden = self.act_fn(gate) * up
+
+            out = self.down_proj(hidden)
+
+        else:  # mlp
+            hidden = self.act_fn(self.up_proj(x))
+            out = self.down_proj(hidden)
+
+        # Optional Waleed clipping
+        if self.is_waleed10:
+            out = self.waleed_beta * torch.tanh(out / self.waleed_beta)
+
+        return out
+
+
+class TinyLlamaDecoderLayer(nn.Module):
+    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LlamaAttention(config=config, layer_idx=layer_idx)
+        self.mlp = TinyLlamaMLP(config)
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None,
+                position_embeddings=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,)
+
+
+def _build_causal_mask(attention_mask, seq_len, dtype, device):
+    min_value = torch.finfo(dtype).min
+    causal = torch.full((seq_len, seq_len), min_value, dtype=dtype, device=device)
+    causal = torch.triu(causal, diagonal=1)
+    causal = causal[None, None, :, :]
+    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 = attention_mask[:, None, None, :].to(device) == 0
+    causal = causal.masked_fill(padding, min_value)
+    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)
+        self.rotary_emb = LlamaRotaryEmbedding(config=config)
+        self.post_init()
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                inputs_embeds=None, return_dict=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
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        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
+    _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=None, attention_mask=None, position_ids=None,
+                inputs_embeds=None, labels=None, return_dict=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,
+        }
+
+# ============================================================================
+# END OF MODEL DEFINITIONS
+# ============================================================================
+
+# Percentiles we report everywhere
+PERCENTILES = [25.0, 50.0, 90.0, 99.0, 99.9]
+
+
+@dataclass
+class TensorStats:
+    mean: float
+    max_abs: float
+    std: float
+    norm: float
+    numel: int
+    min_val: float
+    max_val: float
+    range_val: float
+    p25: float
+    p50: float
+    p90: float
+    p99: float
+    p999: float
+
+    @staticmethod
+    def zero() -> "TensorStats":
+        return TensorStats(
+            mean=0.0, max_abs=0.0, std=0.0, norm=0.0, numel=0,
+            min_val=0.0, max_val=0.0, range_val=0.0,
+            p25=0.0, p50=0.0, p90=0.0, p99=0.0, p999=0.0,
+        )
+
+
+@dataclass
+class ModelResult:
+    model_name: str
+    dataset_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
+    checkpoint_step: int = 0
+
+
+# ============================================================================
+# ACTIVATION HOOK MANAGER — with name normalization & extra hooks
+# ============================================================================
+
+class ActivationHookManager:
+    def __init__(
+        self,
+        store_dtype: torch.dtype = torch.float16,
+        exclude_substrings: Optional[List[str]] = None,
+        max_elements_per_tensor: Optional[int] = 500_000,
+        disk_offload: bool = False,
+        disk_dir: Optional[str] = None,
+        capture_pre_down: bool = True,
+        capture_layer_output: bool = True,
+    ):
+        self.raw_values: Dict[str, List[torch.Tensor]] = defaultdict(list)
+        self.hooks = []
+        self._attention_mask: Optional[torch.Tensor] = None
+        self.store_dtype = store_dtype
+        self.exclude_substrings = exclude_substrings if exclude_substrings is not None else ["lm_head"]
+        self.max_elements_per_tensor = max_elements_per_tensor
+        self._per_batch_quota: Optional[int] = None
+        self._counts: Dict[str, int] = defaultdict(int)
+        self.disk_offload = disk_offload
+        self.disk_dir = disk_dir
+        self._file_handles: Dict[str, "object"] = {}
+        self.capture_pre_down = capture_pre_down
+        self.capture_layer_output = capture_layer_output
+        if self.disk_offload:
+            self.disk_dir = disk_dir or tempfile.mkdtemp(prefix="act_cache_")
+            os.makedirs(self.disk_dir, exist_ok=True)
+
+    @staticmethod
+    def _safe_filename(name: str) -> str:
+        return name.replace("/", "__").replace(".", "_") + ".bin"
+
+    def set_expected_batches(self, num_batches: int):
+        if self.max_elements_per_tensor is not None:
+            self._per_batch_quota = max(1, self.max_elements_per_tensor // max(1, num_batches))
+        else:
+            self._per_batch_quota = None
+
+    def set_attention_mask(self, attention_mask: Optional[torch.Tensor]):
+        self._attention_mask = attention_mask
+
+    def _make_hook(self, name: str, is_pre: bool = False):
+        """Factory for hook functions. If is_pre, captures input; otherwise output."""
+        def hook(module, input_or_output, output=None):
+            # For pre_hook: input_or_output is the input tuple, output is None.
+            # For post_hook: input_or_output is the input tuple (ignored), output is the module output.
+            if is_pre:
+                # Capture the first input tensor (usually the only one)
+                tensor = input_or_output[0] if isinstance(input_or_output, tuple) else input_or_output
+            else:
+                # For post hook, we capture output
+                if isinstance(output, torch.Tensor):
+                    tensor = output
+                elif isinstance(output, tuple) and len(output) > 0 and isinstance(output[0], torch.Tensor):
+                    tensor = output[0]
+                else:
+                    return
+
+            if not torch.is_tensor(tensor):
+                return
+
+            # Apply attention mask if available (only for post hooks that have matching shape)
+            mask = self._attention_mask
+            if not is_pre and mask is not None and tensor.dim() >= 2 and tensor.shape[0] == mask.shape[0] and tensor.shape[1] == mask.shape[1]:
+                flat = tensor[mask.bool()].reshape(-1)
+            else:
+                flat = tensor.reshape(-1)
+
+            if flat.numel() == 0:
+                return
+
+            if self._per_batch_quota is not None and flat.numel() > self._per_batch_quota:
+                idx = torch.randperm(flat.numel(), device=flat.device)[:self._per_batch_quota]
+                sampled = flat[idx]
+            else:
+                sampled = flat
+
+            if self.disk_offload:
+                self._write_to_disk(name, sampled)
+            else:
+                self.raw_values[name].append(sampled.detach().to(self.store_dtype))
+            self._counts[name] += sampled.numel()
+        return hook
+
+    def _write_to_disk(self, name: str, flat: torch.Tensor):
+        arr = flat.detach().to(self.store_dtype).cpu().numpy()
+        fh = self._file_handles.get(name)
+        if fh is None:
+            path = os.path.join(self.disk_dir, self._safe_filename(name))
+            fh = open(path, "ab")
+            self._file_handles[name] = fh
+        arr.tofile(fh)
+
+    def register_hooks(self, model: torch.nn.Module):
+        """Register hooks for MLP activations, pre-down, and layer outputs."""
+        # 1. Standard hooks on gate_proj, up_proj, down_proj (post-activation)
+        for name, module in model.named_modules():
+            if any(x in name for x in ["gate_proj", "up_proj", "down_proj"]):
+                # Normalize: W_g/W_u already handled in canonical names; for down_proj keep as is
+                # For down_proj we also want a pre-hook separately
+                # For post-hook we use the original name (already canonical for gate/up)
+                norm_name = name.replace("W_g", "gate_proj").replace("W_u", "up_proj")
+                hook = module.register_forward_hook(self._make_hook(norm_name, is_pre=False))
+                self.hooks.append(hook)
+
+        # 2. Pre-down_proj hook (captures input to down_proj)
+        if self.capture_pre_down:
+            for name, module in model.named_modules():
+                if "down_proj" in name:
+                    # Name: replace "down_proj" with "pre_down"
+                    pre_name = name.replace("down_proj", "pre_down")
+                    hook = module.register_forward_pre_hook(self._make_hook(pre_name, is_pre=True))
+                    self.hooks.append(hook)
+
+        # 3. Layer output hook (captures final hidden_states from each decoder layer)
+        if self.capture_layer_output:
+            for name, module in model.named_modules():
+                if isinstance(module, TinyLlamaDecoderLayer):
+                    # Use the module's full name, e.g., model.layers.0
+                    out_name = f"{name}.output"
+                    hook = module.register_forward_hook(self._make_hook(out_name, is_pre=False))
+                    self.hooks.append(hook)
+
+    def clear(self):
+        self.raw_values.clear()
+        self._counts.clear()
+        for fh in self._file_handles.values():
+            try:
+                fh.close()
+            except Exception:
+                pass
+        self._file_handles.clear()
+        if self.disk_offload and self.disk_dir:
+            shutil.rmtree(self.disk_dir, ignore_errors=True)
+
+    def remove_hooks(self):
+        for hook in self.hooks:
+            hook.remove()
+        self.hooks.clear()
+
+    @staticmethod
+    def _compute_stats_from_tensor(flat: torch.Tensor) -> TensorStats:
+        flat32 = flat.to(torch.float32)
+        numel = flat32.numel()
+        mean = flat32.mean().item()
+        std = flat32.std(unbiased=False).item()
+        norm = flat32.norm().item()
+        min_val = flat32.min().item()
+        max_val = flat32.max().item()
+        max_abs = max(abs(min_val), abs(max_val))
+
+        q = torch.tensor([p / 100.0 for p in PERCENTILES], device=flat32.device, dtype=torch.float32)
+        try:
+            quantiles = torch.quantile(flat32, q).tolist()
+        except RuntimeError:
+            import numpy as np
+            arr = flat32.cpu().numpy()
+            quantiles = np.percentile(arr, PERCENTILES).tolist()
+
+        p25, p50, p90, p99, p999 = quantiles
+
+        return TensorStats(
+            mean=mean, max_abs=max_abs, std=std, norm=norm, numel=numel,
+            min_val=min_val, max_val=max_val, range_val=max_val - min_val,
+            p25=p25, p50=p50, p90=p90, p99=p99, p999=p999,
+        )
+
+    def get_all_stats(self) -> Tuple[Dict[str, "TensorStats"], "TensorStats"]:
+        if self.disk_offload:
+            return self._get_all_stats_from_disk()
+
+        layer_stats: Dict[str, TensorStats] = {}
+        for name, chunks in self.raw_values.items():
+            if not chunks:
+                continue
+            flat = torch.cat(chunks)
+            layer_stats[name] = self._compute_stats_from_tensor(flat)
+            del chunks, flat
+            self.raw_values[name] = []
+
+        global_stats = TensorStats.zero()
+        return layer_stats, global_stats
+
+    def _get_all_stats_from_disk(self) -> Tuple[Dict[str, "TensorStats"], "TensorStats"]:
+        for fh in self._file_handles.values():
+            try:
+                fh.close()
+            except Exception:
+                pass
+        self._file_handles.clear()
+
+        np_dtype = np.float16 if self.store_dtype == torch.float16 else np.float32
+        layer_stats: Dict[str, TensorStats] = {}
+
+        for name in list(self._counts.keys()):
+            path = os.path.join(self.disk_dir, self._safe_filename(name))
+            if not os.path.exists(path):
+                continue
+            arr = np.fromfile(path, dtype=np_dtype)
+            os.remove(path)
+            if arr.size == 0:
+                continue
+            t = torch.from_numpy(arr.astype(np.float32))
+            layer_stats[name] = self._compute_stats_from_tensor(t)
+            del t, arr
+
+        shutil.rmtree(self.disk_dir, ignore_errors=True)
+        global_stats = TensorStats.zero()
+        return layer_stats, global_stats
+
+
+# ============================================================================
+# LLM ANALYZER — TRAINING‑STYLE TOKENIZATION, FULL DATASET SUPPORT
+# ============================================================================
+
+class LLMAnalyzer:
+    def __init__(
+        self,
+        max_length: int = 512,
+        max_samples: Optional[int] = 1000,
+        batch_size: int = 4,
+        wandb_project: Optional[str] = None,
+        cbt_simple_max_chars: int = 120,
+        min_line_chars: int = 20,
+        activation_store_dtype: torch.dtype = torch.float16,
+        track_activations: bool = True,
+        activation_hook_exclude: Optional[List[str]] = None,
+        max_elements_per_tensor: Optional[int] = 500_000,
+        activation_disk_offload: bool = False,
+        activation_disk_dir: Optional[str] = None,
+        capture_pre_down: bool = True,
+        capture_layer_output: bool = True,
+    ):
+        if not torch.cuda.is_available():
+            raise RuntimeError("This script requires a GPU. CUDA is not available.")
+        self.device = "cuda"
+        self.dtype = torch.float32
+
+        self.max_length = max_length
+        self.max_samples = max_samples
+        self.batch_size = batch_size
+        self.wandb_project = wandb_project or os.environ.get("WANDB_PROJECT", "llm-activation-analysis")
+        self.cbt_simple_max_chars = cbt_simple_max_chars
+        self.min_line_chars = min_line_chars
+        self.activation_store_dtype = activation_store_dtype
+        self.track_activations = track_activations
+        self.activation_hook_exclude = activation_hook_exclude if activation_hook_exclude is not None else ["lm_head"]
+        self.max_elements_per_tensor = max_elements_per_tensor
+        self.activation_disk_offload = activation_disk_offload
+        self.activation_disk_dir = activation_disk_dir
+        self.capture_pre_down = capture_pre_down
+        self.capture_layer_output = capture_layer_output
+        self._cache = {}
+        self._dataset_cache: Dict[str, List[str]] = {}
+
+    # ------------------------------------------------------------------
+    # Dataset loaders – now respect max_samples=None (load all)
+    # ------------------------------------------------------------------
+
+    def _load_tinystories(self, split: str = "validation") -> List[str]:
+        print(f"[Dataset] Loading TinyStories (split='{split}')...")
+        ds = hf_load_dataset("roneneldan/TinyStories", split=split)
+        lines = []
+        for example in ds:
+            text = (example.get("text") or "").strip()
+            if len(text) >= self.min_line_chars:
+                lines.append(text)
+            if self.max_samples is not None and len(lines) >= self.max_samples:
+                break
+        print(f"[Dataset] Loaded {len(lines)} TinyStories examples")
+        return lines
+
+    def _load_simplestories(self, split: str = "test") -> List[str]:
+        print(f"[Dataset] Loading SimpleStories (split='{split}')...")
+        ds = hf_load_dataset("SimpleStories/SimpleStories", split=split)
+        lines = []
+        for example in ds:
+            text = (example.get("story") or "").strip()
+            if len(text) >= self.min_line_chars:
+                lines.append(text)
+            if self.max_samples is not None and len(lines) >= self.max_samples:
+                break
+        print(f"[Dataset] Loaded {len(lines)} SimpleStories examples")
+        return lines
+
+    def _load_cbt_simple(self, split: str = "test") -> List[str]:
+        print(f"[Dataset] Loading CBT-simple (cam-cst/cbt, raw, split='{split}', "
+              f"max_chars={self.cbt_simple_max_chars})...")
+        ds = hf_load_dataset("cam-cst/cbt", "raw", split=split)
+        lines = []
+        for example in ds:
+            content = example.get("content", "")
+            if not content:
+                continue
+            for chunk in content.replace("\r\n", "\n").split("\n"):
+                chunk = chunk.strip()
+                if self.min_line_chars <= len(chunk) <= self.cbt_simple_max_chars:
+                    lines.append(chunk)
+            if self.max_samples is not None and len(lines) >= self.max_samples:
+                break
+        print(f"[Dataset] Loaded {len(lines)} simple CBT sentences")
+        return lines[:self.max_samples] if self.max_samples is not None else lines
+
+    DATASET_LOADERS = {
+        "tinystories": lambda self, split: self._load_tinystories(split or "validation"),
+        "simplestories": lambda self, split: self._load_simplestories(split or "test"),
+        "cbt_simple": lambda self, split: self._load_cbt_simple(split or "test"),
+    }
+
+    def load_dataset(self, dataset_name: str, split: Optional[str] = None) -> List[str]:
+        cache_key = f"{dataset_name}:{split}"
+        if cache_key in self._dataset_cache:
+            return self._dataset_cache[cache_key]
+        if dataset_name not in self.DATASET_LOADERS:
+            raise ValueError(f"Unknown dataset '{dataset_name}'. Options: {list(self.DATASET_LOADERS)}")
+        texts = self.DATASET_LOADERS[dataset_name](self, split)
+        self._dataset_cache[cache_key] = texts
+        return texts
+
+    # ------------------------------------------------------------------
+    # Training‑style block preparation (no padding, no BOS)
+    # ------------------------------------------------------------------
+
+    def _prepare_blocks(self, texts: List[str], tokenizer, max_seq_len: int) -> List[List[int]]:
+        """
+        Tokenize texts exactly as in training:
+          - add_special_tokens=False
+          - append EOS token (tokenizer.eos_token_id)
+          - concatenate all stories into one flat list
+          - split into blocks of size max_seq_len, discarding any remainder
+        Returns a list of token ID lists, each exactly max_seq_len long.
+        """
+        eos_id = tokenizer.eos_token_id  # should be 0 for this tokenizer
+        flat_tokens = []
+        for text in texts:
+            ids = tokenizer.encode(text, add_special_tokens=False)
+            ids.append(eos_id)
+            flat_tokens.extend(ids)
+
+        num_blocks = len(flat_tokens) // max_seq_len
+        blocks = []
+        for i in range(num_blocks):
+            start = i * max_seq_len
+            end = start + max_seq_len
+            blocks.append(flat_tokens[start:end])
+        return blocks
+
+    # ------------------------------------------------------------------
+    # Model loading – reads waleed_beta and powlu_m from config
+    # ------------------------------------------------------------------
+
+    def load_model(self, model_name: str):
+        if model_name in self._cache:
+            return self._cache[model_name]
+
+        print(f"[Loading] {model_name} ...")
+        parts = model_name.split("/")
+        if len(parts) > 2:
+            repo_id = "/".join(parts[:2])
+            subfolder = "/".join(parts[2:])
+        else:
+            repo_id = model_name
+            subfolder = None
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            "w-ahmad/tiny-stories-tokenizer",
+            trust_remote_code=True,
+            padding_side="left"
+        )
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        config = TinyLlamaConfig.from_pretrained(
+            repo_id, subfolder=subfolder, trust_remote_code=True
+        )
+
+        model = TinyLlamaForCausalLM.from_pretrained(
+            repo_id,
+            config=config,
+            subfolder=subfolder,
+            torch_dtype=self.dtype,
+            trust_remote_code=True,
+        )
+        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} | waleed_beta: {config.waleed_beta} | powlu_m: {config.powlu_m}")
+        return tokenizer, model, config, num_params
+
+    # ------------------------------------------------------------------
+    # Discover checkpoints (unchanged)
+    # ------------------------------------------------------------------
+
+    def discover_checkpoints(self) -> Dict[str, List[Tuple[int, str]]]:
+        from huggingface_hub import list_repo_files
+        print(f"🔍 Scanning Hugging Face repo: {REPO_ID}/{BASE_PATH} ...")
+        try:
+            all_files = list_repo_files(REPO_ID)
+        except Exception as e:
+            print(f"❌ Failed to fetch from Hugging Face Hub: {e}")
+            print("💡 Falling back to manual list.")
+            return {
+                "glu-silu-waleed10_run": [(100, "checkpoint-100")],
+            }
+
+        run_checkpoints = defaultdict(list)
+        prefix = BASE_PATH + "/"
+        for file_path in all_files:
+            if not file_path.startswith(prefix):
+                continue
+            rel_path = file_path[len(prefix):]
+            parts = rel_path.split("/")
+            if len(parts) < 2:
+                continue
+            run_folder = parts[0]
+            ckpt_folder = parts[1]
+            if not ckpt_folder.startswith("checkpoint-"):
+                continue
+            if not (file_path.endswith("config.json") or file_path.endswith("model.safetensors")):
+                continue
+            try:
+                ckpt_num = int(ckpt_folder.replace("checkpoint-", ""))
+            except ValueError:
+                continue
+            run_checkpoints[run_folder].append((ckpt_num, ckpt_folder))
+
+        result = {}
+        for run_folder, ckpts in run_checkpoints.items():
+            seen = set()
+            unique = []
+            for ckpt_num, ckpt_folder in sorted(ckpts, key=lambda x: x[0]):
+                if ckpt_num not in seen:
+                    seen.add(ckpt_num)
+                    unique.append((ckpt_num, ckpt_folder))
+            result[run_folder] = unique
+
+        print(f"✅ Found {len(result)} run folders:")
+        for run_folder, ckpts in sorted(result.items()):
+            print(f"   - {run_folder}: {len(ckpts)} checkpoints ({ckpts[0][0]}...{ckpts[-1][0]})")
+        return result
+
+    # ------------------------------------------------------------------
+    # Single checkpoint evaluation
+    # ------------------------------------------------------------------
+
+    def _evaluate_checkpoint(
+        self,
+        model_name: str,
+        dataset_name: str,
+        blocks: List[List[int]],
+    ) -> ModelResult:
+        tokenizer, model, config, num_params = self.load_model(model_name)
+
+        total_loss = 0.0
+        total_tokens = 0
+        num_batches = (len(blocks) + self.batch_size - 1) // self.batch_size
+
+        hook_mgr = None
+        if self.track_activations:
+            hook_mgr = ActivationHookManager(
+                store_dtype=self.activation_store_dtype,
+                exclude_substrings=self.activation_hook_exclude,
+                max_elements_per_tensor=self.max_elements_per_tensor,
+                disk_offload=self.activation_disk_offload,
+                disk_dir=self.activation_disk_dir,
+                capture_pre_down=self.capture_pre_down,
+                capture_layer_output=self.capture_layer_output,
+            )
+            hook_mgr.register_hooks(model)
+            hook_mgr.set_expected_batches(num_batches)
+
+        if hook_mgr is not None:
+            n_hooks = len(hook_mgr.hooks)
+            mode = "disk" if self.activation_disk_offload else "in-memory"
+            if self.max_elements_per_tensor is None:
+                print(f"[Memory] ⚠️  {n_hooks} modules hooked ({mode}) with NO cap.")
+            else:
+                est_mb = (n_hooks * self.max_elements_per_tensor * 2) / (1024 ** 2)
+                where = f"on disk at {hook_mgr.disk_dir}" if self.activation_disk_offload else "in RAM/VRAM"
+                print(f"[Memory] {n_hooks} modules hooked, capped at {self.max_elements_per_tensor:,} elements/tensor (~{est_mb:.0f} MB ceiling, stored {where}).")
+
+        try:
+            desc = f"Eval {model_name.split('/')[-2]}/{model_name.split('/')[-1]}"
+            for batch_idx, i in enumerate(tqdm(range(0, len(blocks), self.batch_size), desc=desc, total=num_batches)):
+                batch_blocks = blocks[i:i+self.batch_size]
+                input_ids = torch.tensor(batch_blocks, dtype=torch.long, device=self.device)
+                attention_mask = None
+                labels = input_ids.clone()
+
+                if hook_mgr is not None:
+                    hook_mgr.set_attention_mask(None)
+
+                with torch.no_grad():
+                    outputs = model(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        labels=labels,
+                    )
+
+                logits = outputs.logits
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+                token_losses = loss_fct(
+                    shift_logits.view(-1, shift_logits.size(-1)),
+                    shift_labels.view(-1)
+                )
+                num_valid_tokens = token_losses.numel()
+                batch_loss = token_losses.sum().item()
+                total_loss += batch_loss
+                total_tokens += num_valid_tokens
+
+                if batch_idx % 20 == 19:
+                    torch.cuda.empty_cache()
+
+            if hook_mgr is not None:
+                hook_mgr.remove_hooks()
+                final_layer_stats, global_stats = hook_mgr.get_all_stats()
+            else:
+                final_layer_stats, global_stats = {}, TensorStats.zero()
+        except Exception:
+            if hook_mgr is not None:
+                hook_mgr.remove_hooks()
+                hook_mgr.clear()
+            raise
+
+        avg_loss = total_loss / max(total_tokens, 1)
+        perplexity = torch.exp(torch.tensor(avg_loss)).item()
+
+        checkpoint_step = 0
+        for part in model_name.split("/"):
+            if part.startswith("checkpoint-"):
+                try:
+                    checkpoint_step = int(part.replace("checkpoint-", ""))
+                except ValueError:
+                    pass
+
+        result = ModelResult(
+            model_name=model_name,
+            dataset_name=dataset_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,
+            checkpoint_step=checkpoint_step,
+        )
+
+        self._cache.pop(model_name, None)
+        del model
+        torch.cuda.empty_cache()
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Logging (updated to include new stats)
+    # ------------------------------------------------------------------
+
+    def _log_checkpoint_to_wandb(self, result: ModelResult, step: int):
+        wandb.log({
+            "final/loss": result.loss,
+            "final/perplexity": result.perplexity,
+            "final/num_tokens": result.num_tokens,
+        }, step=step)
+
+        if not self.track_activations:
+            return
+
+        # Log per-tensor stats
+        for tensor_name, stats in result.layer_acts.items():
+            clean_name = tensor_name.replace(".", "/")
+            # Determine the kind from the name: pre_down, output, or others (gate/up/down)
+            if "pre_down" in tensor_name:
+                prefix = f"train/{clean_name}/pre_down"
+            elif "output" in tensor_name and "layers" in tensor_name:
+                prefix = f"train/{clean_name}/output"
+            else:
+                prefix = f"train/{clean_name}/act"   # existing activation stats
+
+            wandb.log({
+                f"{prefix}/mean": stats.mean,
+                f"{prefix}/max_abs": stats.max_abs,
+                f"{prefix}/std": stats.std,
+                f"{prefix}/norm": stats.norm,
+                f"{prefix}/min": stats.min_val,
+                f"{prefix}/max": stats.max_val,
+                f"{prefix}/range": stats.range_val,
+                f"{prefix}/p25": stats.p25,
+                f"{prefix}/p50": stats.p50,
+                f"{prefix}/p90": stats.p90,
+                f"{prefix}/p99": stats.p99,
+                f"{prefix}/p999": stats.p999,
+            }, step=step)
+
+        # Tables: include all tensors with a column indicating the type
+        table_data = []
+        for tensor_name, stats in sorted(result.layer_acts.items()):
+            # Determine type for table
+            if "pre_down" in tensor_name:
+                ttype = "pre_down"
+            elif "output" in tensor_name and "layers" in tensor_name:
+                ttype = "output"
+            else:
+                ttype = "activation"
+            table_data.append([
+                tensor_name,
+                ttype,
+                stats.mean,
+                stats.max_abs,
+                stats.std,
+                stats.norm,
+                stats.min_val,
+                stats.max_val,
+                stats.range_val,
+                stats.p25,
+                stats.p50,
+                stats.p90,
+                stats.p99,
+                stats.p999,
+                stats.numel,
+            ])
+
+        if table_data:
+            table = wandb.Table(
+                columns=["tensor_name", "type", "mean", "max_abs", "std", "norm",
+                         "min", "max", "range", "p25", "p50", "p90", "p99", "p999", "numel"],
+                data=table_data
+            )
+            wandb.log({"activation_table": table}, step=step)
+
+    # ------------------------------------------------------------------
+    # Run‑folder evaluation – includes waleed_beta and powlu_m in name
+    # ------------------------------------------------------------------
+
+    def _evaluate_run_folder(
+        self,
+        run_folder: str,
+        dataset_name: str,
+        blocks: List[List[int]],
+        checkpoints: List[Tuple[int, str]],
+    ) -> List[ModelResult]:
+        first_model_name = f"{REPO_ID}/{BASE_PATH}/{run_folder}/{checkpoints[0][1]}"
+        _, _, first_config, first_num_params = self.load_model(first_model_name)
+
+        effective_beta = first_config.waleed_beta
+        powlu_m = first_config.powlu_m
+        beta_str = f"_beta{effective_beta}" if effective_beta is not None else ""
+        powlu_str = f"_powlu{powlu_m}" if powlu_m is not None else ""
+
+        run_name = f"{run_folder}__{dataset_name}{beta_str}{powlu_str}"
+
+        wandb.init(
+            project=self.wandb_project,
+            name=run_name,
+            config={
+                "run_folder": run_folder,
+                "dataset": dataset_name,
+                "max_length": self.max_length,
+                "max_samples": self.max_samples,
+                "batch_size": self.batch_size,
+                "dtype": "float32",
+                "num_params": first_num_params,
+                "num_layers": first_config.num_hidden_layers,
+                "hidden_size": first_config.hidden_size,
+                "waleed_beta": effective_beta,
+                "powlu_m": powlu_m,
+                "checkpoints": [c[0] for c in checkpoints],
+                "capture_pre_down": self.capture_pre_down,
+                "capture_layer_output": self.capture_layer_output,
+            },
+            reinit=True
+        )
+
+        self._cache.pop(first_model_name, None)
+        torch.cuda.empty_cache()
+
+        results = []
+        for ckpt_num, ckpt_folder in checkpoints:
+            model_name = f"{REPO_ID}/{BASE_PATH}/{run_folder}/{ckpt_folder}"
+            try:
+                result = self._evaluate_checkpoint(model_name, dataset_name, blocks)
+                self._log_checkpoint_to_wandb(result, step=ckpt_num)
+                results.append(result)
+                print(f"   ✅ checkpoint-{ckpt_num}: loss={result.loss:.4f}  ppl={result.perplexity:.2f}")
+            except Exception as e:
+                print(f"   ❌ checkpoint-{ckpt_num} failed: {e}")
+                import traceback
+                traceback.print_exc()
+                continue
+
+        wandb.finish()
+        return results
+
+    # ------------------------------------------------------------------
+    # Main compute – prepares training‑style blocks once per dataset
+    # ------------------------------------------------------------------
+
+    def compute(
+        self,
+        dataset_names: List[str],
+        run_folders: Optional[List[str]] = None,
+    ) -> Dict[str, Dict[str, List[ModelResult]]]:
+        all_run_folders = self.discover_checkpoints()
+
+        if run_folders is not None:
+            filtered = {}
+            for name in run_folders:
+                if name in all_run_folders:
+                    filtered[name] = all_run_folders[name]
+                else:
+                    print(f"⚠️  Run folder '{name}' not found. Skipping.")
+            all_run_folders = filtered
+            if not all_run_folders:
+                print("❌ No valid run folders to evaluate. Exiting.")
+                return {}
+
+        # Load tokenizer once (same for all datasets)
+        tokenizer = AutoTokenizer.from_pretrained("w-ahmad/tiny-stories-tokenizer", trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        blocks_by_dataset = {}
+        for dataset_name in dataset_names:
+            texts = self.load_dataset(dataset_name)
+            blocks = self._prepare_blocks(texts, tokenizer, self.max_length)
+            blocks_by_dataset[dataset_name] = blocks
+            print(f"[Dataset] {dataset_name} → {len(blocks)} blocks of length {self.max_length}")
+
+        results_by_dataset: Dict[str, Dict[str, List[ModelResult]]] = {
+            name: {} for name in dataset_names
+        }
+
+        for dataset_name in dataset_names:
+            blocks = blocks_by_dataset[dataset_name]
+            for run_folder, checkpoints in sorted(all_run_folders.items()):
+                print(f"\n🚀 Evaluating {run_folder} on {dataset_name} ({len(checkpoints)} checkpoints)...")
+                results = self._evaluate_run_folder(run_folder, dataset_name, blocks, checkpoints)
+                results_by_dataset[dataset_name][run_folder] = results
+
+        return results_by_dataset
+
+    # ------------------------------------------------------------------
+    # Reporting (updated to include new stats in the top‑5 summary)
+    # ------------------------------------------------------------------
+
+    def print_report(self, results_by_dataset: Dict[str, Dict[str, List[ModelResult]]]):
+        for dataset_name, run_results in results_by_dataset.items():
+            print("\n" + "=" * 150)
+            print(f" Dataset: {dataset_name}")
+            print("=" * 150)
+
+            print(f"{'Run Folder':<38} {'Ckpt':>7} {'Loss':>9} {'PPL':>9} {'Tokens':>8}  {'Num Layers':>10}")
+            print("-" * 90)
+
+            for run_folder, results in sorted(run_results.items()):
+                if not results:
+                    continue
+                latest = max(results, key=lambda r: r.checkpoint_step)
+                name = run_folder[:36]
+                print(
+                    f"{name:<38} "
+                    f"{latest.checkpoint_step:>7} "
+                    f"{latest.loss:>9.4f} "
+                    f"{latest.perplexity:>9.2f} "
+                    f"{latest.num_tokens:>8} "
+                    f"{latest.num_layers:>10}"
+                )
+
+        # Show top‑5 per tensor type for the latest checkpoint
+        print("\n[Per-Tensor Max Abs Top 5 (latest checkpoint per run)]")
+        for dataset_name, run_results in results_by_dataset.items():
+            for run_folder, results in sorted(run_results.items()):
+                if not results:
+                    continue
+                latest = max(results, key=lambda r: r.checkpoint_step)
+                name = f"{run_folder} [{dataset_name}]"
+                # Group by type
+                activations = []
+                pre_downs = []
+                outputs = []
+                for tensor_name, stats in latest.layer_acts.items():
+                    if "pre_down" in tensor_name:
+                        pre_downs.append((tensor_name, stats))
+                    elif "output" in tensor_name and "layers" in tensor_name:
+                        outputs.append((tensor_name, stats))
+                    else:
+                        activations.append((tensor_name, stats))
+
+                # Top 5 for each type
+                for group_name, group_list in [
+                    ("Activations", activations),
+                    ("Pre-down", pre_downs),
+                    ("Layer outputs", outputs)
+                ]:
+                    if not group_list:
+                        continue
+                    sorted_group = sorted(group_list, key=lambda x: x[1].max_abs, reverse=True)[:5]
+                    print(f"\n  {name} (ckpt-{latest.checkpoint_step}) – {group_name}:")
+                    for tensor_name, stats in sorted_group:
+                        print(
+                            f"    {tensor_name:<50} max_abs={stats.max_abs:>9.4f}  "
+                            f"mean={stats.mean:>9.4f}  std={stats.std:>9.4f}  norm={stats.norm:>9.4f}"
+                        )
+
+    def export_json(self, results_by_dataset: Dict[str, Dict[str, List[ModelResult]]], path: str):
+        data = []
+        for dataset_name, run_results in results_by_dataset.items():
+            for run_folder, results in run_results.items():
+                for r in results:
+                    entry = {
+                        "run_folder": run_folder,
+                        "model": r.model_name,
+                        "dataset": r.dataset_name,
+                        "checkpoint_step": r.checkpoint_step,
+                        "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}")
+
+
+# ============================================================================
+# DRIVER CODE
+# ============================================================================
+if __name__ == "__main__":
+    # --- HARDCODE WANDB LOGIN (optional, replace with your key) ---
+    # wandb.login(key="wandb_v1_...")
+    # -------------------------------------------------------------
+
+    analyzer = LLMAnalyzer(
+        max_length=512,
+        max_samples=100,           # use None for all samples
+        batch_size=50,
+        wandb_project="RESEARCH-10",
+        activation_disk_offload=False,
+        max_elements_per_tensor=None,
+        activation_hook_exclude=[],
+        track_activations=True,     # set to False to skip all activation collection
+        capture_pre_down=True,      # enable pre-down capture
+        capture_layer_output=True,  # enable layer output capture
+    )
+
+    DATASETS = ["tinystories"]
+    # RUN_FOLDERS = ["glu-powlu-100L_run"]   # uncomment to limit runs
+
+    print(f"\n🚀 Starting evaluation across datasets: {DATASETS}...")
+    results_by_dataset = analyzer.compute(dataset_names=DATASETS)   # add run_folders if needed
+
+    analyzer.print_report(results_by_dataset)
+    analyzer.export_json(results_by_dataset, "evaluation_results_powlu.json")
+
+    print("\n✅ Done! Pre-down and layer output stats are now captured.")

zain/Activation/out/glu-relu-21L_run/checkpoint-100/config.json

@@ -0,0 +1,37 @@
+{
+  "activation": "relu",
+  "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": 21,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "powlu_m": 3.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.16.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096,
+  "waleed_beta": 10.0
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-100/model.safetensors

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

zain/Activation/out/glu-relu-21L_run/checkpoint-100/optimizer.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-100/rng_state.pth

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

zain/Activation/out/glu-relu-21L_run/checkpoint-100/scheduler.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-100/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|>"
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-100/trainer_state.json

@@ -0,0 +1,77 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.003370407819346141,
+  "eval_steps": 100,
+  "global_step": 100,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0006740815638692282,
+      "grad_norm": 1.7578125,
+      "learning_rate": 2.66e-05,
+      "loss": 8.330369567871093,
+      "step": 20
+    },
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.40625,
+      "learning_rate": 5.46e-05,
+      "loss": 8.17730255126953,
+      "step": 40
+    },
+    {
+      "epoch": 0.0020222446916076846,
+      "grad_norm": 1.3359375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.952220916748047,
+      "step": 60
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3125,
+      "learning_rate": 0.0001106,
+      "loss": 7.694232177734375,
+      "step": 80
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "grad_norm": 1.234375,
+      "learning_rate": 0.0001386,
+      "loss": 7.353645324707031,
+      "step": 100
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "eval_loss": 7.14666223526001,
+      "eval_runtime": 8.8883,
+      "eval_samples_per_second": 1071.854,
+      "eval_steps_per_second": 1.125,
+      "step": 100
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 33876973977600.0,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-100/training_args.bin

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/config.json

@@ -0,0 +1,37 @@
+{
+  "activation": "relu",
+  "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": 21,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "powlu_m": 3.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.16.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096,
+  "waleed_beta": 10.0
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/model.safetensors

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/optimizer.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/rng_state.pth

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/scheduler.pt

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

zain/Activation/out/glu-relu-21L_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|>"
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/trainer_state.json

@@ -0,0 +1,464 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.03370407819346141,
+  "eval_steps": 100,
+  "global_step": 1000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0006740815638692282,
+      "grad_norm": 1.7578125,
+      "learning_rate": 2.66e-05,
+      "loss": 8.330369567871093,
+      "step": 20
+    },
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.40625,
+      "learning_rate": 5.46e-05,
+      "loss": 8.17730255126953,
+      "step": 40
+    },
+    {
+      "epoch": 0.0020222446916076846,
+      "grad_norm": 1.3359375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.952220916748047,
+      "step": 60
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3125,
+      "learning_rate": 0.0001106,
+      "loss": 7.694232177734375,
+      "step": 80
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "grad_norm": 1.234375,
+      "learning_rate": 0.0001386,
+      "loss": 7.353645324707031,
+      "step": 100
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "eval_loss": 7.14666223526001,
+      "eval_runtime": 8.8883,
+      "eval_samples_per_second": 1071.854,
+      "eval_steps_per_second": 1.125,
+      "step": 100
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.1796875,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.9598846435546875,
+      "step": 120
+    },
+    {
+      "epoch": 0.0047185709470845974,
+      "grad_norm": 1.1015625,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.568833923339843,
+      "step": 140
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.046875,
+      "learning_rate": 0.0002226,
+      "loss": 6.253362655639648,
+      "step": 160
+    },
+    {
+      "epoch": 0.006066734074823054,
+      "grad_norm": 0.86328125,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.974322509765625,
+      "step": 180
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0002786,
+      "loss": 5.698452377319336,
+      "step": 200
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "eval_loss": 5.558600902557373,
+      "eval_runtime": 9.1031,
+      "eval_samples_per_second": 1046.564,
+      "eval_steps_per_second": 1.099,
+      "step": 200
+    },
+    {
+      "epoch": 0.00741489720256151,
+      "grad_norm": 0.8046875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.451509857177735,
+      "step": 220
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.0078125,
+      "learning_rate": 0.0003346,
+      "loss": 5.163370895385742,
+      "step": 240
+    },
+    {
+      "epoch": 0.008763060330299966,
+      "grad_norm": 1.0859375,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 4.929123306274414,
+      "step": 260
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 0.84765625,
+      "learning_rate": 0.0003906,
+      "loss": 4.702068710327149,
+      "step": 280
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0004186,
+      "loss": 4.501141738891602,
+      "step": 300
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "eval_loss": 4.396740913391113,
+      "eval_runtime": 8.9474,
+      "eval_samples_per_second": 1064.782,
+      "eval_steps_per_second": 1.118,
+      "step": 300
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.0004466,
+      "loss": 4.323627853393555,
+      "step": 320
+    },
+    {
+      "epoch": 0.011459386585776879,
+      "grad_norm": 0.88671875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.204174041748047,
+      "step": 340
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.92578125,
+      "learning_rate": 0.0005026,
+      "loss": 4.052695846557617,
+      "step": 360
+    },
+    {
+      "epoch": 0.012807549713515335,
+      "grad_norm": 0.6484375,
+      "learning_rate": 0.0005306,
+      "loss": 3.947071075439453,
+      "step": 380
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.73828125,
+      "learning_rate": 0.0005586,
+      "loss": 3.855146026611328,
+      "step": 400
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.8139851093292236,
+      "eval_runtime": 8.9743,
+      "eval_samples_per_second": 1061.59,
+      "eval_steps_per_second": 1.114,
+      "step": 400
+    },
+    {
+      "epoch": 0.014155712841253791,
+      "grad_norm": 0.80859375,
+      "learning_rate": 0.0005866,
+      "loss": 3.76041259765625,
+      "step": 420
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 0.77734375,
+      "learning_rate": 0.0006146,
+      "loss": 3.708443450927734,
+      "step": 440
+    },
+    {
+      "epoch": 0.015503875968992248,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006426,
+      "loss": 3.637612152099609,
+      "step": 460
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.556614303588867,
+      "step": 480
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006986,
+      "loss": 3.5122703552246093,
+      "step": 500
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "eval_loss": 3.4916131496429443,
+      "eval_runtime": 8.8876,
+      "eval_samples_per_second": 1071.946,
+      "eval_steps_per_second": 1.125,
+      "step": 500
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0007,
+      "loss": 3.477622222900391,
+      "step": 520
+    },
+    {
+      "epoch": 0.01820020222446916,
+      "grad_norm": 0.75390625,
+      "learning_rate": 0.0007,
+      "loss": 3.4115901947021485,
+      "step": 540
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 3.359669876098633,
+      "step": 560
+    },
+    {
+      "epoch": 0.01954836535220762,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 3.3188705444335938,
+      "step": 580
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2840240478515623,
+      "step": 600
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "eval_loss": 3.2660911083221436,
+      "eval_runtime": 8.9246,
+      "eval_samples_per_second": 1067.504,
+      "eval_steps_per_second": 1.121,
+      "step": 600
+    },
+    {
+      "epoch": 0.020896528479946073,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0007,
+      "loss": 3.223845672607422,
+      "step": 620
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2053298950195312,
+      "step": 640
+    },
+    {
+      "epoch": 0.022244691607684528,
+      "grad_norm": 0.890625,
+      "learning_rate": 0.0007,
+      "loss": 3.1749242782592773,
+      "step": 660
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 3.134219741821289,
+      "step": 680
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0007,
+      "loss": 3.1245954513549803,
+      "step": 700
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "eval_loss": 3.113696575164795,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.373,
+      "eval_steps_per_second": 1.128,
+      "step": 700
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.0987874984741213,
+      "step": 720
+    },
+    {
+      "epoch": 0.02494101786316144,
+      "grad_norm": 0.73046875,
+      "learning_rate": 0.0007,
+      "loss": 3.08369140625,
+      "step": 740
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.03726863861084,
+      "step": 760
+    },
+    {
+      "epoch": 0.0262891809908999,
+      "grad_norm": 0.64453125,
+      "learning_rate": 0.0007,
+      "loss": 2.9970935821533202,
+      "step": 780
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.625,
+      "learning_rate": 0.0007,
+      "loss": 3.006105422973633,
+      "step": 800
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 2.999389886856079,
+      "eval_runtime": 8.984,
+      "eval_samples_per_second": 1060.442,
+      "eval_steps_per_second": 1.113,
+      "step": 800
+    },
+    {
+      "epoch": 0.027637344118638354,
+      "grad_norm": 0.609375,
+      "learning_rate": 0.0007,
+      "loss": 2.9657209396362303,
+      "step": 820
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 2.9664005279541015,
+      "step": 840
+    },
+    {
+      "epoch": 0.028985507246376812,
+      "grad_norm": 0.65625,
+      "learning_rate": 0.0007,
+      "loss": 2.988846206665039,
+      "step": 860
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.66015625,
+      "learning_rate": 0.0007,
+      "loss": 2.9386875152587892,
+      "step": 880
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "grad_norm": 0.7109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8965709686279295,
+      "step": 900
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "eval_loss": 2.9093611240386963,
+      "eval_runtime": 8.9344,
+      "eval_samples_per_second": 1066.332,
+      "eval_steps_per_second": 1.119,
+      "step": 900
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.9101829528808594,
+      "step": 920
+    },
+    {
+      "epoch": 0.031681833501853725,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 2.8905385971069335,
+      "step": 940
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.74609375,
+      "learning_rate": 0.0007,
+      "loss": 2.8552478790283202,
+      "step": 960
+    },
+    {
+      "epoch": 0.03302999662959218,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8467588424682617,
+      "step": 980
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 2.831803894042969,
+      "step": 1000
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "eval_loss": 2.8364417552948,
+      "eval_runtime": 8.9776,
+      "eval_samples_per_second": 1061.195,
+      "eval_steps_per_second": 1.114,
+      "step": 1000
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 338769739776000.0,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1000/training_args.bin

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/config.json

@@ -0,0 +1,37 @@
+{
+  "activation": "relu",
+  "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": 21,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "powlu_m": 3.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.16.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096,
+  "waleed_beta": 10.0
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/model.safetensors

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/optimizer.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/rng_state.pth

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/scheduler.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/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|>"
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/trainer_state.json

@@ -0,0 +1,507 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.03707448601280755,
+  "eval_steps": 100,
+  "global_step": 1100,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0006740815638692282,
+      "grad_norm": 1.7578125,
+      "learning_rate": 2.66e-05,
+      "loss": 8.330369567871093,
+      "step": 20
+    },
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.40625,
+      "learning_rate": 5.46e-05,
+      "loss": 8.17730255126953,
+      "step": 40
+    },
+    {
+      "epoch": 0.0020222446916076846,
+      "grad_norm": 1.3359375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.952220916748047,
+      "step": 60
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3125,
+      "learning_rate": 0.0001106,
+      "loss": 7.694232177734375,
+      "step": 80
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "grad_norm": 1.234375,
+      "learning_rate": 0.0001386,
+      "loss": 7.353645324707031,
+      "step": 100
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "eval_loss": 7.14666223526001,
+      "eval_runtime": 8.8883,
+      "eval_samples_per_second": 1071.854,
+      "eval_steps_per_second": 1.125,
+      "step": 100
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.1796875,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.9598846435546875,
+      "step": 120
+    },
+    {
+      "epoch": 0.0047185709470845974,
+      "grad_norm": 1.1015625,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.568833923339843,
+      "step": 140
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.046875,
+      "learning_rate": 0.0002226,
+      "loss": 6.253362655639648,
+      "step": 160
+    },
+    {
+      "epoch": 0.006066734074823054,
+      "grad_norm": 0.86328125,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.974322509765625,
+      "step": 180
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0002786,
+      "loss": 5.698452377319336,
+      "step": 200
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "eval_loss": 5.558600902557373,
+      "eval_runtime": 9.1031,
+      "eval_samples_per_second": 1046.564,
+      "eval_steps_per_second": 1.099,
+      "step": 200
+    },
+    {
+      "epoch": 0.00741489720256151,
+      "grad_norm": 0.8046875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.451509857177735,
+      "step": 220
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.0078125,
+      "learning_rate": 0.0003346,
+      "loss": 5.163370895385742,
+      "step": 240
+    },
+    {
+      "epoch": 0.008763060330299966,
+      "grad_norm": 1.0859375,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 4.929123306274414,
+      "step": 260
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 0.84765625,
+      "learning_rate": 0.0003906,
+      "loss": 4.702068710327149,
+      "step": 280
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0004186,
+      "loss": 4.501141738891602,
+      "step": 300
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "eval_loss": 4.396740913391113,
+      "eval_runtime": 8.9474,
+      "eval_samples_per_second": 1064.782,
+      "eval_steps_per_second": 1.118,
+      "step": 300
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.0004466,
+      "loss": 4.323627853393555,
+      "step": 320
+    },
+    {
+      "epoch": 0.011459386585776879,
+      "grad_norm": 0.88671875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.204174041748047,
+      "step": 340
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.92578125,
+      "learning_rate": 0.0005026,
+      "loss": 4.052695846557617,
+      "step": 360
+    },
+    {
+      "epoch": 0.012807549713515335,
+      "grad_norm": 0.6484375,
+      "learning_rate": 0.0005306,
+      "loss": 3.947071075439453,
+      "step": 380
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.73828125,
+      "learning_rate": 0.0005586,
+      "loss": 3.855146026611328,
+      "step": 400
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.8139851093292236,
+      "eval_runtime": 8.9743,
+      "eval_samples_per_second": 1061.59,
+      "eval_steps_per_second": 1.114,
+      "step": 400
+    },
+    {
+      "epoch": 0.014155712841253791,
+      "grad_norm": 0.80859375,
+      "learning_rate": 0.0005866,
+      "loss": 3.76041259765625,
+      "step": 420
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 0.77734375,
+      "learning_rate": 0.0006146,
+      "loss": 3.708443450927734,
+      "step": 440
+    },
+    {
+      "epoch": 0.015503875968992248,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006426,
+      "loss": 3.637612152099609,
+      "step": 460
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.556614303588867,
+      "step": 480
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006986,
+      "loss": 3.5122703552246093,
+      "step": 500
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "eval_loss": 3.4916131496429443,
+      "eval_runtime": 8.8876,
+      "eval_samples_per_second": 1071.946,
+      "eval_steps_per_second": 1.125,
+      "step": 500
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0007,
+      "loss": 3.477622222900391,
+      "step": 520
+    },
+    {
+      "epoch": 0.01820020222446916,
+      "grad_norm": 0.75390625,
+      "learning_rate": 0.0007,
+      "loss": 3.4115901947021485,
+      "step": 540
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 3.359669876098633,
+      "step": 560
+    },
+    {
+      "epoch": 0.01954836535220762,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 3.3188705444335938,
+      "step": 580
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2840240478515623,
+      "step": 600
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "eval_loss": 3.2660911083221436,
+      "eval_runtime": 8.9246,
+      "eval_samples_per_second": 1067.504,
+      "eval_steps_per_second": 1.121,
+      "step": 600
+    },
+    {
+      "epoch": 0.020896528479946073,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0007,
+      "loss": 3.223845672607422,
+      "step": 620
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2053298950195312,
+      "step": 640
+    },
+    {
+      "epoch": 0.022244691607684528,
+      "grad_norm": 0.890625,
+      "learning_rate": 0.0007,
+      "loss": 3.1749242782592773,
+      "step": 660
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 3.134219741821289,
+      "step": 680
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0007,
+      "loss": 3.1245954513549803,
+      "step": 700
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "eval_loss": 3.113696575164795,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.373,
+      "eval_steps_per_second": 1.128,
+      "step": 700
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.0987874984741213,
+      "step": 720
+    },
+    {
+      "epoch": 0.02494101786316144,
+      "grad_norm": 0.73046875,
+      "learning_rate": 0.0007,
+      "loss": 3.08369140625,
+      "step": 740
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.03726863861084,
+      "step": 760
+    },
+    {
+      "epoch": 0.0262891809908999,
+      "grad_norm": 0.64453125,
+      "learning_rate": 0.0007,
+      "loss": 2.9970935821533202,
+      "step": 780
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.625,
+      "learning_rate": 0.0007,
+      "loss": 3.006105422973633,
+      "step": 800
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 2.999389886856079,
+      "eval_runtime": 8.984,
+      "eval_samples_per_second": 1060.442,
+      "eval_steps_per_second": 1.113,
+      "step": 800
+    },
+    {
+      "epoch": 0.027637344118638354,
+      "grad_norm": 0.609375,
+      "learning_rate": 0.0007,
+      "loss": 2.9657209396362303,
+      "step": 820
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 2.9664005279541015,
+      "step": 840
+    },
+    {
+      "epoch": 0.028985507246376812,
+      "grad_norm": 0.65625,
+      "learning_rate": 0.0007,
+      "loss": 2.988846206665039,
+      "step": 860
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.66015625,
+      "learning_rate": 0.0007,
+      "loss": 2.9386875152587892,
+      "step": 880
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "grad_norm": 0.7109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8965709686279295,
+      "step": 900
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "eval_loss": 2.9093611240386963,
+      "eval_runtime": 8.9344,
+      "eval_samples_per_second": 1066.332,
+      "eval_steps_per_second": 1.119,
+      "step": 900
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.9101829528808594,
+      "step": 920
+    },
+    {
+      "epoch": 0.031681833501853725,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 2.8905385971069335,
+      "step": 940
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.74609375,
+      "learning_rate": 0.0007,
+      "loss": 2.8552478790283202,
+      "step": 960
+    },
+    {
+      "epoch": 0.03302999662959218,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8467588424682617,
+      "step": 980
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 2.831803894042969,
+      "step": 1000
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "eval_loss": 2.8364417552948,
+      "eval_runtime": 8.9776,
+      "eval_samples_per_second": 1061.195,
+      "eval_steps_per_second": 1.114,
+      "step": 1000
+    },
+    {
+      "epoch": 0.034378159757330634,
+      "grad_norm": 0.640625,
+      "learning_rate": 0.0007,
+      "loss": 2.831756591796875,
+      "step": 1020
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8086965560913084,
+      "step": 1040
+    },
+    {
+      "epoch": 0.03572632288506909,
+      "grad_norm": 0.59375,
+      "learning_rate": 0.0007,
+      "loss": 2.7877662658691404,
+      "step": 1060
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.71875,
+      "learning_rate": 0.0007,
+      "loss": 2.7670337677001955,
+      "step": 1080
+    },
+    {
+      "epoch": 0.03707448601280755,
+      "grad_norm": 0.65234375,
+      "learning_rate": 0.0007,
+      "loss": 2.762655830383301,
+      "step": 1100
+    },
+    {
+      "epoch": 0.03707448601280755,
+      "eval_loss": 2.7699334621429443,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.369,
+      "eval_steps_per_second": 1.128,
+      "step": 1100
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 372646713753600.0,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1100/training_args.bin

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/config.json

@@ -0,0 +1,37 @@
+{
+  "activation": "relu",
+  "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": 21,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "powlu_m": 3.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.16.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096,
+  "waleed_beta": 10.0
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/model.safetensors

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/optimizer.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/rng_state.pth

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/scheduler.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/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|>"
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/trainer_state.json

@@ -0,0 +1,550 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.04044489383215369,
+  "eval_steps": 100,
+  "global_step": 1200,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0006740815638692282,
+      "grad_norm": 1.7578125,
+      "learning_rate": 2.66e-05,
+      "loss": 8.330369567871093,
+      "step": 20
+    },
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.40625,
+      "learning_rate": 5.46e-05,
+      "loss": 8.17730255126953,
+      "step": 40
+    },
+    {
+      "epoch": 0.0020222446916076846,
+      "grad_norm": 1.3359375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.952220916748047,
+      "step": 60
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3125,
+      "learning_rate": 0.0001106,
+      "loss": 7.694232177734375,
+      "step": 80
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "grad_norm": 1.234375,
+      "learning_rate": 0.0001386,
+      "loss": 7.353645324707031,
+      "step": 100
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "eval_loss": 7.14666223526001,
+      "eval_runtime": 8.8883,
+      "eval_samples_per_second": 1071.854,
+      "eval_steps_per_second": 1.125,
+      "step": 100
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.1796875,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.9598846435546875,
+      "step": 120
+    },
+    {
+      "epoch": 0.0047185709470845974,
+      "grad_norm": 1.1015625,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.568833923339843,
+      "step": 140
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.046875,
+      "learning_rate": 0.0002226,
+      "loss": 6.253362655639648,
+      "step": 160
+    },
+    {
+      "epoch": 0.006066734074823054,
+      "grad_norm": 0.86328125,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.974322509765625,
+      "step": 180
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0002786,
+      "loss": 5.698452377319336,
+      "step": 200
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "eval_loss": 5.558600902557373,
+      "eval_runtime": 9.1031,
+      "eval_samples_per_second": 1046.564,
+      "eval_steps_per_second": 1.099,
+      "step": 200
+    },
+    {
+      "epoch": 0.00741489720256151,
+      "grad_norm": 0.8046875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.451509857177735,
+      "step": 220
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.0078125,
+      "learning_rate": 0.0003346,
+      "loss": 5.163370895385742,
+      "step": 240
+    },
+    {
+      "epoch": 0.008763060330299966,
+      "grad_norm": 1.0859375,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 4.929123306274414,
+      "step": 260
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 0.84765625,
+      "learning_rate": 0.0003906,
+      "loss": 4.702068710327149,
+      "step": 280
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0004186,
+      "loss": 4.501141738891602,
+      "step": 300
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "eval_loss": 4.396740913391113,
+      "eval_runtime": 8.9474,
+      "eval_samples_per_second": 1064.782,
+      "eval_steps_per_second": 1.118,
+      "step": 300
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.0004466,
+      "loss": 4.323627853393555,
+      "step": 320
+    },
+    {
+      "epoch": 0.011459386585776879,
+      "grad_norm": 0.88671875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.204174041748047,
+      "step": 340
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.92578125,
+      "learning_rate": 0.0005026,
+      "loss": 4.052695846557617,
+      "step": 360
+    },
+    {
+      "epoch": 0.012807549713515335,
+      "grad_norm": 0.6484375,
+      "learning_rate": 0.0005306,
+      "loss": 3.947071075439453,
+      "step": 380
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.73828125,
+      "learning_rate": 0.0005586,
+      "loss": 3.855146026611328,
+      "step": 400
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.8139851093292236,
+      "eval_runtime": 8.9743,
+      "eval_samples_per_second": 1061.59,
+      "eval_steps_per_second": 1.114,
+      "step": 400
+    },
+    {
+      "epoch": 0.014155712841253791,
+      "grad_norm": 0.80859375,
+      "learning_rate": 0.0005866,
+      "loss": 3.76041259765625,
+      "step": 420
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 0.77734375,
+      "learning_rate": 0.0006146,
+      "loss": 3.708443450927734,
+      "step": 440
+    },
+    {
+      "epoch": 0.015503875968992248,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006426,
+      "loss": 3.637612152099609,
+      "step": 460
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.556614303588867,
+      "step": 480
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006986,
+      "loss": 3.5122703552246093,
+      "step": 500
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "eval_loss": 3.4916131496429443,
+      "eval_runtime": 8.8876,
+      "eval_samples_per_second": 1071.946,
+      "eval_steps_per_second": 1.125,
+      "step": 500
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0007,
+      "loss": 3.477622222900391,
+      "step": 520
+    },
+    {
+      "epoch": 0.01820020222446916,
+      "grad_norm": 0.75390625,
+      "learning_rate": 0.0007,
+      "loss": 3.4115901947021485,
+      "step": 540
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 3.359669876098633,
+      "step": 560
+    },
+    {
+      "epoch": 0.01954836535220762,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 3.3188705444335938,
+      "step": 580
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2840240478515623,
+      "step": 600
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "eval_loss": 3.2660911083221436,
+      "eval_runtime": 8.9246,
+      "eval_samples_per_second": 1067.504,
+      "eval_steps_per_second": 1.121,
+      "step": 600
+    },
+    {
+      "epoch": 0.020896528479946073,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0007,
+      "loss": 3.223845672607422,
+      "step": 620
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2053298950195312,
+      "step": 640
+    },
+    {
+      "epoch": 0.022244691607684528,
+      "grad_norm": 0.890625,
+      "learning_rate": 0.0007,
+      "loss": 3.1749242782592773,
+      "step": 660
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 3.134219741821289,
+      "step": 680
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0007,
+      "loss": 3.1245954513549803,
+      "step": 700
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "eval_loss": 3.113696575164795,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.373,
+      "eval_steps_per_second": 1.128,
+      "step": 700
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.0987874984741213,
+      "step": 720
+    },
+    {
+      "epoch": 0.02494101786316144,
+      "grad_norm": 0.73046875,
+      "learning_rate": 0.0007,
+      "loss": 3.08369140625,
+      "step": 740
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.03726863861084,
+      "step": 760
+    },
+    {
+      "epoch": 0.0262891809908999,
+      "grad_norm": 0.64453125,
+      "learning_rate": 0.0007,
+      "loss": 2.9970935821533202,
+      "step": 780
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.625,
+      "learning_rate": 0.0007,
+      "loss": 3.006105422973633,
+      "step": 800
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 2.999389886856079,
+      "eval_runtime": 8.984,
+      "eval_samples_per_second": 1060.442,
+      "eval_steps_per_second": 1.113,
+      "step": 800
+    },
+    {
+      "epoch": 0.027637344118638354,
+      "grad_norm": 0.609375,
+      "learning_rate": 0.0007,
+      "loss": 2.9657209396362303,
+      "step": 820
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 2.9664005279541015,
+      "step": 840
+    },
+    {
+      "epoch": 0.028985507246376812,
+      "grad_norm": 0.65625,
+      "learning_rate": 0.0007,
+      "loss": 2.988846206665039,
+      "step": 860
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.66015625,
+      "learning_rate": 0.0007,
+      "loss": 2.9386875152587892,
+      "step": 880
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "grad_norm": 0.7109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8965709686279295,
+      "step": 900
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "eval_loss": 2.9093611240386963,
+      "eval_runtime": 8.9344,
+      "eval_samples_per_second": 1066.332,
+      "eval_steps_per_second": 1.119,
+      "step": 900
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.9101829528808594,
+      "step": 920
+    },
+    {
+      "epoch": 0.031681833501853725,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 2.8905385971069335,
+      "step": 940
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.74609375,
+      "learning_rate": 0.0007,
+      "loss": 2.8552478790283202,
+      "step": 960
+    },
+    {
+      "epoch": 0.03302999662959218,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8467588424682617,
+      "step": 980
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 2.831803894042969,
+      "step": 1000
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "eval_loss": 2.8364417552948,
+      "eval_runtime": 8.9776,
+      "eval_samples_per_second": 1061.195,
+      "eval_steps_per_second": 1.114,
+      "step": 1000
+    },
+    {
+      "epoch": 0.034378159757330634,
+      "grad_norm": 0.640625,
+      "learning_rate": 0.0007,
+      "loss": 2.831756591796875,
+      "step": 1020
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8086965560913084,
+      "step": 1040
+    },
+    {
+      "epoch": 0.03572632288506909,
+      "grad_norm": 0.59375,
+      "learning_rate": 0.0007,
+      "loss": 2.7877662658691404,
+      "step": 1060
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.71875,
+      "learning_rate": 0.0007,
+      "loss": 2.7670337677001955,
+      "step": 1080
+    },
+    {
+      "epoch": 0.03707448601280755,
+      "grad_norm": 0.65234375,
+      "learning_rate": 0.0007,
+      "loss": 2.762655830383301,
+      "step": 1100
+    },
+    {
+      "epoch": 0.03707448601280755,
+      "eval_loss": 2.7699334621429443,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.369,
+      "eval_steps_per_second": 1.128,
+      "step": 1100
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.640625,
+      "learning_rate": 0.0007,
+      "loss": 2.762282371520996,
+      "step": 1120
+    },
+    {
+      "epoch": 0.03842264914054601,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0007,
+      "loss": 2.734586715698242,
+      "step": 1140
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.59765625,
+      "learning_rate": 0.0007,
+      "loss": 2.72774658203125,
+      "step": 1160
+    },
+    {
+      "epoch": 0.03977081226828446,
+      "grad_norm": 0.65625,
+      "learning_rate": 0.0007,
+      "loss": 2.7388858795166016,
+      "step": 1180
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.625,
+      "learning_rate": 0.0007,
+      "loss": 2.7111089706420897,
+      "step": 1200
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "eval_loss": 2.7126073837280273,
+      "eval_runtime": 8.8945,
+      "eval_samples_per_second": 1071.112,
+      "eval_steps_per_second": 1.124,
+      "step": 1200
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 406523687731200.0,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1200/training_args.bin

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/config.json

@@ -0,0 +1,37 @@
+{
+  "activation": "relu",
+  "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": 21,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "powlu_m": 3.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.16.0.dev0",
+  "use_cache": false,
+  "vocab_size": 4096,
+  "waleed_beta": 10.0
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/model.safetensors

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/optimizer.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/rng_state.pth

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/scheduler.pt

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

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/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|>"
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/trainer_state.json

@@ -0,0 +1,593 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.043815301651499834,
+  "eval_steps": 100,
+  "global_step": 1300,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0006740815638692282,
+      "grad_norm": 1.7578125,
+      "learning_rate": 2.66e-05,
+      "loss": 8.330369567871093,
+      "step": 20
+    },
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 1.40625,
+      "learning_rate": 5.46e-05,
+      "loss": 8.17730255126953,
+      "step": 40
+    },
+    {
+      "epoch": 0.0020222446916076846,
+      "grad_norm": 1.3359375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.952220916748047,
+      "step": 60
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3125,
+      "learning_rate": 0.0001106,
+      "loss": 7.694232177734375,
+      "step": 80
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "grad_norm": 1.234375,
+      "learning_rate": 0.0001386,
+      "loss": 7.353645324707031,
+      "step": 100
+    },
+    {
+      "epoch": 0.003370407819346141,
+      "eval_loss": 7.14666223526001,
+      "eval_runtime": 8.8883,
+      "eval_samples_per_second": 1071.854,
+      "eval_steps_per_second": 1.125,
+      "step": 100
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.1796875,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.9598846435546875,
+      "step": 120
+    },
+    {
+      "epoch": 0.0047185709470845974,
+      "grad_norm": 1.1015625,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.568833923339843,
+      "step": 140
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.046875,
+      "learning_rate": 0.0002226,
+      "loss": 6.253362655639648,
+      "step": 160
+    },
+    {
+      "epoch": 0.006066734074823054,
+      "grad_norm": 0.86328125,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.974322509765625,
+      "step": 180
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0002786,
+      "loss": 5.698452377319336,
+      "step": 200
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "eval_loss": 5.558600902557373,
+      "eval_runtime": 9.1031,
+      "eval_samples_per_second": 1046.564,
+      "eval_steps_per_second": 1.099,
+      "step": 200
+    },
+    {
+      "epoch": 0.00741489720256151,
+      "grad_norm": 0.8046875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.451509857177735,
+      "step": 220
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.0078125,
+      "learning_rate": 0.0003346,
+      "loss": 5.163370895385742,
+      "step": 240
+    },
+    {
+      "epoch": 0.008763060330299966,
+      "grad_norm": 1.0859375,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 4.929123306274414,
+      "step": 260
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 0.84765625,
+      "learning_rate": 0.0003906,
+      "loss": 4.702068710327149,
+      "step": 280
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0004186,
+      "loss": 4.501141738891602,
+      "step": 300
+    },
+    {
+      "epoch": 0.010111223458038422,
+      "eval_loss": 4.396740913391113,
+      "eval_runtime": 8.9474,
+      "eval_samples_per_second": 1064.782,
+      "eval_steps_per_second": 1.118,
+      "step": 300
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.0004466,
+      "loss": 4.323627853393555,
+      "step": 320
+    },
+    {
+      "epoch": 0.011459386585776879,
+      "grad_norm": 0.88671875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.204174041748047,
+      "step": 340
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.92578125,
+      "learning_rate": 0.0005026,
+      "loss": 4.052695846557617,
+      "step": 360
+    },
+    {
+      "epoch": 0.012807549713515335,
+      "grad_norm": 0.6484375,
+      "learning_rate": 0.0005306,
+      "loss": 3.947071075439453,
+      "step": 380
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.73828125,
+      "learning_rate": 0.0005586,
+      "loss": 3.855146026611328,
+      "step": 400
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "eval_loss": 3.8139851093292236,
+      "eval_runtime": 8.9743,
+      "eval_samples_per_second": 1061.59,
+      "eval_steps_per_second": 1.114,
+      "step": 400
+    },
+    {
+      "epoch": 0.014155712841253791,
+      "grad_norm": 0.80859375,
+      "learning_rate": 0.0005866,
+      "loss": 3.76041259765625,
+      "step": 420
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 0.77734375,
+      "learning_rate": 0.0006146,
+      "loss": 3.708443450927734,
+      "step": 440
+    },
+    {
+      "epoch": 0.015503875968992248,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006426,
+      "loss": 3.637612152099609,
+      "step": 460
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.556614303588867,
+      "step": 480
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0006986,
+      "loss": 3.5122703552246093,
+      "step": 500
+    },
+    {
+      "epoch": 0.016852039096730706,
+      "eval_loss": 3.4916131496429443,
+      "eval_runtime": 8.8876,
+      "eval_samples_per_second": 1071.946,
+      "eval_steps_per_second": 1.125,
+      "step": 500
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.80078125,
+      "learning_rate": 0.0007,
+      "loss": 3.477622222900391,
+      "step": 520
+    },
+    {
+      "epoch": 0.01820020222446916,
+      "grad_norm": 0.75390625,
+      "learning_rate": 0.0007,
+      "loss": 3.4115901947021485,
+      "step": 540
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 3.359669876098633,
+      "step": 560
+    },
+    {
+      "epoch": 0.01954836535220762,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 3.3188705444335938,
+      "step": 580
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2840240478515623,
+      "step": 600
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "eval_loss": 3.2660911083221436,
+      "eval_runtime": 8.9246,
+      "eval_samples_per_second": 1067.504,
+      "eval_steps_per_second": 1.121,
+      "step": 600
+    },
+    {
+      "epoch": 0.020896528479946073,
+      "grad_norm": 0.7421875,
+      "learning_rate": 0.0007,
+      "loss": 3.223845672607422,
+      "step": 620
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 3.2053298950195312,
+      "step": 640
+    },
+    {
+      "epoch": 0.022244691607684528,
+      "grad_norm": 0.890625,
+      "learning_rate": 0.0007,
+      "loss": 3.1749242782592773,
+      "step": 660
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 3.134219741821289,
+      "step": 680
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "grad_norm": 0.6640625,
+      "learning_rate": 0.0007,
+      "loss": 3.1245954513549803,
+      "step": 700
+    },
+    {
+      "epoch": 0.023592854735422986,
+      "eval_loss": 3.113696575164795,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.373,
+      "eval_steps_per_second": 1.128,
+      "step": 700
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.0987874984741213,
+      "step": 720
+    },
+    {
+      "epoch": 0.02494101786316144,
+      "grad_norm": 0.73046875,
+      "learning_rate": 0.0007,
+      "loss": 3.08369140625,
+      "step": 740
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.62890625,
+      "learning_rate": 0.0007,
+      "loss": 3.03726863861084,
+      "step": 760
+    },
+    {
+      "epoch": 0.0262891809908999,
+      "grad_norm": 0.64453125,
+      "learning_rate": 0.0007,
+      "loss": 2.9970935821533202,
+      "step": 780
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.625,
+      "learning_rate": 0.0007,
+      "loss": 3.006105422973633,
+      "step": 800
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "eval_loss": 2.999389886856079,
+      "eval_runtime": 8.984,
+      "eval_samples_per_second": 1060.442,
+      "eval_steps_per_second": 1.113,
+      "step": 800
+    },
+    {
+      "epoch": 0.027637344118638354,
+      "grad_norm": 0.609375,
+      "learning_rate": 0.0007,
+      "loss": 2.9657209396362303,
+      "step": 820
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.67578125,
+      "learning_rate": 0.0007,
+      "loss": 2.9664005279541015,
+      "step": 840
+    },
+    {
+      "epoch": 0.028985507246376812,
+      "grad_norm": 0.65625,
+      "learning_rate": 0.0007,
+      "loss": 2.988846206665039,
+      "step": 860
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.66015625,
+      "learning_rate": 0.0007,
+      "loss": 2.9386875152587892,
+      "step": 880
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "grad_norm": 0.7109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8965709686279295,
+      "step": 900
+    },
+    {
+      "epoch": 0.030333670374115267,
+      "eval_loss": 2.9093611240386963,
+      "eval_runtime": 8.9344,
+      "eval_samples_per_second": 1066.332,
+      "eval_steps_per_second": 1.119,
+      "step": 900
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.9101829528808594,
+      "step": 920
+    },
+    {
+      "epoch": 0.031681833501853725,
+      "grad_norm": 0.5859375,
+      "learning_rate": 0.0007,
+      "loss": 2.8905385971069335,
+      "step": 940
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.74609375,
+      "learning_rate": 0.0007,
+      "loss": 2.8552478790283202,
+      "step": 960
+    },
+    {
+      "epoch": 0.03302999662959218,
+      "grad_norm": 0.62109375,
+      "learning_rate": 0.0007,
+      "loss": 2.8467588424682617,
+      "step": 980
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0007,
+      "loss": 2.831803894042969,
+      "step": 1000
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "eval_loss": 2.8364417552948,
+      "eval_runtime": 8.9776,
+      "eval_samples_per_second": 1061.195,
+      "eval_steps_per_second": 1.114,
+      "step": 1000
+    },
+    {
+      "epoch": 0.034378159757330634,
+      "grad_norm": 0.640625,
+      "learning_rate": 0.0007,
+      "loss": 2.831756591796875,
+      "step": 1020
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8086965560913084,
+      "step": 1040
+    },
+    {
+      "epoch": 0.03572632288506909,
+      "grad_norm": 0.59375,
+      "learning_rate": 0.0007,
+      "loss": 2.7877662658691404,
+      "step": 1060
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.71875,
+      "learning_rate": 0.0007,
+      "loss": 2.7670337677001955,
+      "step": 1080
+    },
+    {
+      "epoch": 0.03707448601280755,
+      "grad_norm": 0.65234375,
+      "learning_rate": 0.0007,
+      "loss": 2.762655830383301,
+      "step": 1100
+    },
+    {
+      "epoch": 0.03707448601280755,
+      "eval_loss": 2.7699334621429443,
+      "eval_runtime": 8.8675,
+      "eval_samples_per_second": 1074.369,
+      "eval_steps_per_second": 1.128,
+      "step": 1100
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.640625,
+      "learning_rate": 0.0007,
+      "loss": 2.762282371520996,
+      "step": 1120
+    },
+    {
+      "epoch": 0.03842264914054601,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0007,
+      "loss": 2.734586715698242,
+      "step": 1140
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.59765625,
+      "learning_rate": 0.0007,
+      "loss": 2.72774658203125,
+      "step": 1160
+    },
+    {
+      "epoch": 0.03977081226828446,
+      "grad_norm": 0.65625,
+      "learning_rate": 0.0007,
+      "loss": 2.7388858795166016,
+      "step": 1180
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.625,
+      "learning_rate": 0.0007,
+      "loss": 2.7111089706420897,
+      "step": 1200
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "eval_loss": 2.7126073837280273,
+      "eval_runtime": 8.8945,
+      "eval_samples_per_second": 1071.112,
+      "eval_steps_per_second": 1.124,
+      "step": 1200
+    },
+    {
+      "epoch": 0.04111897539602292,
+      "grad_norm": 0.6328125,
+      "learning_rate": 0.0007,
+      "loss": 2.681635856628418,
+      "step": 1220
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.609375,
+      "learning_rate": 0.0007,
+      "loss": 2.699861526489258,
+      "step": 1240
+    },
+    {
+      "epoch": 0.042467138523761376,
+      "grad_norm": 0.7109375,
+      "learning_rate": 0.0007,
+      "loss": 2.704066848754883,
+      "step": 1260
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.6328125,
+      "learning_rate": 0.0007,
+      "loss": 2.675613784790039,
+      "step": 1280
+    },
+    {
+      "epoch": 0.043815301651499834,
+      "grad_norm": 0.6171875,
+      "learning_rate": 0.0007,
+      "loss": 2.6731496810913087,
+      "step": 1300
+    },
+    {
+      "epoch": 0.043815301651499834,
+      "eval_loss": 2.6630420684814453,
+      "eval_runtime": 8.9026,
+      "eval_samples_per_second": 1070.134,
+      "eval_steps_per_second": 1.123,
+      "step": 1300
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 440400661708800.0,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

zain/Activation/out/glu-relu-21L_run/checkpoint-1300/training_args.bin

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