[jinyouzhi]

Stage 1 of #9081: the guard (option 3), which @LiRunGuo and @Achyuthan-S agreed should come first.

Problem

The TP/DP groups in deepspeed.utils.groups are process-wide, and the first engine creates them. When a second engine is built in the same process:

• Different autotp_size: _init_tp_mesh_device returns early because _TENSOR_MODEL_PARALLEL_GROUP is already set. An engine that asks for TP4 after a TP2 engine silently becomes a TP2 model (reproducer in [RFC] Cleanup: move process-wide parallel-group globals to per-engine groups — follow-up to #8241 / #8851 #9081).
• No AutoTP and no mpu (pointed out by @Achyuthan-S on [RFC] Cleanup: move process-wide parallel-group globals to per-engine groups — follow-up to #8241 / #8851 #9081): this engine never reaches _init_tp_mesh_device. The earlier AutoTP engine left groups.mpu = groups, so the new engine takes the TP2 engine's data parallel group, which covers only part of the ranks. On master it then fails in _configure_distributed_model with AttributeError: 'NoneType' object has no attribute 'get_model_parallel_rank'.

Change

• groups._check_existing_tp_groups() compares the requested TP size, and the DP size when one is given, with the existing groups. It raises a RuntimeError that names both layouts and points to [RFC] Cleanup: move process-wide parallel-group globals to per-engine groups — follow-up to #8241 / #8851 #9081.
• _init_tp_mesh_device calls it before its early return. This covers the engine path, tp_model_init (init_utils.py) and TpTrainingManager (tp_manager.py).
• DeepSpeedEngine.__init__ calls it with size 1 for an engine without AutoTP that passes neither mpu nor mesh_param, when TP groups already exist.
• Engines with the same layout keep sharing the groups, as before.
• The AutoTP tutorial lists this as limitation 5.

Not covered

These are left for stage 2 of #9081, as the RFC lays out:

• An explicit mpu replaces groups.mpu. Every later reader of groups then sees the newest engine's groups. [rollout] In-memory relayout between sharded layouts (two-stage relay, pipelined execution) and ZeRO-3 training-state move #9074 relies on this: it passes an mpu to every engine. Making this an error would break that approach, and it is not a misconfiguration of the new engine.
• An engine without an mpu after one with a user mpu (the DP case @LiRunGuo reported on [RFC] Cleanup: move process-wide parallel-group globals to per-engine groups — follow-up to #8241 / #8851 #9081). The second engine takes the first one's DP group. Megatron-style code may rely on passing the mpu once ("For subsequent use, just query groups"), so a check here would be a behavior change. I would rather fix it with per-engine groups than turn it into an error.
• DeepSpeedExamples#1021, which clears the TP/DP globals before each initialize. The new engine then builds its own groups, so the check does not fire.

Duplicate check

I ran gh pr list --search "9081 in:body", and searched for _init_tp_mesh_device and AutoTP guard PRs. The only related open PR, #8801, changes the bwc.py fallbacks and is unrelated. On #9081 nobody else has taken the guard. @Achyuthan-S plans to move the stage3.py universal-checkpoint reader in a later stage.

Tests

Hardware: 4× Intel Arc Pro B70 (XPU, oneCCL), torch 2.14.0.dev20260804+xpu, on upstream master 5dc44e8.

CUDA_VISIBLE_DEVICES=0,1,2,3 pytest -q tests/unit/v1/autotp/test_autotp_multiple_models.py tests/unit/utils/test_groups.py
14 passed

pytest -q tests/unit/v1/autotp tests/unit/utils/test_groups.py tests/unit/v1/moe/test_autoep_autotp_runtime.py
36 passed, 84 skipped, 1 failed

• The skips are the AutoTP tests that call skip_on_device() on XPU, plus H100-only tests.
• The single failure was TestGroupCacheTeardown::test_nonmember_destroy_preserves_cached_groups: a worker was killed by signal 6 during teardown. That test does not reach the changed code. It passed when rerun alone, three times on this branch and once on master.
• New tests (TestAutoTPEnginesShareProcessGroups, 4 ranks):
  • TP2 then TP2 shares the groups;
  • TP2 then TP4 raises, where master silently built a TP2 engine;
  • TP2 then a non-AutoTP engine raises, where master failed with the AttributeError above.
  • Both raising cases fail on master.

Compatibility matrix: ZeRO stage 0, bf16 off, AutoTP 1/2/4, no offload, on XPU. Not run on CUDA. The check runs only at initialize and does not touch training steps, so there is no effect on numerics, throughput or memory.

This PR was prepared with AI assistance; the tests above were run on the hardware listed.