Auto upload zain 2026-08-14T19:47:24.937839

.gitattributes

@@ -41,3 +41,4 @@ zain/Activation/wandb/run-20260812_225525-hf77resg/run-hf77resg.wandb filter=lfs
 zain/Activation/wandb/run-20260812_230915-vvjf0upl/run-vvjf0upl.wandb filter=lfs diff=lfs merge=lfs -text
 zain/Activation/wandb/run-20260812_232326-8u0k1nti/run-8u0k1nti.wandb filter=lfs diff=lfs merge=lfs -text
 zain/Activation/wandb/run-20260812_233704-tpo5j00e/run-tpo5j00e.wandb filter=lfs diff=lfs merge=lfs -text
+zain/Activation/wandb/run-20260814_193907-kk1xbqht/run-kk1xbqht.wandb filter=lfs diff=lfs merge=lfs -text

zain/Activation/exp.py

@@ -1,8 +1,6 @@
-"""
-Tiny Llama GLU Lab — Consolidated training library.
-One file: model definition, activation registry, stability monitoring,
-time tracking, dataset builder, and Trainer factory.
-"""
+# =====================================================================
+# exp.py – FULL FILE, ALL FIXES INCLUDED (hook requires_grad check)
+# =====================================================================
 
 import math
 import os
@@ -32,6 +30,7 @@ from transformers.models.llama.modeling_llama import (
 )
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from datasets import load_dataset
+from huggingface_hub import snapshot_download
 
 
 # =============================================================================
@@ -39,7 +38,6 @@ from datasets import load_dataset
 # =============================================================================
 
 class GLUActivationRegistry:
-    """Own every gating activation you test. Add new variants in one line."""
     _registry: Dict[str, Callable[[torch.Tensor], torch.Tensor]] = {}
 
     @classmethod
@@ -54,7 +52,6 @@ class GLUActivationRegistry:
             )
         return cls._registry[name]
 
-
 # Built-ins
 GLUActivationRegistry.register("silu", nn.functional.silu)
 GLUActivationRegistry.register("swish", nn.functional.silu)
@@ -64,8 +61,6 @@ GLUActivationRegistry.register("sigmoid", torch.sigmoid)
 GLUActivationRegistry.register("tanh", torch.tanh)
 GLUActivationRegistry.register("softplus", nn.functional.softplus)
 GLUActivationRegistry.register("linear", lambda x: x)
-
-# Custom
 GLUActivationRegistry.register("s10", lambda x: x * x * torch.sigmoid(x))
 GLUActivationRegistry.register("w1a", lambda x: x * torch.tanh(x))
 
@@ -75,13 +70,6 @@ GLUActivationRegistry.register("w1a", lambda x: x * torch.tanh(x))
 # =============================================================================
 
 class TinyLlamaConfig(LlamaConfig):
-    """
-    Exact Llama config plus two fields:
-      - mlp_type: "glu" or "mlp" (standard Transformer MLP)
-      - activation: name of the activation function to use inside the MLP block.
-      - waleed_beta: β_cap for the waleed10 post‑clip (default 10.0)
-    Enforces pure MHA by requiring num_key_value_heads == num_attention_heads.
-    """
     model_type = "tiny_llama"
 
     def __init__(
@@ -103,20 +91,13 @@ class TinyLlamaConfig(LlamaConfig):
 
 
 # =============================================================================
-# 3. MODEL
+# 3. MODEL (with override hook)
 # =============================================================================
 
 class TinyLlamaMLP(nn.Module):
-    """
-    Unified MLP block supporting:
-      - Standard GLU:   down_proj( act(gate_proj(x)) * up_proj(x) )
-      - Standard MLP:   down_proj( act(up_proj(x)) )
-      - Gated variants (situglu, waleed) with internal tanh scaling.
-      - waleed10 / silu-waleed10 with post‑down‑proj clipping.
-
-    All projections are always named the same way, ensuring consistent
-    tensor logging regardless of activation.
-    """
+    override_active = False
+    override_value = -100.0
+
     def __init__(self, config: TinyLlamaConfig):
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -125,27 +106,25 @@ class TinyLlamaMLP(nn.Module):
         self.activation_name = config.activation
         self.waleed_beta = getattr(config, "waleed_beta", 10.0)
 
-        # Effective intermediate size (scaled for MLP)
         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} for parameter parity.")
+            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
 
-        # Always define the projections (names consistent across variants)
         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:  # mlp
+        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 function (for standard variants)
+        # Activation handling
         if self.mlp_type == "glu":
             if self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
                 self.act_fn = None
@@ -155,12 +134,9 @@ class TinyLlamaMLP(nn.Module):
                 self.act_fn = GLUActivationRegistry.get("silu")
             else:
                 self.act_fn = GLUActivationRegistry.get(self.activation_name)
-        else:  # mlp
+        else:
             if self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
-                raise ValueError(
-                    f"Activation '{self.activation_name}' requires a gated architecture (GLU). "
-                    f"Please use mlp_type='glu'."
-                )
+                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":
@@ -168,7 +144,6 @@ class TinyLlamaMLP(nn.Module):
             else:
                 self.act_fn = GLUActivationRegistry.get(self.activation_name)
 
-        # Determine beta values for situglu/waleed variants
         if self.activation_name in ("situglu_low", "waleedglu_low"):
             self.beta1 = 2.5
             self.beta2 = 4.0
@@ -176,7 +151,6 @@ class TinyLlamaMLP(nn.Module):
             self.beta1 = 4.0
             self.beta2 = 25.0
 
-        # Flags for special handling
         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")
@@ -206,9 +180,17 @@ class TinyLlamaMLP(nn.Module):
         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__()
@@ -216,23 +198,12 @@ class TinyLlamaDecoderLayer(nn.Module):
         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
-        )
-        # Residual stream gateways (zero params, hookable)
+        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: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        **kwargs,
-    ):
-        # Attention sub-layer
+    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)
@@ -245,7 +216,6 @@ class TinyLlamaDecoderLayer(nn.Module):
         hidden_states = residual + attn_out
         hidden_states = self.residual_post_attn(hidden_states)
 
