DeepSpeed
ae1dba54 - Fix Python 3.9 import-time TypeError in AutoEP ep_router (#8119)

Commit
3 days ago
Fix Python 3.9 import-time TypeError in AutoEP ep_router (#8119) ## What Fixes #8102. On Python 3.9, every `deepspeed.initialize()` call crashes at import time with: ``` TypeError: unsupported operand type(s) for |: 'type' and 'NoneType' ``` ## Root cause `deepspeed/moe/ep_router.py` uses PEP 604 unions in eagerly evaluated annotation positions (`num_expert_groups: int | None` and `num_limited_groups: int | None` in `TokenChoiceTopKRouter.__init__`, `expert_bias: torch.Tensor | None = None` in `forward`). Def-signature annotations are evaluated at function-definition time, and an AST scan of the whole package shows `ep_router.py` is the only module that uses this syntax without `from __future__ import annotations`; every sibling AutoEP module (`auto_ep.py`, `auto_ep_layer.py`, `auto_ep_config.py`, `ep_repack.py`, the presets) already has the future import, so their annotations are deferred and harmless. The crash escapes DeepSpeed's own safety net: `engine.py` wraps `from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer` (which imports `ep_router`) in `except ImportError` inside `_configure_distributed_model`, but the exception raised is `TypeError`, not `ImportError`, so it propagates out of every `deepspeed.initialize()` on Python 3.9. Note the issue body points at `auto_ep.py` / `auto_ep_presets/base.py`; those files already defer annotations, `ep_router.py` is the actual offender. `setup.py` still advertises Python 3.8+ support, so this restores importability rather than changing the syntax: the fix adds the same `from __future__ import annotations` line the rest of AutoEP uses. No runtime behavior change. ## Reproduction Verified on real Python 3.9.6: executing current-master `ep_router.py` raises exactly the reported `TypeError`; with the future import added it imports cleanly. ## Test Adds `TestPy39AnnotationSafety::test_autoep_import_chain_defers_pep604_annotations` to `tests/unit/v1/moe/test_autoep_unit.py`: it AST-scans every module in the AutoEP import chain and fails, naming the module and line numbers, if any of them evaluates a PEP 604 union at import time without deferring annotations. It is version-independent (fails on the bug even when CI runs 3.10+). Fails before this fix (`deepspeed.moe.ep_router ... lines [53, 54, 137]`), passes after. Full `test_autoep_unit.py` suite: 60 passed, 1 skipped. `pre-commit run` clean on both changed files. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Author
Parents
Loading