Open xrsrke opened 10 months ago
APIs
from pipegoose.nn.expert_parallel import ExpertParallel, ExpertLoss parallel_context = ParallelContext.from_torch(expert_parallel_size=8) mlp = CustomExpert() router = CustomRouter() noise_policy = CustomNoisePolicy() loss_func = nn.CrossEntropy() model = ExpertParallel( model, expert=mlp, router=router, noise_policy=noise_policy, enable_tensor_parallelism=True, parallel_context=parallel_context, ).parallelize() loss_func = ExpertLoss(loss_func, aux_weight=0.1)
TODOs
[x] Top-1, Top-2 router
[x] ExpertParallel (turn a š¤ transformers to a MoE automatically)
ExpertParallel
transformers
[ ] Does expert embedding need to multiply its corresponding router probability?
[ ] Make ExpertParallel work with data parallelism
[ ] Optionally apply tensor parallelism to an expert layer
[ ] Make ExpertParallel work in pipeline parallelism
[ ] Make ExpertParallel work with ZeRO-1
[x] Loss function (include aux and z loss)
[ ] Move inputs to target expert device
Engineering Reading
MoE Reading
APIs
TODOs
[x] Top-1, Top-2 router
[x]
ExpertParallel
(turn a š¤transformers
to a MoE automatically)[ ] Does expert embedding need to multiply its corresponding router probability?
[ ] Make
ExpertParallel
work with data parallelism[ ] Optionally apply tensor parallelism to an expert layer
[ ] Make
ExpertParallel
work in pipeline parallelism[ ] Make
ExpertParallel
work with ZeRO-1[x] Loss function (include aux and z loss)
[ ] Move inputs to target expert device
Engineering Reading
MoE Reading