|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
|
import json |
|
|
import math |
|
|
from dataclasses import dataclass, asdict |
|
|
from pathlib import Path |
|
|
from typing import Any, Dict, List, Optional |
|
|
|
|
|
import torch |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class RunningStat: |
|
|
count: int = 0 |
|
|
sum: float = 0.0 |
|
|
sumsq: float = 0.0 |
|
|
min: Optional[float] = None |
|
|
max: Optional[float] = None |
|
|
zero_count: int = 0 |
|
|
nan_count: int = 0 |
|
|
inf_count: int = 0 |
|
|
|
|
|
def update_from_tensor(self, t: torch.Tensor): |
|
|
with torch.no_grad(): |
|
|
nan_mask = torch.isnan(t) |
|
|
inf_mask = torch.isinf(t) |
|
|
self.nan_count += int(nan_mask.sum().item()) |
|
|
self.inf_count += int(inf_mask.sum().item()) |
|
|
t = torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0) |
|
|
|
|
|
self.zero_count += int((t == 0).sum().item()) |
|
|
|
|
|
tf = t.float() |
|
|
self.sum += float(tf.sum().item()) |
|
|
self.sumsq += float((tf * tf).sum().item()) |
|
|
self.count += t.numel() |
|
|
|
|
|
t_min = float(t.min().item()) |
|
|
t_max = float(t.max().item()) |
|
|
if self.min is None or t_min < self.min: |
|
|
self.min = t_min |
|
|
if self.max is None or t_max > self.max: |
|
|
self.max = t_max |
|
|
|
|
|
@property |
|
|
def mean(self) -> Optional[float]: |
|
|
if self.count == 0: |
|
|
return None |
|
|
return self.sum / self.count |
|
|
|
|
|
@property |
|
|
def var(self) -> Optional[float]: |
|
|
if self.count == 0: |
|
|
return None |
|
|
m = self.mean |
|
|
return max(0.0, self.sumsq / self.count - (m * m)) |
|
|
|
|
|
@property |
|
|
def std(self) -> Optional[float]: |
|
|
v = self.var |
|
|
if v is None: |
|
|
return None |
|
|
return math.sqrt(v) |
|
|
|
|
|
def to_dict(self) -> Dict[str, Any]: |
|
|
d = asdict(self) |
|
|
d["mean"] = self.mean |
|
|
d["std"] = self.std |
|
|
return d |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class TokenRMSStat: |
|
|
count: int = 0 |
|
|
sum: float = 0.0 |
|
|
sumsq: float = 0.0 |
|
|
|
|
|
def update_from_tensor(self, t: torch.Tensor): |
|
|
with torch.no_grad(): |
|
|
if t.ndim == 1: |
|
|
feats = t.unsqueeze(0) |
|
|
else: |
|
|
feats = t.view(-1, t.shape[-1]) |
|
|
rms = feats.float().pow(2).mean(dim=-1).sqrt() |
|
|
rms = torch.nan_to_num(rms, nan=0.0, posinf=0.0, neginf=0.0) |
|
|
self.count += int(rms.numel()) |
|
|
self.sum += float(rms.sum().item()) |
|
|
self.sumsq += float((rms * rms).sum().item()) |
|
|
|
|
|
@property |
|
|
def mean(self) -> Optional[float]: |
|
|
if self.count == 0: |
|
|
return None |
|
|
return self.sum / self.count |
|
|
|
|
|
@property |
|
|
def var(self) -> Optional[float]: |
|
|
if self.count == 0: |
|
|
return None |
|
|
m = self.mean |
|
|
return max(0.0, self.sumsq / self.count - (m * m)) |
|
|
|
|
|
@property |
|
|
def std(self) -> Optional[float]: |
|
|
v = self.var |
|
|
if v is None: |
|
|
return None |
|
|
return math.sqrt(v) |
|
|
|
|
|
def to_dict(self) -> Dict[str, Any]: |
|
|
return { |
|
|
"count": self.count, |
|
|
"mean": self.mean, |
|
|
"std": self.std, |
|
|
} |
|
|
|
|
|
|
|
|
class ActivationMonitor: |
|
|
def __init__(self, use_tensorboard: bool = False, tb_dir: Optional[str] = None): |
|
|
self.stats: Dict[str, RunningStat] = {} |
|
|
self.token_rms: Dict[str, TokenRMSStat] = {} |
|
|
self.use_tensorboard = use_tensorboard |
|
|
self.tb = None |
|
|
self._global_step = 0 |
|
|
if self.use_tensorboard and tb_dir is not None: |
|
|
try: |
|
|
from torch.utils.tensorboard import SummaryWriter |
|
|
|
|
|
self.tb = SummaryWriter(log_dir=tb_dir) |
|
|
except Exception as e: |
|
|
print(f"TensorBoard not available: {e}") |
|
|
|
|
|
def _get_stat(self, name: str) -> RunningStat: |
|
|
if name not in self.stats: |
|
|
self.stats[name] = RunningStat() |
|
|
return self.stats[name] |
|
|
|
|
|
def _get_token_rms(self, name: str) -> TokenRMSStat: |
|
|
if name not in self.token_rms: |
|
|
self.token_rms[name] = TokenRMSStat() |
|
|
return self.token_rms[name] |
|
|
|
|
|
def hook(self, name: str): |
|
|
def _hook(module, inputs, output): |
|
|
with torch.no_grad(): |
|
|
t = output |
|
|
if isinstance(t, tuple): |
|
|
t = t[0] |
|
|
if not isinstance(t, torch.Tensor): |
|
|
return |
|
|
self._get_stat(name).update_from_tensor(t) |
|
|
self._get_token_rms(name).update_from_tensor(t) |
|
|
|
|
|
if self.tb is not None and (self._global_step % 10 == 0): |
|
|
rs = self.stats[name] |
|
|
tr = self.token_rms[name] |
|
|
if rs.count > 0: |
|
|
self.tb.add_scalar( |
|
|
f"{name}/mean", rs.mean, self._global_step |
|
|
) |
|
|
if rs.std is not None: |
|
|
self.tb.add_scalar( |
|
|
f"{name}/std", rs.std, self._global_step |
|
|
) |
|
|
self.tb.add_scalar( |
|
|
f"{name}/zero_frac", |
|
|
rs.zero_count / max(1, rs.count), |
|
|
self._global_step, |
|
|
) |
|
|
if tr.count > 0 and tr.mean is not None: |
|
|
self.tb.add_scalar( |
|
|
f"{name}/token_rms_mean", |
|
|
tr.mean, |
|
|
self._global_step, |
|
|
) |
|
|
return |
|
|
|
|
|
return _hook |
|
|
|
|
|
def step(self): |
|
|
self._global_step += 1 |
|
|
|
|
|
def close(self): |
|
|
if self.tb is not None: |
|
|
self.tb.flush() |
|
|
self.tb.close() |
|
|
|
|
|
def to_dict(self) -> Dict[str, Any]: |
|
|
out: Dict[str, Any] = {} |
|
|
for k in sorted(self.stats.keys()): |
|
|
out[k] = { |
|
|
"global": self.stats[k].to_dict(), |
|
|
"token_rms": self.token_rms[k].to_dict(), |
|
|
} |
|
|
return out |
|
|
|
|
|
|
|
|
def find_modules_to_hook( |
|
|
model: torch.nn.Module, patterns: List[str] |
|
|
) -> List[str]: |
|
|
names: List[str] = [] |
|
|
for name, _ in model.named_modules(): |
|
|
lname = name.lower() |
|
|
if not lname.startswith("model.layers."): |
|
|
continue |
|
|
for p in patterns: |
|
|
if p in lname: |
|
|
names.append(name) |
|
|
break |
|
|
return sorted(list(set(names))) |
|
|
|
|
|
|
|
|
def compute_attention_entropy( |
|
|
model: AutoModelForCausalLM, |
|
|
tok: AutoTokenizer, |
|
|
prompts: List[str], |
|
|
max_length: int, |
|
|
input_device: torch.device, |
|
|
) -> Dict[int, float]: |
|
|
prev = getattr(model.config, "output_attentions", False) |
|
|
model.config.output_attentions = True |
|
|
|
|
|
with torch.inference_mode(): |
|
|
enc = tok( |
|
|
prompts, |
|
|
return_tensors="pt", |
|
|
padding=True, |
|
|
truncation=True, |
|
|
max_length=max_length, |
|
|
) |
|
|
for k in enc: |
|
|
enc[k] = enc[k].to(input_device) |
|
|
out = model(**enc, output_attentions=True, use_cache=False) |
|
|
atts = out.attentions |
|
|
entropies: Dict[int, float] = {} |
|
|
for i, att in enumerate(atts): |
|
|
probs = att.float().clamp_min(1e-12) |
|
|
ent = -(probs * probs.log()).sum(dim=-1) |
|
|
ent_mean = float(ent.mean().item()) |
|
|
entropies[i] = ent_mean |
|
|
|
|
|
model.config.output_attentions = prev |
|
|
return entropies |
|
|
|
|
|
|
|
|
def load_prompts( |
|
|
prompts: Optional[str], prompts_file: Optional[str] |
|
|
) -> List[str]: |
|
|
lines: List[str] = [] |
|
|
if prompts_file: |
|
|
with open(prompts_file, "r", encoding="utf-8") as f: |
|
|
for line in f: |
|
|
s = line.strip("\n") |
|
|
if s: |
|
|
lines.append(s) |
|
|
if prompts: |
|
|
for s in prompts.split("\n"): |
|
|
s = s.strip() |
|
|
if s: |
|
|
lines.append(s) |
|
|
if not lines: |
|
|
lines = [ |
|
|
"Hello! Briefly introduce yourself.", |
|
|
"Explain the concept of attention in transformers.", |
|
|
"List three use cases for large language models.", |
|
|
] |
|
|
return lines |
|
|
|
|
|
|
|
|
def main(): |
|
|
ap = argparse.ArgumentParser( |
|
|
description="Activation statistics monitor for HF CausalLM models." |
|
|
) |
|
|
ap.add_argument("--model", type=str, required=True, help="Model path or HF ID.") |
|
|
ap.add_argument("--prompts", type=str) |
|
|
ap.add_argument("--prompts_file", type=str) |
|
|
ap.add_argument("--max_length", type=int, default=256) |
|
|
ap.add_argument("--batch_size", type=int, default=4) |
|
|
ap.add_argument( |
|
|
"--dtype", |
|
|
type=str, |
|
|
default="bfloat16", |
|
|
choices=["bfloat16", "float16", "float32"], |
|
|
) |
|
|
ap.add_argument("--device_map", type=str, default="auto") |
|
|
ap.add_argument( |
|
|
"--patterns", |
|
|
type=str, |
|
|
default=( |
|
|
"q_proj,k_proj,v_proj,o_proj,mlp.up_proj,mlp.gate_proj," |
|
|
"mlp.down_proj,layernorm,norm" |
|
|
), |
|
|
) |
|
|
ap.add_argument("--save_json", type=str) |
|
|
ap.add_argument("--tensorboard_dir", type=str) |
|
|
ap.add_argument("--attention_entropy", action="store_true") |
|
|
args = ap.parse_args() |
|
|
|
|
|
dtype_map = { |
|
|
"bfloat16": torch.bfloat16, |
|
|
"float16": torch.float16, |
|
|
"float32": torch.float32, |
|
|
} |
|
|
torch_dtype = dtype_map[args.dtype] |
|
|
|
|
|
print(f"Loading tokenizer/model: {args.model}") |
|
|
tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
args.model, |
|
|
torch_dtype=torch_dtype, |
|
|
trust_remote_code=True, |
|
|
device_map=args.device_map, |
|
|
) |
|
|
model.eval() |
|
|
|
|
|
embed_device = model.get_input_embeddings().weight.device |
|
|
print(f"Sending inputs to: {embed_device}") |
|
|
|
|
|
patterns = [p.strip().lower() for p in args.patterns.split(",") if p.strip()] |
|
|
to_hook = find_modules_to_hook(model, patterns) |
|
|
|
|
|
mon = ActivationMonitor( |
|
|
use_tensorboard=args.tensorboard_dir is not None, |
|
|
tb_dir=args.tensorboard_dir, |
|
|
) |
|
|
handles = [] |
|
|
for name, module in model.named_modules(): |
|
|
if name in to_hook: |
|
|
handles.append(module.register_forward_hook(mon.hook(name))) |
|
|
print(f"Registered hooks on {len(handles)} modules.") |
|
|
|
|
|
prompts = load_prompts(args.prompts, args.prompts_file) |
|
|
|
|
|
with torch.inference_mode(): |
|
|
i = 0 |
|
|
while i < len(prompts): |
|
|
batch_prompts = prompts[i : i + args.batch_size] |
|
|
i += args.batch_size |
|
|
enc = tok( |
|
|
batch_prompts, |
|
|
return_tensors="pt", |
|
|
padding=True, |
|
|
truncation=True, |
|
|
max_length=args.max_length, |
|
|
) |
|
|
for k in enc: |
|
|
enc[k] = enc[k].to(embed_device) |
|
|
_ = model(**enc, use_cache=False) |
|
|
mon.step() |
|
|
|
|
|
attn_entropy: Dict[int, float] = {} |
|
|
if args.attention_entropy: |
|
|
subset = prompts[: min(len(prompts), args.batch_size)] |
|
|
attn_entropy = compute_attention_entropy( |
|
|
model, tok, subset, args.max_length, embed_device |
|
|
) |
|
|
|
|
|
for h in handles: |
|
|
h.remove() |
|
|
mon.close() |
|
|
|
|
|
stats = mon.to_dict() |
|
|
if args.attention_entropy: |
|
|
stats["_attention_entropy"] = attn_entropy |
|
|
|
|
|
print("\nActivation summary (top 10 by token_rms mean):") |
|
|
ranked = sorted( |
|
|
[ |
|
|
(name, d["token_rms"]["mean"] or 0.0) |
|
|
for name, d in stats.items() |
|
|
if name != "_attention_entropy" |
|
|
], |
|
|
key=lambda x: x[1], |
|
|
reverse=True, |
|
|
)[:10] |
|
|
for name, rms_mean in ranked: |
|
|
g = stats[name]["global"] |
|
|
zero_frac = g.get("zero_count", 0) / max(1, g.get("count", 1)) |
|
|
print( |
|
|
f"- {name}: token_rms_mean={rms_mean:.4f}, " |
|
|
f"mean={g.get('mean'):.4f} std={g.get('std'):.4f} " |
|
|
f"min={g.get('min'):.4f} max={g.get('max'):.4f} " |
|
|
f"zero_frac={zero_frac:.4f}" |
|
|
) |
|
|
|
|
|
if args.save_json: |
|
|
out_path = Path(args.save_json) |
|
|
out_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
with open(out_path, "w") as f: |
|
|
json.dump(stats, f, indent=2) |
|
|
print(f"\nSaved stats JSON to: {out_path}") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
