EleutherAI / EleutherAI/bergson
Support gradient tracking for MoE models with fused-parameter experts/router
- Dominant language
- Python
- Stars
- 86
- Forks
- 29
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 37
Description
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...
We could also add support for VAEs etc.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in bergson/gradients.py at LayerAdapter.supported_modules (around lines 237-262) and review the documented limitations referenced in the issue. Compare the fused-parameter layouts of GPT-OSS, Mixtral, Qwen-MoE, and OLMoE, including per-token routing and sparsity, then define completion as correct expert and router attribution in per-example gradients for those models.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100