Spaces:
Running on L40S
Running on L40S
File size: 8,428 Bytes
9f818c5 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""
MoE Specialization Callback
============================
Monitors whether MoE experts are developing distinct, stable roles over training.
A well-trained MoE should have experts that specialize — each processing a different
kind of input — rather than a few generalist experts doing everything while the rest
idle.
Expert Co-activation Rate
-------------------------
If two experts frequently fire together on the same token (both in the top-K
selected), they are likely learning redundant representations. Ideally experts
specialize on non-overlapping token types, so co-activation should stay close
to the chance baseline of K/N (e.g. 8/128 ≈ 0.0625 for the 235B model).
For each layer and each unique expert pair (i, j), we compute:
CoAct(i, j) = N_{i,j} / N_i
where N_{i,j} = number of tokens where both i and j were selected, and N_i =
total tokens routed to expert i. We then summarize across all pairs as max and
mean. A rising mean_coact, especially well above the chance baseline, signals
that the router is collapsing onto a small correlated cluster of experts.
Buffer ownership
----------------
coactivation_counts is reset here (in compute_moe_coactivation_metrics).
Per-expert token counts are derived from coactivation_counts itself
(row_sum + col_sum) / (K-1), so this callback is fully independent of
ExpertHeatmap's reset cycle for total_tokens_per_expert.
"""
import torch
import wandb
from torch.distributed.tensor import DTensor, Partial
from cosmos_framework.callbacks.every_n import EveryN
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.trainer import ImaginaireTrainer
from cosmos_framework.utils import distributed
from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock
def _get_device_mesh(vfm: torch.nn.Module):
weight = vfm.language_model.model.layers[0].self_attn.q_proj.weight
return weight.device_mesh if isinstance(weight, DTensor) else None
def _allreduce_dtensor(t: torch.Tensor, device_mesh) -> torch.Tensor:
"""Sum-reduce a local tensor across all FSDP ranks and return the global tensor."""
return DTensor.from_local(
t,
device_mesh=device_mesh,
placements=[Partial()] * device_mesh.ndim,
).full_tensor()
def compute_moe_coactivation_metrics(vfm: torch.nn.Module) -> dict[str, dict]:
"""
Compute per-layer Expert Co-activation metrics for both towers.
For each unique expert pair (i < j) in the upper triangle of the N×N
coactivation matrix, computes:
CoAct(i, j) = N_{i,j} / N_i
where N_{i,j} is the count of tokens where both i and j were in the top-K,
and N_i is the total token count for expert i (the row expert, i.e. the
lower-indexed expert in the pair).
N_i is derived directly from the co-activation matrix rather than from
the shared total_tokens_per_expert buffer, so this metric is independent
of ExpertHeatmap's reset cycle. Each token routed to expert i contributes
to (K-1) co-activation pairs, so N_i = (row_sum_i + col_sum_i) / (K-1).
High co-activation relative to the chance baseline (K/N) indicates that
certain expert pairs are systematically selected together — a sign of
redundancy rather than specialization.
Returns a dict: tower -> {
"layer_indices": list[int] — actual model layer positions
"max_coact": Tensor[num_moe_layers] — worst pair per layer
"mean_coact": Tensor[num_moe_layers] — average over all pairs
"chance_baseline": float — K/N, same for all layers (reference)
}
"""
with torch.no_grad():
device_mesh = _get_device_mesh(vfm)
if device_mesh is None:
return {}
results: dict[str, dict] = {}
for tower in ["und", "gen"]:
layer_indices, max_coacts, mean_coacts, chance_baselines = [], [], [], []
num_layers = len(vfm.language_model.model.layers)
for layer_idx in range(num_layers):
layer = vfm.language_model.model.layers[layer_idx]
mlp = layer.mlp if tower == "und" else getattr(layer, "mlp_moe_gen", None)
if not isinstance(mlp, Qwen3VLMoeTextSparseMoeBlock):
continue
coact_counts = _allreduce_dtensor(mlp.get_coactivation_counts(reset=True), device_mesh) # [N, N]
n = mlp.num_experts
k = mlp.top_k
# Derive per-expert token counts directly from the co-activation
# matrix so we don't depend on ExpertHeatmap's reset cycle.
# Each token that routes to expert i contributes (K-1) entries
# across row i and column i of the upper-triangle matrix.
tokens_per_expert = (coact_counts.sum(dim=1) + coact_counts.sum(dim=0)).float() / (k - 1)
mask = torch.triu(torch.ones(n, n, dtype=torch.bool, device=coact_counts.device), diagonal=1)
# CoAct(i, j) = N_{i,j} / N_i — normalise by how often expert i fires overall.
denom = tokens_per_expert.unsqueeze(1).clamp(min=1) # [N, 1]
coact_rates = (coact_counts.float() / denom)[mask] # [N*(N-1)/2]
layer_indices.append(layer_idx)
max_coacts.append(coact_rates.max())
mean_coacts.append(coact_rates.mean())
# Chance baseline = probability two randomly-chosen top-K slots land on the
# same pair under uniform routing = K/N. Constant across layers and steps,
# logged once per tower as a reference line.
chance_baselines.append(k / n)
if layer_indices:
results[tower] = {
"layer_indices": layer_indices,
"max_coact": torch.stack(max_coacts),
"mean_coact": torch.stack(mean_coacts),
"chance_baseline": chance_baselines[0], # same value for all layers
}
return results
class MoESpecializationCallback(EveryN):
"""
Logs per-layer MoE specialization metrics to W&B every N training steps.
What it captures
----------------
Whether MoE experts are developing distinct routing identities:
Expert Co-activation (logged every N steps)
- mean_coact / max_coact per layer: how often expert pairs fire together
relative to the chance_baseline (K/N). Values well above the baseline
suggest the router is selecting a redundant cluster of experts rather
than a diverse set.
W&B layout
----------
moe_specialization/coact_chance_baseline/<tower> — flat reference (K/N)
moe_specialization/max_coact/<tower>/layer_NNN|mean|max
moe_specialization/mean_coact/<tower>/layer_NNN|mean|max
Args:
every_n (int): Logging interval in training steps.
"""
def __init__(self, every_n: int = 100):
super().__init__(every_n=every_n)
def every_n_impl(
self,
trainer: ImaginaireTrainer,
model: ImaginaireModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor],
loss: torch.Tensor,
iteration: int,
) -> None:
vfm = model.net
coact_results = compute_moe_coactivation_metrics(vfm)
if not (distributed.is_rank0() and wandb.run):
return
log_dict: dict[str, float] = {}
for tower, tower_metrics in coact_results.items():
layer_indices = tower_metrics.pop("layer_indices")
chance_baseline = tower_metrics.pop("chance_baseline")
log_dict[f"moe_specialization/coact_chance_baseline/{tower}"] = chance_baseline
for metric_name, values in tower_metrics.items():
for layer_idx, val in zip(layer_indices, values):
log_dict[f"moe_specialization/{metric_name}/{tower}/layer_{layer_idx:03d}"] = val.item()
log_dict[f"moe_specialization/{metric_name}/{tower}/mean"] = values.mean().item()
log_dict[f"moe_specialization/{metric_name}/{tower}/max"] = values.max().item()
wandb.log(log_dict, step=iteration)
|