Auto upload zain 2026-08-13T21:15:52.507316

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+zain/Activation/wandb/run-20260813_211521-bdgno22l/run-bdgno22l.wandb filter=lfs diff=lfs merge=lfs -text

zain/Activation/README.md

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

zain/Activation/exp.py

@@ -0,0 +1,921 @@
+# =====================================================================
+# exp.py – FULL FILE, ALL FIXES INCLUDED (hook requires_grad check)
+# =====================================================================
+
+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,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.mlp_type = mlp_type
+        self.activation = activation
+        self.waleed_beta = waleed_beta
+        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 override hook)
+# =============================================================================
+
+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)
+
+        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 handling
+        if self.mlp_type == "glu":
+            if 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:
+            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)
+
+        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
+
+        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")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.mlp_type == "glu":
+            gate = self.gate_proj(x)
+            up = self.up_proj(x)
+
+            if 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
+            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)
+
+        if self.is_waleed10:
+            out = self.waleed_beta * torch.tanh(out / self.waleed_beta)
+
+        # FIX: only register hook if the output tensor requires gradients
+        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
+# ----------------------------------------------------------------------------
+
+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
+# =============================================================================
+
+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 (FIXED)
+# =============================================================================
+
+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
+# =============================================================================
+
+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
+# =============================================================================
+
+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
+# =============================================================================
+
+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
+# =============================================================================
+
+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/llm_analyzer_wandb.py

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

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/config.json

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

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/model.safetensors

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

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/optimizer.pt

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

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/rng_state.pth

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

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/scheduler.pt

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

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/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-linear-100L_trash_run/checkpoint-2600/trainer_state.json