-        # MLP sub-layer
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
@@ -254,24 +224,14 @@ class TinyLlamaDecoderLayer(nn.Module):
         return (hidden_states,)
 
 
-# ----------------------------------------------------------------------------
-# ATTENTION MASK
-# ----------------------------------------------------------------------------
-def _build_causal_mask(
-    attention_mask: Optional[torch.Tensor],
-    seq_len: int,
-    dtype: torch.dtype,
-    device: torch.device,
-) -> torch.Tensor:
+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
@@ -289,48 +249,27 @@ class TinyLlamaModel(LlamaPreTrainedModel):
         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.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([TinyLlamaDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
         self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = LlamaRotaryEmbedding(config=config)
         self.post_init()
 
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs,
-    ):
+    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)
-
+            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
-        )
-
+        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,
@@ -339,7 +278,6 @@ class TinyLlamaModel(LlamaPreTrainedModel):
                 position_embeddings=position_embeddings,
             )
             hidden_states = layer_outputs[0]
-
         hidden_states = self.norm(hidden_states)
         if not return_dict:
             return (hidden_states,)
@@ -367,50 +305,25 @@ class TinyLlamaForCausalLM(LlamaPreTrainedModel):
     def get_output_embeddings(self):
         return self.lm_head
 
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs,
-    ):
+    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,
-        )
+        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)
-            )
-
+            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,
-        )
+        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
-    ):
+    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")
@@ -428,57 +341,141 @@ class TinyLlamaForCausalLM(LlamaPreTrainedModel):
 
 
 # =============================================================================
-# 4. MONITORING ENGINE
+# 4. HF CHECKPOINT FETCHER
 # =============================================================================
 
-class StatsEngine:
-    """Compute unified signature for any tensor."""
+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]:
+    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")
-            )
-
-            t_min = tensor.min().item()
-            t_max = tensor.max().item()
-
+            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_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": t_min,
-                "max": t_max,
-                "range": t_max - t_min,
+                "min": tensor.min().item(),
+                "max": tensor.max().item(),
+                "range": tensor.max().item() - tensor.min().item(),
             }
 
 
 class StepAccumulator:
-    """Stores per-tensor entries, aggregates to layer/global scope."""
-
     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)
-        )
+        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]:
@@ -488,21 +485,12 @@ class StepAccumulator:
         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
+        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
-
+        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,
@@ -524,27 +512,15 @@ class StepAccumulator:
             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
-        )
+        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
-        )
-
+        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,
@@ -561,27 +537,17 @@ class StepAccumulator:
         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 + ".")
-        }
+        entries = {k: v for k, v in self.tensors.items() if k.startswith(layer_prefix + ".")}
         return self._aggregate(entries)
 
 
 class HookRegistry:
-    """Attach and throttle forward/backward hooks."""
-
     def __init__(self, model: nn.Module):
         self.model = model
-        self.handles: List[torch.utils.hooks.RemovableHandle] = []
+        self.handles = []
         self.active = False
 
-    def attach_forward(
-        self,
-        module_patterns: List[str],
-        accumulator: StepAccumulator,
-        user_limit: float,
-        dtype_ratio: float,
-    ):
+    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(
@@ -589,51 +555,35 @@ class HookRegistry:
                 )
                 self.handles.append(h)
 
-    def attach_backward(
-        self, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
-    ):
+    def attach_backward(self, param_patterns, accumulator, user_limit, dtype_ratio):
         for name, param in self.model.named_parameters():
-            if param.requires_grad:
-                h = param.register_hook(
-                    self._make_backward_hook(
-                        f"grad.{name}", accumulator, user_limit, dtype_ratio
-                    )
-                )
-                self.handles.append(h)
+            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: str, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
-    ):
+    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_dict = out
-                out = out_dict.get("last_hidden_state")
-                if out is None:
-                    out = next(
-                        (v for v in out_dict.values() if torch.is_tensor(v)), None
-                    )
-            elif isinstance(out, (tuple, list)):
-                out = out[0] if len(out) > 0 else None
-
+                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: str, accumulator: StepAccumulator, user_limit: float, dtype_ratio: float
-    ):
+    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):
@@ -646,13 +596,12 @@ class HookRegistry:
 
 
 class StabilityMonitorCallback(TrainerCallback):
-    """Full stability instrumentation: grad / param / act statistics."""
-
     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,
@@ -660,14 +609,11 @@ class StabilityMonitorCallback(TrainerCallback):
     ):
         self.model = model
         self.monitor_every_n_steps = monitor_every_n_steps
-        self.module_patterns = module_patterns or [".*mlp.*", ".*self_attn.*", ".*residual.*"]
+        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.log_scope = log_scope or {"global": True, "per_layer": True, "per_tensor": False}
         self.monitor_during_eval = monitor_during_eval
 
         self.accumulator = StepAccumulator()
@@ -679,12 +625,14 @@ class StabilityMonitorCallback(TrainerCallback):
             self.dtype_ratio,
         )
         self.hooks.attach_backward(
-            self.accumulator, self.user_limits["grad"], self.dtype_ratio
+            self.param_patterns,
+            self.accumulator,
+            self.user_limits["grad"],
+            self.dtype_ratio,
         )
+        self.pending_metrics = None
 
-        self.pending_metrics: Optional[Dict[str, float]] = None
-
-    def _should_monitor(self, state) -> bool:
+    def _should_monitor(self, state):
         return state.global_step % self.monitor_every_n_steps == 0
 
     def on_step_begin(self, args, state, control, **kwargs):
@@ -696,16 +644,14 @@ class StabilityMonitorCallback(TrainerCallback):
         if not self.hooks.active:
             return
         for name, param in self.model.named_parameters():
-            stats = StatsEngine.compute(
-                param.data, self.user_limits["param"], self.dtype_ratio
-            )
+            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()
 
-    @staticmethod
-    def _kind_of(name: str) -> str:
+    def _kind_of(self, name: str) -> str:
         if name.startswith("act."):
             return "act"
         if name.startswith("grad."):
@@ -714,8 +660,7 @@ class StabilityMonitorCallback(TrainerCallback):
             return "param"
         return "other"
 
