[dist] fix: place Ulysses inside CP on device mesh

.agents/knowledge/constraints.md

@@ -48,7 +48,7 @@ Core entry points:
    - Gradient clipping: `veomni/distributed/fsdp2/clip_grad_norm.py` — handles DTensor grads and ExtraParallel param groups.
 
 6. **Device mesh initialization (`init_parallel_state_from_config()`)**
-   - Builds a global `DeviceMesh` with named dimensions: `pp`, `dp_replicate`, `dp_shard`, `ulysses`, `cp`, `tp` (each included only if size > 1).
+   - Builds a global `DeviceMesh` with named dimensions: `pp`, `dp_replicate`, `dp_shard`, `cp`, `ulysses`, `tp` (`dp_shard` is always included; other dimensions only when size > 1). Ulysses follows CP so its ranks are adjacent when both axes are used.
    - Flattens subviews for common usage: `dp` (all data-parallel), `sp` (ulysses+cp), `dp_shard_sp` (FSDP shard × SP), `dp_sp` (for loss/grad sync across SP+DP).
    - For each ExtraParallel name (e.g. `ep`), builds a `[para_size × para_fsdp_size]` submesh via `init_para_mesh_matrix()`.
 

tests/parallel/test_parallel_state_registry.py

@@ -12,15 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""``_init_parallel_state`` registry behaviour, including ``name=None``.
+"""``_init_parallel_state`` registry and CPU/Gloo mesh behavior."""
 
-Single-rank gloo, so this runs anywhere. The registry, the topology cache and
-the ambient state are process globals, hence the teardown.
-"""
+from datetime import timedelta
 
 import pytest
 import torch.distributed as dist
+import torch.multiprocessing as mp
 
+from veomni.distributed import parallel_state as parallel_state_module
 from veomni.distributed.parallel_state import (
     _init_parallel_state,
     clear_parallel_state,
@@ -67,3 +67,40 @@ def test_a_later_named_init_adopts_the_cached_anonymous_state(single_rank_group)
 
     assert named is anonymous
     assert get_parallel_state_by_name("base") is anonymous
+
+
+def _check_mesh_rank_layout(rank, rendezvous):
+    dist.init_process_group(
+        "gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=4, timeout=timedelta(seconds=60)
+    )
+    original_init_device_mesh = parallel_state_module.init_device_mesh
+    captured_mesh = {}
+
+    def capture_real_mesh(*args, **kwargs):
+        mesh = original_init_device_mesh(*args, **kwargs)
+        captured_mesh["mesh"] = mesh
+        return mesh
+
+    parallel_state_module.init_device_mesh = capture_real_mesh
+    try:
+        # The real mesh and its flattened views are built before the existing
+        # CP+Ulysses guard rejects this unsupported configuration.
+        with pytest.raises(NotImplementedError, match="cannot be combined with Ulysses yet"):
+            _init_parallel_state(cp_size=2, ulysses_size=2, device_type="cpu", name=None)
+
+        mesh = captured_mesh["mesh"]
+        assert mesh.mesh_dim_names == ("dp_shard", "cp", "ulysses")
+        assert mesh.mesh.tolist() == [[[0, 1], [2, 3]]]
+        assert dist.get_process_group_ranks(mesh.get_group("ulysses")) == [rank // 2 * 2, rank // 2 * 2 + 1]
+        assert dist.get_process_group_ranks(mesh.get_group("cp")) == [rank % 2, rank % 2 + 2]
+        for name in ("sp", "dp_sp", "dp_shard_sp"):
+            assert dist.get_process_group_ranks(mesh.get_group(name)) == [0, 1, 2, 3]
+    finally:
+        parallel_state_module.init_device_mesh = original_init_device_mesh
+        dist.destroy_process_group()
+        clear_parallel_state()
+
+
+@pytest.mark.skipif(not dist.is_gloo_available(), reason="requires the CPU Gloo backend")
+def test_mesh_places_ulysses_inside_cp_without_enabling_combined_parallelism(tmp_path):
+    mp.spawn(_check_mesh_rank_layout, args=(str(tmp_path / "rendezvous"),), nprocs=4, join=True)

veomni/distributed/parallel_state.py

@@ -609,11 +609,13 @@ def _init_parallel_state(
     extra_parallel_fsdp_device_mesh = {f"{para_name}": None for para_name in extra_parallel_names}
     extra_parallel_flat_device_mesh = {f"{para_name}": None for para_name in extra_parallel_names}
 
+    # With TP disabled, the last mesh dimension has stride-1 ranks. Keep
+    # Ulysses inside CP so its groups can stay within a node when CP spans nodes.
     mesh_shape = []
     mesh_dim_names = []
     for d, dim_name in zip(
-        [pp_size, dp_replicate_size, dp_shard_size, ulysses_size, cp_size, tp_size],
-        ["pp", "dp_replicate", "dp_shard", "ulysses", "cp", "tp"],
+        [pp_size, dp_replicate_size, dp_shard_size, cp_size, ulysses_size, tp_size],
+        ["pp", "dp_replicate", "dp_shard", "cp", "ulysses", "tp"],
     ):
         if d > 1 or dim_name in ["dp_shard"]:
             mesh_shape.append(d)
@@ -641,14 +643,14 @@ def _init_parallel_state(
         dp_mesh_dim_names.append("dp_shard")
         dp_shard_sp_mesh_dim_names.append("dp_shard")
         dp_sp_mesh_dim_names.append("dp_shard")
-    if ulysses_size > 1:
-        dp_shard_sp_mesh_dim_names.append("ulysses")
-        sp_mesh_dim_names.append("ulysses")
-        dp_sp_mesh_dim_names.append("ulysses")
     if cp_size > 1:
         dp_shard_sp_mesh_dim_names.append("cp")
         sp_mesh_dim_names.append("cp")
         dp_sp_mesh_dim_names.append("cp")
+    if ulysses_size > 1:
+        dp_shard_sp_mesh_dim_names.append("ulysses")
+        sp_mesh_dim_names.append("ulysses")
+        dp_sp_mesh_dim_names.append("ulysses")
 
     if dp_mesh_dim_names != []:
         device_mesh[tuple(dp_mesh_dim_names)]._flatten(mesh_dim_name="dp")