MoEs in src/ and proper load balancing losses

[haeggee]

Copy from PR, huggingface#192, issue huggingface#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:

• Not much more code in src (the moe.py file in models, some if statements inside the main model, the forwarding of additional aux loss value)

• With the config setup, a standard dense model is simply one with num_experts=1. For >1, the dropless MoE is used

• Working for all parallelisms without errors

• With latest commit, pipeline parallel also has correct logging: the load balancing losses are computed inside the MLPs and passed forward + summed cumulatively

• z-loss

Questions and input needed would be on:

• Is there a more elegant way to forward the aux losses?
• Inside the pipeline engine, it's currently hardcoded to check for "load_balancing_loss" for the backward pass (just like it was before with "loss"). Would you want to make this more generic, i.e., loop over all keys in the dictionary to store them? Note that the two losses are separate (and not added up before) for logging.

Please feel free to give feedback or request other changes! I want to make sure it's clean and correct.

For Swiss AI: need to create a conversion script for MoEs. Probably makes sense to add in another PR @ischlag ?

[haeggee]

@TJ-Solergibert suggested that we could make the MoE models a fully new model file (just like llama.py) and also have a new config, in order to keep things clean, separate, and this wouldn't break the current configs/conversion etc.

The downsides of this would be that we add a lot of copied code since the model definition is essentially the same, apart from which MLP definition to use and the extended forward of aux_losses. Any thoughts on this?

[TJ-Solergibert]

Hi! Yes, I suggest to keep the model apart with its own model definition and config. This is a common practice in HuggingFace transformers, check the mistral model comments. I think that having separate models will help us maintain the project in the future and will be something we'll do sooner or later once we add new models.

For the aux_loss, as it's just for logging purposes, we will try to reduce/gather it at the end of each iteration instead of sending them through all pipeline stages.

[ischlag]

hey @AleHD I believe you discussed it with @TJ-Solergibert, is this ready to merge?
@TJ-Solergibert I assigned you as the reviewer for this PR.