From 1aa250e7cbc9bc7a6b5b10f049ea79f2befa9943 Mon Sep 17 00:00:00 2001 From: kohsheent Date: Wed, 29 Jul 2026 18:00:45 +0530 Subject: [PATCH] feat: track MoE models with fused-parameter experts and routers Closes #318 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 discover_targets skips both without saying so. On gpt-oss-20b that leaves 5.8% of the parameters attributed: attention and lm_head only. --track_moe_experts attaches an ExpertLinear per expert projection (model.layers.0.mlp.experts.expert_3.gate_up_proj) and tracks the router in place. gpt-oss-20b reaches 97.2%, Mixtral-8x7B 99.7%. An expert only receives the tokens routed to it, which is ragged across examples, while per-example gradients need [N, S, ...]. Each expert's rows are padded into a rectangular grid and the expert GEMM runs there, so _compute_gradient is untouched: padding rows carry zero activations and receive zero output gradient, contributing nothing to the weight gradient, the bias gradient or any Hessian factor. Routers need a different hook. Every family scores its logits inside its own forward, so register_full_backward_hook reports grad_output[0] is None and the gradient never crosses the module boundary; the collector taps the logits tensor instead. Off by default. Tracking replaces the fused experts' forward with a per-expert loop, which measured 3-11x on collection in CPU eager (GPU unmeasured, and is where losing the fused grouped-matmul kernel will cost most). When tracking is off and a model has fused experts, the collector warns rather than skipping silently, since silence was the original problem. Expansion is reversible and leaves state_dict() untouched, so a model is unaffected once collection ends. Index size grows with layers * experts * 2, ~50 to ~1,600 modules on gpt-oss-20b. --projection_target global absorbs that entirely; --filter_modules narrows it. Tested against per-sample autograd backward passes on gpt-oss, Mixtral, Qwen3-MoE, OLMoE and DeepSeek-V3, spanning both weight orientations, gated and non-gated experts and both router styles, plus a locally declared has_gate=False module for the up_proj path. Also covered: expansion does not change the model's outputs and reverses exactly, works from a model or its base_model, EK-FAC factors use each expert's routed-row mask, and the opt-out wins over an already-expanded model. The CUDA bf16/fp16 tests are written but have not been run. Not covered: llama4 and longcat_flash fuse experts without the use_experts_implementation contract detection relies on; --attribute_tokens is rejected alongside fused experts, since under top-k routing one token feeds several experts and its gradient rows cannot line up with token positions; and --optimizer_state still derives no normalizers for fused expert parameters. --- README.md | 12 +- .../precompute_checkpoints.py | 1 + bergson/collection.py | 1 + bergson/collector/collector.py | 112 +++- bergson/config/config.py | 10 + bergson/gradients.py | 34 +- bergson/hessians/hessian_approximations.py | 2 + bergson/hessians/kfac.py | 4 +- bergson/hessians/shampoo.py | 4 +- bergson/hessians/tkfac.py | 4 +- bergson/huggingface.py | 6 + bergson/moe.py | 244 ++++++++ bergson/utils/batch_size.py | 1 + docs/index.rst | 1 + tests/test_moe.py | 537 ++++++++++++++++++ 15 files changed, 941 insertions(+), 32 deletions(-) create mode 100644 bergson/moe.py create mode 100644 tests/test_moe.py diff --git a/README.md b/README.md index cbe033b3..fdbaeb0d 100644 --- a/README.md +++ b/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 diff --git a/bergson/approx_unrolling/precompute_checkpoints.py b/bergson/approx_unrolling/precompute_checkpoints.py index a61ef7a3..a802ce3c 100644 --- a/bergson/approx_unrolling/precompute_checkpoints.py +++ b/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( diff --git a/bergson/collection.py b/bergson/collection.py index 0aaf712b..385ec4a5 100644 --- a/bergson/collection.py +++ b/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( diff --git a/bergson/collector/collector.py b/bergson/collector/collector.py index f0aa3b8f..95ee0a1f 100644 --- a/bergson/collector/collector.py +++ b/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: diff --git a/bergson/config/config.py b/bergson/config/config.py index 0dce422e..833d3b7a 100644 --- a/bergson/config/config.py +++ b/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.""" diff --git a/bergson/gradients.py b/bergson/gradients.py index 175cfb2d..26aa058d 100644 --- a/bergson/gradients.py +++ b/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) diff --git a/bergson/hessians/hessian_approximations.py b/bergson/hessians/hessian_approximations.py index cf4e8f28..2d3e8f7a 100644 --- a/bergson/hessians/hessian_approximations.py +++ b/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: diff --git a/bergson/hessians/kfac.py b/bergson/hessians/kfac.py index 7128295e..9fd6d9a4 100644 --- a/bergson/hessians/kfac.py +++ b/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] diff --git a/bergson/hessians/shampoo.py b/bergson/hessians/shampoo.py index 26d2112c..e73b9a9d 100644 --- a/bergson/hessians/shampoo.py +++ b/bergson/hessians/shampoo.py @@ -44,7 +44,7 @@ def setup(self) -> None: def forward_hook(self, module: nn.Module, a: Tensor) -> None: """Compute activation covariance: A^T @ A.""" - 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 @@ -64,7 +64,7 @@ def backward_hook(self, module: nn.Module, g: Tensor) -> None: name = assert_type(str, module._name) S_shampoo_po = self.S_shampoo_dict[name] A_shampoo_ki = self.A_shampoo_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] diff --git a/bergson/hessians/tkfac.py b/bergson/hessians/tkfac.py index bef150c1..79ad57e4 100644 --- a/bergson/hessians/tkfac.py +++ b/bergson/hessians/tkfac.py @@ -45,7 +45,7 @@ def setup(self) -> None: def forward_hook(self, module: nn.Module, a: Tensor) -> None: """Compute activation covariance: A^T @ A.""" - 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 @@ -65,7 +65,7 @@ def backward_hook(self, module: nn.Module, g: Tensor) -> None: name = assert_type(str, module._name) S_tcov_po = self.S_tcov_dict[name] A_tcov_ki = self.A_tcov_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] diff --git a/bergson/huggingface.py b/bergson/huggingface.py index fe1ad194..96f9edc9 100644 --- a/bergson/huggingface.py +++ b/bergson/huggingface.py @@ -43,6 +43,7 @@ def __init__( use_optimizer_state: bool = True, scale_by_lr: bool = True, track_order: bool = False, + track_moe_experts: bool = False, ): """ Args: @@ -62,6 +63,9 @@ def __init__( that step's learning rate. Set False for ``g / sqrt(v)``, which matches ``load_from_optimizer``. track_order: Whether to record the shuffled order of training data. + track_moe_experts: Track fused-parameter MoE experts. Off by + default: enabling it replaces the experts' forward for the whole + training run, which costs throughput. attention_cfgs: Information used to split matrix-valued parameters into per-head matrices before down projection. """ @@ -80,6 +84,7 @@ def __init__( self.use_optimizer_state = use_optimizer_state self.scale_by_lr = scale_by_lr self.order: list[dict] | None = [] if track_order else None + self.track_moe_experts = track_moe_experts self.mod_grads = {} self.batch_indices: Tensor | None = None @@ -127,6 +132,7 @@ def on_train_begin( target_modules=target_modules, attention_cfgs=self.attention_cfgs, dtype=self.torch_dtype, + track_moe_experts=self.track_moe_experts, ) self.grad_sizes = { name: math.prod(s) for name, s in self.collector.shapes().items() diff --git a/bergson/moe.py b/bergson/moe.py new file mode 100644 index 00000000..cb835d7a --- /dev/null +++ b/bergson/moe.py @@ -0,0 +1,244 @@ +"""Gradient tracking for MoE layers with fused expert parameters. + +In transformers 5.x an MoE layer stores its experts as one 3D ``nn.Parameter`` +and its router as a 2D one, so neither is an ``nn.Linear`` and the collectors +skip both. +:func:`expand_moe` attaches an :class:`ExpertLinear` per expert projection and +swaps in :func:`bergson_experts_forward`, which pads each expert's routed tokens +into ``[N, L, ...]`` so the collectors see their usual layout. Padding rows get +zero gradient, so gradient computation is unchanged. + +Detection tests for the attributes ``use_experts_implementation`` sets. +``llama4`` and ``longcat_flash`` fuse experts without it and are not covered. +""" + +import types + +import torch +import torch.nn as nn +from torch import Tensor + +from bergson.utils.utils import assert_type + +_EXPERTS_ATTRS = ("num_experts", "has_gate", "has_bias", "is_transposed") + +# An experts module only sees [N*S, hidden], so a pre-hook on the enclosing block +# leaves the batch size in _BATCH_ATTR. _EXPANSION_ATTR holds what expand_moe +# added, for restore_moe to undo. +_BATCH_ATTR = "_bergson_num_examples" +_EXPANSION_ATTR = "_bergson_expansion" +# A tracked router points back at its experts module, whose batch size it shares. +_ROUTER_ATTR = "_bergson_router_experts" + + +class ExpertLinear(nn.Module): + """One expert's slice of a fused MoE parameter, as a linear layer. + + ``weight`` and ``bias`` are unregistered views, so only ``named_modules()`` + grows. ``weight`` keeps the parent's orientation, which + ``LayerAdapter.weight_transposed`` reports. + """ + + def __init__(self, experts: nn.Module, weight_name: str, expert_idx: int): + super().__init__() + # Not setattr: registering the parent makes named_modules() recurse. + self.__dict__["_experts"] = experts + self.weight_name = weight_name + self.expert_idx = expert_idx + self.transposed = bool(experts.is_transposed) # type: ignore[attr-defined] + + rows, cols = self.weight.shape + self.in_features, self.out_features = ( + (rows, cols) if self.transposed else (cols, rows) + ) + + @property + def weight(self) -> Tensor: + return getattr(self._experts, self.weight_name)[self.expert_idx] + + @property + def bias(self) -> Tensor | None: + bias = getattr(self._experts, f"{self.weight_name}_bias", None) + return None if bias is None else bias[self.expert_idx] + + def forward(self, x: Tensor) -> Tensor: + weight = self.weight + out = x @ weight if self.transposed else x @ weight.mT + bias = self.bias + return out if bias is None else out + bias + + +def is_fused_experts(module: nn.Module) -> bool: + """Whether ``module`` is a fused-parameter MoE experts module.""" + down = getattr(module, "down_proj", None) + return ( + all(hasattr(module, attr) for attr in _EXPERTS_ATTRS) + and isinstance(down, nn.Parameter) + and down.ndim == 3 + ) + + +def is_fused_router(module: nn.Module, num_experts: int) -> bool: + """A bare-parameter router: 2D [num_experts, hidden] weight. Routers that are + already linear modules (Llama4's subclasses nn.Linear) are left alone, since + annotating one would overwrite its own attributes.""" + w = getattr(module, "weight", None) + return ( + isinstance(w, nn.Parameter) + and w.ndim == 2 + and w.shape[0] == num_experts + and not hasattr(module, "in_features") + ) + + +def tracked_router(module: nn.Module) -> nn.Module | None: + """The experts module a tracked router belongs to, or None.""" + return getattr(module, _ROUTER_ATTR, None) + + +def router_batch_size(router: nn.Module) -> int: + """Batch size recorded for a tracked router's block this forward.""" + return getattr(getattr(router, _ROUTER_ATTR), _BATCH_ATTR) + + +def weight_names(experts: nn.Module) -> tuple[str, str]: + return ("gate_up_proj" if experts.has_gate else "up_proj", "down_proj") # type: ignore[attr-defined] + + +def _grid(token_idx: Tensor, num_examples: int, seq_len: int): + """Row -> (example, column) in a padded grid, plus the [N, L] mask of real + rows. Rows arrive grouped by example, from a row-major torch.where.""" + example = token_idx // seq_len + counts = torch.bincount(example, minlength=num_examples) + starts = counts.cumsum(0) - counts + slot = torch.arange(token_idx.numel(), device=token_idx.device) - starts[example] + + # An expert with no tokens still needs one column, or its backward hook never + # fires and Builder comes up a key short. + mask = torch.zeros( + num_examples, + max(int(counts.max()), 1), + dtype=torch.bool, + device=token_idx.device, + ) + mask[example, slot] = True + return example, slot, mask + + +def bergson_experts_forward( + self: nn.Module, + hidden_states: Tensor, + top_k_index: Tensor, + top_k_weights: Tensor, +) -> Tensor: + """Per-expert MoE forward, matching the implementations in transformers but + padded so the collectors see [N, S, ...]. Experts with no tokens still run.""" + num_tokens = hidden_states.shape[0] + num_examples = getattr(self, _BATCH_ATTR) + assert ( + num_tokens % num_examples == 0 + ), f"{num_tokens} tokens do not divide into {num_examples} examples" + seq_len = num_tokens // num_examples + + top_k_index = top_k_index.reshape(num_tokens, -1) + top_k_weights = top_k_weights.reshape(num_tokens, -1) + up_name, down_name = weight_names(self) + out = torch.zeros_like(hidden_states) + + for expert_idx in range(self.num_experts): # type: ignore[attr-defined] + # Expert-parallel sentinels get zero routing weight, so never matching + # them is what we want. + token_idx, k_slot = torch.where(top_k_index == expert_idx) + example, slot, mask = _grid(token_idx, num_examples, seq_len) + + shims = getattr(self, f"expert_{expert_idx}") + for shim in shims.children(): + shim._row_mask = mask # collectors mask on this, not [N, S] + + a = hidden_states.new_zeros(*mask.shape, hidden_states.shape[-1]) + a[example, slot] = hidden_states[token_idx] + weights = top_k_weights.new_zeros(mask.shape) + weights[example, slot] = top_k_weights[token_idx, k_slot] + + h = getattr(shims, up_name)(a) + h = self._apply_gate(h) if self.has_gate else self.act_fn(h) # type: ignore[attr-defined] + h = getattr(shims, down_name)(h) + h = h * weights.unsqueeze(-1) + + out = out.index_add(0, token_idx, h[example, slot].to(out.dtype)) + + return out + + +def _record_batch_size(experts: nn.Module): + def hook(module: nn.Module, args: tuple): + x = args[0] + assert x.ndim == 3, f"{type(module).__name__} input is not [N, S, hidden]" + setattr(experts, _BATCH_ATTR, x.shape[0]) + + return hook + + +def expand_moe(model: nn.Module) -> list[str]: + """Expose fused MoE experts to collection, in place. Undo with restore_moe. + + Safe to call twice, and from either a model or its base_model: the undo state + hangs off the experts module, not off whichever root was passed. + """ + for name, experts in list(model.named_modules()): + if not is_fused_experts(experts) or hasattr(experts, _EXPANSION_ATTR): + continue + + containers = [] + for expert_idx in range(experts.num_experts): # type: ignore[attr-defined] + container = nn.Module() + for weight_name in weight_names(experts): + container.add_module( + weight_name, ExpertLinear(experts, weight_name, expert_idx) + ) + child = f"expert_{expert_idx}" + experts.add_module(child, container) + containers.append(child) + + setattr(experts, _BATCH_ATTR, 1) + experts.forward = types.MethodType(bergson_experts_forward, experts) + block = model.get_submodule(name.rpartition(".")[0]) + + # The router keeps its own F.linear; it only needs the metadata + # LayerAdapter reads, and a way to reach the batch size. + num_experts = int(experts.num_experts) # type: ignore[arg-type] + routers = [c for c in block.children() if is_fused_router(c, num_experts)] + for router in routers: + out_features, in_features = assert_type(Tensor, router.weight).shape + setattr(router, "out_features", int(out_features)) + setattr(router, "in_features", int(in_features)) + # Via __dict__: setattr would register the experts module as a child + # of the router, dragging its parameters into the router's path. + router.__dict__[_ROUTER_ATTR] = experts + + handle = block.register_forward_pre_hook(_record_batch_size(experts)) + setattr(experts, _EXPANSION_ATTR, (containers, routers, handle)) + + return [ + n + for n, m in model.named_modules() + if isinstance(m, ExpertLinear) or tracked_router(m) is not None + ] + + +def restore_moe(model: nn.Module) -> None: + """Undo expand_moe, from any root containing the experts.""" + for experts in list(model.modules()): + if not hasattr(experts, _EXPANSION_ATTR): + continue + + containers, routers, handle = getattr(experts, _EXPANSION_ATTR) + handle.remove() + for container in containers: + delattr(experts, container) + for router in routers: + for attr in ("out_features", "in_features", _ROUTER_ATTR): + delattr(router, attr) + del experts.forward + delattr(experts, _BATCH_ATTR) + delattr(experts, _EXPANSION_ATTR) diff --git a/bergson/utils/batch_size.py b/bergson/utils/batch_size.py index 063b1c82..a9cdb32b 100644 --- a/bergson/utils/batch_size.py +++ b/bergson/utils/batch_size.py @@ -109,6 +109,7 @@ def maybe_auto_batch_size( target_modules=target_modules, data=ds, # type: ignore skip_index=True, + track_moe_experts=cfg.track_moe_experts, ), starting_batch_size=cfg.token_batch_size, ) diff --git a/docs/index.rst b/docs/index.rst index fd7d7b5d..28573ac5 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -13,6 +13,7 @@ We provide options for analyzing models and datasets at any scale or level of gr * Parallelize Bergson operations across multiple GPUs or nodes. * Load gradients with or without their module-wise structure. * Split attention module gradients by head. +* Attribute Mixture-of-Experts models whose experts and router are fused ``nn.Parameter``\ s, per expert (opt-in via ``--track_moe_experts``). Installation ------------ diff --git a/tests/test_moe.py b/tests/test_moe.py new file mode 100644 index 00000000..77857e0c --- /dev/null +++ b/tests/test_moe.py @@ -0,0 +1,537 @@ +"""Gradient collection for MoE models with fused expert and router parameters. + +gpt-oss and Mixtral are the reference pair: transposed vs not, biased vs not, +interleaved vs concatenated gate. Models are built from configs in-process, so +the suite stays offline. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from datasets import Dataset +from transformers import ( + AutoModelForCausalLM, + DeepseekV3Config, + GptOssConfig, + MixtralConfig, + OlmoeConfig, + Qwen3MoeConfig, +) +from transformers.integrations import use_experts_implementation + +from bergson.collector.collector import HookCollectorBase +from bergson.collector.gradient_collectors import GradientCollector +from bergson.config import IndexConfig +from bergson.gradients import AdamNormalizer, GradientProcessor, LayerAdapter +from bergson.hessians.kfac import CovarianceCollector +from bergson.moe import ExpertLinear, expand_moe, restore_moe, tracked_router + +FAMILIES = ("gpt_oss", "mixtral") +NUM_EXPERTS = 4 +TOP_K = 2 +SEQ_LEN = 7 + +# include_bias is a no-op on Mixtral, which has no biased module at all, while +# every module gpt-oss tracks carries one. +BIAS_CASES = [("gpt_oss", False), ("gpt_oss", True), ("mixtral", False)] + +SHARED = dict( + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=64, + max_position_embeddings=64, + num_experts_per_tok=TOP_K, +) + +# DeepSeek-V3 brings a sigmoid grouped-top-k router with a score-correction bias +# and a shared expert; the other two vary the block wiring. +CONFIGS = { + "gpt_oss": lambda: GptOssConfig( + intermediate_size=16, num_local_experts=NUM_EXPERTS, head_dim=8, **SHARED + ), + "mixtral": lambda: MixtralConfig( + intermediate_size=16, num_local_experts=NUM_EXPERTS, **SHARED + ), + "deepseek_v3": lambda: DeepseekV3Config( + intermediate_size=32, + moe_intermediate_size=16, + n_routed_experts=NUM_EXPERTS, + n_group=1, + topk_group=1, + first_k_dense_replace=0, + n_shared_experts=1, + qk_rope_head_dim=8, + qk_nope_head_dim=8, + v_head_dim=8, + kv_lora_rank=8, + q_lora_rank=8, + **{**SHARED, "num_key_value_heads": 4}, + ), + "qwen3_moe": lambda: Qwen3MoeConfig( + intermediate_size=32, + moe_intermediate_size=16, + num_experts=NUM_EXPERTS, + head_dim=8, + **SHARED, + ), + "olmoe": lambda: OlmoeConfig( + intermediate_size=16, num_experts=NUM_EXPERTS, **SHARED + ), +} + + +def build_model(family: str) -> nn.Module: + torch.manual_seed(0) + model = AutoModelForCausalLM.from_config(CONFIGS[family](), dtype=torch.float32) + model.eval() + return model + + +def moe_names(target_info) -> list[str]: + """Expert and router modules among the discovered targets.""" + return [ + n + for n in target_info + if "experts.expert_" in n or n.rpartition(".")[2] in ("gate", "router") + ] + + +def make_collector(model, *, processor=None, track_moe_experts=True, **cfg_kwargs): + return GradientCollector( + model=model.base_model, + cfg=IndexConfig(run_path="/tmp/bergson-moe-test", **cfg_kwargs), + data=Dataset.from_dict({"input_ids": [[1] * SEQ_LEN]}), + processor=processor or GradientProcessor(), + skip_index=True, + track_moe_experts=track_moe_experts, + ) + + +def backward_pass(model, x: torch.Tensor) -> None: + model.zero_grad() + (model(input_ids=x).logits ** 2).sum().backward() + + +def autograd_gradient(base, name: str, include_bias: bool) -> torch.Tensor: + """``name``'s weight gradient from the last backward, read off the *fused* + parameter and oriented to [out, in]. An independent reference, not a + restatement of the collector's arithmetic.""" + layer = base.get_submodule(name) + if isinstance(layer, ExpertLinear): + grad = getattr(layer._experts, layer.weight_name).grad[layer.expert_idx] + if LayerAdapter.weight_transposed(layer): + grad = grad.T + fused_bias = getattr(layer._experts, f"{layer.weight_name}_bias", None) + if include_bias and fused_bias is not None: + grad = torch.cat([grad, fused_bias.grad[layer.expert_idx, :, None]], dim=1) + else: + grad = layer.weight.grad + if include_bias and getattr(layer, "bias", None) is not None: + grad = torch.cat([grad, layer.bias.grad[:, None]], dim=1) + return grad.flatten() + + +def assert_matches_autograd( + base, + names, + collected, + run_example, + *, + include_bias=False, + normalizers=None, + tol=None, +): + """Check each collected per-example gradient against a per-sample backward. + + ``run_example(i)`` must leave the gradients holding example ``i``'s alone. + With ``normalizers``, the reference is normalized as the collector would, + catching an expert paired with the wrong normalizer or orientation. + """ + atol, rtol = tol or (1e-5, 1e-4) + for example in range(len(next(iter(collected.values())))): + run_example(example) + for name in names: + expected = autograd_gradient(base, name, include_bias) + if normalizers is not None: + n = normalizers[name] + expected = n.normalize_weight( + expected.view(n.weight_avg_sq.shape).clone() + ).flatten() + torch.testing.assert_close( + collected[name][example].float(), + expected.float(), + atol=atol, + rtol=rtol, + msg=f"{name}, example {example}", + ) + assert max(g.abs().max() for g in collected.values()) > 0, "all-zero gradients" + + +def collect_and_compare(model, batch_size: int, include_bias: bool, tol=None): + model.requires_grad_(True) + x = torch.randint(0, 64, (batch_size, SEQ_LEN), device=model.device) + + collector = make_collector( + model, processor=GradientProcessor(include_bias=include_bias) + ) + names = moe_names(collector.target_info) + assert names, "no fused MoE experts or routers discovered" + + with collector: + backward_pass(model, x) + collected = {k: v.clone() for k, v in collector.mod_grads.items()} + + # Builder concatenates every module in shapes(), so none may be missing, + # including experts that happened to receive no tokens. + assert set(collected) == set(collector.shapes()) + + assert_matches_autograd( + model.base_model, + names, + collected, + lambda i: backward_pass(model, x[i : i + 1]), + include_bias=include_bias, + tol=tol, + ) + + +@pytest.mark.parametrize("batch_size", [1, 3]) +@pytest.mark.parametrize("family,include_bias", BIAS_CASES) +def test_per_example_gradients_match_autograd(family, batch_size, include_bias): + collect_and_compare(build_model(family), batch_size, include_bias) + + +@pytest.mark.parametrize("family", ["deepseek_v3", "qwen3_moe", "olmoe"]) +def test_other_families_match_autograd(family): + """Detection generalizes past the reference pair.""" + collect_and_compare(build_model(family), batch_size=2, include_bias=False) + + +@pytest.mark.parametrize("family", FAMILIES) +def test_expansion_covers_every_fused_parameter(family): + """Expansion adds exactly the fused expert and router weights.""" + base = build_model(family).base_model + + def tracked() -> int: + total = 0 + for name in HookCollectorBase.discover_targets(base): + layer = base.get_submodule(name) + total += getattr(layer, LayerAdapter.in_attr(layer)) * getattr( + layer, LayerAdapter.out_attr(layer) + ) + return total + + before = tracked() + expand_moe(base) + after = tracked() + + experts = sum(p.numel() for p in base.parameters() if p.ndim == 3) + routers = sum( + m.weight.numel() for m in base.modules() if tracked_router(m) is not None + ) + assert experts and routers + assert after - before == experts + routers + + total = sum(p.numel() for p in base.parameters()) + assert before / total < 0.5 < after / total + + +@pytest.mark.parametrize("family", FAMILIES) +def test_expansion_is_transparent_and_reversible(family): + """Expansion changes what is visible, not what the model computes.""" + model = build_model(family) + x = torch.randint(0, 64, (3, SEQ_LEN)) + with torch.no_grad(): + reference = model(input_ids=x).logits.clone() + + parameters = {n for n, _ in model.named_parameters()} + state_dict = set(model.state_dict()) + modules = set(dict(model.named_modules())) + + added = expand_moe(model) + assert len(added) == 2 * (2 * NUM_EXPERTS + 1) # two projections + a router + assert expand_moe(model) == added, "expansion should be idempotent" + + with torch.no_grad(): + torch.testing.assert_close(model(input_ids=x).logits, reference) + assert {n for n, _ in model.named_parameters()} == parameters + assert set(model.state_dict()) == state_dict + assert modules < set(dict(model.named_modules())) + + restore_moe(model) + assert set(dict(model.named_modules())) == modules + with torch.no_grad(): + torch.testing.assert_close(model(input_ids=x).logits, reference) + + +@pytest.mark.parametrize("family", FAMILIES) +def test_expansion_state_is_root_independent(family): + """Collection runs on base_model while callers hold model, so undo state + keyed to one root would leave the other thinking nothing was expanded.""" + model = build_model(family) + roots = ((model, model.base_model), (model.base_model, model)) + for expand_root, other_root in roots: + added = expand_moe(expand_root) + assert any(isinstance(m, ExpertLinear) for m in model.modules()) + # Names come back relative to the root asked, so compare counts. + assert len(expand_moe(other_root)) == len(added) + + restore_moe(other_root) + assert not any(isinstance(m, ExpertLinear) for m in model.modules()) + assert not any(tracked_router(m) is not None for m in model.modules()) + + +def test_experts_are_untracked_by_default_and_warn(): + """Tracking is opt-in, and skipping experts is announced, not silent.""" + model = build_model("gpt_oss") + with pytest.warns(UserWarning, match="fused MoE expert modules are not"): + collector = GradientCollector( + model=model.base_model, + cfg=IndexConfig(run_path="/tmp/bergson-moe-test"), + data=Dataset.from_dict({"input_ids": [[1] * SEQ_LEN]}), + processor=GradientProcessor(), + skip_index=True, + ) + assert moe_names(collector.target_info) == [] + + +def test_opt_out_is_authoritative_over_an_expanded_model(): + model = build_model("gpt_oss") + expand_moe(model.base_model) + collector = make_collector(model, track_moe_experts=False) + assert moe_names(collector.target_info) == [] + + +@pytest.mark.parametrize("family,include_bias", BIAS_CASES) +def test_projection_shapes_and_determinism(family, include_bias): + """Every MoE module yields a deterministic [N, p*p] block.""" + model = build_model(family) + p = 4 + x = torch.randint(0, 64, (2, SEQ_LEN)) + collector = make_collector( + model, processor=GradientProcessor(projection_dim=p, include_bias=include_bias) + ) + names = moe_names(collector.target_info) + + grads = [] + for _ in range(2): + with collector: + backward_pass(model, x) + grads.append({n: collector.mod_grads[n].clone() for n in names}) + + for name in names: + assert collector.shapes()[name] == torch.Size((p, p)) + assert grads[0][name].shape == (2, p * p) + torch.testing.assert_close(grads[0][name], grads[1][name]) + assert max(g.abs().max() for g in grads[0].values()) > 0 + + +def test_global_projection_absorbs_experts(): + """projection_target='global' sums all modules into one vector per example, + so expert tracking costs nothing extra in the index.""" + model = build_model("gpt_oss") + p = 16 + collector = make_collector( + model, processor=GradientProcessor(projection_dim=p, projection_target="global") + ) + assert moe_names(collector.target_info) + assert collector.shapes() == {"gradients": torch.Size((p,))} + + with collector: + backward_pass(model, torch.randint(0, 64, (3, SEQ_LEN))) + + grads = collector.mod_grads["gradients"] + assert grads.shape == (3, p) + assert torch.isfinite(grads).all() and grads.abs().max() > 0 + + +def test_shapes_match_collected_widths(): + """shapes() is what Builder sizes the index from. Run on gpt-oss, where a + bias column widens every gradient.""" + model = build_model("gpt_oss") + collector = make_collector(model, processor=GradientProcessor(include_bias=True)) + with collector: + backward_pass(model, torch.randint(0, 64, (2, SEQ_LEN))) + + for name, shape in collector.shapes().items(): + assert collector.mod_grads[name].shape == (2, math.prod(shape)), name + + +def test_attribute_tokens_rejected(): + """Under top-k routing one token feeds several experts, so per-token rows + cannot line up with token positions.""" + collector = make_collector(build_model("gpt_oss"), attribute_tokens=True) + with pytest.raises(ValueError, match="attribute_tokens is incompatible"): + collector.__enter__() + + +@pytest.mark.parametrize("family", FAMILIES) +def test_ekfac_covariance_over_experts(family, tmp_path): + """EK-FAC factors use each expert's routed-row mask, not the batch mask.""" + model = build_model(family) + x = torch.randint(0, 64, (2, SEQ_LEN)) + collector = CovarianceCollector( + model=model.base_model, + processor=GradientProcessor(), + dtype=torch.float32, + path=str(tmp_path), + track_moe_experts=True, + ) + names = moe_names(collector.target_info) + assert names, "no MoE modules discovered, the checks below would be vacuous" + + with collector.with_batch(torch.ones(x.shape, dtype=torch.bool)): + backward_pass(model, x) + + for name in names: + layer = model.base_model.get_submodule(name) + i = getattr(layer, LayerAdapter.in_attr(layer)) + o = getattr(layer, LayerAdapter.out_attr(layer)) + assert collector.A_cov_dict[name].shape[-1] == i, name + assert collector.S_cov_dict[name].shape[-1] == o, name + + +@pytest.mark.parametrize("family", FAMILIES) +def test_normalized_gradients_match_autograd(family): + """Per-expert normalizers reach the right expert in the right orientation.""" + model = build_model(family) + model.requires_grad_(True) + base = model.base_model + x = torch.randint(0, 64, (2, SEQ_LEN)) + + collector = make_collector(model) + names = moe_names(collector.target_info) + + torch.manual_seed(1) + normalizers = {} + for name in names: + layer = base.get_submodule(name) + shape = ( + getattr(layer, LayerAdapter.out_attr(layer)), + getattr(layer, LayerAdapter.in_attr(layer)), + ) + normalizers[name] = AdamNormalizer(torch.rand(shape) + 0.1) + collector.processor = GradientProcessor(normalizers=normalizers) + + with collector: + backward_pass(model, x) + collected = {k: v.clone() for k, v in collector.mod_grads.items()} + + assert_matches_autograd( + base, + names, + collected, + lambda i: backward_pass(model, x[i : i + 1]), + normalizers=normalizers, + ) + + +@use_experts_implementation(has_gate=False) +class _NonGatedExperts(nn.Module): + """A non-gated fused experts module, as nemotron_h declares one: up_proj + instead of gate_up_proj, a plain activation instead of gating. No released + model small enough to build offline has that pair. No forward, because the + collector always substitutes bergson_experts_forward.""" + + def __init__(self, config, hidden: int, intermediate: int): + super().__init__() + self.num_experts = NUM_EXPERTS + self.up_proj = nn.Parameter( + torch.randn(NUM_EXPERTS, intermediate, hidden) * 0.1 + ) + self.down_proj = nn.Parameter( + torch.randn(NUM_EXPERTS, hidden, intermediate) * 0.1 + ) + self.act_fn = nn.GELU() + + +class _NonGatedMoEBlock(nn.Module): + def __init__(self, hidden: int = 16, intermediate: int = 8): + super().__init__() + self.hidden = hidden + self.gate = nn.Module() + self.gate.num_experts = NUM_EXPERTS + self.gate.weight = nn.Parameter(torch.randn(NUM_EXPERTS, hidden) * 0.1) + self.gate.forward = self._route # type: ignore[method-assign] + self.experts = _NonGatedExperts( + SimpleNamespace(_experts_implementation="eager"), hidden, intermediate + ) + + @property + def device(self): + return next(self.parameters()).device + + def _route(self, hidden_states): + logits = nn.functional.linear(hidden_states, self.gate.weight) + weights, index = torch.topk(logits.softmax(-1), TOP_K, dim=-1) + return logits, weights, index + + def forward(self, x): + batch, seq, hidden = x.shape + flat = x.reshape(-1, hidden) + _, weights, index = self.gate(flat) + return self.experts(flat, index, weights).reshape(batch, seq, hidden) + + +def test_non_gated_experts_gradients_match_autograd(): + """The up_proj path is collected as correctly as the gated one.""" + torch.manual_seed(0) + block = _NonGatedMoEBlock() + x = torch.randn(3, SEQ_LEN, block.hidden) + + collector = GradientCollector( + model=block, + cfg=IndexConfig(run_path="/tmp/bergson-moe-test"), + data=Dataset.from_dict({"input_ids": [[1] * SEQ_LEN] * 3}), + processor=GradientProcessor(), + skip_index=True, + track_moe_experts=True, + ) + names = moe_names(collector.target_info) + assert len(names) == 2 * NUM_EXPERTS + 1 + assert all( + collector.model.get_submodule(n).weight_name == "up_proj" + for n in names + if n.endswith("up_proj") + ) + + def run(example: int) -> None: + block.zero_grad() + (block(x[example : example + 1]) ** 2).sum().backward() + + with collector: + block.zero_grad() + (block(x) ** 2).sum().backward() + collected = {k: v.clone() for k, v in collector.mod_grads.items()} + + assert_matches_autograd(block, names, collected, run) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("family", FAMILIES) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_gpu_gradients_and_transparency(family, dtype): + """The padded grid, its scatter and the index_add all run in the model's + dtype, so reduced precision is where a dtype mistake surfaces. On GPU the + unexpanded model also dispatches to a fused grouped-matmul kernel, so this + is the only check of fidelity to that kernel rather than to eager.""" + model = build_model(family).to(device="cuda", dtype=dtype) + x = torch.randint(0, 64, (3, SEQ_LEN), device="cuda") + + with torch.no_grad(): + reference = model(input_ids=x).logits.clone() + expand_moe(model) + with torch.no_grad(): + torch.testing.assert_close( + model(input_ids=x).logits, reference, atol=5e-2, rtol=5e-2 + ) + restore_moe(model) + + # Loose: bf16 carries ~3 decimal digits, and the collector reassociates the + # sum differently from a per-sample backward. + collect_and_compare(model, batch_size=3, include_bias=False, tol=(5e-2, 5e-2))