Bergson currently only tracks nn.Linear, HF Conv1D, and nn.Conv{1,2,3}d modules (see LayerAdapter.supported_modules in bergson/gradients.py:238, known limitations section of focs). Modern MoE models with fused-parameter experts/routers (bare nn.Parameters rather than nn.Linear layers — e.g. gpt-oss, Mixtral, Qwen-MoE, OLMoE in transformers 5.x) are silently skipped.
We should add gradient-collection support for these fused MoE layouts so experts and routers are attributed correctly.
Notes:
- LayerAdapter (bergson/gradients.py:237-262) is the extension point.
- Should cover GPT-OSS, Mixtral, Qwen-MoE, and OLMoE.-
- Account for per-token expert routing/sparsity when computing per-example gradients.
- Possibly the simplest solution involves letting the user manually specify bare 2D nn.Parameters, or find the 2D nn.Parameters...
Bergson currently only tracks nn.Linear, HF Conv1D, and nn.Conv{1,2,3}d modules (see LayerAdapter.supported_modules in bergson/gradients.py:238, known limitations section of focs). Modern MoE models with fused-parameter experts/routers (bare nn.Parameters rather than nn.Linear layers — e.g. gpt-oss, Mixtral, Qwen-MoE, OLMoE in transformers 5.x) are silently skipped.
We should add gradient-collection support for these fused MoE layouts so experts and routers are attributed correctly.
Notes: