Auto upload zain 2026-08-12T20:48:11.035594

zain/Activation/exp.py

@@ -146,11 +146,8 @@ class TinyLlamaMLP(nn.Module):
         self.down_proj = nn.Linear(effective_intermediate, self.hidden_size, bias=False)
 
         # Activation function (for standard variants)
-        # For gated variants (situglu, waleed) we handle them in forward, but we still need a placeholder.
-        # For waleed10/silu-waleed10 we need linear or silu.
         if self.mlp_type == "glu":
             if self.activation_name in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
-                # These will be handled in forward; no act_fn needed.
                 self.act_fn = None
             elif self.activation_name == "waleed10":
                 self.act_fn = GLUActivationRegistry.get("linear")
@@ -175,7 +172,7 @@ class TinyLlamaMLP(nn.Module):
         if self.activation_name in ("situglu_low", "waleedglu_low"):
             self.beta1 = 2.5
             self.beta2 = 4.0
-        else:  # original situglu or waleed
+        else:
             self.beta1 = 4.0
             self.beta2 = 25.0
 
@@ -183,7 +180,7 @@ class TinyLlamaMLP(nn.Module):
         self.is_situglu = self.activation_name in ("situglu", "situglu_low")
         self.is_waleed = self.activation_name in ("waleed", "waleedglu_low")
         self.is_waleed10 = self.activation_name in ("waleed10", "silu-waleed10")
-        self.has_sigmoid_gate = self.activation_name.startswith("situglu")  # sigmoid in gate
+        self.has_sigmoid_gate = self.activation_name.startswith("situglu")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.mlp_type == "glu":
@@ -191,7 +188,6 @@ class TinyLlamaMLP(nn.Module):
             up = self.up_proj(x)
 
             if self.is_situglu or self.is_waleed:
-                # Gated variants with tanh scaling
                 if self.has_sigmoid_gate:
                     gate = self.beta1 * torch.tanh(gate / self.beta1) * torch.sigmoid(gate)
                 else:
@@ -199,7 +195,6 @@ class TinyLlamaMLP(nn.Module):
                 up = self.beta2 * torch.tanh(up / self.beta2)
                 hidden = gate * up
             else:
-                # Standard GLU (activation applied to gate)
                 hidden = self.act_fn(gate) * up
 
             out = self.down_proj(hidden)
@@ -208,7 +203,6 @@ class TinyLlamaMLP(nn.Module):
             hidden = self.act_fn(self.up_proj(x))
             out = self.down_proj(hidden)
 
-        # Post‑clip for waleed10 variants
         if self.is_waleed10:
             out = self.waleed_beta * torch.tanh(out / self.waleed_beta)
 
@@ -225,11 +219,7 @@ class TinyLlamaDecoderLayer(nn.Module):
         self.post_attention_layernorm = LlamaRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
-        # ---------------------------------------------------------------------
-        # NEW: zero-parameter Identity gateways for residual-stream logging.
-        # These expose the residual tensor as named modules so the hook
-        # registry can capture them with pattern ".*residual.*".
-        # ---------------------------------------------------------------------
+        # Residual stream gateways (zero params, hookable)
         self.residual_pre_attn = nn.Identity()
         self.residual_post_attn = nn.Identity()
         self.residual_post_mlp = nn.Identity()
@@ -242,7 +232,7 @@ class TinyLlamaDecoderLayer(nn.Module):
         position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         **kwargs,
     ):
-        # --- Attention sub-layer ---
+        # Attention sub-layer
         residual = hidden_states
         hidden_states = self.residual_pre_attn(hidden_states)
         hidden_states = self.input_layernorm(hidden_states)
@@ -255,7 +245,7 @@ class TinyLlamaDecoderLayer(nn.Module):
         hidden_states = residual + attn_out
         hidden_states = self.residual_post_attn(hidden_states)
 
-        # --- MLP sub-layer ---
+        # MLP sub-layer
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
@@ -265,7 +255,7 @@ class TinyLlamaDecoderLayer(nn.Module):
 
 
 # ----------------------------------------------------------------------------
