Commit 2ba1c7d
committed
Update base for Update on "[MoE][2/n]Move EP setup from trainer to config registry and add model_registry params"
## Summary
Remove the `apply_ep()` call from `Trainer.__post_init__` and move EP-aware token dispatcher configuration to model config time.
Previously, model configs always created a `LocalTokenDispatcher.Config` (EP=1 default), then `Trainer.__post_init__` would immediately replace it with the correct EP dispatcher via `apply_ep()`. This two-phase setup split EP concerns between the model config layer and the trainer, making it harder to reason about what dispatcher a model actually uses.
Now:
- `model_registry()` accepts `attn_backend` and `moe_comm_backend` params across all model families (llama3, llama4, deepseek_v3, gpt_oss, qwen3), so config registries can build the right dispatcher config upfront
- `apply_ep()` and `find_pad_multiple` imports are removed from `trainer.py`
- `ep_degree` is moved from `LocalTokenDispatcher.Config` (base) to `AllToAllTokenDispatcher.Config` and `DeepEPTokenDispatcher.Config` where it's actually used
- Stale comment referencing removed `ExpertSequenceParallel` class is updated
## Test Plan
Existing unit tests: `python -m pytest tests/unit_tests/test_expert_parallel.py -x`
[ghstack-poisoned]1 parent 27420b9 commit 2ba1c7d
0 file changed
0 commit comments