-    @staticmethod
-    def _strip_kind(name: str) -> str:
+    def _strip_kind(self, name: str) -> str:
         if name.startswith("act."):
             return name[4:]
         if name.startswith(("grad.", "param.")):
@@ -723,10 +668,9 @@ class StabilityMonitorCallback(TrainerCallback):
         return name
 
     def _build_metrics(self, scope: str = "train") -> Dict[str, float]:
-        metrics: Dict[str, float] = {}
-
+        metrics = {}
         if self.log_scope.get("global", True):
-            by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
+            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():
@@ -744,7 +688,7 @@ class StabilityMonitorCallback(TrainerCallback):
                         prefix = ".".join(parts[: i + 2])
                         layer_prefixes.add(prefix)
             for prefix in layer_prefixes:
-                by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
+                by_kind = {}
                 for k, v in self.accumulator.tensors.items():
                     clean = self._strip_kind(k)
                     if clean.startswith(prefix + ".") or clean == prefix:
@@ -764,12 +708,17 @@ class StabilityMonitorCallback(TrainerCallback):
                     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):
@@ -779,9 +728,9 @@ class StabilityMonitorCallback(TrainerCallback):
             self.accumulator.clear()
             self.hooks.set_active(True)
             for name, param in self.model.named_parameters():
-                stats = StatsEngine.compute(
-                    param.data, self.user_limits["param"], self.dtype_ratio
-                )
+                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")
 
@@ -790,14 +739,16 @@ class StabilityMonitorCallback(TrainerCallback):
         self.accumulator.clear()
 
 
-class TimeTrackerCallback(TrainerCallback):
-    """Precise training & eval timing with remaining-time estimates."""
+# =============================================================================
+# 7. TIME TRACKER
+# =============================================================================
 
+class TimeTrackerCallback(TrainerCallback):
     def __init__(self):
-        self.step_start: Optional[float] = None
-        self.epoch_start: Optional[float] = None
+        self.step_start = None
+        self.epoch_start = None
         self.total_train_time = 0.0
-        self.step_times: List[float] = []
+        self.step_times = []
 
     def on_epoch_begin(self, args, state, control, **kwargs):
         self.epoch_start = time.perf_counter()
@@ -826,13 +777,12 @@ class TimeTrackerCallback(TrainerCallback):
             remaining = (state.max_steps - state.global_step) * avg
             logs["train/estimated_remaining_minutes"] = remaining / 60.0
 
-    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
-        pass
 
+# =============================================================================
+# 8. METRICS LOGGER
+# =============================================================================
 
 class MetricsLoggerCallback(TrainerCallback):
-    """Persist every logged dict as JSONL in the output dir."""
-
     def __init__(self, output_dir: str):
         self.output_dir = Path(output_dir)
         self.output_dir.mkdir(parents=True, exist_ok=True)
@@ -841,113 +791,13 @@ class MetricsLoggerCallback(TrainerCallback):
     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,
-        }
+        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")
 
 
 # =============================================================================
-# CONTAMINATION CALLBACK
-# =============================================================================
-
-class ContaminationCallback(TrainerCallback):
-    """
-    Intentionally corrupt input_ids and labels for a window of steps.
-    Modes: "shift" (+1 mod vocab), "random" (uniform random IDs).
-    """
-    def __init__(
-        self,
-        vocab_size: int,
-        enabled: bool = False,
-        start_step: int = 0,
-        duration_steps: int = 0,
-        mode: str = "shift",
-        fraction: float = 1.0,
-        seed: Optional[int] = None,
-    ):
-        self.vocab_size = vocab_size
-        self.enabled = enabled
-        self.start_step = start_step
-        self.duration_steps = duration_steps
-        self.mode = mode
-        self.fraction = fraction
-        self.seed = seed
-        self.generator = torch.Generator()
-        if seed is not None:
-            self.generator.manual_seed(seed)
-        self._active = False
-
-    def _should_corrupt(self, state) -> bool:
-        if not self.enabled:
-            return False
-        step = state.global_step
-        return self.start_step <= step < self.start_step + self.duration_steps
-
-    def on_step_begin(self, args, state, control, **kwargs):
-        if not self._should_corrupt(state):
-            return
-        batch = kwargs.get("inputs")
-        if batch is None or not isinstance(batch, dict):
-            return
-        input_ids = batch.get("input_ids")
-        labels = batch.get("labels")
-        attention_mask = batch.get("attention_mask")
-        if input_ids is None or labels is None:
-            return
-
-        device = input_ids.device
-        batch_size, seq_len = input_ids.shape
-
-        if self.fraction >= 1.0:
-            corrupt_mask = torch.ones((batch_size, seq_len), dtype=torch.bool, device=device)
-        elif self.fraction <= 0.0:
-            return
-        else:
-            rand = torch.rand((batch_size, seq_len), generator=self.generator, device=device)
-            corrupt_mask = rand < self.fraction
-
-        if attention_mask is not None:
-            padding_mask = attention_mask == 0
-            corrupt_mask = corrupt_mask & (~padding_mask)
-
-        if not corrupt_mask.any():
-            return
-
-        if self.mode == "shift":
-            input_ids_corrupt = (input_ids + 1) % self.vocab_size
-            labels_corrupt = (labels + 1) % self.vocab_size
-        elif self.mode == "random":
-            random_ids = torch.randint(
-                0, self.vocab_size, input_ids.shape,
-                generator=self.generator, device=device
-            )
-            input_ids_corrupt = random_ids
-            labels_corrupt = random_ids
-        else:
-            raise ValueError(f"Unknown contamination mode: {self.mode}")
-
-        input_ids.masked_scatter_(corrupt_mask, input_ids_corrupt[corrupt_mask])
-        labels.masked_scatter_(corrupt_mask, labels_corrupt[corrupt_mask])
-        batch["input_ids"] = input_ids
-        batch["labels"] = labels
-
-        if not self._active:
-            print(f"[Contamination] Started at step {state.global_step} for {self.duration_steps} steps (mode={self.mode})")
-            self._active = True
-
-    def on_step_end(self, args, state, control, **kwargs):
-        if self._active and state.global_step >= self.start_step + self.duration_steps:
-            print(f"[Contamination] Ended at step {state.global_step}")
-            self._active = False
-
-
-# =============================================================================
-# 5. DATA & TRAINER FACTORY
+# 9. DATA & TRAINER FACTORY
 # =============================================================================
 
 def build_dataset(
@@ -957,11 +807,9 @@ def build_dataset(
     dataset_name: str = "roneneldan/TinyStories",
     max_samples: Optional[int] = None,
 ):
-    """Concatenate and chunk TinyStories for causal LM."""
     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))))
