[`distributed`] Add dense and expert device mesh (#48857)
* Add dense and expert device mesh views with a MeshManager
`initialize_distributed_mesh` now builds two named views of the same ranks:
`(pp, fsdp, tp)` for dense layers and `(pp, efsdp, ep)` for experts, both
keeping size-one axes so callers select dimensions by name. `MeshManager`
routes `ep`/`efsdp` lookups to the expert view and everything else to the
dense view. `DistributedConfig` gains `ep_size` (defaults to `tp_size` when
`enable_expert_parallel=True`) and `efsdp_size`, with size validation.
Model execution is unchanged: expert sharding and FSDP still use the `tp`
and `fsdp` axes, and loading rejects `ep_size != tp_size` until the
all-to-all dispatcher lands.
* cleaning
* clean
* update
* clean doc
* Use ep_size in the Kimi K2.5 expert parallel example
* Enable expert parallelism in the Mega MoE example
* warn once
* Add dense and expert device mesh views with a MeshManager
`initialize_distributed_mesh` now builds two named views of the same ranks:
`(pp, fsdp, tp)` for dense layers and `(pp, efsdp, ep)` for experts, both
keeping size-one axes so callers select dimensions by name. `MeshManager`
routes `ep`/`efsdp` lookups to the expert view and everything else to the
dense view. `DistributedConfig` gains `ep_size` (defaults to `tp_size` when
`enable_expert_parallel=True`) and `efsdp_size`, with size validation.
Model execution is unchanged: expert sharding and FSDP still use the `tp`
and `fsdp` axes, and loading rejects `ep_size != tp_size` until the
all-to-all dispatcher lands.
* cleaning
* update
* clean doc
* Update src/transformers/distributed/mixin.py
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
* style: ruff import block formatting in distributed config
* add more comments
* Add dense and expert device mesh views with a MeshManager
`initialize_distributed_mesh` now builds two named views of the same ranks:
`(pp, fsdp, tp)` for dense layers and `(pp, efsdp, ep)` for experts, both
keeping size-one axes so callers select dimensions by name. `MeshManager`
routes `ep`/`efsdp` lookups to the expert view and everything else to the
dense view. `DistributedConfig` gains `ep_size` (defaults to `tp_size` when
`enable_expert_parallel=True`) and `efsdp_size`, with size validation.
Model execution is unchanged: expert sharding and FSDP still use the `tp`
and `fsdp` axes, and loading rejects `ep_size != tp_size` until the
all-to-all dispatcher lands.
* cleaning
* clean
* update
* clean doc
* Use ep_size in the Kimi K2.5 expert parallel example
* Enable expert parallelism in the Mega MoE example
* warn once
* Add dense and expert device mesh views with a MeshManager
`initialize_distributed_mesh` now builds two named views of the same ranks:
`(pp, fsdp, tp)` for dense layers and `(pp, efsdp, ep)` for experts, both
keeping size-one axes so callers select dimensions by name. `MeshManager`
routes `ep`/`efsdp` lookups to the expert view and everything else to the
dense view. `DistributedConfig` gains `ep_size` (defaults to `tp_size` when
`enable_expert_parallel=True`) and `efsdp_size`, with size validation.
Model execution is unchanged: expert sharding and FSDP still use the `tp`
and `fsdp` axes, and loading rejects `ep_size != tp_size` until the
all-to-all dispatcher lands.
* cleaning
* update
* clean doc
* Update src/transformers/distributed/mixin.py
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
* style: ruff import block formatting in distributed config
* add more comments
* fix
* typo
* clearer
* clean
* Update src/transformers/distributed/utils.py
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
* Update src/transformers/distributed/configuration_utils.py
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
* linting
* fix comment
---------
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
