File size: 8,210 Bytes
fc3971e | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 | """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
|