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