@@ -0,0 +1,960 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.17526120660599934,
+  "eval_steps": 2498,
+  "global_step": 2600,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 2.125,
+      "learning_rate": 2.66e-05,
+      "loss": 8.29943389892578,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3046875,
+      "learning_rate": 5.46e-05,
+      "loss": 8.036395263671874,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.2734375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.834572601318359,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.265625,
+      "learning_rate": 0.0001106,
+      "loss": 7.5718994140625,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 1.1953125,
+      "learning_rate": 0.0001386,
+      "loss": 7.237230682373047,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.1484375,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.864912414550782,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 1.0234375,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.509800720214844,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 1.0234375,
+      "learning_rate": 0.0002226,
+      "loss": 6.20410041809082,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.875,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.940516662597656,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.8203125,
+      "learning_rate": 0.0002786,
+      "loss": 5.676180267333985,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 0.8828125,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.412886428833008,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 0.81640625,
+      "learning_rate": 0.0003346,
+      "loss": 5.135355377197266,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 1.1796875,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 4.882264709472656,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.94921875,
+      "learning_rate": 0.0003906,
+      "loss": 4.6449634552001955,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.046875,
+      "learning_rate": 0.0004186,
+      "loss": 4.446089553833008,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 0.87890625,
+      "learning_rate": 0.0004466,
+      "loss": 4.2602802276611325,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.625,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.108924484252929,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.82421875,
+      "learning_rate": 0.0005026,
+      "loss": 3.984059143066406,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.78515625,
+      "learning_rate": 0.0005306,
+      "loss": 3.864672088623047,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.6796875,
+      "learning_rate": 0.0005586,
+      "loss": 3.741139602661133,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.63671875,
+      "learning_rate": 0.0005866,
+      "loss": 3.6620559692382812,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.578125,
+      "learning_rate": 0.0006146,
+      "loss": 3.604226303100586,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.484375,
+      "learning_rate": 0.0006426,
+      "loss": 3.5065567016601564,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.60546875,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.437101364135742,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.57421875,
+      "learning_rate": 0.0006986,
+      "loss": 3.3709640502929688,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.484375,
+      "learning_rate": 0.0007,
+      "loss": 3.3115055084228517,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.578125,
+      "learning_rate": 0.0007,
+      "loss": 3.240653228759766,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.48828125,
+      "learning_rate": 0.0007,
+      "loss": 3.1956838607788085,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 3.1403549194335936,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.484375,
+      "learning_rate": 0.0007,
+      "loss": 3.109465980529785,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 3.052653121948242,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.48046875,
+      "learning_rate": 0.0007,
+      "loss": 3.0361297607421873,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 2.984110450744629,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 0.51171875,
+      "learning_rate": 0.0007,
+      "loss": 2.961138153076172,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.922393226623535,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.416015625,
+      "learning_rate": 0.0007,
+      "loss": 2.892241287231445,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.48828125,
+      "learning_rate": 0.0007,
+      "loss": 2.858970069885254,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.486328125,
+      "learning_rate": 0.0007,
+      "loss": 2.8428062438964843,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.5234375,
+      "learning_rate": 0.0007,
+      "loss": 2.837088203430176,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.7979549407958983,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.46875,
+      "learning_rate": 0.0007,
+      "loss": 2.787800979614258,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.4296875,
+      "learning_rate": 0.0007,
+      "loss": 2.7603782653808593,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.44140625,
+      "learning_rate": 0.0007,
+      "loss": 2.7198511123657227,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.439453125,
+      "learning_rate": 0.0007,
+      "loss": 2.6977243423461914,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.46484375,
+      "learning_rate": 0.0007,
+      "loss": 2.6839799880981445,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.661431312561035,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.423828125,
+      "learning_rate": 0.0007,
+      "loss": 2.6492240905761717,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 2.6231040954589844,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.4609375,
+      "learning_rate": 0.0007,
+      "loss": 2.610037422180176,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.455078125,
+      "learning_rate": 0.0007,
+      "loss": 2.6107507705688477,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06875631951466127,
+      "grad_norm": 0.42578125,
+      "learning_rate": 0.0007,
+      "loss": 2.5831680297851562,
+      "step": 1020
+    },
+    {
+      "epoch": 0.07010448264239973,
+      "grad_norm": 0.453125,
+      "learning_rate": 0.0007,
+      "loss": 2.5671245574951174,
+      "step": 1040
+    },
+    {
+      "epoch": 0.07145264577013818,
+      "grad_norm": 0.4453125,
+      "learning_rate": 0.0007,
+      "loss": 2.5532295227050783,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07280080889787664,
+      "grad_norm": 0.4453125,
+      "learning_rate": 0.0007,
+      "loss": 2.535228729248047,
+      "step": 1080
+    },
+    {
+      "epoch": 0.0741489720256151,
+      "grad_norm": 0.431640625,
+      "learning_rate": 0.0007,
+      "loss": 2.5336929321289063,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07549713515335356,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.4990976333618162,
+      "step": 1120
+    },
+    {
+      "epoch": 0.07684529828109202,
+      "grad_norm": 0.4296875,
+      "learning_rate": 0.0007,
+      "loss": 2.4878835678100586,
+      "step": 1140
+    },
+    {
+      "epoch": 0.07819346140883048,
+      "grad_norm": 0.4296875,
+      "learning_rate": 0.0007,
+      "loss": 2.4860113143920897,
+      "step": 1160
+    },
+    {
+      "epoch": 0.07954162453656892,
+      "grad_norm": 0.4296875,
+      "learning_rate": 0.0007,
+      "loss": 2.4814546585083006,
+      "step": 1180
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "grad_norm": 0.474609375,
+      "learning_rate": 0.0007,
+      "loss": 2.4764120101928713,
+      "step": 1200
+    },
+    {
+      "epoch": 0.08223795079204584,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.4622081756591796,
+      "step": 1220
+    },
+    {
+      "epoch": 0.0835861139197843,
+      "grad_norm": 0.3984375,
+      "learning_rate": 0.0007,
+      "loss": 2.450253486633301,
+      "step": 1240
+    },
+    {
+      "epoch": 0.08493427704752275,
+      "grad_norm": 0.427734375,
+      "learning_rate": 0.0007,
+      "loss": 2.4399486541748048,
+      "step": 1260
+    },
+    {
+      "epoch": 0.08628244017526121,
+      "grad_norm": 0.45703125,
+      "learning_rate": 0.0007,
+      "loss": 2.4211536407470704,
+      "step": 1280
+    },
+    {
+      "epoch": 0.08763060330299967,
+      "grad_norm": 0.44921875,
+      "learning_rate": 0.0007,
+      "loss": 2.413456916809082,
+      "step": 1300
+    },
+    {
+      "epoch": 0.08897876643073811,
+      "grad_norm": 0.408203125,
+      "learning_rate": 0.0007,
+      "loss": 2.405606269836426,
+      "step": 1320
+    },
+    {
+      "epoch": 0.09032692955847657,
+      "grad_norm": 0.474609375,
+      "learning_rate": 0.0007,
+      "loss": 2.3934694290161134,
+      "step": 1340
+    },
+    {
+      "epoch": 0.09167509268621503,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.393663024902344,
+      "step": 1360
+    },
+    {
+      "epoch": 0.09302325581395349,
+      "grad_norm": 0.404296875,
+      "learning_rate": 0.0007,
+      "loss": 2.3733840942382813,
+      "step": 1380
+    },
+    {
+      "epoch": 0.09437141894169195,
+      "grad_norm": 0.419921875,
+      "learning_rate": 0.0007,
+      "loss": 2.3658527374267577,
+      "step": 1400
+    },
+    {
+      "epoch": 0.0957195820694304,
+      "grad_norm": 0.451171875,
+      "learning_rate": 0.0007,
+      "loss": 2.3782094955444335,
+      "step": 1420
+    },
+    {
+      "epoch": 0.09706774519716886,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0007,
+      "loss": 2.357101249694824,
+      "step": 1440
+    },
+    {
+      "epoch": 0.09841590832490732,
+      "grad_norm": 0.44140625,
+      "learning_rate": 0.0007,
+      "loss": 2.3440685272216797,
+      "step": 1460
+    },
+    {
+      "epoch": 0.09976407145264576,
+      "grad_norm": 0.40625,
+      "learning_rate": 0.0007,
+      "loss": 2.33266544342041,
+      "step": 1480
+    },
+    {
+      "epoch": 0.10111223458038422,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.318921661376953,
+      "step": 1500
+    },
+    {
+      "epoch": 0.10246039770812268,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.3219444274902346,
+      "step": 1520
+    },
+    {
+      "epoch": 0.10380856083586114,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0007,
+      "loss": 2.312142181396484,
+      "step": 1540
+    },
+    {
+      "epoch": 0.1051567239635996,
+      "grad_norm": 0.400390625,
+      "learning_rate": 0.0007,
+      "loss": 2.300857162475586,
+      "step": 1560
+    },
+    {
+      "epoch": 0.10650488709133805,
+      "grad_norm": 0.41015625,
+      "learning_rate": 0.0007,
+      "loss": 2.3189367294311523,
+      "step": 1580
+    },
+    {
+      "epoch": 0.10785305021907651,
+      "grad_norm": 0.42578125,
+      "learning_rate": 0.0007,
+      "loss": 2.2867937088012695,
+      "step": 1600
+    },
+    {
+      "epoch": 0.10920121334681497,
+      "grad_norm": 0.392578125,
+      "learning_rate": 0.0007,
+      "loss": 2.286478614807129,
+      "step": 1620
+    },
+    {
+      "epoch": 0.11054937647455342,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0007,
+      "loss": 2.299053192138672,
+      "step": 1640
+    },
+    {
+      "epoch": 0.11189753960229187,
+      "grad_norm": 0.404296875,
+      "learning_rate": 0.0007,
+      "loss": 2.2822343826293947,
+      "step": 1660
+    },
+    {
+      "epoch": 0.11324570273003033,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0007,
+      "loss": 2.264908218383789,
+      "step": 1680
+    },
+    {
+      "epoch": 0.11459386585776879,
+      "grad_norm": 0.416015625,
+      "learning_rate": 0.0007,
+      "loss": 2.277639389038086,
+      "step": 1700
+    },
+    {
+      "epoch": 0.11594202898550725,
+      "grad_norm": 0.46484375,
+      "learning_rate": 0.0007,
+      "loss": 2.2545700073242188,
+      "step": 1720
+    },
+    {
+      "epoch": 0.1172901921132457,
+      "grad_norm": 0.40625,
+      "learning_rate": 0.0007,
+      "loss": 2.2448637008666994,
+      "step": 1740
+    },
+    {
+      "epoch": 0.11863835524098416,
+      "grad_norm": 0.4296875,
+      "learning_rate": 0.0007,
+      "loss": 2.250792694091797,
+      "step": 1760
+    },
+    {
+      "epoch": 0.11998651836872262,
+      "grad_norm": 0.41796875,
+      "learning_rate": 0.0007,
+      "loss": 2.264188766479492,
+      "step": 1780
+    },
+    {
+      "epoch": 0.12133468149646107,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0007,
+      "loss": 2.2614953994750975,
+      "step": 1800
+    },
+    {
+      "epoch": 0.12268284462419952,
+      "grad_norm": 0.408203125,
+      "learning_rate": 0.0007,
+      "loss": 2.252177619934082,
+      "step": 1820
+    },
+    {
+      "epoch": 0.12403100775193798,
+      "grad_norm": 0.390625,
+      "learning_rate": 0.0007,
+      "loss": 2.219951057434082,
+      "step": 1840
+    },
+    {
+      "epoch": 0.12537917087967643,
+      "grad_norm": 0.400390625,
+      "learning_rate": 0.0007,
+      "loss": 2.219348907470703,
+      "step": 1860
+    },
+    {
+      "epoch": 0.1267273340074149,
+      "grad_norm": 0.4140625,
+      "learning_rate": 0.0007,
+      "loss": 2.2108036041259767,
+      "step": 1880
+    },
+    {
+      "epoch": 0.12807549713515334,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0007,
+      "loss": 2.219270706176758,
+      "step": 1900
+    },
+    {
+      "epoch": 0.12942366026289182,
+      "grad_norm": 0.392578125,
+      "learning_rate": 0.0007,
+      "loss": 2.1969385147094727,
+      "step": 1920
+    },
+    {
+      "epoch": 0.13077182339063026,
+      "grad_norm": 0.42578125,
+      "learning_rate": 0.0007,
+      "loss": 2.207957077026367,
+      "step": 1940
+    },
+    {
+      "epoch": 0.13211998651836873,
+      "grad_norm": 0.416015625,
+      "learning_rate": 0.0007,
+      "loss": 2.1822214126586914,
+      "step": 1960
+    },
+    {
+      "epoch": 0.13346814964610718,
+      "grad_norm": 0.427734375,
+      "learning_rate": 0.0007,
+      "loss": 2.20617618560791,
+      "step": 1980
+    },
+    {
+      "epoch": 0.13481631277384565,
+      "grad_norm": 0.384765625,
+      "learning_rate": 0.0007,
+      "loss": 2.179953193664551,
+      "step": 2000
+    },
+    {
+      "epoch": 0.1361644759015841,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0007,
+      "loss": 2.190639877319336,
+      "step": 2020
+    },
+    {
+      "epoch": 0.13751263902932254,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.1713678359985353,
+      "step": 2040
+    },
+    {
+      "epoch": 0.138860802157061,
+      "grad_norm": 0.396484375,
+      "learning_rate": 0.0007,
+      "loss": 2.189349365234375,
+      "step": 2060
+    },
+    {
+      "epoch": 0.14020896528479945,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0007,
+      "loss": 2.1666955947875977,
+      "step": 2080
+    },
+    {
+      "epoch": 0.14155712841253792,
+      "grad_norm": 0.376953125,
+      "learning_rate": 0.0007,
+      "loss": 2.1569326400756834,
+      "step": 2100
+    },
+    {
+      "epoch": 0.14290529154027637,
+      "grad_norm": 0.3828125,
+      "learning_rate": 0.0007,
+      "loss": 2.166208839416504,
+      "step": 2120
+    },
+    {
+      "epoch": 0.14425345466801484,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0007,
+      "loss": 2.166676712036133,
+      "step": 2140
+    },
+    {
+      "epoch": 0.14560161779575329,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0007,
+      "loss": 2.152173614501953,
+      "step": 2160
+    },
+    {
+      "epoch": 0.14694978092349173,
+      "grad_norm": 0.39453125,
+      "learning_rate": 0.0007,
+      "loss": 2.162281608581543,
+      "step": 2180
+    },
+    {
+      "epoch": 0.1482979440512302,
+      "grad_norm": 0.384765625,
+      "learning_rate": 0.0007,
+      "loss": 2.1356754302978516,
+      "step": 2200
+    },
+    {
+      "epoch": 0.14964610717896865,
+      "grad_norm": 0.40625,
+      "learning_rate": 0.0007,
+      "loss": 2.1436878204345704,
+      "step": 2220
+    },
+    {
+      "epoch": 0.15099427030670712,
+      "grad_norm": 0.419921875,
+      "learning_rate": 0.0007,
+      "loss": 2.1559894561767576,
+      "step": 2240
+    },
+    {
+      "epoch": 0.15234243343444556,
+      "grad_norm": 0.404296875,
+      "learning_rate": 0.0007,
+      "loss": 2.1330844879150392,
+      "step": 2260
+    },
+    {
+      "epoch": 0.15369059656218403,
+      "grad_norm": 0.37890625,
+      "learning_rate": 0.0007,
+      "loss": 2.1436931610107424,
+      "step": 2280
+    },
+    {
+      "epoch": 0.15503875968992248,
+      "grad_norm": 0.39453125,
+      "learning_rate": 0.0007,
+      "loss": 2.12701358795166,
+      "step": 2300
+    },
+    {
+      "epoch": 0.15638692281766095,
+      "grad_norm": 0.412109375,
+      "learning_rate": 0.0007,
+      "loss": 2.135152053833008,
+      "step": 2320
+    },
+    {
+      "epoch": 0.1577350859453994,
+      "grad_norm": 0.37109375,
+      "learning_rate": 0.0007,
+      "loss": 2.1244712829589845,
+      "step": 2340
+    },
+    {
+      "epoch": 0.15908324907313784,
+      "grad_norm": 0.38671875,
+      "learning_rate": 0.0007,
+      "loss": 2.127437973022461,
+      "step": 2360
+    },
+    {
+      "epoch": 0.1604314122008763,
+      "grad_norm": 0.388671875,
+      "learning_rate": 0.0007,
+      "loss": 2.1295385360717773,
+      "step": 2380
+    },
+    {
+      "epoch": 0.16177957532861476,
+      "grad_norm": 0.40625,
+      "learning_rate": 0.0007,
+      "loss": 2.1162397384643556,
+      "step": 2400
+    },
+    {
+      "epoch": 0.16312773845635323,
+      "grad_norm": 0.40234375,
+      "learning_rate": 0.0007,
+      "loss": 2.118353843688965,
+      "step": 2420
+    },
+    {
+      "epoch": 0.16447590158409167,
+      "grad_norm": 0.380859375,
+      "learning_rate": 0.0007,
+      "loss": 2.104597473144531,
+      "step": 2440
+    },
+    {
+      "epoch": 0.16582406471183014,
+      "grad_norm": 0.3828125,
+      "learning_rate": 0.0007,
+      "loss": 2.1171491622924803,
+      "step": 2460
+    },
+    {
+      "epoch": 0.1671722278395686,
+      "grad_norm": 0.392578125,
+      "learning_rate": 0.0007,
+      "loss": 2.1171087265014648,
+      "step": 2480
+    },
+    {
+      "epoch": 0.1683855746545332,
+      "eval_loss": 2.1098570823669434,
+      "eval_runtime": 15.3278,
+      "eval_samples_per_second": 621.551,
+      "eval_steps_per_second": 0.652,
+      "step": 2498
+    },
+    {
+      "epoch": 0.16852039096730703,
+      "grad_norm": 0.37890625,
+      "learning_rate": 0.0007,
+      "loss": 2.100436973571777,
+      "step": 2500
+    },
+    {
+      "epoch": 0.16852039096730703,
+      "eval_loss": 2.1097066402435303,
+      "eval_runtime": 15.3325,
+      "eval_samples_per_second": 621.361,
+      "eval_steps_per_second": 0.652,
+      "step": 2500
+    },
+    {
+      "epoch": 0.1698685540950455,
+      "grad_norm": 26607616.0,
+      "learning_rate": 3.4200000000000005e-05,
+      "loss": 2.1015687942504884,
+      "step": 2520
+    },
+    {
+      "epoch": 0.17121671722278395,
+      "grad_norm": 31064064.0,
+      "learning_rate": 7.020000000000001e-05,
+      "loss": 2.1284088134765624,
+      "step": 2540
+    },
+    {
+      "epoch": 0.17256488035052242,
+      "grad_norm": 222298112.0,
+      "learning_rate": 9e-05,
+      "loss": 4.244554901123047,
+      "step": 2560
+    },
+    {
+      "epoch": 0.17391304347826086,
+      "grad_norm": 488636416.0,
+      "learning_rate": 9e-05,
+      "loss": 7.165976715087891,
+      "step": 2580
+    },
+    {
+      "epoch": 0.17526120660599934,
+      "grad_norm": 884998144.0,
+      "learning_rate": 9e-05,
+      "loss": 7.48333740234375,
+      "step": 2600
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2700,
+  "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": 8388337886822400.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

zain/Activation/out/glu-linear-100L_trash_run/checkpoint-2600/training_args.bin

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

zain/Activation/sweep.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""Sweep explicit GLU / MLP variants with identical data and hyperparameters."""
+
+import argparse
+import copy
+import json
+import re
+import time
+import os
+from pathlib import Path
+
+import yaml
+import wandb
+import torch
+from transformers import AutoTokenizer, set_seed
+from exp import (
+    TinyLlamaConfig,
+    TinyLlamaForCausalLM,
+    build_dataset,
+    create_trainer,
+    fetch_latest_checkpoint_from_hub,
+    ResumeFreezeOverrideCallback,
+)
+
+
+def format_param_count(total_params: int) -> str:
+    if total_params >= 1e9:
+        return f"{total_params / 1e9:.1f}B"
+    else:
+        return f"{total_params / 1e6:.1f}M"
+
+
+def parse_variant(variant: str):
+    parts = variant.split('-')
+    if len(parts) < 2:
+        raise ValueError(f"Invalid variant format: '{variant}'. Expected: <glu|mlp>-<activation>[-<layers>L]")
+    prefix = parts[0]
+    if prefix not in ('glu', 'mlp'):
+        raise ValueError(f"Invalid prefix: '{prefix}'. Must be 'glu' or 'mlp'.")
+    last = parts[-1]
+    if last.endswith('L') and last[:-1].isdigit():
+        layers = int(last[:-1])
+        activation = '-'.join(parts[1:-1])
+    else:
+        layers = None
+        activation = '-'.join(parts[1:])
+    if not activation:
+        raise ValueError(f"Missing activation name in variant: '{variant}'")
+    return prefix, activation, layers
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="Base YAML config")
+    parser.add_argument("--variants", nargs="+", required=True, help="List of variants")
+    parser.add_argument("--push", action="store_true")
+    args = parser.parse_args()
+
+    with open(args.config) as f:
+        base = yaml.safe_load(f)
+
+    seed = base.get("training", {}).get("seed", 42)
+    set_seed(seed)
+
+    wandb_project = base.get("training", {}).get("wandb_project")
+    if wandb_project:
+        os.environ["WANDB_PROJECT"] = wandb_project
+        print(f"[WandB] Project locked to: {wandb_project}")
+
+    tok_name = base["model"].get("tokenizer_name", "meta-llama/Llama-2-7b-hf")
+    tokenizer = AutoTokenizer.from_pretrained(tok_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    msl = base["model"].get("max_position_embeddings", 512)
+    train_ds = build_dataset(tokenizer, max_seq_len=msl, split="train", max_samples=None)
+    eval_ds = build_dataset(tokenizer, max_seq_len=msl, split="validation", max_samples=None)
+
+    results = []
+
+    for variant in args.variants:
+        prefix, act, layers = parse_variant(variant)
+
+        if prefix == "mlp" and act in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
+            raise ValueError(
+                f"Activation '{act}' requires GLU. Please use 'glu-{act}'."
+            )
+
+        cfg = copy.deepcopy(base)
+        cfg["model"]["mlp_type"] = prefix
+        cfg["model"]["activation"] = act
+        if layers is not None:
+            cfg["model"]["num_hidden_layers"] = layers
+
+        actual_layers = cfg["model"]["num_hidden_layers"]
+        variant_label = f"{prefix}-{act}-{actual_layers}L"
+
+        # Output directory: add _trash if override enabled
+        go = cfg.get("gradient_override", {})
+        if go.get("enabled", False):
+            base_out = Path(cfg["training"]["output_dir"]).parent
+            variant_name = variant_label + "_trash_run"
+        else:
+            base_out = Path(cfg["training"]["output_dir"]).parent
+            variant_name = variant_label + "_run"
+
+        out_dir = base_out / variant_name
+        cfg["training"]["output_dir"] = str(out_dir)
+
+        # Checkpoint fetching
+        checkpoint_path = None
+        resume_cfg = cfg.get("training", {})
+        if resume_cfg.get("resume_from_hub", False):
+            hub_cfg = cfg.get("hub", {})
+            repo_id = hub_cfg.get("repo_id")
+            subpath = hub_cfg.get("subpath", "")
+            if not repo_id:
+                raise ValueError("hub.repo_id must be set when resume_from_hub is true")
+            checkpoint_step = resume_cfg.get("checkpoint_step")
+            print(f"[Resume] Fetching {variant_label} from HF Hub: {repo_id}/{subpath}/{variant_label}_run")
+            checkpoint_path = fetch_latest_checkpoint_from_hub(
+                repo_id=repo_id,
+                subpath=subpath,
+                variant=variant_label,
+                checkpoint_step=checkpoint_step,
+            )
+            print(f"[Resume] Downloaded to: {checkpoint_path}")
+        elif resume_cfg.get("resume_from"):
+            print("[Warning] resume_from set to a specific path; all variants will try to use the same checkpoint.")
+            checkpoint_path = resume_cfg.get("resume_from")
+
+        set_seed(seed)
+
+        print(f"\n{'='*60}\n>>> Variant: {variant_label} | Out: {out_dir}\n{'='*60}")
+
+        config = TinyLlamaConfig(**cfg["model"])
+        model = TinyLlamaForCausalLM(config)
+        model = model.to(torch.bfloat16)
+
+        total_params = sum(p.numel() for p in model.parameters())
+        param_str = format_param_count(total_params)
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        run_name = f"LM-{variant_label}-{param_str}-{timestamp}"
+        cfg["training"]["run_name"] = run_name
+
+        hub_id_base = cfg["training"].get("hub_model_id", "tiny-llama-lab")
+        cfg["training"]["hub_model_id"] = f"{hub_id_base}-{variant_label}"
+
+        os.environ.pop("WANDB_RUN_ID", None)
+
+        trainer = create_trainer(model, tokenizer, cfg, train_ds, eval_ds)
+
+        # ---- CRITICAL FIX: manually set trainer on the callback ----
+        freeze_mlp = cfg["training"].get("freeze_mlp", False)
+        override_enabled = go.get("enabled", False)
+        if freeze_mlp or override_enabled:
+            for cb in trainer.callback_handler.callbacks:
+                if isinstance(cb, ResumeFreezeOverrideCallback):
+                    cb._trainer = trainer
+                    break
+
+            # Also delete optimizer.pt from the checkpoint so the Trainer doesn't try to load it
+            if checkpoint_path is not None and os.path.isdir(checkpoint_path):
+                opt_path = os.path.join(checkpoint_path, "optimizer.pt")
+                if os.path.exists(opt_path):
+                    os.remove(opt_path)
+                    print("[Freeze] Removed optimizer.pt from checkpoint to avoid state mismatch.")
+
+        # ---- NOW TRAIN ----
+        try:
+            trainer.train(resume_from_checkpoint=checkpoint_path)
+            metrics = trainer.evaluate()
+            results.append({
+                "variant": variant_label,
+                "eval_loss": metrics.get("eval_loss"),
+                "out": str(out_dir),
+                "run_name": run_name,
+                "status": "success",
+            })
+            trainer.save_model(str(out_dir))
+            if args.push or cfg["training"].get("push_to_hub", False):
+                trainer.push_to_hub()
+            print(f">>> FINISHED {variant_label} successfully")
+
+        except Exception as e:
+            print(f"!!! VARIANT {variant_label} FAILED with: {e}")
+            import traceback
+            traceback.print_exc()
+            results.append({
+                "variant": variant_label,
+                "error": str(e),
+                "out": str(out_dir),
+                "status": "failed",
+            })
+            # Continue with next variant
+            continue
+
+        finally:
+            wandb.finish()
+            torch.cuda.empty_cache()   # free memory before next variant
+
+    summary = Path(base["training"]["output_dir"]).parent / "sweep_summary.json"
+    summary.write_text(json.dumps(results, indent=2))
+    print("\nSweep complete:")
+    for r in results:
+        if r.get("status") == "success":
+            print(f"  {r['variant']:20s}  eval_loss={r['eval_loss']:.4f}")
+        else:
+            print(f"  {r['variant']:20s}  FAILED: {r.get('error', 'unknown')}")
+
+    # Return non-zero exit code if any variant failed
+    if any(r.get("status") == "failed" for r in results):
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

zain/Activation/train.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+"""Train one TinyLlama variant with optional resume from HF Hub."""
+
+import argparse
+import yaml
+import os
+import torch
+from pathlib import Path
+
+from transformers import AutoTokenizer, set_seed
+from exp import (
+    TinyLlamaConfig,
+    TinyLlamaForCausalLM,
+    build_dataset,
+    create_trainer,
+    fetch_latest_checkpoint_from_hub,
+    ResumeFreezeOverrideCallback,
+)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="Path to YAML config")
+    parser.add_argument("--variant", required=True, help="e.g., glu-silu-100L")
+    parser.add_argument("--resume_from", help="Local checkpoint path (overrides config)")
+    parser.add_argument("--push", action="store_true", help="Push final model to HF Hub")
+    args = parser.parse_args()
+
+    with open(args.config) as f:
+        cfg = yaml.safe_load(f)
+
+    seed = cfg.get("training", {}).get("seed", 42)
+    set_seed(seed)
+
+    wandb_project = cfg.get("training", {}).get("wandb_project")
+    if wandb_project:
+        os.environ["WANDB_PROJECT"] = wandb_project
+        print(f"[WandB] Project locked to: {wandb_project}")
+
+    # Determine checkpoint path
+    checkpoint_path = None
+    if args.resume_from:
+        checkpoint_path = args.resume_from
+        print(f"[Resume] Using CLI-provided local path: {checkpoint_path}")
+    else:
+        resume_config = cfg.get("training", {})
+        if resume_config.get("resume_from_hub", False):
+            hub_cfg = cfg.get("hub", {})
+            repo_id = hub_cfg.get("repo_id")
+            subpath = hub_cfg.get("subpath", "")
+            if not repo_id:
+                raise ValueError("hub.repo_id must be set when resume_from_hub is true")
+            checkpoint_step = resume_config.get("checkpoint_step")
+            print(f"[Resume] Fetching from HF Hub: {repo_id}/{subpath}/{args.variant}_run")
+            checkpoint_path = fetch_latest_checkpoint_from_hub(
+                repo_id=repo_id,
+                subpath=subpath,
+                variant=args.variant,
+                checkpoint_step=checkpoint_step,
+            )
+            print(f"[Resume] Downloaded to: {checkpoint_path}")
+        elif resume_config.get("resume_from"):
+            checkpoint_path = resume_config.get("resume_from")
+            print(f"[Resume] Using config-provided local path: {checkpoint_path}")
+
+    # Build model & tokenizer
+    model_cfg = cfg["model"]
+    train_cfg = cfg.get("training", {})
+
+    # Output dir with _trash if override enabled
+    go = cfg.get("gradient_override", {})
+    if go.get("enabled", False):
+        base_out = Path(train_cfg.get("output_dir", "./out"))
+        variant_name = args.variant + "_trash"
+        output_dir = base_out / variant_name
+        train_cfg["output_dir"] = str(output_dir)
+        print(f"[Override] Output directory set to: {output_dir}")
+
+    tok_name = model_cfg.pop("tokenizer_name", "meta-llama/Llama-2-7b-hf")
+    tokenizer = AutoTokenizer.from_pretrained(tok_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    tiny_config = TinyLlamaConfig(**model_cfg)
+    model = TinyLlamaForCausalLM(tiny_config)
+    model = model.to(torch.bfloat16)
+
+    n_params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"Model: {n_params:.2f}M params | MLP type: {tiny_config.mlp_type} | Activation: {tiny_config.activation}")
+
+    msl = model_cfg.get("max_position_embeddings", 512)
+    train_ds = build_dataset(tokenizer, max_seq_len=msl, split="train", max_samples=None)
+    eval_ds = build_dataset(tokenizer, max_seq_len=msl, split="validation", max_samples=None)
+
+    trainer = create_trainer(model, tokenizer, cfg, train_ds, eval_ds)
+
+    # Manually set trainer on callback if freezing/override is enabled
+    freeze_mlp = train_cfg.get("freeze_mlp", False)
+    override_enabled = go.get("enabled", False)
+    if freeze_mlp or override_enabled:
+        for cb in trainer.callback_handler.callbacks:
+            if isinstance(cb, ResumeFreezeOverrideCallback):
+                cb._trainer = trainer
+                break
+        # Delete optimizer.pt to avoid state mismatch
+        if checkpoint_path is not None and os.path.isdir(checkpoint_path):
+            opt_path = os.path.join(checkpoint_path, "optimizer.pt")
+            if os.path.exists(opt_path):
+                os.remove(opt_path)
+                print("[Freeze] Removed optimizer.pt from checkpoint.")
+
+    trainer.train(resume_from_checkpoint=checkpoint_path)
+
+    out = train_cfg.get("output_dir", "./out")
+    trainer.save_model(out)
+    if args.push or train_cfg.get("push_to_hub", False):
+        trainer.push_to_hub()
+    print(f"Done. Artifacts in {out}")
+
+
+if __name__ == "__main__":
+    main()

zain/Activation/wandb/debug-internal.log

@@ -0,0 +1,33 @@
+{"time":"2026-08-13T21:15:21.457059839Z","level":"INFO","msg":"wandb-core"}
+{"time":"2026-08-13T21:15:21.457218835Z","level":"INFO","msg":"stream: starting","core version":"0.28.1"}
+{"time":"2026-08-13T21:15:21.71665521Z","level":"INFO","msg":"stream: created new stream","id":"bdgno22l"}
+{"time":"2026-08-13T21:15:21.716724793Z","level":"INFO","msg":"handler: started"}
+{"time":"2026-08-13T21:15:21.716834966Z","level":"INFO","msg":"stream: started"}
+{"time":"2026-08-13T21:15:21.716845775Z","level":"INFO","msg":"writer: started","stream_id":"bdgno22l"}
+{"time":"2026-08-13T21:15:21.716864355Z","level":"INFO","msg":"sender: started"}
+{"time":"2026-08-13T21:15:22.087946715Z","level":"INFO","msg":"filestream: sending request","total_files":1,"console_offset":0,"console_lines":1}
+{"time":"2026-08-13T21:15:22.187941633Z","level":"INFO","msg":"filestream: request sent","status":"200 OK"}
+{"time":"2026-08-13T21:15:33.902773632Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":2617}
+{"time":"2026-08-13T21:15:33.902804288Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":1}
+{"time":"2026-08-13T21:15:33.903189126Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":2650}
+{"time":"2026-08-13T21:15:33.903223176Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":1}
+{"time":"2026-08-13T21:15:33.905188752Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":3405}
+{"time":"2026-08-13T21:15:33.905644326Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":109}
+{"time":"2026-08-13T21:15:33.905737177Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":3547}
+{"time":"2026-08-13T21:15:33.905849936Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":11}
+{"time":"2026-08-13T21:15:33.907613871Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":4244}
+{"time":"2026-08-13T21:15:33.914808342Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":2917}
+{"time":"2026-08-13T21:15:33.91681943Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":7726}
+{"time":"2026-08-13T21:15:33.917089153Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":83}
+{"time":"2026-08-13T21:15:33.917182957Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":7844}
+{"time":"2026-08-13T21:15:33.91727044Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":13}
+{"time":"2026-08-13T21:15:33.919229048Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":8663}
+{"time":"2026-08-13T21:15:33.919463562Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":61}
+{"time":"2026-08-13T21:15:33.92636457Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":9853}
+{"time":"2026-08-13T21:15:33.926563588Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":8}
+{"time":"2026-08-13T21:15:33.92776116Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":10322}
+{"time":"2026-08-13T21:15:33.927881025Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":18}
+{"time":"2026-08-13T21:15:37.112579744Z","level":"INFO","msg":"filestream: sending request","total_files":4,"history_offset":0,"history_lines":2,"events_offset":0,"events_lines":1,"console_offset":1,"console_lines":5,"uploaded_len":2}
+{"time":"2026-08-13T21:15:37.764424515Z","level":"INFO","msg":"filestream: request sent","status":"200 OK"}
+{"time":"2026-08-13T21:15:52.106707323Z","level":"INFO","msg":"filestream: sending request","total_files":4,"history_offset":2,"history_lines":3,"events_offset":1,"events_lines":2,"console_offset":3,"console_lines":1}
+{"time":"2026-08-13T21:15:53.006564044Z","level":"INFO","msg":"filestream: request sent","status":"200 OK"}

zain/Activation/wandb/debug.log

@@ -0,0 +1,23 @@
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_setup.py:_flush():81] Current SDK version is 0.28.1
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_setup.py:_flush():81] Configure stats pid to 52372
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:setup_run_log_directory():729] Logging user logs to /mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug.log
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:setup_run_log_directory():730] Logging internal logs to /mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug-internal.log
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():772] calling init triggers
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():777] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():820] starting backend
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():826] Connected to an existing wandb-core service via WANDB_SERVICE
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():835] sending inform_init request
+2026-08-13 21:15:21,717 INFO    MainThread:52372 [wandb_init.py:init():840] backend started and connected
+2026-08-13 21:15:21,720 INFO    MainThread:52372 [wandb_init.py:init():910] updated telemetry
+2026-08-13 21:15:21,727 INFO    MainThread:52372 [wandb_init.py:init():933] communicating run to backend with 90.0 second timeout
+2026-08-13 21:15:22,001 INFO    MainThread:52372 [wandb_init.py:init():978] starting run threads in backend
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_console_start():2621] atexit reg
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_redirect():2471] redirect: wrap_raw
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_redirect():2540] Wrapping output streams.
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_redirect():2563] Redirects installed.
+2026-08-13 21:15:22,077 INFO    MainThread:52372 [wandb_init.py:init():1016] run started, returning control to user process
+2026-08-13 21:15:22,078 INFO    MainThread:52372 [wandb_run.py:_config_callback():1346] config_cb None None {'transformers_version': '5.16.0.dev0', 'architectures': None, 'output_hidden_states': False, 'return_dict': True, 'dtype': None, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'id2label': {0: 'LABEL_0', 1: 'LABEL_1'}, 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, 'problem_type': None, 'vocab_size': 4096, 'hidden_size': 128, 'intermediate_size': 256, 'num_hidden_layers': 100, 'num_attention_heads': 4, 'num_key_value_heads': 4, 'hidden_act': 'silu', 'max_position_embeddings': 512, 'initializer_range': 0.02, 'rms_norm_eps': 1e-06, 'use_cache': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2, 'pretraining_tp': 1, 'tie_word_embeddings': True, 'rope_parameters': {'rope_theta': 10000.0, 'rope_type': 'default'}, 'attention_bias': False, 'attention_dropout': 0.0, 'mlp_bias': False, 'head_dim': 32, '_name_or_path': '', 'tokenizer_name': 'w-ahmad/tiny-stories-tokenizer', 'mlp_type': 'glu', 'activation': 'linear', 'waleed_beta': 10.0, 'model_type': 'tiny_llama', 'output_attentions': False, 'output_dir': 'out/glu-linear-100L_trash_run', 'per_device_train_batch_size': 64, 'num_train_epochs': 1, 'max_steps': 2700, 'learning_rate': 9e-05, 'lr_scheduler_type': 'constant_with_warmup', 'lr_scheduler_kwargs': None, 'warmup_steps': 50, 'optim': 'adamw_torch_fused', 'optim_args': None, 'weight_decay': 0.0, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_epsilon': 1e-08, 'optim_target_modules': None, 'gradient_accumulation_steps': 1, 'average_tokens_across_devices': True, 'max_grad_norm': 1.0, 'label_smoothing_factor': 0.0, 'bf16': True, 'fp16': False, 'bf16_full_eval': False, 'fp16_full_eval': False, 'tf32': None, 'gradient_checkpointing': False, 'gradient_checkpointing_kwargs': None, 'torch_compile': False, 'torch_compile_backend': None, 'torch_compile_mode': None, 'use_liger_kernel': False, 'liger_kernel_config': None, 'neftune_noise_alpha': None, 'torch_empty_cache_steps': None, 'auto_find_batch_size': False, 'logging_strategy': 'steps', 'logging_steps': 20, 'logging_first_step': False, 'log_on_each_node': True, 'logging_nan_inf_filter': True, 'include_num_input_tokens_seen': 'no', 'log_level': 'passive', 'log_level_replica': 'warning', 'disable_tqdm': False, 'report_to': ['wandb'], 'run_name': 'LM-glu-linear-100L-16.9M-20260813-211520', 'project': 'huggingface', 'trackio_space_id': None, 'trackio_bucket_id': None, 'trackio_static_space_id': None, 'eval_strategy': 'steps', 'eval_steps': 2498, 'eval_delay': 0, 'per_device_eval_batch_size': 1024, 'prediction_loss_only': False, 'eval_on_start': False, 'eval_do_concat_batches': True, 'eval_use_gather_object': False, 'eval_accumulation_steps': None, 'include_for_metrics': [], 'batch_eval_metrics': False, 'save_only_model': False, 'save_strategy': 'steps', 'save_steps': 100, 'save_on_each_node': False, 'save_total_limit': None, 'enable_jit_checkpoint': False, 'push_to_hub': False, 'hub_token': '<HUB_TOKEN>', 'hub_private_repo': None, 'hub_model_id': 'w-ahmad/6L-glu-linear-100L', 'hub_strategy': 'every_save', 'hub_always_push': False, 'hub_revision': None, 'load_best_model_at_end': False, 'metric_for_best_model': None, 'greater_is_better': None, 'ignore_data_skip': False, 'restore_callback_states_from_checkpoint': False, 'full_determinism': False, 'seed': 42, 'data_seed': 42, 'use_cpu': False, 'accelerator_config': {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}, 'parallelism_config': None, 'dataloader_drop_last': False, 'dataloader_num_workers': 0, 'dataloader_pin_memory': True, 'dataloader_persistent_workers': False, 'dataloader_prefetch_factor': None, 'dataloader_multiprocessing_context': None, 'dataloader_in_order': True, 'remove_unused_columns': False, 'label_names': None, 'train_sampling_strategy': 'random', 'length_column_name': 'length', 'ddp_find_unused_parameters': None, 'ddp_bucket_cap_mb': None, 'ddp_broadcast_buffers': None, 'ddp_static_graph': None, 'ddp_backend': None, 'ddp_timeout': 1800, 'fsdp': None, 'fsdp_config': None, 'deepspeed': None, 'debug': [], 'skip_memory_metrics': True, 'do_train': False, 'do_eval': True, 'do_predict': False, 'resume_from_checkpoint': None, 'local_rank': -1}
+2026-08-13 21:15:22,082 INFO    MainThread:52372 [wandb_config.py:__setitem__():155] [no run ID] config set model/num_parameters = 16934016 - <bound method Run._config_callback of <wandb.sdk.wandb_run.Run object at 0x15235c33c510>>
+2026-08-13 21:15:22,082 INFO    MainThread:52372 [wandb_run.py:_config_callback():1346] config_cb model/num_parameters 16934016 None

zain/Activation/wandb/run-20260813_211521-bdgno22l/files/output.log

@@ -0,0 +1,14 @@
+[Override] Activated with value = -10.0
+  0%|          | 0/2700 [00:00<?, ?it/s][transformers] `use_return_dict` is deprecated! Use `return_dict` instead!
+[INFO] Causal mask (float with -inf) applied to all attention layers.
+ 96%|█████████▋| 2600/2700 [00:27<00:00, 102.99it/s][transformers] TinyLlamaForCausalLM has generative capabilities, as `prepare_inputs_for_generation` is explicitly defined. However, it doesn't directly inherit from `GenerationMixin`. From 👉v4.50👈 onwards, `PreTrainedModel` will NOT inherit from `GenerationMixin`, and this model will lose the ability to call `generate` and other related functions.
+{'loss': '2.102', 'grad_norm': '2.661e+07', 'learning_rate': '3.42e-05', 'epoch': '0.1699', 'train/total_time_seconds': '5.379', 'train/time_per_step_avg': '0.2689', 'train/epoch_time_elapsed': '6.585', 'train/estimated_remaining_minutes': '0.006403', 'train/tensor_act_model_layers_0_residual_pre_attn/norm': '187.1', 'train/tensor_act_model_layers_0_residual_pre_attn/mean': '0.001579', 'train/tensor_act_model_layers_0_residual_pre_attn/std': '0.09131', 'train/tensor_act_model_layers_0_residual_pre_attn/max_abs': '0.2734', 'train/tensor_act_model_layers_0_residual_pre_attn/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_residual_pre_attn/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_residual_pre_attn/min': '-0.2734', 'train/tensor_act_model_layers_0_residual_pre_attn/max': '0.2734', 'train/tensor_act_model_layers_0_residual_pre_attn/range': '0.5469', 'train/tensor_act_model_layers_0_residual_post_attn/norm': '188.5', 'train/tensor_act_model_layers_0_residual_post_attn/mean': '0.001442', 'train/tensor_act_model_layers_0_residual_post_attn/std': '0.09229', 'train/tensor_act_model_layers_0_residual_post_attn/max_abs': '0.3574', 'train/tensor_act_model_layers_0_residual_post_attn/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_residual_post_attn/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_residual_post_attn/min': '-0.3574', 'train/tensor_act_model_layers_0_residual_post_attn/max': '0.3418', 'train/tensor_act_model_layers_0_residual_post_attn/range': '0.6992', 'train/tensor_act_model_layers_0_mlp_gate_proj/norm': '1041', 'train/tensor_act_model_layers_0_mlp_gate_proj/mean': '0.0003281', 'train/tensor_act_model_layers_0_mlp_gate_proj/std': '0.3594', 'train/tensor_act_model_layers_0_mlp_gate_proj/max_abs': '1.828', 'train/tensor_act_model_layers_0_mlp_gate_proj/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp_gate_proj/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp_gate_proj/min': '-1.719', 'train/tensor_act_model_layers_0_mlp_gate_proj/max': '1.828', 'train/tensor_act_model_layers_0_mlp_gate_proj/range': '3.547', 'train/tensor_act_model_layers_0_mlp_up_proj/norm': '1049', 'train/tensor_act_model_layers_0_mlp_up_proj/mean': '-0.00238', 'train/tensor_act_model_layers_0_mlp_up_proj/std': '0.3613', 'train/tensor_act_model_layers_0_mlp_up_proj/max_abs': '2.125', 'train/tensor_act_model_layers_0_mlp_up_proj/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp_up_proj/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp_up_proj/min': '-1.992', 'train/tensor_act_model_layers_0_mlp_up_proj/max': '2.125', 'train/tensor_act_model_layers_0_mlp_up_proj/range': '4.117', 'train/tensor_act_model_layers_0_mlp_down_proj/norm': '200.9', 'train/tensor_act_model_layers_0_mlp_down_proj/mean': '0.0004768', 'train/tensor_act_model_layers_0_mlp_down_proj/std': '0.09814', 'train/tensor_act_model_layers_0_mlp_down_proj/max_abs': '0.5273', 'train/tensor_act_model_layers_0_mlp_down_proj/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp_down_proj/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp_down_proj/min': '-0.4902', 'train/tensor_act_model_layers_0_mlp_down_proj/max': '0.5273', 'train/tensor_act_model_layers_0_mlp_down_proj/range': '1.018', 'train/tensor_act_model_layers_0_mlp/norm': '200.9', 'train/tensor_act_model_layers_0_mlp/mean': '0.0004768', 'train/tensor_act_model_layers_0_mlp/std': '0.09814', 'train/tensor_act_model_layers_0_mlp/max_abs': '0.5273', 'train/tensor_act_model_layers_0_mlp/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp/min': '-0.4902', 'train/tensor_act_model_layers_0_mlp/max': '0.5273', 'train/tensor_act_model_layers_0_mlp/range': '1.018', 'train/tensor_act_model_layers_0_residual_post_mlp/norm': '291.2', 'train/tensor_act_model_layers_0_residual_post_mlp/mean': '0.001923', 'train/tensor_act_model_layers_0_residual_post_mlp/std': '0.1426', 'train/tensor_act_model_layers_0_residual_post_mlp/max_abs':
+{'loss': '2.128', 'grad_norm': '3.106e+07', 'learning_rate': '7.02e-05', 'epoch': '0.1712', 'train/total_time_seconds': '9.549', 'train/time_per_step_avg': '0.2387', 'train/epoch_time_elapsed': '11.79', 'train/estimated_remaining_minutes': '0.01003'}
+{'loss': '4.245', 'grad_norm': '2.223e+08', 'learning_rate': '9e-05', 'epoch': '0.1726', 'train/total_time_seconds': '14.12', 'train/time_per_step_avg': '0.2354', 'train/epoch_time_elapsed': '17.43', 'train/estimated_remaining_minutes': '0.01287', 'train/tensor_act_model_layers_0_residual_pre_attn/norm': '187.2', 'train/tensor_act_model_layers_0_residual_pre_attn/mean': '0.001549', 'train/tensor_act_model_layers_0_residual_pre_attn/std': '0.09131', 'train/tensor_act_model_layers_0_residual_pre_attn/max_abs': '0.2734', 'train/tensor_act_model_layers_0_residual_pre_attn/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_residual_pre_attn/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_residual_pre_attn/min': '-0.2734', 'train/tensor_act_model_layers_0_residual_pre_attn/max': '0.2734', 'train/tensor_act_model_layers_0_residual_pre_attn/range': '0.5469', 'train/tensor_act_model_layers_0_residual_post_attn/norm': '188.6', 'train/tensor_act_model_layers_0_residual_post_attn/mean': '0.001312', 'train/tensor_act_model_layers_0_residual_post_attn/std': '0.09229', 'train/tensor_act_model_layers_0_residual_post_attn/max_abs': '0.3477', 'train/tensor_act_model_layers_0_residual_post_attn/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_residual_post_attn/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_residual_post_attn/min': '-0.3477', 'train/tensor_act_model_layers_0_residual_post_attn/max': '0.3477', 'train/tensor_act_model_layers_0_residual_post_attn/range': '0.6953', 'train/tensor_act_model_layers_0_mlp_gate_proj/norm': '1087', 'train/tensor_act_model_layers_0_mlp_gate_proj/mean': '-0.000843', 'train/tensor_act_model_layers_0_mlp_gate_proj/std': '0.375', 'train/tensor_act_model_layers_0_mlp_gate_proj/max_abs': '1.867', 'train/tensor_act_model_layers_0_mlp_gate_proj/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp_gate_proj/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp_gate_proj/min': '-1.867', 'train/tensor_act_model_layers_0_mlp_gate_proj/max': '1.852', 'train/tensor_act_model_layers_0_mlp_gate_proj/range': '3.719', 'train/tensor_act_model_layers_0_mlp_up_proj/norm': '1094', 'train/tensor_act_model_layers_0_mlp_up_proj/mean': '-0.001305', 'train/tensor_act_model_layers_0_mlp_up_proj/std': '0.377', 'train/tensor_act_model_layers_0_mlp_up_proj/max_abs': '2.266', 'train/tensor_act_model_layers_0_mlp_up_proj/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp_up_proj/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp_up_proj/min': '-1.984', 'train/tensor_act_model_layers_0_mlp_up_proj/max': '2.266', 'train/tensor_act_model_layers_0_mlp_up_proj/range': '4.25', 'train/tensor_act_model_layers_0_mlp_down_proj/norm': '233.8', 'train/tensor_act_model_layers_0_mlp_down_proj/mean': '0.02856', 'train/tensor_act_model_layers_0_mlp_down_proj/std': '0.1104', 'train/tensor_act_model_layers_0_mlp_down_proj/max_abs': '0.7109', 'train/tensor_act_model_layers_0_mlp_down_proj/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp_down_proj/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp_down_proj/min': '-0.5039', 'train/tensor_act_model_layers_0_mlp_down_proj/max': '0.7109', 'train/tensor_act_model_layers_0_mlp_down_proj/range': '1.215', 'train/tensor_act_model_layers_0_mlp/norm': '233.8', 'train/tensor_act_model_layers_0_mlp/mean': '0.02856', 'train/tensor_act_model_layers_0_mlp/std': '0.1104', 'train/tensor_act_model_layers_0_mlp/max_abs': '0.7109', 'train/tensor_act_model_layers_0_mlp/frac_near_dtype_limit': '0', 'train/tensor_act_model_layers_0_mlp/frac_near_user_limit': '0', 'train/tensor_act_model_layers_0_mlp/min': '-0.5039', 'train/tensor_act_model_layers_0_mlp/max': '0.7109', 'train/tensor_act_model_layers_0_mlp/range': '1.215', 'train/tensor_act_model_layers_0_residual_post_mlp/norm': '315', 'train/tensor_act_model_layers_0_residual_post_mlp/mean': '0.02991', 'train/tensor_act_model_layers_0_residual_post_mlp/std': '0.1504', 'train/tensor_act_model_layers_0_residual_post_mlp/max_abs': '0.7734', 'tra
+{'loss': '7.166', 'grad_norm': '4.886e+08', 'learning_rate': '9e-05', 'epoch': '0.1739', 'train/total_time_seconds': '18.38', 'train/time_per_step_avg': '0.2297', 'train/epoch_time_elapsed': '22.75', 'train/estimated_remaining_minutes': '0.01424'}
+{'loss': '7.483', 'grad_norm': '8.85e+08', 'learning_rate': '9e-05', 'epoch': '0.1753', 'train/total_time_seconds': '22.58', 'train/time_per_step_avg': '0.2258', 'train/epoch_time_elapsed': '27.92', 'train/estimated_remaining_minutes': '0.01447'}
+  - If you're using `trust_remote_code=True`, you can get rid of this warning by loading the model with an auto class. See https://huggingface.co/docs/transformers/en/model_doc/auto#auto-classes
+  - If you are the owner of the model architecture code, please modify your model class such that it inherits from `GenerationMixin` (after `PreTrainedModel`, otherwise you'll get an exception).
+  - If you are not the owner of the model architecture class, please contact the model code owner to update it.
+Writing model shards: 100%|██████████| 1/1 [00:00<00:00, 12.56it/s]
+Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

zain/Activation/wandb/run-20260813_211521-bdgno22l/files/requirements.txt

@@ -0,0 +1,149 @@
+asttokens==3.0.1
+comm==0.2.3
+debugpy==1.8.21
+decorator==5.3.1
+executing==2.2.1
+nest-asyncio==1.6.0
+parso==0.8.7
+platformdirs==4.11.0
+psutil==7.2.2
+ptyprocess==0.7.0
+pure_eval==0.2.3
+Pygments==2.20.0
+pyzmq==27.1.0
+setuptools==83.0.0
+six==1.17.0
+tornado==6.5.7
+traitlets==5.15.0
+fsspec==2026.4.0
+wcwidth==0.8.2
+ipython_pygments_lexers==1.1.1
+jedi==0.20.0
+jupyter_core==5.9.1
+matplotlib-inline==0.2.2
+pexpect==4.9.0
+prompt_toolkit==3.0.53
+python-dateutil==2.9.0.post0
+stack_data==0.6.3
+wheel==0.47.0
+jupyter_client==8.9.1
+pip==26.1.2
+ipython==9.15.0
+ipykernel==7.2.0
+threadpoolctl==3.6.0
+pyparsing==3.3.2
+typing_extensions==4.15.0
+Jinja2==3.1.6
+narwhals==2.24.0
+kiwisolver==1.5.0
+joblib==1.5.3
+fonttools==4.63.0
+cycler==0.12.1
+scipy==1.17.1
+pandas==3.0.5
+contourpy==1.3.3
+scikit-learn==1.9.0
+matplotlib==3.11.1
+urllib3==2.7.0
+tqdm==4.70.0
+idna==3.18
+charset-normalizer==3.4.9
+certifi==2026.7.22
+requests==2.34.2
+seaborn==0.13.2
+uv==0.12.0
+shellingham==1.5.4
+mpmath==1.3.0
+attrs==26.1.0
+hf-xet==1.5.2
+nvidia-nccl-cu12==2.21.5
+MarkupSafe==3.0.3
+regex==2026.7.19
+importlib_metadata==9.0.0
+httpcore==1.0.9
+annotated-doc==0.0.5
+multidict==6.7.1
+aiohttp==3.14.3
+aiosignal==1.4.0
+xxhash==3.8.1
+aiohappyeyeballs==2.7.1
+mdurl==0.1.2
+cuda-toolkit==13.0.3.0
+networkx==3.6.1
+PyYAML==6.0.3
+nvidia-cufile==1.15.1.6
+typer==0.27.0
+torchaudio==2.6.0+cu124
+rich==15.0.0
+nvidia-cufft-cu12==11.2.1.3
+h11==0.16.0
+dill==0.4.1
+cuda-pathfinder==1.6.0
+filelock==3.29.0
+nvidia-nvtx-cu12==12.4.127
+httpx==0.28.1
+anyio==4.14.2
+numpy==2.4.4
+yarl==1.24.5
+click==8.4.2
+triton==3.2.0
+frozenlist==1.8.0
+zipp==4.1.0
+propcache==0.5.2
+tokenizers==0.22.2
+markdown-it-py==4.2.0
+nvidia-cuda-runtime==13.0.96
+cuda-bindings==13.3.1
+nvidia-cuda-cupti==13.0.85
+torch==2.6.0+cu124
+multiprocess==0.70.19
+pillow==12.2.0
+transformers==5.16.0.dev0
+wandb==0.28.1
+nvidia-curand==10.4.0.35
+sympy==1.13.1
+nvidia-cusparse==12.6.3.3
+nvidia-cuda-nvrtc==13.0.88
+typing-inspection==0.4.2
+nvidia-cusolver==12.0.4.66
+nvidia-cufft==12.0.0.61
+nvidia-cudnn-cu13==9.20.0.48
+nvidia-cublas==13.1.1.3
+pyarrow==25.0.0
+evaluate==0.4.6
+diffusers==0.39.0
+pydantic==2.13.4
+annotated-types==0.8.0
+protobuf==7.35.1
+sentry-sdk==2.66.1
+einops==0.8.2
+packaging==26.2
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusparselt-cu12==0.6.2
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cuda-runtime-cu12==12.4.127
+torchvision==0.21.0+cu124
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cudnn-cu12==9.1.0.70
+huggingface_hub==1.26.0
+datasets==5.0.1
+safetensors==0.8.0
+accelerate==1.14.0
+pydantic_core==2.46.4
+ninja==1.13.0
+autocommand==2.2.2
+backports.tarfile==1.2.0
+importlib_metadata==8.7.1
+jaraco.text==4.0.0
+jaraco.context==6.1.0
+jaraco.functools==4.4.0
+more-itertools==10.8.0
+packaging==26.0
+platformdirs==4.4.0
+tomli==2.4.0
+wheel==0.46.3
+zipp==3.23.0

zain/Activation/wandb/run-20260813_211521-bdgno22l/files/wandb-metadata.json

@@ -0,0 +1,102 @@
+{
+  "os":  "Linux-5.15.0-126-generic-x86_64-with-glibc2.35",
+  "python":  "CPython 3.11.15",
+  "startedAt":  "2026-08-13T21:15:21.454482Z",
+  "args":  [
+    "--config",
+    "/mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation/configs/baseline100L.yaml",
+    "--variants",
+    "glu-linear-100L",
+    "glu-silu-100L",
+    "glu-silu-waleed10-100L",
+    "glu-situglu-100L",
+    "glu-situglu_low-100L",
+    "glu-waleed-100L",
+    "glu-waleed10-100L",
+    "glu-waleedglu_low-100L"
+  ],
+  "program":  "/mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation/sweep.py",
+  "codePath":  "sweep.py",
+  "codePathLocal":  "sweep.py",
+  "git":  {
+    "remote":  "https://github.com/w-ahmad1a10/Activation.git",
+    "commit":  "8381885396c6fb3cbe0304e3b3c5498dda2c7440"
+  },
+  "email":  "deepnevro@gmail.com",
+  "root":  "/mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation",
+  "host":  "deeplens-k3s-node1",
+  "executable":  "/mnt/data/zainulabideen/zain-exp/notebooks/my_env/bin/python",
+  "cpu_count":  112,
+  "cpu_count_logical":  224,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "1560765693952",
+      "used":  "710564417536"
+    }
+  },
+  "memory":  {
+    "total":  "2164089937920"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-39c684a5-fde6-83d7-1663-0859795881ae"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-68012e5a-38b6-b643-0ca6-62fb66720bf3"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-132944c4-b689-2b5f-89a4-d730401677ab"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-2df386cc-6d26-d0e2-7a2d-a057b0d95864"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-bfa16575-1d94-1aa2-4537-2c93433f42ef"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-bc6c3e3c-9b90-09ca-c034-774961847c54"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-00a441e1-7c95-e7d6-4c35-43d6b291aea9"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-1c4d29a2-4647-6fce-d8fc-0c5ecfbbd6ea"
+    }
+  ],
+  "cudaVersion":  "12.4",
+  "writerId":  "3doo1fos9k4ofmuy6vlh53x3p7tpfmqy"
+}

zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug-core.log

@@ -0,0 +1,100 @@
+{"time":"2026-08-13T20:50:56.88837736Z","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmp_1eb1gnf/port-4013524.txt","pid":4013524,"detached":false,"idle-timeout":600000000000,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false}
+{"time":"2026-08-13T20:50:56.889974484Z","level":"INFO","msg":"server: will exit if parent process dies","ppid":4013524}
+{"time":"2026-08-13T20:50:56.889954218Z","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-4013524-4015748-1172316019/socket","Net":"unix"}}
+{"time":"2026-08-13T20:50:57.063325317Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"}
+{"time":"2026-08-13T20:51:59.231462609Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"2(@)"}
+{"time":"2026-08-13T20:51:59.35547125Z","level":"INFO","msg":"handleInformInit: received","streamId":"qudyw6my","id":"2(@)"}
+{"time":"2026-08-13T20:51:59.618336359Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"qudyw6my","id":"2(@)"}
+{"time":"2026-08-13T20:52:05.07357923Z","level":"INFO","msg":"connection: cancelling request","id":"2(@)","requestId":"67oxstahnbwk"}
+{"time":"2026-08-13T20:53:24.810084125Z","level":"INFO","msg":"connection: cancelling request","id":"2(@)","requestId":"67oxstahnbwk"}
+{"time":"2026-08-13T20:53:26.647793248Z","level":"INFO","msg":"connection: cancelling request","id":"2(@)","requestId":"67oxstahnbwk"}
+{"time":"2026-08-13T20:53:26.906580178Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"qudyw6my","id":"2(@)"}
+{"time":"2026-08-13T20:53:26.907531442Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"qudyw6my","id":"2(@)"}
+{"time":"2026-08-13T20:53:28.4626298Z","level":"INFO","msg":"connection: closing","id":"2(@)"}
+{"time":"2026-08-13T20:53:28.462723867Z","level":"INFO","msg":"connection: closed successfully","id":"2(@)"}
+{"time":"2026-08-13T20:53:28.462642414Z","level":"INFO","msg":"processOutgoingData: finished","id":"2(@)"}
+{"time":"2026-08-13T20:53:28.462736968Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"2(@)"}
+{"time":"2026-08-13T20:58:26.484371032Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"3(@)"}
+{"time":"2026-08-13T20:58:26.565457961Z","level":"INFO","msg":"handleInformInit: received","streamId":"k88lj9m1","id":"3(@)"}
+{"time":"2026-08-13T20:58:26.825939415Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"k88lj9m1","id":"3(@)"}
+{"time":"2026-08-13T20:58:32.214683089Z","level":"INFO","msg":"connection: cancelling request","id":"3(@)","requestId":"aja3hhk7rh8b"}
+{"time":"2026-08-13T20:59:51.636693819Z","level":"INFO","msg":"connection: cancelling request","id":"3(@)","requestId":"aja3hhk7rh8b"}
+{"time":"2026-08-13T20:59:53.52715723Z","level":"INFO","msg":"connection: cancelling request","id":"3(@)","requestId":"aja3hhk7rh8b"}
+{"time":"2026-08-13T20:59:53.844075007Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"k88lj9m1","id":"3(@)"}
+{"time":"2026-08-13T20:59:53.844796525Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"k88lj9m1","id":"3(@)"}
+{"time":"2026-08-13T20:59:55.642797778Z","level":"INFO","msg":"processOutgoingData: finished","id":"3(@)"}
+{"time":"2026-08-13T20:59:55.6427907Z","level":"INFO","msg":"connection: closing","id":"3(@)"}
+{"time":"2026-08-13T20:59:55.642884054Z","level":"INFO","msg":"connection: closed successfully","id":"3(@)"}
+{"time":"2026-08-13T20:59:55.642888633Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"3(@)"}
+{"time":"2026-08-13T21:10:30.327897889Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"4(@)"}
+{"time":"2026-08-13T21:10:30.430510211Z","level":"INFO","msg":"handleInformInit: received","streamId":"idtw6vfo","id":"4(@)"}
+{"time":"2026-08-13T21:10:30.674038153Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"idtw6vfo","id":"4(@)"}
+{"time":"2026-08-13T21:10:33.084198753Z","level":"INFO","msg":"connection: cancelling request","id":"4(@)","requestId":"y5arwk2qvfef"}
+{"time":"2026-08-13T21:10:33.613412909Z","level":"INFO","msg":"connection: cancelling request","id":"4(@)","requestId":"y5arwk2qvfef"}
+{"time":"2026-08-13T21:10:33.614423646Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"idtw6vfo","id":"4(@)"}
+{"time":"2026-08-13T21:10:33.615071285Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"idtw6vfo","id":"4(@)"}
+{"time":"2026-08-13T21:10:35.274863492Z","level":"INFO","msg":"connection: closing","id":"4(@)"}
+{"time":"2026-08-13T21:10:35.274947383Z","level":"INFO","msg":"connection: closed successfully","id":"4(@)"}
+{"time":"2026-08-13T21:10:35.274864147Z","level":"INFO","msg":"processOutgoingData: finished","id":"4(@)"}
+{"time":"2026-08-13T21:10:35.274956775Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"4(@)"}
+{"time":"2026-08-13T21:10:59.004852712Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"5(@)"}
+{"time":"2026-08-13T21:10:59.090409085Z","level":"INFO","msg":"handleInformInit: received","streamId":"cvqont09","id":"5(@)"}
+{"time":"2026-08-13T21:10:59.354009138Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"cvqont09","id":"5(@)"}
+{"time":"2026-08-13T21:11:04.733635221Z","level":"INFO","msg":"connection: cancelling request","id":"5(@)","requestId":"3mnjy6uk50sj"}
+{"time":"2026-08-13T21:11:24.240496285Z","level":"INFO","msg":"connection: cancelling request","id":"5(@)","requestId":"3mnjy6uk50sj"}
+{"time":"2026-08-13T21:11:24.943824437Z","level":"INFO","msg":"connection: cancelling request","id":"5(@)","requestId":"3mnjy6uk50sj"}
+{"time":"2026-08-13T21:11:24.945036572Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"cvqont09","id":"5(@)"}
+{"time":"2026-08-13T21:11:24.945565304Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"cvqont09","id":"5(@)"}
+{"time":"2026-08-13T21:11:26.576272285Z","level":"INFO","msg":"connection: closing","id":"5(@)"}
+{"time":"2026-08-13T21:11:26.57629102Z","level":"INFO","msg":"processOutgoingData: finished","id":"5(@)"}
+{"time":"2026-08-13T21:11:26.576355678Z","level":"INFO","msg":"connection: closed successfully","id":"5(@)"}
+{"time":"2026-08-13T21:11:26.576364137Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"5(@)"}
+{"time":"2026-08-13T21:11:43.037962092Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"6(@)"}
+{"time":"2026-08-13T21:11:43.126951891Z","level":"INFO","msg":"handleInformInit: received","streamId":"olcsup4b","id":"6(@)"}
+{"time":"2026-08-13T21:11:43.384950035Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"olcsup4b","id":"6(@)"}
+{"time":"2026-08-13T21:11:48.74325104Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"0eevcmnde8cb"}
+{"time":"2026-08-13T21:12:13.594157711Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"0eevcmnde8cb"}
+{"time":"2026-08-13T21:12:14.147755589Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"0eevcmnde8cb"}
+{"time":"2026-08-13T21:12:14.149406248Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"olcsup4b","id":"6(@)"}
+{"time":"2026-08-13T21:12:14.150092264Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"olcsup4b","id":"6(@)"}
+{"time":"2026-08-13T21:12:15.149367854Z","level":"INFO","msg":"handleInformInit: received","streamId":"o4os2eko","id":"6(@)"}
+{"time":"2026-08-13T21:12:15.419403518Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"o4os2eko","id":"6(@)"}
+{"time":"2026-08-13T21:12:20.903171066Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"lmwiaqx3lznd"}
+{"time":"2026-08-13T21:12:31.561337484Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"lmwiaqx3lznd"}
+{"time":"2026-08-13T21:12:32.158848875Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"lmwiaqx3lznd"}
+{"time":"2026-08-13T21:12:32.160185701Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"o4os2eko","id":"6(@)"}
+{"time":"2026-08-13T21:12:32.160843182Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"o4os2eko","id":"6(@)"}
+{"time":"2026-08-13T21:12:33.26906436Z","level":"INFO","msg":"handleInformInit: received","streamId":"69zdkrji","id":"6(@)"}
+{"time":"2026-08-13T21:12:33.527818889Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"69zdkrji","id":"6(@)"}
+{"time":"2026-08-13T21:12:38.885554227Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"btjdkox503n3"}
+{"time":"2026-08-13T21:13:03.69502194Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"btjdkox503n3"}
+{"time":"2026-08-13T21:13:04.210578234Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"btjdkox503n3"}
+{"time":"2026-08-13T21:13:04.211977276Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"69zdkrji","id":"6(@)"}
+{"time":"2026-08-13T21:13:04.212517763Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"69zdkrji","id":"6(@)"}
+{"time":"2026-08-13T21:13:05.519021818Z","level":"INFO","msg":"handleInformInit: received","streamId":"5yyvte9z","id":"6(@)"}
+{"time":"2026-08-13T21:13:05.781547687Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"5yyvte9z","id":"6(@)"}
+{"time":"2026-08-13T21:13:11.147370358Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"8ljb7gkyfrgm"}
+{"time":"2026-08-13T21:13:38.248378546Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"8ljb7gkyfrgm"}
+{"time":"2026-08-13T21:13:38.969603918Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"8ljb7gkyfrgm"}
+{"time":"2026-08-13T21:13:38.970854991Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"5yyvte9z","id":"6(@)"}
+{"time":"2026-08-13T21:13:38.971319805Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"5yyvte9z","id":"6(@)"}
+{"time":"2026-08-13T21:13:40.112984858Z","level":"INFO","msg":"handleInformInit: received","streamId":"etbhfswt","id":"6(@)"}
+{"time":"2026-08-13T21:13:40.374066248Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"etbhfswt","id":"6(@)"}
+{"time":"2026-08-13T21:13:45.77023074Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"y1z3ogdatulo"}
+{"time":"2026-08-13T21:14:12.750151439Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"y1z3ogdatulo"}
+{"time":"2026-08-13T21:14:13.738994621Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"y1z3ogdatulo"}
+{"time":"2026-08-13T21:14:13.739800429Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"etbhfswt","id":"6(@)"}
+{"time":"2026-08-13T21:14:13.74128316Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"etbhfswt","id":"6(@)"}
+{"time":"2026-08-13T21:14:33.07143742Z","level":"INFO","msg":"handleInformInit: received","streamId":"vouvkqbm","id":"6(@)"}
+{"time":"2026-08-13T21:14:33.432731087Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"vouvkqbm","id":"6(@)"}
+{"time":"2026-08-13T21:14:38.883290826Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"m1t6sbtmcgac"}
+{"time":"2026-08-13T21:15:04.878189732Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"m1t6sbtmcgac"}
+{"time":"2026-08-13T21:15:05.414751942Z","level":"INFO","msg":"connection: cancelling request","id":"6(@)","requestId":"m1t6sbtmcgac"}
+{"time":"2026-08-13T21:15:07.242042844Z","level":"INFO","msg":"connection: closing","id":"6(@)"}
+{"time":"2026-08-13T21:15:07.24213917Z","level":"INFO","msg":"connection: closed successfully","id":"6(@)"}
+{"time":"2026-08-13T21:15:07.242050444Z","level":"INFO","msg":"processOutgoingData: finished","id":"6(@)"}
+{"time":"2026-08-13T21:15:07.242164062Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"6(@)"}
+{"time":"2026-08-13T21:15:21.224912597Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"7(@)"}
+{"time":"2026-08-13T21:15:21.456878232Z","level":"INFO","msg":"handleInformInit: received","streamId":"bdgno22l","id":"7(@)"}
+{"time":"2026-08-13T21:15:21.716846245Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"bdgno22l","id":"7(@)"}
+{"time":"2026-08-13T21:15:27.078470483Z","level":"INFO","msg":"connection: cancelling request","id":"7(@)","requestId":"7nz3jdwc0fg3"}

zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug-internal.log

@@ -0,0 +1,33 @@
+{"time":"2026-08-13T21:15:21.457059839Z","level":"INFO","msg":"wandb-core"}
+{"time":"2026-08-13T21:15:21.457218835Z","level":"INFO","msg":"stream: starting","core version":"0.28.1"}
+{"time":"2026-08-13T21:15:21.71665521Z","level":"INFO","msg":"stream: created new stream","id":"bdgno22l"}
+{"time":"2026-08-13T21:15:21.716724793Z","level":"INFO","msg":"handler: started"}
+{"time":"2026-08-13T21:15:21.716834966Z","level":"INFO","msg":"stream: started"}
+{"time":"2026-08-13T21:15:21.716845775Z","level":"INFO","msg":"writer: started","stream_id":"bdgno22l"}
+{"time":"2026-08-13T21:15:21.716864355Z","level":"INFO","msg":"sender: started"}
+{"time":"2026-08-13T21:15:22.087946715Z","level":"INFO","msg":"filestream: sending request","total_files":1,"console_offset":0,"console_lines":1}
+{"time":"2026-08-13T21:15:22.187941633Z","level":"INFO","msg":"filestream: request sent","status":"200 OK"}
+{"time":"2026-08-13T21:15:33.902773632Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":2617}
+{"time":"2026-08-13T21:15:33.902804288Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":1}
+{"time":"2026-08-13T21:15:33.903189126Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":2650}
+{"time":"2026-08-13T21:15:33.903223176Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":1}
+{"time":"2026-08-13T21:15:33.905188752Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":3405}
+{"time":"2026-08-13T21:15:33.905644326Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":109}
+{"time":"2026-08-13T21:15:33.905737177Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":3547}
+{"time":"2026-08-13T21:15:33.905849936Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":11}
+{"time":"2026-08-13T21:15:33.907613871Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":4244}
+{"time":"2026-08-13T21:15:33.914808342Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":2917}
+{"time":"2026-08-13T21:15:33.91681943Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":7726}
+{"time":"2026-08-13T21:15:33.917089153Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":83}
+{"time":"2026-08-13T21:15:33.917182957Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":7844}
+{"time":"2026-08-13T21:15:33.91727044Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":13}
+{"time":"2026-08-13T21:15:33.919229048Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":8663}
+{"time":"2026-08-13T21:15:33.919463562Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":61}
+{"time":"2026-08-13T21:15:33.92636457Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":9853}
+{"time":"2026-08-13T21:15:33.926563588Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":8}
+{"time":"2026-08-13T21:15:33.92776116Z","level":"INFO","msg":"flowcontrol: backed up, offloading to disk","recordNumber":10322}
+{"time":"2026-08-13T21:15:33.927881025Z","level":"INFO","msg":"flowcontrol: unblocked","totalOffloaded":18}
+{"time":"2026-08-13T21:15:37.112579744Z","level":"INFO","msg":"filestream: sending request","total_files":4,"history_offset":0,"history_lines":2,"events_offset":0,"events_lines":1,"console_offset":1,"console_lines":5,"uploaded_len":2}
+{"time":"2026-08-13T21:15:37.764424515Z","level":"INFO","msg":"filestream: request sent","status":"200 OK"}
+{"time":"2026-08-13T21:15:52.106707323Z","level":"INFO","msg":"filestream: sending request","total_files":4,"history_offset":2,"history_lines":3,"events_offset":1,"events_lines":2,"console_offset":3,"console_lines":1}
+{"time":"2026-08-13T21:15:53.006564044Z","level":"INFO","msg":"filestream: request sent","status":"200 OK"}

zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug.log

@@ -0,0 +1,23 @@
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_setup.py:_flush():81] Current SDK version is 0.28.1
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_setup.py:_flush():81] Configure stats pid to 52372
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_setup.py:_flush():81] Loading settings from environment variables
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:setup_run_log_directory():729] Logging user logs to /mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug.log
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:setup_run_log_directory():730] Logging internal logs to /mnt/data/zainulabideen/zain-exp/notebooks/zain/Activation/wandb/run-20260813_211521-bdgno22l/logs/debug-internal.log
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():772] calling init triggers
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():777] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():820] starting backend
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():826] Connected to an existing wandb-core service via WANDB_SERVICE
+2026-08-13 21:15:21,455 INFO    MainThread:52372 [wandb_init.py:init():835] sending inform_init request
+2026-08-13 21:15:21,717 INFO    MainThread:52372 [wandb_init.py:init():840] backend started and connected
+2026-08-13 21:15:21,720 INFO    MainThread:52372 [wandb_init.py:init():910] updated telemetry
+2026-08-13 21:15:21,727 INFO    MainThread:52372 [wandb_init.py:init():933] communicating run to backend with 90.0 second timeout
+2026-08-13 21:15:22,001 INFO    MainThread:52372 [wandb_init.py:init():978] starting run threads in backend
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_console_start():2621] atexit reg
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_redirect():2471] redirect: wrap_raw
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_redirect():2540] Wrapping output streams.
+2026-08-13 21:15:22,074 INFO    MainThread:52372 [wandb_run.py:_redirect():2563] Redirects installed.
+2026-08-13 21:15:22,077 INFO    MainThread:52372 [wandb_init.py:init():1016] run started, returning control to user process
+2026-08-13 21:15:22,078 INFO    MainThread:52372 [wandb_run.py:_config_callback():1346] config_cb None None {'transformers_version': '5.16.0.dev0', 'architectures': None, 'output_hidden_states': False, 'return_dict': True, 'dtype': None, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'id2label': {0: 'LABEL_0', 1: 'LABEL_1'}, 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, 'problem_type': None, 'vocab_size': 4096, 'hidden_size': 128, 'intermediate_size': 256, 'num_hidden_layers': 100, 'num_attention_heads': 4, 'num_key_value_heads': 4, 'hidden_act': 'silu', 'max_position_embeddings': 512, 'initializer_range': 0.02, 'rms_norm_eps': 1e-06, 'use_cache': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2, 'pretraining_tp': 1, 'tie_word_embeddings': True, 'rope_parameters': {'rope_theta': 10000.0, 'rope_type': 'default'}, 'attention_bias': False, 'attention_dropout': 0.0, 'mlp_bias': False, 'head_dim': 32, '_name_or_path': '', 'tokenizer_name': 'w-ahmad/tiny-stories-tokenizer', 'mlp_type': 'glu', 'activation': 'linear', 'waleed_beta': 10.0, 'model_type': 'tiny_llama', 'output_attentions': False, 'output_dir': 'out/glu-linear-100L_trash_run', 'per_device_train_batch_size': 64, 'num_train_epochs': 1, 'max_steps': 2700, 'learning_rate': 9e-05, 'lr_scheduler_type': 'constant_with_warmup', 'lr_scheduler_kwargs': None, 'warmup_steps': 50, 'optim': 'adamw_torch_fused', 'optim_args': None, 'weight_decay': 0.0, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_epsilon': 1e-08, 'optim_target_modules': None, 'gradient_accumulation_steps': 1, 'average_tokens_across_devices': True, 'max_grad_norm': 1.0, 'label_smoothing_factor': 0.0, 'bf16': True, 'fp16': False, 'bf16_full_eval': False, 'fp16_full_eval': False, 'tf32': None, 'gradient_checkpointing': False, 'gradient_checkpointing_kwargs': None, 'torch_compile': False, 'torch_compile_backend': None, 'torch_compile_mode': None, 'use_liger_kernel': False, 'liger_kernel_config': None, 'neftune_noise_alpha': None, 'torch_empty_cache_steps': None, 'auto_find_batch_size': False, 'logging_strategy': 'steps', 'logging_steps': 20, 'logging_first_step': False, 'log_on_each_node': True, 'logging_nan_inf_filter': True, 'include_num_input_tokens_seen': 'no', 'log_level': 'passive', 'log_level_replica': 'warning', 'disable_tqdm': False, 'report_to': ['wandb'], 'run_name': 'LM-glu-linear-100L-16.9M-20260813-211520', 'project': 'huggingface', 'trackio_space_id': None, 'trackio_bucket_id': None, 'trackio_static_space_id': None, 'eval_strategy': 'steps', 'eval_steps': 2498, 'eval_delay': 0, 'per_device_eval_batch_size': 1024, 'prediction_loss_only': False, 'eval_on_start': False, 'eval_do_concat_batches': True, 'eval_use_gather_object': False, 'eval_accumulation_steps': None, 'include_for_metrics': [], 'batch_eval_metrics': False, 'save_only_model': False, 'save_strategy': 'steps', 'save_steps': 100, 'save_on_each_node': False, 'save_total_limit': None, 'enable_jit_checkpoint': False, 'push_to_hub': False, 'hub_token': '<HUB_TOKEN>', 'hub_private_repo': None, 'hub_model_id': 'w-ahmad/6L-glu-linear-100L', 'hub_strategy': 'every_save', 'hub_always_push': False, 'hub_revision': None, 'load_best_model_at_end': False, 'metric_for_best_model': None, 'greater_is_better': None, 'ignore_data_skip': False, 'restore_callback_states_from_checkpoint': False, 'full_determinism': False, 'seed': 42, 'data_seed': 42, 'use_cpu': False, 'accelerator_config': {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}, 'parallelism_config': None, 'dataloader_drop_last': False, 'dataloader_num_workers': 0, 'dataloader_pin_memory': True, 'dataloader_persistent_workers': False, 'dataloader_prefetch_factor': None, 'dataloader_multiprocessing_context': None, 'dataloader_in_order': True, 'remove_unused_columns': False, 'label_names': None, 'train_sampling_strategy': 'random', 'length_column_name': 'length', 'ddp_find_unused_parameters': None, 'ddp_bucket_cap_mb': None, 'ddp_broadcast_buffers': None, 'ddp_static_graph': None, 'ddp_backend': None, 'ddp_timeout': 1800, 'fsdp': None, 'fsdp_config': None, 'deepspeed': None, 'debug': [], 'skip_memory_metrics': True, 'do_train': False, 'do_eval': True, 'do_predict': False, 'resume_from_checkpoint': None, 'local_rank': -1}
+2026-08-13 21:15:22,082 INFO    MainThread:52372 [wandb_config.py:__setitem__():155] [no run ID] config set model/num_parameters = 16934016 - <bound method Run._config_callback of <wandb.sdk.wandb_run.Run object at 0x15235c33c510>>
+2026-08-13 21:15:22,082 INFO    MainThread:52372 [wandb_run.py:_config_callback():1346] config_cb model/num_parameters 16934016 None

zain/Activation/wandb/run-20260813_211521-bdgno22l/run-bdgno22l.wandb

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