DeepSpeed

084e5f6e - Validate the AutoSP SP/DP mesh when registering process groups (#8674)

Validate the AutoSP SP/DP mesh when registering process groups (#8674) ## What breaks AutoSP keeps its SP/DP process groups in a process-global registry (`deepspeed/compile/custom_ops/sp_dp_registry.py`). The compiled graph reads `sp_size()` and `get_group()` from it **at runtime**, inside the `autosp::all_to_all` op and the sharding offsets. ### 1. A different mesh silently reuses the stale groups `populate_registry()` returns early once any mesh is registered: ```python if GROUP_REGISTRY.get('is_reg', False): return ``` If a second AutoSP compilation in the same process (e.g. another engine) requests a different `sequence_parallel_size`, it keeps the first mesh without warning. The graph is then sharded and exchanged with the old SP size, while the engine's own data-parallel groups follow the new config. On 4 ranks, `populate_registry(2, 2)` followed by `populate_registry(4, 1)` gives, on master: ``` rank0: requested sp=4, registry sp_size()=2 group size=2 rank1: requested sp=4, registry sp_size()=2 group size=2 rank2: requested sp=4, registry sp_size()=2 group size=2 rank3: requested sp=4, registry sp_size()=2 group size=2 ``` Rebuilding the groups is not an option either, since graphs already compiled for the first mesh would then pick up the new SP size at runtime. So a changed mesh now raises a `RuntimeError`. Re-registering the same mesh still reuses the existing groups. ### 2. A mesh smaller than the world size fails later with a `KeyError` `apply_autosp()` only asserted `sp_size * dp_size <= world_size`. With a smaller mesh, ranks outside it have no SP group. On 4 ranks with `(sp, dp) = (2, 1)`: ``` rank0: group size=2 rank1: group size=2 rank2: KeyError: 1 rank3: KeyError: 1 ``` `populate_registry()` now requires `sp_size * dp_size == world_size` and raises a `ValueError` before creating any group. The check was moved from `apply_autosp()` into `populate_registry()` so that every caller gets it. The engine derives `dp_size = world_size // sp_size` in `extract_mesh_size()`, so case 2 is only reachable when `apply_autosp()` / `populate_registry()` are called directly. Case 1 is reachable through the engine. ## Tests In `tests/unit/v1/compile/test_compile_autosp.py`: - `TestAutoSPMeshValidation::test_rejects_changed_mesh`: an already-registered `(2, 2)` mesh rejects `(4, 1)` and keeps its SP/DP sizes. - `TestAutoSPMeshValidation::test_rejects_mesh_not_covering_world[2-1, 1-2]`: a mesh that does not cover a world of 4 raises before any process group is created. - `TestAutoSPMeshRegistry::test_same_mesh_is_reused` (4 ranks, real process groups): registering the same mesh twice keeps the same group objects. The rejection tests are plain pytest tests with `get_world_size` mocked, not `DistributedTest`s. When `pytest.raises` fails inside a `DistributedTest` worker it raises a `BaseException`, which hangs the worker pool instead of reporting the failure. | | master | this PR | |---|---|---| | `test_rejects_changed_mesh` | FAILED | PASSED | | `test_rejects_mesh_not_covering_world[2-1]` | FAILED | PASSED | | `test_rejects_mesh_not_covering_world[1-2]` | FAILED | PASSED | | `test_same_mesh_is_reused` | PASSED | PASSED | The rest of the IR-level tests in `test_compile_autosp.py` pass as well (26 passed). --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>