File size: 3,690 Bytes
d5147ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Per-step MoE expert call frequency (layer × loop). Rank0 JSONL, no collectives."""
from __future__ import annotations
import json
import os
from collections import defaultdict
from pathlib import Path
from typing import Optional
from torch import Tensor, nn
from models.transformer import TransformerBlock
_ACTIVE: Optional["ExpertFreqLogger"] = None
def get_active_expert_logger() -> Optional["ExpertFreqLogger"]:
return _ACTIVE
def set_active_expert_logger(logger: Optional["ExpertFreqLogger"]) -> None:
global _ACTIVE
_ACTIVE = logger
class ExpertFreqLogger:
"""Accumulate routing counts during forward; flush one JSONL line per step."""
def __init__(self, model: nn.Module, out_path: str | Path, *, every: int = 1) -> None:
self.every = max(1, int(every))
self.out_path = Path(out_path)
self.out_path.parent.mkdir(parents=True, exist_ok=True)
self._moe_tag_by_id: dict[int, str] = {}
self._sum: dict[tuple[str, str, int], float] = defaultdict(float)
self._cycle = ""
self._discover(model)
def _discover(self, model: nn.Module) -> None:
backbone = getattr(model, "model", model)
for level_name, short in (("H_level", "H"), ("L_level", "L")):
level = getattr(backbone, level_name, None)
if level is None:
continue
core = getattr(level, "core", None)
if core is None or not hasattr(core, "layers"):
continue
for i, block in enumerate(core.layers):
if not isinstance(block, TransformerBlock):
continue
if getattr(block, "ffn_type", "") == "moe":
self._moe_tag_by_id[id(block.mlp)] = f"{short}/L{i:02d}"
if self._moe_tag_by_id:
return
layers = getattr(backbone, "layers", None)
if layers is None:
return
for i, block in enumerate(layers):
if not isinstance(block, TransformerBlock):
continue
if getattr(block, "ffn_type", "") == "moe":
self._moe_tag_by_id[id(block.mlp)] = f"H/L{i:02d}"
def begin_step(self) -> None:
self._sum.clear()
self._cycle = ""
def set_cycle(self, level: str, idx: int) -> None:
self._cycle = f"c{level}{int(idx):02d}"
def record(self, moe_module: nn.Module, load_counts: Tensor) -> None:
tag = self._moe_tag_by_id.get(id(moe_module))
if tag is None:
return
counts = load_counts.detach()
if counts.device.type != "cpu":
counts = counts.float().cpu()
else:
counts = counts.float()
cyc = self._cycle
for eid in range(int(counts.numel())):
self._sum[(tag, cyc, eid)] += float(counts[eid].item())
def should_flush(self, step: int) -> bool:
return step > 0 and step % self.every == 0
def flush(self, step: int) -> None:
if not self.should_flush(step) or not self._sum:
self._sum.clear()
return
layers: dict[str, dict[str, list[float]]] = {}
n_exp = 0
for (tag, cyc, eid), v in self._sum.items():
n_exp = max(n_exp, eid + 1)
bucket = layers.setdefault(tag, {})
row = bucket.setdefault(cyc or "cH00", [])
while len(row) <= eid:
row.append(0.0)
row[eid] += float(v)
rec = {"step": int(step), "n_experts": n_exp, "layers": layers}
with open(self.out_path, "a", encoding="utf-8") as f:
f.write(json.dumps(rec, separators=(",", ":")) + "\n")
self._sum.clear()
|