Signed-off-by: Moshe Island <[email protected]>

Currently MoE uses Megatron-DeepSpeed APIs to get tensor-parallel info (rank,
world_size, group).

In order to enable MoE for PipelineModule, modify to use backward-compatible
methods that can access either Megatron, DeepSpeed Topology or Old Megatron
APIs.

Since MoE is not part of deepspeed runtime, move backward compatible methods
to deepseed.utils and modify imports as required.

Signed-off-by: Moshe Island <[email protected]>

Signed-off-by: Moshe Island <[email protected]>

Currently, only "total_loss" is displayed.
If model has additional losses (e.g. MoE loss), display them as well.
Similar to "total loss", additional losses are displayed for the full batch
after mean reduced across DP ranks.

Signed-off-by: Moshe Island <[email protected]>

Signed-off-by: Moshe Island <[email protected]>

Currently, when using no-drop tokens, we calculate locally the capacity and
then all-reduce(op=max) on world group.

This fails when using pipeline parallel (with micro batches), since different
stage workers are handling different model layers (or at warmup, where first
stage workers are processing while last stage workers are idle).

Fix it by running the all-reduce on the expert group.

Signed-off-by: Moshe Island <[email protected]>

This commit enhances expert group creation for both modes:
- DP + PP + EP
- DP + TP + PP + EP

Signed-off-by: Moshe Island <[email protected]>

When using MoE with MoE-TP disabled, use pipeline parallel group to max or sum
MoE gradients.

This also fixes the behavior for following configuration:
No pipeline, TP enabled, MoE TP disabled.

Signed-off-by: Moshe Island <[email protected]>

Signed-off-by: Moshe Island <[email protected]>