Skip to content

Commit 2ba1c7d

Browse files
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

File tree

    0 commit comments

    Comments
     (0)