-# ATTENTION MASK – float mask with 0.0 / -inf (works with all backends)
+# ATTENTION MASK
 # ----------------------------------------------------------------------------
 def _build_causal_mask(
     attention_mask: Optional[torch.Tensor],
@@ -273,17 +263,10 @@ def _build_causal_mask(
     dtype: torch.dtype,
     device: torch.device,
 ) -> torch.Tensor:
-    """
-    Build a 4D float attention mask for scaled_dot_product_attention.
-    - 0.0 where attention is allowed
-    - -inf where it is masked (causal future + padding)
-    """
     min_value = torch.finfo(dtype).min
-
-    # Causal mask: upper triangle (future) = -inf
     causal = torch.full((seq_len, seq_len), fill_value=min_value, dtype=dtype, device=device)
     causal = torch.triu(causal, diagonal=1)
-    causal = causal[None, None, :, :]  # (1, 1, seq_len, seq_len)
+    causal = causal[None, None, :, :]
 
     if attention_mask is None:
         batch_size = 1
@@ -291,15 +274,11 @@ def _build_causal_mask(
 
     batch_size = attention_mask.shape[0]
     causal = causal.expand(batch_size, 1, seq_len, seq_len).clone()
-
-    # Padding: where attention_mask == 0, set to -inf
-    padding = attention_mask[:, None, None, :].to(device) == 0  # (batch, 1, 1, seq_len)
+    padding = attention_mask[:, None, None, :].to(device) == 0
     causal = causal.masked_fill(padding, min_value)
-
     return causal
 
 
-# Global flag to print mask message only once
 _MASK_PRINTED = False
 
 
@@ -330,10 +309,7 @@ class TinyLlamaModel(LlamaPreTrainedModel):
         **kwargs,
     ):
         global _MASK_PRINTED
-
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
+        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)
 
@@ -346,7 +322,6 @@ class TinyLlamaModel(LlamaPreTrainedModel):
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
-        # Build float causal + padding mask (print only once)
         seq_len = hidden_states.shape[1]
         causal_mask = _build_causal_mask(
             attention_mask, seq_len, hidden_states.dtype, hidden_states.device
@@ -402,9 +377,7 @@ class TinyLlamaForCausalLM(LlamaPreTrainedModel):
         return_dict: Optional[bool] = None,
         **kwargs,
     ):
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
+        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,
@@ -459,7 +432,7 @@ class TinyLlamaForCausalLM(LlamaPreTrainedModel):
 # =============================================================================
 
 class StatsEngine:
-    """Compute the unified signature for any tensor (now with range + percentiles)."""
+    """Compute unified signature for any tensor."""
 
     @staticmethod
     def compute(
@@ -474,66 +447,30 @@ class StatsEngine:
                 else float("inf")
             )
 
-            # --- Base statistics ---
-            norm = tensor.norm(2).item()
-            mean = tensor.mean().item()
-            std = tensor.std().item()
-            max_abs = abs_t.max().item()
-
-            # --- Exact min / max / range ---
             t_min = tensor.min().item()
             t_max = tensor.max().item()
-            t_range = t_max - t_min
-
-            # --- Exact percentiles (float32 CPU for dtype safety) ---
-            p01 = p25 = p50 = p75 = p90 = p95 = p99 = 0.0
-            try:
-                flat_f32 = tensor.detach().reshape(-1).to(torch.float32).cpu()
-                if flat_f32.numel() > 0:
-                    q_vals = torch.quantile(
-                        flat_f32,
-                        torch.tensor(
-                            [0.01, 0.25, 0.5, 0.75, 0.90, 0.95, 0.99],
-                            dtype=torch.float32,
-                        ),
-                    )
-                    p01, p25, p50, p75, p90, p95, p99 = (v.item() for v in q_vals)
-            except Exception:
-                pass  # If quantile fails (very unlikely), leave as 0.0
 
             return {
-                "norm": norm,
-                "mean": mean,
-                "std": std,
-                "max_abs": max_abs,
+                "norm": tensor.norm(2).item(),
+                "mean": tensor.mean().item(),
+                "std": tensor.std().item(),
+                "max_abs": abs_t.max().item(),
                 "frac_near_dtype_limit": (
                     (abs_t > dtype_limit).float().mean().item()
                     if not math.isinf(dtype_limit)
                     else 0.0
                 ),
                 "frac_near_user_limit": (abs_t > user_limit).float().mean().item(),
-                # --- NEW: distributional tail metrics (exact per-tensor) ---
                 "min": t_min,
                 "max": t_max,
-                "range": t_range,
-                "p01": p01,
-                "p25": p25,
-                "p50": p50,
-                "p75": p75,
-                "p90": p90,
-                "p95": p95,
-                "p99": p99,
+                "range": t_max - t_min,
             }
 
 
 class StepAccumulator:
-    """
-    Stores per-tensor entries, then aggregates to layer-scope or global-scope
-    using exact population formulas (no tensor retention).
-    """
+    """Stores per-tensor entries, aggregates to layer/global scope."""
 
     def __init__(self):
-        # name -> {numel, norm, mean, std, max_abs, frac_near_dtype_limit, frac_near_user_limit, min, max, range, p01..p99}
         self.tensors: Dict[str, Dict[str, float]] = {}
 
     def add(self, name: str, numel: int, stats: Dict[str, float]):
@@ -563,18 +500,8 @@ class StepAccumulator:
             a["frac_near_user_limit"] * a["numel"] + b["frac_near_user_limit"] * b["numel"]
         ) / total_n
 
-        # Exact min/max across micro-batches (gradient accumulation)
         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")))
-        t_range = t_max - t_min
-
-        # Percentiles: weighted average across micro-batches (best effort)
-        def _wpct(key: str) -> float:
-            av = a.get(key, 0.0)
-            bv = b.get(key, 0.0)
-            if av == 0.0 and bv == 0.0:
-                return 0.0
-            return (av * a["numel"] + bv * b["numel"]) / total_n
 
         return {
             "numel": total_n,
@@ -586,14 +513,7 @@ class StepAccumulator:
             "frac_near_user_limit": frac_user,
             "min": t_min,
             "max": t_max,
-            "range": t_range,
-            "p01": _wpct("p01"),
-            "p25": _wpct("p25"),
-            "p50": _wpct("p50"),
-            "p75": _wpct("p75"),
-            "p90": _wpct("p90"),
-            "p95": _wpct("p95"),
-            "p99": _wpct("p99"),
+            "range": t_max - t_min,
         }
 
     def clear(self):
@@ -605,19 +525,14 @@ class StepAccumulator:
         numels = [e["numel"] for e in entries.values()]
         total_n = sum(numels)
 
-        # L2 norm
         norm = math.sqrt(sum(e["norm"] ** 2 for e in entries.values()))
-        # Max abs
         max_abs = max(e["max_abs"] for e in entries.values())
-        # Weighted mean
         mean = sum(e["mean"] * e["numel"] for e in entries.values()) / total_n
-        # Pooled std: sqrt( E[σ² + μ²] - μ_global² )
         ex2 = (
             sum(e["numel"] * (e["std"] ** 2 + e["mean"] ** 2) for e in entries.values())
             / total_n
         )
         std = math.sqrt(max(0.0, ex2 - mean ** 2))
-        # Weighted fractions
         frac_dtype = (
             sum(e["frac_near_dtype_limit"] * e["numel"] for e in entries.values())
             / total_n
@@ -627,10 +542,8 @@ class StepAccumulator:
             / total_n
         )
 
-        # Exact bounds across all tensors in this scope
         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())
-        t_range = t_max - t_min
 
         return {
             "norm": norm,
@@ -641,9 +554,7 @@ class StepAccumulator:
             "frac_near_user_limit": frac_user,
             "min": t_min,
             "max": t_max,
-            "range": t_range,
-            # NOTE: percentiles intentionally omitted from aggregated scopes.
-            # They are only meaningful at per-tensor scope.
+            "range": t_max - t_min,
         }
 
     def get_global_stats(self) -> Dict[str, float]:
@@ -696,10 +607,6 @@ class HookRegistry:
         def hook(module, inp, out):
             if not self.active:
                 return
-
-            # Modules can return a tensor, a tuple (take first item), or a
-            # dict (e.g. TinyLlamaModel returns {"last_hidden_state": ...}).
-            # Pull out the first real tensor we find; skip cleanly if none.
             if isinstance(out, dict):
                 out_dict = out
                 out = out_dict.get("last_hidden_state")
@@ -739,10 +646,7 @@ class HookRegistry:
 
 
 class StabilityMonitorCallback(TrainerCallback):
