| """Attach / detach TENP collectors on SparseMoEMLP layers of a DiT.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from typing import Iterable |
|
|
| import torch |
|
|
| from inference.model.dit.dit_module import SparseMoEMLP |
|
|
| from .collector import ( |
| DEFAULT_PERTURB_RATIO, |
| DEFAULT_PERTURB_REPEATS, |
| DEFAULT_TOKEN_CACHE_PER_BIN, |
| LayerCollector, |
| ) |
| from .temporal import DEFAULT_NUM_BINS |
|
|
| COLLECT_DIR_ENV = "MAGI_TENP_COLLECT_DIR" |
| COLLECT_LAYERS_ENV = "MAGI_TENP_LAYERS" |
| COLLECT_BINS_ENV = "MAGI_TENP_NUM_BINS" |
| COLLECT_SENS_ENV = "MAGI_TENP_SENSITIVITY" |
| COLLECT_RHO_ENV = "MAGI_TENP_PERTURB_RATIO" |
| COLLECT_REPEAT_ENV = "MAGI_TENP_PERTURB_REPEATS" |
| COLLECT_GRAM_ENV = "MAGI_TENP_GRAM" |
| COLLECT_CACHE_ENV = "MAGI_TENP_TOKEN_CACHE_PER_BIN" |
|
|
|
|
| def iter_sparse_moe(model) -> Iterable[tuple[int, SparseMoEMLP]]: |
| layers = model.block.layers |
| for layer in layers: |
| mlp = getattr(layer, "mlp", None) |
| if isinstance(mlp, SparseMoEMLP): |
| yield int(mlp.layer_idx), mlp |
|
|
|
|
| def _env_int(name: str, default: int) -> int: |
| raw = os.getenv(name, "").strip() |
| if not raw: |
| return int(default) |
| return int(raw) |
|
|
|
|
| def _env_float(name: str, default: float) -> float: |
| raw = os.getenv(name, "").strip() |
| if not raw: |
| return float(default) |
| return float(raw) |
|
|
|
|
| def _env_bool(name: str, default: bool) -> bool: |
| raw = os.getenv(name, "").strip().lower() |
| if not raw: |
| return bool(default) |
| return raw not in ("0", "false", "no", "off") |
|
|
|
|
| def attach_collectors( |
| model, |
| layers: set[int] | None = None, |
| device: torch.device | None = None, |
| num_bins: int | None = None, |
| collect_sensitivity: bool | None = None, |
| perturb_ratio: float | None = None, |
| perturb_repeats: int | None = None, |
| collect_gram: bool | None = None, |
| token_cache_per_bin: int | None = None, |
| ) -> dict[int, LayerCollector]: |
| collectors: dict[int, LayerCollector] = {} |
| bins = DEFAULT_NUM_BINS if num_bins is None else int(num_bins) |
| sens = True if collect_sensitivity is None else bool(collect_sensitivity) |
| rho = DEFAULT_PERTURB_RATIO if perturb_ratio is None else float(perturb_ratio) |
| repeats = DEFAULT_PERTURB_REPEATS if perturb_repeats is None else int(perturb_repeats) |
| gram = True if collect_gram is None else bool(collect_gram) |
| cache_n = ( |
| DEFAULT_TOKEN_CACHE_PER_BIN |
| if token_cache_per_bin is None |
| else max(int(token_cache_per_bin), 0) |
| ) |
| for layer_idx, mlp in iter_sparse_moe(model): |
| if layers is not None and int(layer_idx) not in layers: |
| continue |
| d_int = int(mlp.experts[0].intermediate_size) |
| dev = device or mlp.experts[0].down_proj.weight.device |
| col = LayerCollector( |
| d_int, |
| mlp.num_experts, |
| torch.device(dev), |
| mlp=mlp, |
| num_bins=bins, |
| collect_sensitivity=sens, |
| perturb_ratio=rho, |
| perturb_repeats=repeats, |
| collect_gram=gram, |
| token_cache_per_bin=cache_n, |
| ) |
| mlp.set_mid_observer(col) |
| collectors[int(layer_idx)] = col |
| return collectors |
|
|
|
|
| def set_collectors_timestep( |
| model, |
| timestep, |
| step_index: int | None = None, |
| num_steps: int | None = None, |
| ) -> None: |
| from inference.model.dit.dit_module import set_sparse_moe_diffusion_timestep |
|
|
| set_sparse_moe_diffusion_timestep(model, timestep, step_index=step_index, num_steps=num_steps) |
| collectors = getattr(model, "_tenp_collectors", None) |
| if not collectors: |
| return |
| if torch.is_tensor(timestep): |
| t = float(timestep.detach().float().reshape(-1)[0].item()) |
| else: |
| t = float(timestep) |
| for col in collectors.values(): |
| col.current_timestep = t |
| col.current_step = None if step_index is None else int(step_index) |
| col.num_steps = None if num_steps is None else int(num_steps) |
|
|
|
|
| def detach_collectors(model) -> None: |
| for _layer_idx, mlp in iter_sparse_moe(model): |
| mlp.set_mid_observer(None) |
|
|
|
|
| def parse_layer_set(raw: str | None) -> set[int] | None: |
| if raw is None or not str(raw).strip(): |
| return None |
| return {int(x) for x in str(raw).split(",") if x.strip()} |
|
|
|
|
| def save_collectors(path: str, collectors: dict[int, LayerCollector]) -> None: |
| payload = { |
| "version": 2, |
| "num_bins": next(iter(collectors.values())).num_bins if collectors else DEFAULT_NUM_BINS, |
| "layers": {str(k): col.state_dict() for k, col in collectors.items()}, |
| } |
| torch.save(payload, path) |
|
|
|
|
| def load_collector_payload(path: str, map_location="cpu") -> dict: |
| return torch.load(path, map_location=map_location, weights_only=False) |
|
|
|
|
| def collector_token_summary(collectors: dict[int, LayerCollector]) -> dict: |
| payload = { |
| "layers": {str(k): col.state_dict() for k, col in collectors.items()}, |
| } |
| return summary_from_payload(payload) |
|
|
|
|
| def summary_from_payload(payload: dict) -> dict: |
| layers = [] |
| for layer_s in sorted(payload["layers"], key=lambda x: int(x)): |
| layer_sd = payload["layers"][layer_s] |
| experts = [] |
| audio_primary = False |
| for eid_s in sorted(layer_sd["stats"], key=lambda x: int(x)): |
| st = layer_sd["stats"][eid_s] |
| n_mod = [int(x) for x in st["n_tokens_mod"].tolist()] |
| total = int(sum(n_mod)) |
| p_a = (n_mod[1] / total) if total else 0.0 |
| if p_a > 0.5: |
| audio_primary = True |
| rec = { |
| "expert_id": int(eid_s), |
| "n_tokens": int(st["n_tokens"]), |
| "n_mod": {"V": n_mod[0], "A": n_mod[1], "T": n_mod[2]}, |
| "p_A": round(p_a, 4), |
| } |
| n_bin = st.get("n_tokens_bin") |
| if n_bin is not None: |
| rec["n_tokens_bin"] = [int(x) for x in n_bin.tolist()] |
| experts.append(rec) |
| layers.append( |
| { |
| "layer": int(layer_s), |
| "experts": experts, |
| "audio_primary_any": audio_primary, |
| } |
| ) |
| return { |
| "audio_primary_discount": False, |
| "note": "Audio-primary p_mid discount stays off. n_A in text-only infer is not real audio evidence.", |
| "layers": layers, |
| } |
|
|
|
|
| def maybe_attach_from_env(model) -> dict[int, LayerCollector] | None: |
| from .unstructured import maybe_attach_capture |
|
|
| maybe_attach_capture(model) |
| out_dir = os.getenv(COLLECT_DIR_ENV, "").strip() |
| if not out_dir: |
| return None |
| layers = parse_layer_set(os.getenv(COLLECT_LAYERS_ENV, "")) |
| collectors = attach_collectors( |
| model, |
| layers=layers, |
| num_bins=_env_int(COLLECT_BINS_ENV, DEFAULT_NUM_BINS), |
| collect_sensitivity=_env_bool(COLLECT_SENS_ENV, True), |
| perturb_ratio=_env_float(COLLECT_RHO_ENV, DEFAULT_PERTURB_RATIO), |
| perturb_repeats=_env_int(COLLECT_REPEAT_ENV, DEFAULT_PERTURB_REPEATS), |
| collect_gram=_env_bool(COLLECT_GRAM_ENV, False), |
| token_cache_per_bin=_env_int(COLLECT_CACHE_ENV, 0), |
| ) |
| model._tenp_collectors = collectors |
| model._tenp_collect_dir = out_dir |
| layer_ids = sorted(collectors) |
| print( |
| f"TENP collect dir={out_dir} layers={layer_ids or 'none'} " |
| f"bins={_env_int(COLLECT_BINS_ENV, DEFAULT_NUM_BINS)}" |
| ) |
| return collectors |
|
|
|
|
| def maybe_save_from_env(model) -> str | None: |
| from .unstructured import maybe_prune_captured |
|
|
| maybe_prune_captured(model) |
| collectors = getattr(model, "_tenp_collectors", None) |
| out_dir = getattr(model, "_tenp_collect_dir", "") or os.getenv(COLLECT_DIR_ENV, "").strip() |
| if not collectors or not out_dir: |
| return None |
| os.makedirs(out_dir, exist_ok=True) |
| stats_path = os.path.join(out_dir, "tenp_stats.pt") |
| save_collectors(stats_path, collectors) |
| summary = collector_token_summary(collectors) |
| summary_path = os.path.join(out_dir, "tenp_collect_summary.json") |
| with open(summary_path, "w", encoding="utf-8") as f: |
| json.dump(summary, f, indent=2) |
| f.write("\n") |
| detach_collectors(model) |
| print(f"TENP wrote {stats_path} and {summary_path}") |
| return stats_path |
|
|
