| """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() |
|
|