-    """
-    Full stability instrumentation: grad / param / act statistics
-    at global, per-layer, and per-tensor scope.
-    """
+    """Full stability instrumentation: grad / param / act statistics."""
 
     def __init__(
         self,
@@ -756,7 +660,6 @@ class StabilityMonitorCallback(TrainerCallback):
     ):
         self.model = model
         self.monitor_every_n_steps = monitor_every_n_steps
-        # NEW: default patterns now include residual gateways
         self.module_patterns = module_patterns or [".*mlp.*", ".*self_attn.*", ".*residual.*"]
         self.user_limits = user_limits or {"grad": 1.0, "param": 100.0, "act": 50.0}
         self.dtype_ratio = dtype_proximity_ratio
@@ -792,8 +695,6 @@ class StabilityMonitorCallback(TrainerCallback):
     def on_step_end(self, args, state, control, **kwargs):
         if not self.hooks.active:
             return
-
-        # Parameter stats (post-optimizer step)
         for name, param in self.model.named_parameters():
             stats = StatsEngine.compute(
                 param.data, self.user_limits["param"], self.dtype_ratio
@@ -805,7 +706,6 @@ class StabilityMonitorCallback(TrainerCallback):
 
     @staticmethod
     def _kind_of(name: str) -> str:
-        """Classify a tensor key by its source: activation, gradient, or parameter."""
         if name.startswith("act."):
             return "act"
         if name.startswith("grad."):
@@ -825,18 +725,15 @@ class StabilityMonitorCallback(TrainerCallback):
     def _build_metrics(self, scope: str = "train") -> Dict[str, float]:
         metrics: Dict[str, float] = {}
 
-        # --- Global (split by kind: act / grad / param — never pooled together) ---
         if self.log_scope.get("global", True):
             by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
             for k, v in self.accumulator.tensors.items():
                 by_kind.setdefault(self._kind_of(k), {})[k] = v
-
             for kind, entries in by_kind.items():
                 stats = self.accumulator._aggregate(entries)
                 for kk, vv in stats.items():
                     metrics[f"{scope}/global/{kind}/{kk}"] = vv
 
-        # --- Per-layer (group by model.layers.{i}, split by kind) ---
         if self.log_scope.get("per_layer", True):
             layer_prefixes = set()
             for name in self.accumulator.tensors:
@@ -846,14 +743,12 @@ class StabilityMonitorCallback(TrainerCallback):
                     if p == "layers" and i + 1 < len(parts):
                         prefix = ".".join(parts[: i + 2])
                         layer_prefixes.add(prefix)
-
             for prefix in layer_prefixes:
                 by_kind: Dict[str, Dict[str, Dict[str, float]]] = {}
                 for k, v in self.accumulator.tensors.items():
                     clean = self._strip_kind(k)
                     if clean.startswith(prefix + ".") or clean == prefix:
                         by_kind.setdefault(self._kind_of(k), {})[k] = v
-
                 safe = prefix.replace(".", "_")
                 for kind, entries in by_kind.items():
                     if not entries:
@@ -862,7 +757,6 @@ class StabilityMonitorCallback(TrainerCallback):
                     for kk, vv in stats.items():
                         metrics[f"{scope}/layer_{safe}/{kind}/{kk}"] = vv
 
-        # --- Per-tensor (exact percentiles live here) ---
         if self.log_scope.get("per_tensor", False):
             for name, stats in self.accumulator.tensors.items():
                 safe = name.replace(".", "_")
@@ -879,19 +773,9 @@ class StabilityMonitorCallback(TrainerCallback):
             self.pending_metrics = None
 
     def on_prediction_step(self, args, state, control, **kwargs):
-        """Fires once per eval/predict batch. Trainer.evaluate() calls this
-        for every batch in the eval loop, then calls self.log(output.metrics)
-        (which dispatches on_log to every callback, including the wandb/
-        tensorboard reporting callbacks) BEFORE on_evaluate() runs. So to get
-        eval-time stats into that same on_log dispatch, we have to build
-        pending_metrics here, not in on_evaluate — by the time on_evaluate
-        fires, self.log() has already happened and it's too late.
-        """
         if not self.monitor_during_eval:
             return
         if not self.hooks.active:
-            # First batch of this eval pass: start a fresh accumulation and
-            # snapshot parameter stats once (they don't change during eval).
             self.accumulator.clear()
             self.hooks.set_active(True)
             for name, param in self.model.named_parameters():
@@ -931,16 +815,12 @@ class TimeTrackerCallback(TrainerCallback):
     def on_log(self, args, state, control, logs=None, **kwargs):
         if logs is None:
             return
-
         logs["train/total_time_seconds"] = self.total_train_time
-
         if self.step_times:
             recent = self.step_times[-100:]
             logs["train/time_per_step_avg"] = sum(recent) / len(recent)
-
         if self.epoch_start is not None:
             logs["train/epoch_time_elapsed"] = time.perf_counter() - self.epoch_start
-
         if state.max_steps and state.global_step > 0:
             avg = self.total_train_time / state.global_step
             remaining = (state.max_steps - state.global_step) * avg
@@ -972,16 +852,13 @@ class MetricsLoggerCallback(TrainerCallback):
 
 
 # =============================================================================
-# NEW: CONTAMINATION CALLBACK
+# CONTAMINATION CALLBACK
 # =============================================================================
 
 class ContaminationCallback(TrainerCallback):
     """
-    Intentionally corrupt input_ids and labels for a window of steps,
-    causing a controlled loss spike. Supports two modes:
-      - "shift": every token ID is incremented by 1 (mod vocab_size)
-      - "random": token IDs are replaced with uniform random IDs
-    Fraction controls what proportion of tokens (per sequence) are corrupted.
+    Intentionally corrupt input_ids and labels for a window of steps.
+    Modes: "shift" (+1 mod vocab), "random" (uniform random IDs).
     """
     def __init__(
         self,
@@ -1000,12 +877,9 @@ class ContaminationCallback(TrainerCallback):
         self.mode = mode
         self.fraction = fraction
         self.seed = seed
-
-        # Create a dedicated generator for reproducibility
         self.generator = torch.Generator()
         if seed is not None:
             self.generator.manual_seed(seed)
-
         self._active = False
 
     def _should_corrupt(self, state) -> bool:
@@ -1017,58 +891,37 @@ class ContaminationCallback(TrainerCallback):
     def on_step_begin(self, args, state, control, **kwargs):
         if not self._should_corrupt(state):
             return
-
-        # The batch is passed in kwargs under key "inputs" (Trainer convention).
-        # We also need to grab the model's device to generate tensors on the same device.
         batch = kwargs.get("inputs")
-        if batch is None:
+        if batch is None or not isinstance(batch, dict):
             return
-
-        # We need the device – get it from the model or from the input tensors.
-        # We can access the model through the trainer? Not directly here.
-        # But we can infer device from batch tensors.
-        if 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
 
-        # Build mask of positions to corrupt (based on fraction)
-        # We generate a mask of shape (batch_size, seq_len) with True for positions to corrupt.
         if self.fraction >= 1.0:
             corrupt_mask = torch.ones((batch_size, seq_len), dtype=torch.bool, device=device)
         elif self.fraction <= 0.0:
-            return  # nothing to corrupt
+            return
         else:
-            # Generate random floats and threshold
             rand = torch.rand((batch_size, seq_len), generator=self.generator, device=device)
             corrupt_mask = rand < self.fraction
 
-        # If attention_mask is present, exclude padding positions from corruption
         if attention_mask is not None:
-            # attention_mask is 1 for real tokens, 0 for padding
             padding_mask = attention_mask == 0
             corrupt_mask = corrupt_mask & (~padding_mask)
 
         if not corrupt_mask.any():
             return
 
-        # Apply corruption
         if self.mode == "shift":
-            # Increment by 1 modulo vocab_size
-            # We need to handle vocab_size properly (IDs go 0..vocab_size-1)
             input_ids_corrupt = (input_ids + 1) % self.vocab_size
             labels_corrupt = (labels + 1) % self.vocab_size
         elif self.mode == "random":
-            # Replace with uniform random IDs
-            # Create random tensor of same shape
             random_ids = torch.randint(
                 0, self.vocab_size, input_ids.shape,
                 generator=self.generator, device=device
@@ -1078,21 +931,16 @@ class ContaminationCallback(TrainerCallback):
         else:
             raise ValueError(f"Unknown contamination mode: {self.mode}")
 
-        # Apply only where mask is True
         input_ids.masked_scatter_(corrupt_mask, input_ids_corrupt[corrupt_mask])
         labels.masked_scatter_(corrupt_mask, labels_corrupt[corrupt_mask])
-
-        # Reassign into batch (mutated in-place)
         batch["input_ids"] = input_ids
         batch["labels"] = labels
 
-        # Log a one‑time message when contamination starts
         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 we just passed the end of the window, de‑activate
         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
@@ -1109,9 +957,8 @@ def build_dataset(
     dataset_name: str = "roneneldan/TinyStories",
     max_samples: Optional[int] = None,
 ):
-    """Concatenate and chunk TinyStories for causal LM. Fast path with multiprocessing."""
+    """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})")
@@ -1197,11 +1044,8 @@ def create_trainer(
         remove_unused_columns=False,
     )
 
-    callbacks = [
-        TimeTrackerCallback(),
-    ]
+    callbacks = [TimeTrackerCallback()]
 
-    # --- Contamination callback (optional) ---
     cc = config.get("contamination", {})
     if cc.get("enabled", False):
         vocab_size = model.config.vocab_size
@@ -1223,14 +1067,9 @@ def create_trainer(
                 model=model,
                 monitor_every_n_steps=mc.get("monitor_every_n_steps", 10),
                 module_patterns=mc.get("module_patterns", [".*mlp.*", ".*self_attn.*", ".*residual.*"]),
-                user_limits=mc.get(
-                    "user_limits", {"grad": 1.0, "param": 100.0, "act": 50.0}
-                ),
+                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},
-                ),
+                log_scope=mc.get("log_scope", {"global": True, "per_layer": True, "per_tensor": False}),
                 monitor_during_eval=mc.get("monitor_during_eval", False),
             )
         )
@@ -1248,10 +1087,8 @@ def create_trainer(
         callbacks=callbacks,
     )
 
-    # Move reporting integrations to the end
     try:
         from transformers.integrations import get_reporting_integration_callbacks
-
         reporting_types = tuple(get_reporting_integration_callbacks(args.report_to))
     except Exception:
         reporting_types = ()

zain/Activation/sweep.py

@@ -39,14 +39,13 @@ def parse_variant(variant: str):
     if prefix not in ('glu', 'mlp'):
         raise ValueError(f"Invalid prefix: '{prefix}'. Must be 'glu' or 'mlp'.")
 
-    # The last part might be layers like "10L"
     last = parts[-1]
     if last.endswith('L') and last[:-1].isdigit():
         layers = int(last[:-1])
-        activation = '-'.join(parts[1:-1])   # everything between prefix and layers
+        activation = '-'.join(parts[1:-1])
     else:
         layers = None
-        activation = '-'.join(parts[1:])     # everything after prefix
+        activation = '-'.join(parts[1:])
 
     if not activation:
         raise ValueError(f"Missing activation name in variant: '{variant}'")
@@ -84,32 +83,20 @@ def main():
         tokenizer.pad_token = tokenizer.eos_token
 
     msl = base["model"].get("max_position_embeddings", 512)
-    train_ds = build_dataset(
-        tokenizer,
-        max_seq_len=msl,
-        split="train",
-        max_samples=None
-    )
-    eval_ds = build_dataset(
-        tokenizer,
-        max_seq_len=msl,
-        split="validation",
-        max_samples=None
-    )
+    train_ds = build_dataset(tokenizer, max_seq_len=msl, split="train", max_samples=None)
+    eval_ds = build_dataset(tokenizer, max_seq_len=msl, split="validation", max_samples=None)
 
     results = []
 
     for variant in args.variants:
         prefix, act, layers = parse_variant(variant)
 
-        # Validation: mlp-situglu and mlp-waleed are banned
         if prefix == "mlp" and act in ("situglu", "waleed", "situglu_low", "waleedglu_low"):
             raise ValueError(
                 f"Activation '{act}' requires a gated architecture (GLU). "
                 f"Please use 'glu-{act}' instead."
             )
 
-        # Build config overrides
         cfg = copy.deepcopy(base)
         cfg["model"]["mlp_type"] = prefix
         cfg["model"]["activation"] = act
@@ -123,16 +110,13 @@ def main():
         actual_layers = cfg["model"]["num_hidden_layers"]
         variant_label = f"{prefix}-{act}-{actual_layers}L"
 
-        # Unique output directory
         out_dir = Path(cfg["training"]["output_dir"]).parent / f"{variant_label}_run"
         cfg["training"]["output_dir"] = str(out_dir)
 
-        # Re-seed for reproducibility across variants
         set_seed(seed)
 
         print(f"\n{'='*60}\n>>> Variant: {variant_label} | Out: {out_dir}\n{'='*60}")
 
-        # Instantiate model
         config = TinyLlamaConfig(**cfg["model"])
         model = TinyLlamaForCausalLM(config)
         model = model.to(torch.bfloat16)
@@ -143,11 +127,9 @@ def main():
         run_name = f"LM-{variant_label}-{param_str}-{timestamp}"
         cfg["training"]["run_name"] = run_name
 
-        # Also update hub_model_id to include variant and layer count
         hub_id_base = cfg["training"].get("hub_model_id", "tiny-llama-lab")
         cfg["training"]["hub_model_id"] = f"{hub_id_base}-{variant_label}"
 
-        # Ensure a fresh WandB run – remove any global WANDB_RUN_ID
         os.environ.pop("WANDB_RUN_ID", None)
 
         trainer = create_trainer(model, tokenizer, cfg, train_ds, eval_ds)
@@ -156,7 +138,7 @@ def main():
             trainer.train()
             metrics = trainer.evaluate()
             results.append({
-                "variant": variant_label,   # now always includes layer count
+                "variant": variant_label,
                 "eval_loss": metrics.get("eval_loss"),
                 "out": str(out_dir),
                 "run_name": run_name,
@@ -165,10 +147,8 @@ def main():
             if args.push or cfg["training"].get("push_to_hub", False):
                 trainer.push_to_hub()
         finally:
-            # Explicitly finish WandB run to avoid re‑using the same run
             wandb.finish()
 
-    # Save summary
     summary = Path(base["training"]["output_dir"]).parent / "sweep_summary.json"
     summary.write_text(json.dumps(results, indent=2))
     print("\nSweep complete:")

zain/Activation/train.py

@@ -18,12 +18,11 @@ def main():
     with open(args.config) as f:
         cfg = yaml.safe_load(f)
 
-    # Explicit seed before any randomness
     seed = cfg.get("training", {}).get("seed", 42)
     set_seed(seed)
 
     # -------------------------------------------------------------------------
-    # NEW: Force run into a specific WandB project (read from config).
+    # LOCK WANDB PROJECT: read from config, force into environment.
     # -------------------------------------------------------------------------
     wandb_project = cfg.get("training", {}).get("wandb_project")
     if wandb_project:
@@ -34,13 +33,11 @@ def main():
     model_cfg = cfg["model"]
     train_cfg = cfg.get("training", {})
 
-    # Tokenizer
     tok_name = model_cfg.pop("tokenizer_name", "meta-llama/Llama-2-7b-hf")
     tokenizer = AutoTokenizer.from_pretrained(tok_name)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    # Model
     tiny_config = TinyLlamaConfig(**model_cfg)
     model = TinyLlamaForCausalLM(tiny_config)
     model = model.to(torch.bfloat16)
@@ -48,26 +45,13 @@ def main():
     n_params = sum(p.numel() for p in model.parameters()) / 1e6
     print(f"Model: {n_params:.2f}M params | MLP type: {tiny_config.mlp_type} | Activation: {tiny_config.activation}")
 
-    # Data
     msl = model_cfg.get("max_position_embeddings", 512)
-    train_ds = build_dataset(
-        tokenizer,
-        max_seq_len=msl,
-        split="train",
-        max_samples=None
-    )
-    eval_ds = build_dataset(
-        tokenizer,
-        max_seq_len=msl,
-        split="validation",
-        max_samples=None
-    )
+    train_ds = build_dataset(tokenizer, max_seq_len=msl, split="train", max_samples=None)
+    eval_ds = build_dataset(tokenizer, max_seq_len=msl, split="validation", max_samples=None)
 
-    # Train
     trainer = create_trainer(model, tokenizer, cfg, train_ds, eval_ds)
     trainer.train()
 
-    # Save & push
     out = train_cfg.get("output_dir", "./out")
     trainer.save_model(out)
     if args.push or train_cfg.get("push_to_hub", False):