Move MoE Implementation into src/, add Load Balancing Losses#192

Open

haeggee wants to merge 16 commits intohuggingface:mainfrom swiss-ai:moe

+1,025-83

Commits on Apr 9, 2024

Implemented wandb entity configuration

AleHD
authored and
haeggee
committed

Commits on Apr 14, 2024

fix pipeline bug with moe load balancing loss; correct compute estimation includes moes
haeggee
committed

Commits on May 2, 2024

proper swiglu for mlp in expert parallel
haeggee
committed

Commits on Jun 5, 2024

Merge remote-tracking branch 'upstream/main' into moe
haeggee
committed

Commits on Jun 6, 2024

add nn.linear module in init parametrization (for moe router)
haeggee
committed

Commits on Jun 13, 2024

refactor: aux loss computed in forward of each mlp -> enables correct logs bc loss forwarded through layers
haeggee
committed

Commits on Jul 2, 2024

Commits on Jul 3, 2024