-        print(f"[Dataset] Using first {len(ds)} samples for training (max_samples={max_samples})")
 
     def tokenize(examples):
         out = tokenizer(examples["text"], add_special_tokens=False)
@@ -971,34 +819,20 @@ def build_dataset(
             out["attention_mask"] = [mask + [1] for mask in out["attention_mask"]]
         return out
 
-    tokenized = ds.map(
-        tokenize,
-        batched=True,
-        num_proc=4,
-        remove_columns=ds.column_names,
-        desc=f"Tokenizing {split}",
-    )
+    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()
-        }
+        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)]
+            k: [t[i:i + max_seq_len] for i in range(0, total_length, max_seq_len)]
             for k, t in concatenated.items()
         }
         result["labels"] = result["input_ids"].copy()
         return result
 
-    return tokenized.map(
-        group_texts,
-        batched=True,
-        batch_size=10000,
-        num_proc=4,
-        desc=f"Chunking {split}",
-    )
+    return tokenized.map(group_texts, batched=True, batch_size=10000, num_proc=4)
 
 
 def create_trainer(
@@ -1008,15 +842,17 @@ def create_trainer(
     train_dataset,
     eval_dataset=None,
 ):
-    """Assemble HF Trainer with all custom callbacks."""
     tc = config.get("training", {})
     mc = config.get("monitor", {})
+    go = config.get("gradient_override", {})
 
-    run_name = tc.get("run_name", None)
+    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=run_name,
+        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),
@@ -1046,27 +882,22 @@ def create_trainer(
 
     callbacks = [TimeTrackerCallback()]
 
-    cc = config.get("contamination", {})
-    if cc.get("enabled", False):
-        vocab_size = model.config.vocab_size
-        callbacks.append(
-            ContaminationCallback(
-                vocab_size=vocab_size,
-                enabled=True,
-                start_step=cc.get("start_step", 0),
-                duration_steps=cc.get("duration_steps", 0),
-                mode=cc.get("mode", "shift"),
-                fraction=cc.get("fraction", 1.0),
-                seed=cc.get("seed", None),
-            )
-        )
+    # ---- 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=mc.get("module_patterns", [".*mlp.*", ".*self_attn.*", ".*residual.*"]),
+                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}),
@@ -1087,16 +918,4 @@ def create_trainer(
         callbacks=callbacks,
     )
 
-    try:
-        from transformers.integrations import get_reporting_integration_callbacks
-        reporting_types = tuple(get_reporting_integration_callbacks(args.report_to))
-    except Exception:
-        reporting_types = ()
-
-    if reporting_types:
-        handler = trainer.callback_handler
-        reporting_cbs = [cb for cb in handler.callbacks if isinstance(cb, reporting_types)]
-        other_cbs = [cb for cb in handler.callbacks if not isinstance(cb, reporting_types)]
-        handler.callbacks = other_cbs + reporting_cbs
-
     return trainer

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

@@ -0,0 +1,36 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 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": 4
+}

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,69 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.006740815638692282,
+  "eval_steps": 2498,
+  "global_step": 100,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 2.09375,
+      "learning_rate": 2.66e-05,
+      "loss": 8.298905181884766,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3359375,
+      "learning_rate": 5.46e-05,
+      "loss": 8.045733642578124,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.2734375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.839435577392578,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.265625,
+      "learning_rate": 0.0001106,
+      "loss": 7.574758148193359,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 1.1953125,
+      "learning_rate": 0.0001386,
+      "loss": 7.235897064208984,
+      "step": 100
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 322628380262400.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

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

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

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

@@ -0,0 +1,36 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 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": 4
+}

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

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

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

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

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

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

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

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

zain/Activation/out/glu-relu-100L_run/checkpoint-1000/tokenizer_config.json

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

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

@@ -0,0 +1,384 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.06740815638692282,
+  "eval_steps": 2498,
+  "global_step": 1000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 2.09375,
+      "learning_rate": 2.66e-05,
+      "loss": 8.298905181884766,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3359375,
+      "learning_rate": 5.46e-05,
+      "loss": 8.045733642578124,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.2734375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.839435577392578,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.265625,
+      "learning_rate": 0.0001106,
+      "loss": 7.574758148193359,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 1.1953125,
+      "learning_rate": 0.0001386,
+      "loss": 7.235897064208984,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.140625,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.862369537353516,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 1.0234375,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.508391571044922,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 1.078125,
+      "learning_rate": 0.0002226,
+      "loss": 6.215537261962891,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.76171875,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.984259414672851,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.82421875,
+      "learning_rate": 0.0002786,
+      "loss": 5.756989669799805,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.531437683105469,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 1.1171875,
+      "learning_rate": 0.0003346,
+      "loss": 5.294247817993164,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 5.055661392211914,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.94921875,
+      "learning_rate": 0.0003906,
+      "loss": 4.797793197631836,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.421875,
+      "learning_rate": 0.0004186,
+      "loss": 4.594622802734375,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.2265625,
+      "learning_rate": 0.0004466,
+      "loss": 4.39941635131836,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5546875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.24058723449707,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.83984375,
+      "learning_rate": 0.0005026,
+      "loss": 4.102418899536133,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.70703125,
+      "learning_rate": 0.0005306,
+      "loss": 3.975328063964844,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.6171875,
+      "learning_rate": 0.0005586,
+      "loss": 3.839755630493164,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0005866,
+      "loss": 3.7486804962158202,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0006146,
+      "loss": 3.6831081390380858,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.875,
+      "learning_rate": 0.0006426,
+      "loss": 3.5874557495117188,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.5174354553222655,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.66796875,
+      "learning_rate": 0.0006986,
+      "loss": 3.439300537109375,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.54296875,
+      "learning_rate": 0.0007,
+      "loss": 3.386067581176758,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.6015625,
+      "learning_rate": 0.0007,
+      "loss": 3.3094718933105467,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.55859375,
+      "learning_rate": 0.0007,
+      "loss": 3.2609573364257813,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0007,
+      "loss": 3.2046344757080076,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.5234375,
+      "learning_rate": 0.0007,
+      "loss": 3.17300968170166,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.578125,
+      "learning_rate": 0.0007,
+      "loss": 3.1147342681884767,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 3.0973543167114257,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.453125,
+      "learning_rate": 0.0007,
+      "loss": 3.043070602416992,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 0.455078125,
+      "learning_rate": 0.0007,
+      "loss": 3.0213077545166014,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 0.419921875,
+      "learning_rate": 0.0007,
+      "loss": 2.977329063415527,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.490234375,
+      "learning_rate": 0.0007,
+      "loss": 2.942974090576172,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.48046875,
+      "learning_rate": 0.0007,
+      "loss": 2.9064186096191404,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.890872001647949,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8978702545166017,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8463741302490235,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.47265625,
+      "learning_rate": 0.0007,
+      "loss": 2.8323719024658205,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8007469177246094,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.7573240280151365,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.7341793060302733,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0007,
+      "loss": 2.7180759429931642,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.5625,
+      "learning_rate": 0.0007,
+      "loss": 2.692394256591797,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.678949546813965,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.651218223571777,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.632099723815918,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.6305675506591797,
+      "step": 1000
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 3226283802624000.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

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

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

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

