feat: track MoE models with fused-parameter experts and routers #318

README.md

@@ -262,13 +262,17 @@ Not all models are affected — run `bergson test_model_configuration` before en
 
 See `benchmarks/` for scripts to reproduce and generate benchmarks on your own hardware.
 
-# Known limitations
+# Mixture-of-Experts models
 
-## MoE fused-parameter experts are not attributed
+In `transformers` 5.x an MoE layer stores its experts as one fused 3D `nn.Parameter` and its router as a 2D one, so neither is an `nn.Linear` and both are skipped. On gpt-oss-20b that leaves 5.8% of the parameters attributed.
 
-Bergson only tracks `nn.Linear`, HF `Conv1D`, and `nn.Conv{1,2,3}d` modules. In modern MoE models (e.g. gpt-oss, Mixtral, Qwen-MoE, OLMoE in `transformers` 5.x), the experts and router are fused bare `nn.Parameter`s rather than `nn.Linear` layers, so they are silently skipped — only attention projections and `lm_head` are tracked (~1-2% of the model's parameters).
+`--track_moe_experts` exposes each expert projection as its own module (`model.layers.0.mlp.experts.expert_3.gate_up_proj`) and tracks the router in place, taking gpt-oss-20b to 97.2% and Mixtral-8x7B to 99.7%. Each expert's routed tokens are gathered into a padded grid first, so per-example gradients respect top-k sparsity. Layouts with one `nn.Linear` per expert are unaffected.
 
-Legacy MoE layouts that implement each expert as a separate `nn.Linear` are fully tracked.
+Tracking is off by default: it replaces the experts' forward with a per-expert loop, which cost 3-11x on collection in CPU eager measurements (GPU is unmeasured, and is where losing the fused grouped-matmul kernel matters most). Without the flag, bergson warns that a model's experts are being skipped rather than staying silent about it.
+
+Tracked modules grow by `layers * experts * 2`, from ~50 to ~1,600 on gpt-oss-20b, and the per-module index grows with them. `--projection_target global` sums all modules into one vector per example and so absorbs the growth; `--filter_modules '*.experts.*'` drops the experts again.
+
+Not covered: `llama4` and `longcat_flash` fuse expert parameters without the `use_experts_implementation` contract detection relies on. `--attribute_tokens` is rejected alongside fused experts, since under top-k routing one token contributes to several experts and its gradient rows cannot line up with token positions.
 
 # Development
 

bergson/approx_unrolling/precompute_checkpoints.py

@@ -120,6 +120,7 @@ def _lambda_worker(
         output_subdir=output_subdir,
         filter_modules=index_cfg.filter_modules,
         dtype=convert_precision_to_torch(hessian_cfg.hessian_dtype),
+        track_moe_experts=index_cfg.track_moe_experts,
     )
 
     computer = CollectorComputer(

bergson/collection.py

@@ -33,6 +33,7 @@ def collect_gradients(
         preprocess_cfg=preprocess_cfg or PreprocessConfig(),
         attention_cfgs=attention_cfgs or {},
         filter_modules=cfg.filter_modules,
+        track_moe_experts=cfg.track_moe_experts,
     )
 
     computer = CollectorComputer(

bergson/collector/collector.py

@@ -2,6 +2,7 @@
 import hashlib
 import math
 import os
+import warnings
 from abc import ABC, abstractmethod
 from contextlib import ContextDecorator, nullcontext
 from dataclasses import astuple, dataclass, field
@@ -36,6 +37,14 @@
     GradientProcessor,
     LayerAdapter,
 )
+from bergson.moe import (
+    ExpertLinear,
+    expand_moe,
+    is_fused_experts,
+    restore_moe,
+    router_batch_size,
+    tracked_router,
+)
 from bergson.utils.logger import get_logger
 from bergson.utils.peft import set_peft_enabled
 from bergson.utils.utils import assert_type
@@ -100,6 +109,10 @@ class HookCollectorBase(ContextDecorator, ABC):
     save_dtype: torch.dtype = torch.float32
     """Dtype gradients are cast to on the way out. Set in subclass ``setup()``."""
 
+    track_moe_experts: bool = False
+    """Expose fused-parameter MoE experts to collection, via
+    :func:`bergson.moe.expand_moe`. Off by default: see ``IndexConfig``."""
+
     logger = get_logger("HookCollectorBase", level="INFO")
 
     def __post_init__(
@@ -112,6 +125,15 @@ def __post_init__(
         self._fwd_hooks: list[RemovableHandle] = []
         self._bwd_hooks: list[RemovableHandle] = []
 
+        # Before discovery: expansion is what makes fused experts trackable at
+        # all. Restore on the way down too, so the flag wins over an expansion
+        # something else already did.
+        if self.track_moe_experts:
+            expand_moe(self.model)
+        else:
+            restore_moe(self.model)
+            self._warn_if_moe_untracked()
+
         # Discover target modules using the static method
         self.target_info = self.discover_targets(
             self.model,
@@ -130,6 +152,18 @@ def __post_init__(
         # Allow subclasses to perform custom initialization
         self.setup()
 
+    def _warn_if_moe_untracked(self) -> None:
+        """Skipping fused experts used to be silent; say so instead."""
+        skipped = sum(1 for m in self.model.modules() if is_fused_experts(m))
+        if skipped:
+            warnings.warn(
+                f"{skipped} fused MoE expert modules are not being attributed, "
+                f"leaving only attention and lm_head tracked. Pass "
+                f"--track_moe_experts to include them, at a cost in collection "
+                f"speed; see the README.",
+                stacklevel=2,
+            )
+
     @staticmethod
     def discover_targets(
         model: nn.Module,
@@ -158,7 +192,7 @@ def discover_targets(
         """
         target_info = {}
         for name, layer in model.named_modules():
-            if not isinstance(layer, LayerAdapter.supported_modules):
+            if not LayerAdapter.is_supported(layer):
                 continue
 
             if target_modules is not None and name not in target_modules:
@@ -432,8 +466,25 @@ def with_batch(self, collection_mask: Tensor | None = None) -> "HookCollectorBas
         self._current_collection_mask = collection_mask
         return self
 
+    def collection_mask(self, module: nn.Module) -> Tensor | None:
+        """The mask selecting ``module``'s gradient-carrying positions. A fused
+        MoE expert sees its own ``[N, L]`` grid of routed tokens, not the batch's
+        ``[N, S]`` positions."""
+        row_mask = getattr(module, "_row_mask", None)
+        return self._current_collection_mask if row_mask is None else row_mask
+
     def __enter__(self):
         """Register forward and backward hooks on all target modules."""
+        if self.attribute_tokens and any(
+            isinstance(self.model.get_submodule(n), ExpertLinear)
+            for n in self.target_info
+        ):
+            raise ValueError(
+                "attribute_tokens is incompatible with fused MoE experts: under "
+                "top-k routing one token feeds several experts, so its gradient "
+                "rows cannot line up with token positions."
+            )
+
         for name in self.target_info:
             layer = self.model.get_submodule(name)
 
@@ -442,30 +493,58 @@ def __enter__(self):
             layer._collect_bias = self.target_info[name][2]  # type: ignore[attr-defined]
 
             # Register hooks
-            fwd_hook = layer.register_forward_hook(self._process_input)
-            self._fwd_hooks.append(fwd_hook)
-
-            bwd_hook = layer.register_full_backward_hook(self._process_grad)
-            self._bwd_hooks.append(bwd_hook)
+            if tracked_router(layer) is not None:
+                self._fwd_hooks.append(layer.register_forward_hook(self._tap_router))
+            else:
+                self._fwd_hooks.append(layer.register_forward_hook(self._process_input))
+                self._bwd_hooks.append(
+                    layer.register_full_backward_hook(self._process_grad)
+                )
 
         return self
 
     def _process_input(self, module: nn.Module, inp: tuple, _):
         """Internal forward hook that extracts input and delegates to subclass."""
-        x = inp[0].detach()
+        x = self._unflatten(module, inp[0].detach())
         assert x.ndim == 3, f"Expected input of shape [N, S, I], got {x.shape}"
 
         self.forward_hook(module, x)
 
+    def _unflatten(self, module: nn.Module, x: Tensor) -> Tensor:
+        """A router sees hidden states already flattened to [N*S, D]."""
+        if x.ndim == 2 and tracked_router(module) is not None:
+            return x.view(router_batch_size(module), -1, x.shape[-1])
+        return x
+
+    def _tap_router(self, module: nn.Module, inp: tuple, out):
+        """A router scores its logits inside its own forward, so grad_output on a
+        module backward hook is None. Hook the logits tensor instead."""
+        self._process_input(module, inp, out)
+        o = getattr(module, LayerAdapter.out_attr(module))
+        logits = next(
+            t
+            for t in (out if isinstance(out, tuple) else (out,))
+            if isinstance(t, Tensor) and t.requires_grad and t.shape[-1] == o
+        )
+        logits.register_hook(lambda g: self._dispatch_grad(module, g))
+
     def _process_grad(self, module: nn.Module, _, grad_out):
         """Internal backward hook that extracts gradient and delegates to subclass."""
         # Sanity checks
-        assert isinstance(module, LayerAdapter.supported_modules), (
+        assert LayerAdapter.is_supported(module), (
             f"Expected a module of type {LayerAdapter.supported_modules}, "
             f"got {type(module)}"
         )
 
-        g = grad_out[0].detach()  # [N, S, O]
+        self._dispatch_grad(module, grad_out[0])
+
+    def _dispatch_grad(self, module: nn.Module, grad: Tensor):
+        """Hand an output gradient to the subclass as [N, S, O]."""
+        g = self._unflatten(module, grad.detach())
+
+        a = getattr(module, "_inputs", None)
+        if isinstance(a, Tensor) and a.dtype != g.dtype:
+            g = g.to(a.dtype)  # deepseek and glm4 routers upcast before F.linear
 
         self.backward_hook(module, g)
         if hasattr(module, "_inputs"):
@@ -475,10 +554,9 @@ def __exit__(self, exc_type, exc, tb):
         """Clean up hooks and allow subclass cleanup."""
         # Clean up temporary attributes
         for layer in self.model.modules():
-            if hasattr(layer, "_inputs"):
-                del layer._inputs
-            if hasattr(layer, "_name"):
-                del layer._name
+            for attr in ("_inputs", "_name", "_row_mask"):
+                if hasattr(layer, attr):
+                    delattr(layer, attr)
 
         # Remove all registered hooks
         for h in self._fwd_hooks:
@@ -631,7 +709,7 @@ def _compute_gradient(self, module: nn.Module, g: Float[Tensor, "N S O"]) -> Ten
                     P = self.double_sided_projection(name, P, g, p, o, i)
 
                 P = P.flatten(2)  # [N, S, grad_dim]
-                P = P[self._current_collection_mask]  # [total_valid, grad_dim]
+                P = P[self.collection_mask(module)]  # [total_valid, grad_dim]
             else:
                 P = g.mT @ a  # [N,O,S] @ [N,S,I] → [N,O,I]
 
@@ -680,7 +758,7 @@ def _compute_gradient(self, module: nn.Module, g: Float[Tensor, "N S O"]) -> Ten
                     # [N, S, O/p, 1] * [N, S, 1, I/q] → [N, S, O/p, I/q]
                     P = g.unsqueeze(-1) * a.unsqueeze(-2)
                 P = P.flatten(2)  # [N, S, grad_dim]
-                P = P[self._current_collection_mask]  # [total_valid, grad_dim]
+                P = P[self.collection_mask(module)]  # [total_valid, grad_dim]
             else:
                 if bias_grad is not None and p is not None:
                     P = self.double_sided_projection_with_bias(
@@ -717,7 +795,7 @@ def _compute_gradient(self, module: nn.Module, g: Float[Tensor, "N S O"]) -> Ten
                 )
                 if self.attribute_tokens:
                     P = P.flatten(2)  # [N, S, grad_dim]
-                    P = P[self._current_collection_mask]  # [total_valid, grad_dim]
+                    P = P[self.collection_mask(module)]  # [total_valid, grad_dim]
             else:
                 # a was already projected in forward if p is set;
                 # project g individually
@@ -733,7 +811,7 @@ def _compute_gradient(self, module: nn.Module, g: Float[Tensor, "N S O"]) -> Ten
                     if bias_grad is not None:
                         P = torch.cat([P, bias_grad.unsqueeze(-1)], dim=-1)
                     P = P.flatten(2)  # [N, S, grad_dim]
-                    P = P[self._current_collection_mask]  # [total_valid, grad_dim]
+                    P = P[self.collection_mask(module)]  # [total_valid, grad_dim]
                 else:
                     P = g.mT @ a  # [N, O/p, I/p]
                     if bias_grad is not None:

bergson/config/config.py

@@ -586,6 +586,16 @@ class IndexConfig(AttributionConfig, Serializable):
     """Whether to compute per-token gradients instead of per-example.
     Incompatible with reduce mode."""
 
+    track_moe_experts: bool = False
+    """Track MoE layers storing their experts as one fused ``nn.Parameter`` and
+    their router as a 2D one (gpt-oss, Mixtral, Qwen-MoE, OLMoE, DeepSeek-V3).
+    Without this only attention and ``lm_head`` are attributed, a few percent of
+    such a model, and bergson warns when it finds untracked experts.
+
+    Off by default: tracking replaces the experts' forward with a per-expert
+    loop, measured at 3-11x on collection, and grows the per-module index by
+    ``layers * experts * 2``. Incompatible with ``attribute_tokens``."""
+
     modules: list[str] = field(default_factory=list)
     """Modules to use for the query. If empty, all modules will be used."""
 

bergson/gradients.py

@@ -9,6 +9,8 @@
 from torch import Tensor
 from transformers.pytorch_utils import Conv1D as HFConv1D
 
+from bergson.moe import ExpertLinear, tracked_router
+
 NORMALIZER_TYPES: dict[str, type["Normalizer"]] = {}
 
 
@@ -240,36 +242,58 @@ def save(self, path: Path):
 
 
 class LayerAdapter:
-    supported_modules = (nn.Linear, HFConv1D, nn.Conv1d, nn.Conv2d, nn.Conv3d)
+    supported_modules = (
+        nn.Linear,
+        HFConv1D,
+        nn.Conv1d,
+        nn.Conv2d,
+        nn.Conv3d,
+        ExpertLinear,
+    )
+
+    @staticmethod
+    def is_supported(layer: nn.Module) -> bool:
+        """Whether gradients can be collected for ``layer``. Wider than an
+        isinstance check: a tracked MoE router keeps its own class."""
+        return (
+            isinstance(layer, LayerAdapter.supported_modules)
+            or tracked_router(layer) is not None
+        )
 
     @staticmethod
     def in_attr(layer: nn.Module) -> str:
         match layer:
-            case nn.Linear():
+            case nn.Linear() | ExpertLinear():
                 return "in_features"
             case HFConv1D():
                 return "nx"
             case nn.Conv1d() | nn.Conv2d() | nn.Conv3d():
                 return "in_channels"
+            case _ if tracked_router(layer) is not None:
+                return "in_features"
             case _:
                 raise ValueError(f"Unsupported layer type: {type(layer)}")
 
     @staticmethod
     def out_attr(layer: nn.Module) -> str:
         match layer:
-            case nn.Linear():
+            case nn.Linear() | ExpertLinear():
                 return "out_features"
             case HFConv1D():
                 return "nf"
             case nn.Conv1d() | nn.Conv2d() | nn.Conv3d():
                 return "out_channels"
+            case _ if tracked_router(layer) is not None:
+                return "out_features"
             case _:
                 raise ValueError(f"Unsupported layer type: {type(layer)}")
 
     @staticmethod
     def weight_transposed(layer: nn.Module) -> bool:
-        """Whether the layer stores its weight ``[in, out]`` (HF Conv1D)
-        rather than ``[out, in]``."""
+        """Whether the layer stores its weight ``[in, out]`` (HF Conv1D, and
+        MoE experts under ``is_transposed``) rather than ``[out, in]``."""
+        if isinstance(layer, ExpertLinear):
+            return layer.transposed
         return isinstance(layer, HFConv1D)
 
 

bergson/hessians/hessian_approximations.py

@@ -164,6 +164,7 @@ def hessian_worker(
             target_modules=target_modules,
             attention_cfgs=attention_cfgs,
             filter_modules=index_cfg.filter_modules,
+            track_moe_experts=index_cfg.track_moe_experts,
         )
         computer = CollectorComputer(
             model=model,  # type: ignore
@@ -253,6 +254,7 @@ def collect_hessians(
         "filter_modules": index_cfg.filter_modules,
         "processor": GradientProcessor(include_bias=index_cfg.include_bias),
         "dtype": hessian_dtype,
+        "track_moe_experts": index_cfg.track_moe_experts,
     }
     desc = f"Approximating Hessians with {hessian_cfg.method}"
     if ev_correction:

bergson/hessians/kfac.py

@@ -44,7 +44,7 @@ def forward_hook(self, module: nn.Module, a: Tensor) -> None:
         """Compute activation covariance: A^T @ A."""
         name = assert_type(str, module._name)
         A_cov_ki = self.A_cov_dict[name]
-        mask = self._current_collection_mask
+        mask = self.collection_mask(module)
         assert mask is not None, "Collection mask not set for forward hook."
 
         # a: [N, S, I], collection mask: [N, S] -> select gradient-carrying positions
@@ -75,7 +75,7 @@ def backward_hook(self, module: nn.Module, g: Tensor) -> None:
         """Compute gradient covariance: G^T @ G."""
         name = assert_type(str, module._name)
         S_cov_po = self.S_cov_dict[name]
-        mask = self._current_collection_mask
+        mask = self.collection_mask(module)
 
         # g: [N, S, O], mask: [N, S] -> select gradient-carrying positions
         g_bo = g[mask].to(self.dtype)  # [num_valid, O]