Move MoE Implementation into src/, add Load Balancing Losses #192
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
WIP #159
Hey everyone :)
As mentioned in the issue, I think it's worth it to move the MoE code into source, so I'm working on it here. MoEs have become a common/standard architecture choice for LLMs now. With this PR, it would be possible to run MoEs via the pip installed nanotron -- no need to be inside examples/ folder. Also, auxiliary losses are essential for training MoEs at scale (stability, load balancing & max utilization!).
Points I've tried to make sure:
moe.py
file in models, some if statements inside the main model, the forwarding of additional aux loss value)Still TODO:
Questions and input needed would be on:
Please feel free to give feedback or request other changes! I want to make sure it's clean and correct.