@@ -0,0 +1,36 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 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": 4
+}

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

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

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

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

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

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

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

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

zain/Activation/out/glu-relu-100L_run/checkpoint-1100/tokenizer_config.json

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

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

@@ -0,0 +1,419 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.0741489720256151,
+  "eval_steps": 2498,
+  "global_step": 1100,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 2.09375,
+      "learning_rate": 2.66e-05,
+      "loss": 8.298905181884766,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3359375,
+      "learning_rate": 5.46e-05,
+      "loss": 8.045733642578124,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.2734375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.839435577392578,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.265625,
+      "learning_rate": 0.0001106,
+      "loss": 7.574758148193359,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 1.1953125,
+      "learning_rate": 0.0001386,
+      "loss": 7.235897064208984,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.140625,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.862369537353516,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 1.0234375,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.508391571044922,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 1.078125,
+      "learning_rate": 0.0002226,
+      "loss": 6.215537261962891,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.76171875,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.984259414672851,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.82421875,
+      "learning_rate": 0.0002786,
+      "loss": 5.756989669799805,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.531437683105469,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 1.1171875,
+      "learning_rate": 0.0003346,
+      "loss": 5.294247817993164,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 5.055661392211914,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.94921875,
+      "learning_rate": 0.0003906,
+      "loss": 4.797793197631836,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.421875,
+      "learning_rate": 0.0004186,
+      "loss": 4.594622802734375,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.2265625,
+      "learning_rate": 0.0004466,
+      "loss": 4.39941635131836,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5546875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.24058723449707,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.83984375,
+      "learning_rate": 0.0005026,
+      "loss": 4.102418899536133,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.70703125,
+      "learning_rate": 0.0005306,
+      "loss": 3.975328063964844,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.6171875,
+      "learning_rate": 0.0005586,
+      "loss": 3.839755630493164,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0005866,
+      "loss": 3.7486804962158202,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0006146,
+      "loss": 3.6831081390380858,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.875,
+      "learning_rate": 0.0006426,
+      "loss": 3.5874557495117188,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.5174354553222655,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.66796875,
+      "learning_rate": 0.0006986,
+      "loss": 3.439300537109375,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.54296875,
+      "learning_rate": 0.0007,
+      "loss": 3.386067581176758,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.6015625,
+      "learning_rate": 0.0007,
+      "loss": 3.3094718933105467,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.55859375,
+      "learning_rate": 0.0007,
+      "loss": 3.2609573364257813,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0007,
+      "loss": 3.2046344757080076,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.5234375,
+      "learning_rate": 0.0007,
+      "loss": 3.17300968170166,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.578125,
+      "learning_rate": 0.0007,
+      "loss": 3.1147342681884767,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 3.0973543167114257,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.453125,
+      "learning_rate": 0.0007,
+      "loss": 3.043070602416992,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 0.455078125,
+      "learning_rate": 0.0007,
+      "loss": 3.0213077545166014,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 0.419921875,
+      "learning_rate": 0.0007,
+      "loss": 2.977329063415527,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.490234375,
+      "learning_rate": 0.0007,
+      "loss": 2.942974090576172,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.48046875,
+      "learning_rate": 0.0007,
+      "loss": 2.9064186096191404,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.890872001647949,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8978702545166017,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8463741302490235,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.47265625,
+      "learning_rate": 0.0007,
+      "loss": 2.8323719024658205,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8007469177246094,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.7573240280151365,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.7341793060302733,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0007,
+      "loss": 2.7180759429931642,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.5625,
+      "learning_rate": 0.0007,
+      "loss": 2.692394256591797,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.678949546813965,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.651218223571777,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.632099723815918,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.6305675506591797,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06875631951466127,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.5994905471801757,
+      "step": 1020
+    },
+    {
+      "epoch": 0.07010448264239973,
+      "grad_norm": 0.474609375,
+      "learning_rate": 0.0007,
+      "loss": 2.5834341049194336,
+      "step": 1040
+    },
+    {
+      "epoch": 0.07145264577013818,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.567459487915039,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07280080889787664,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.54827823638916,
+      "step": 1080
+    },
+    {
+      "epoch": 0.0741489720256151,
+      "grad_norm": 0.478515625,
+      "learning_rate": 0.0007,
+      "loss": 2.5431196212768556,
+      "step": 1100
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 3548912182886400.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

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

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

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

@@ -0,0 +1,36 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 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": 4
+}

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

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

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

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

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

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

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

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

zain/Activation/out/glu-relu-100L_run/checkpoint-1200/tokenizer_config.json

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

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

@@ -0,0 +1,454 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.08088978766430738,
+  "eval_steps": 2498,
+  "global_step": 1200,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 2.09375,
+      "learning_rate": 2.66e-05,
+      "loss": 8.298905181884766,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3359375,
+      "learning_rate": 5.46e-05,
+      "loss": 8.045733642578124,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.2734375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.839435577392578,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.265625,
+      "learning_rate": 0.0001106,
+      "loss": 7.574758148193359,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 1.1953125,
+      "learning_rate": 0.0001386,
+      "loss": 7.235897064208984,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.140625,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.862369537353516,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 1.0234375,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.508391571044922,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 1.078125,
+      "learning_rate": 0.0002226,
+      "loss": 6.215537261962891,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.76171875,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.984259414672851,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.82421875,
+      "learning_rate": 0.0002786,
+      "loss": 5.756989669799805,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.531437683105469,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 1.1171875,
+      "learning_rate": 0.0003346,
+      "loss": 5.294247817993164,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 5.055661392211914,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.94921875,
+      "learning_rate": 0.0003906,
+      "loss": 4.797793197631836,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.421875,
+      "learning_rate": 0.0004186,
+      "loss": 4.594622802734375,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.2265625,
+      "learning_rate": 0.0004466,
+      "loss": 4.39941635131836,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5546875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.24058723449707,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.83984375,
+      "learning_rate": 0.0005026,
+      "loss": 4.102418899536133,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.70703125,
+      "learning_rate": 0.0005306,
+      "loss": 3.975328063964844,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.6171875,
+      "learning_rate": 0.0005586,
+      "loss": 3.839755630493164,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0005866,
+      "loss": 3.7486804962158202,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0006146,
+      "loss": 3.6831081390380858,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.875,
+      "learning_rate": 0.0006426,
+      "loss": 3.5874557495117188,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.5174354553222655,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.66796875,
+      "learning_rate": 0.0006986,
+      "loss": 3.439300537109375,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.54296875,
+      "learning_rate": 0.0007,
+      "loss": 3.386067581176758,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.6015625,
+      "learning_rate": 0.0007,
+      "loss": 3.3094718933105467,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.55859375,
+      "learning_rate": 0.0007,
+      "loss": 3.2609573364257813,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0007,
+      "loss": 3.2046344757080076,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.5234375,
+      "learning_rate": 0.0007,
+      "loss": 3.17300968170166,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.578125,
+      "learning_rate": 0.0007,
+      "loss": 3.1147342681884767,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 3.0973543167114257,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.453125,
+      "learning_rate": 0.0007,
+      "loss": 3.043070602416992,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 0.455078125,
+      "learning_rate": 0.0007,
+      "loss": 3.0213077545166014,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 0.419921875,
+      "learning_rate": 0.0007,
+      "loss": 2.977329063415527,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.490234375,
+      "learning_rate": 0.0007,
+      "loss": 2.942974090576172,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.48046875,
+      "learning_rate": 0.0007,
+      "loss": 2.9064186096191404,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.890872001647949,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8978702545166017,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8463741302490235,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.47265625,
+      "learning_rate": 0.0007,
+      "loss": 2.8323719024658205,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8007469177246094,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.7573240280151365,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.7341793060302733,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0007,
+      "loss": 2.7180759429931642,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.5625,
+      "learning_rate": 0.0007,
+      "loss": 2.692394256591797,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.678949546813965,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.651218223571777,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.632099723815918,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.6305675506591797,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06875631951466127,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.5994905471801757,
+      "step": 1020
+    },
+    {
+      "epoch": 0.07010448264239973,
+      "grad_norm": 0.474609375,
+      "learning_rate": 0.0007,
+      "loss": 2.5834341049194336,
+      "step": 1040
+    },
+    {
+      "epoch": 0.07145264577013818,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.567459487915039,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07280080889787664,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.54827823638916,
+      "step": 1080
+    },
+    {
+      "epoch": 0.0741489720256151,
+      "grad_norm": 0.478515625,
+      "learning_rate": 0.0007,
+      "loss": 2.5431196212768556,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07549713515335356,
+      "grad_norm": 0.4453125,
+      "learning_rate": 0.0007,
+      "loss": 2.5049072265625,
+      "step": 1120
+    },
+    {
+      "epoch": 0.07684529828109202,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 2.4941583633422852,
+      "step": 1140
+    },
+    {
+      "epoch": 0.07819346140883048,
+      "grad_norm": 0.46875,
+      "learning_rate": 0.0007,
+      "loss": 2.4895328521728515,
+      "step": 1160
+    },
+    {
+      "epoch": 0.07954162453656892,
+      "grad_norm": 0.48828125,
+      "learning_rate": 0.0007,
+      "loss": 2.4841281890869142,
+      "step": 1180
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "grad_norm": 0.46484375,
+      "learning_rate": 0.0007,
+      "loss": 2.475014495849609,
+      "step": 1200
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 3871540563148800.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

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

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

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

@@ -0,0 +1,36 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 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": 4
+}

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

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

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

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

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

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

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

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

zain/Activation/out/glu-relu-100L_run/checkpoint-1300/tokenizer_config.json

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

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

@@ -0,0 +1,489 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.08763060330299967,
+  "eval_steps": 2498,
+  "global_step": 1300,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0013481631277384564,
+      "grad_norm": 2.09375,
+      "learning_rate": 2.66e-05,
+      "loss": 8.298905181884766,
+      "step": 20
+    },
+    {
+      "epoch": 0.002696326255476913,
+      "grad_norm": 1.3359375,
+      "learning_rate": 5.46e-05,
+      "loss": 8.045733642578124,
+      "step": 40
+    },
+    {
+      "epoch": 0.004044489383215369,
+      "grad_norm": 1.2734375,
+      "learning_rate": 8.259999999999999e-05,
+      "loss": 7.839435577392578,
+      "step": 60
+    },
+    {
+      "epoch": 0.005392652510953826,
+      "grad_norm": 1.265625,
+      "learning_rate": 0.0001106,
+      "loss": 7.574758148193359,
+      "step": 80
+    },
+    {
+      "epoch": 0.006740815638692282,
+      "grad_norm": 1.1953125,
+      "learning_rate": 0.0001386,
+      "loss": 7.235897064208984,
+      "step": 100
+    },
+    {
+      "epoch": 0.008088978766430738,
+      "grad_norm": 1.140625,
+      "learning_rate": 0.00016659999999999998,
+      "loss": 6.862369537353516,
+      "step": 120
+    },
+    {
+      "epoch": 0.009437141894169195,
+      "grad_norm": 1.0234375,
+      "learning_rate": 0.00019460000000000001,
+      "loss": 6.508391571044922,
+      "step": 140
+    },
+    {
+      "epoch": 0.010785305021907651,
+      "grad_norm": 1.078125,
+      "learning_rate": 0.0002226,
+      "loss": 6.215537261962891,
+      "step": 160
+    },
+    {
+      "epoch": 0.012133468149646108,
+      "grad_norm": 0.76171875,
+      "learning_rate": 0.00025059999999999997,
+      "loss": 5.984259414672851,
+      "step": 180
+    },
+    {
+      "epoch": 0.013481631277384564,
+      "grad_norm": 0.82421875,
+      "learning_rate": 0.0002786,
+      "loss": 5.756989669799805,
+      "step": 200
+    },
+    {
+      "epoch": 0.01482979440512302,
+      "grad_norm": 1.171875,
+      "learning_rate": 0.00030659999999999997,
+      "loss": 5.531437683105469,
+      "step": 220
+    },
+    {
+      "epoch": 0.016177957532861477,
+      "grad_norm": 1.1171875,
+      "learning_rate": 0.0003346,
+      "loss": 5.294247817993164,
+      "step": 240
+    },
+    {
+      "epoch": 0.01752612066059993,
+      "grad_norm": 0.9453125,
+      "learning_rate": 0.00036260000000000003,
+      "loss": 5.055661392211914,
+      "step": 260
+    },
+    {
+      "epoch": 0.01887428378833839,
+      "grad_norm": 0.94921875,
+      "learning_rate": 0.0003906,
+      "loss": 4.797793197631836,
+      "step": 280
+    },
+    {
+      "epoch": 0.020222446916076844,
+      "grad_norm": 1.421875,
+      "learning_rate": 0.0004186,
+      "loss": 4.594622802734375,
+      "step": 300
+    },
+    {
+      "epoch": 0.021570610043815303,
+      "grad_norm": 1.2265625,
+      "learning_rate": 0.0004466,
+      "loss": 4.39941635131836,
+      "step": 320
+    },
+    {
+      "epoch": 0.022918773171553757,
+      "grad_norm": 0.5546875,
+      "learning_rate": 0.00047460000000000004,
+      "loss": 4.24058723449707,
+      "step": 340
+    },
+    {
+      "epoch": 0.024266936299292215,
+      "grad_norm": 0.83984375,
+      "learning_rate": 0.0005026,
+      "loss": 4.102418899536133,
+      "step": 360
+    },
+    {
+      "epoch": 0.02561509942703067,
+      "grad_norm": 0.70703125,
+      "learning_rate": 0.0005306,
+      "loss": 3.975328063964844,
+      "step": 380
+    },
+    {
+      "epoch": 0.026963262554769128,
+      "grad_norm": 0.6171875,
+      "learning_rate": 0.0005586,
+      "loss": 3.839755630493164,
+      "step": 400
+    },
+    {
+      "epoch": 0.028311425682507583,
+      "grad_norm": 0.91796875,
+      "learning_rate": 0.0005866,
+      "loss": 3.7486804962158202,
+      "step": 420
+    },
+    {
+      "epoch": 0.02965958881024604,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0006146,
+      "loss": 3.6831081390380858,
+      "step": 440
+    },
+    {
+      "epoch": 0.031007751937984496,
+      "grad_norm": 0.875,
+      "learning_rate": 0.0006426,
+      "loss": 3.5874557495117188,
+      "step": 460
+    },
+    {
+      "epoch": 0.032355915065722954,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0006705999999999999,
+      "loss": 3.5174354553222655,
+      "step": 480
+    },
+    {
+      "epoch": 0.03370407819346141,
+      "grad_norm": 0.66796875,
+      "learning_rate": 0.0006986,
+      "loss": 3.439300537109375,
+      "step": 500
+    },
+    {
+      "epoch": 0.03505224132119986,
+      "grad_norm": 0.54296875,
+      "learning_rate": 0.0007,
+      "loss": 3.386067581176758,
+      "step": 520
+    },
+    {
+      "epoch": 0.03640040444893832,
+      "grad_norm": 0.6015625,
+      "learning_rate": 0.0007,
+      "loss": 3.3094718933105467,
+      "step": 540
+    },
+    {
+      "epoch": 0.03774856757667678,
+      "grad_norm": 0.55859375,
+      "learning_rate": 0.0007,
+      "loss": 3.2609573364257813,
+      "step": 560
+    },
+    {
+      "epoch": 0.03909673070441524,
+      "grad_norm": 0.61328125,
+      "learning_rate": 0.0007,
+      "loss": 3.2046344757080076,
+      "step": 580
+    },
+    {
+      "epoch": 0.04044489383215369,
+      "grad_norm": 0.5234375,
+      "learning_rate": 0.0007,
+      "loss": 3.17300968170166,
+      "step": 600
+    },
+    {
+      "epoch": 0.04179305695989215,
+      "grad_norm": 0.578125,
+      "learning_rate": 0.0007,
+      "loss": 3.1147342681884767,
+      "step": 620
+    },
+    {
+      "epoch": 0.043141220087630605,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 3.0973543167114257,
+      "step": 640
+    },
+    {
+      "epoch": 0.044489383215369056,
+      "grad_norm": 0.453125,
+      "learning_rate": 0.0007,
+      "loss": 3.043070602416992,
+      "step": 660
+    },
+    {
+      "epoch": 0.045837546343107514,
+      "grad_norm": 0.455078125,
+      "learning_rate": 0.0007,
+      "loss": 3.0213077545166014,
+      "step": 680
+    },
+    {
+      "epoch": 0.04718570947084597,
+      "grad_norm": 0.419921875,
+      "learning_rate": 0.0007,
+      "loss": 2.977329063415527,
+      "step": 700
+    },
+    {
+      "epoch": 0.04853387259858443,
+      "grad_norm": 0.490234375,
+      "learning_rate": 0.0007,
+      "loss": 2.942974090576172,
+      "step": 720
+    },
+    {
+      "epoch": 0.04988203572632288,
+      "grad_norm": 0.48046875,
+      "learning_rate": 0.0007,
+      "loss": 2.9064186096191404,
+      "step": 740
+    },
+    {
+      "epoch": 0.05123019885406134,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.890872001647949,
+      "step": 760
+    },
+    {
+      "epoch": 0.0525783619817998,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8978702545166017,
+      "step": 780
+    },
+    {
+      "epoch": 0.053926525109538256,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.8463741302490235,
+      "step": 800
+    },
+    {
+      "epoch": 0.05527468823727671,
+      "grad_norm": 0.47265625,
+      "learning_rate": 0.0007,
+      "loss": 2.8323719024658205,
+      "step": 820
+    },
+    {
+      "epoch": 0.056622851365015166,
+      "grad_norm": 0.5,
+      "learning_rate": 0.0007,
+      "loss": 2.8007469177246094,
+      "step": 840
+    },
+    {
+      "epoch": 0.057971014492753624,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.7573240280151365,
+      "step": 860
+    },
+    {
+      "epoch": 0.05931917762049208,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.7341793060302733,
+      "step": 880
+    },
+    {
+      "epoch": 0.06066734074823053,
+      "grad_norm": 0.546875,
+      "learning_rate": 0.0007,
+      "loss": 2.7180759429931642,
+      "step": 900
+    },
+    {
+      "epoch": 0.06201550387596899,
+      "grad_norm": 0.5625,
+      "learning_rate": 0.0007,
+      "loss": 2.692394256591797,
+      "step": 920
+    },
+    {
+      "epoch": 0.06336366700370745,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.678949546813965,
+      "step": 940
+    },
+    {
+      "epoch": 0.06471183013144591,
+      "grad_norm": 0.466796875,
+      "learning_rate": 0.0007,
+      "loss": 2.651218223571777,
+      "step": 960
+    },
+    {
+      "epoch": 0.06605999325918437,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.632099723815918,
+      "step": 980
+    },
+    {
+      "epoch": 0.06740815638692282,
+      "grad_norm": 0.470703125,
+      "learning_rate": 0.0007,
+      "loss": 2.6305675506591797,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06875631951466127,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.5994905471801757,
+      "step": 1020
+    },
+    {
+      "epoch": 0.07010448264239973,
+      "grad_norm": 0.474609375,
+      "learning_rate": 0.0007,
+      "loss": 2.5834341049194336,
+      "step": 1040
+    },
+    {
+      "epoch": 0.07145264577013818,
+      "grad_norm": 0.498046875,
+      "learning_rate": 0.0007,
+      "loss": 2.567459487915039,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07280080889787664,
+      "grad_norm": 0.462890625,
+      "learning_rate": 0.0007,
+      "loss": 2.54827823638916,
+      "step": 1080
+    },
+    {
+      "epoch": 0.0741489720256151,
+      "grad_norm": 0.478515625,
+      "learning_rate": 0.0007,
+      "loss": 2.5431196212768556,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07549713515335356,
+      "grad_norm": 0.4453125,
+      "learning_rate": 0.0007,
+      "loss": 2.5049072265625,
+      "step": 1120
+    },
+    {
+      "epoch": 0.07684529828109202,
+      "grad_norm": 0.4765625,
+      "learning_rate": 0.0007,
+      "loss": 2.4941583633422852,
+      "step": 1140
+    },
+    {
+      "epoch": 0.07819346140883048,
+      "grad_norm": 0.46875,
+      "learning_rate": 0.0007,
+      "loss": 2.4895328521728515,
+      "step": 1160
+    },
+    {
+      "epoch": 0.07954162453656892,
+      "grad_norm": 0.48828125,
+      "learning_rate": 0.0007,
+      "loss": 2.4841281890869142,
+      "step": 1180
+    },
+    {
+      "epoch": 0.08088978766430738,
+      "grad_norm": 0.46484375,
+      "learning_rate": 0.0007,
+      "loss": 2.475014495849609,
+      "step": 1200
+    },
+    {
+      "epoch": 0.08223795079204584,
+      "grad_norm": 0.43359375,
+      "learning_rate": 0.0007,
+      "loss": 2.459096908569336,
+      "step": 1220
+    },
+    {
+      "epoch": 0.0835861139197843,
+      "grad_norm": 0.443359375,
+      "learning_rate": 0.0007,
+      "loss": 2.4455087661743162,
+      "step": 1240
+    },
+    {
+      "epoch": 0.08493427704752275,
+      "grad_norm": 0.48828125,
+      "learning_rate": 0.0007,
+      "loss": 2.432474899291992,
+      "step": 1260
+    },
+    {
+      "epoch": 0.08628244017526121,
+      "grad_norm": 0.427734375,
+      "learning_rate": 0.0007,
+      "loss": 2.4141357421875,
+      "step": 1280
+    },
+    {
+      "epoch": 0.08763060330299967,
+      "grad_norm": 0.46875,
+      "learning_rate": 0.0007,
+      "loss": 2.4043970108032227,
+      "step": 1300
+    }
+  ],
+  "logging_steps": 20,
+  "max_steps": 2500,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 4194168943411200.0,
+  "train_batch_size": 64,
+  "trial_name": null,
+  "trial_params": null
+}

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

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

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

@@ -0,0 +1,36 @@
+{
+  "activation": "relu",
+  "architectures": [
+    "TinyLlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "eos_token_id": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 128,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 512,
+  "mlp_bias": false,
+  "mlp_type": "glu",
+  "model_type": "tiny_llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 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": 4
+}

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

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