From 2f2a0f156b49f52c35b9207afa01e4c909a9bf19 Mon Sep 17 00:00:00 2001 From: "chenzeren.1" Date: Tue, 29 Sep 2026 14:54:38 +0800 Subject: [PATCH 1/2] [dist] fix: place Ulysses inside CP on device mesh --- .agents/knowledge/constraints.md | 2 +- .../parallel/test_parallel_state_registry.py | 39 +++++++++++++++++++ veomni/distributed/parallel_state.py | 14 ++++--- 3 files changed, 48 insertions(+), 7 deletions(-) diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md index 6aded36641..df70451c4b 100644 --- a/.agents/knowledge/constraints.md +++ b/.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()`. diff --git a/tests/parallel/test_parallel_state_registry.py b/tests/parallel/test_parallel_state_registry.py index ebb1ce04ca..b72847247b 100644 --- a/tests/parallel/test_parallel_state_registry.py +++ b/tests/parallel/test_parallel_state_registry.py @@ -21,6 +21,7 @@ import pytest import torch.distributed as dist +from veomni.distributed import parallel_state as parallel_state_module from veomni.distributed.parallel_state import ( _init_parallel_state, clear_parallel_state, @@ -67,3 +68,41 @@ 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 test_mesh_places_ulysses_inside_cp_without_enabling_combined_parallelism(monkeypatch): + class RecordingMesh: + def __init__(self): + self.flat_views = {} + self.view = None + + def __getitem__(self, dims): + self.view = dims + return self + + def _flatten(self, mesh_dim_name): + self.flat_views[mesh_dim_name] = self.view + return self + + mesh = RecordingMesh() + mesh_args = {} + + def record_mesh(*, device_type, mesh_shape, mesh_dim_names): + mesh_args.update(device_type=device_type, mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names) + return mesh + + monkeypatch.setattr(parallel_state_module, "init_device_mesh", record_mesh) + + # The layout can be checked before ParallelState's current CP+Ulysses gate. + # Keep that gate intact: this test only pins the future USP mesh topology. + 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) + + assert mesh_args == { + "device_type": "cpu", + "mesh_shape": (1, 2, 2), + "mesh_dim_names": ("dp_shard", "cp", "ulysses"), + } + assert mesh.flat_views["dp_shard_sp"] == ("dp_shard", "cp", "ulysses") + assert mesh.flat_views["dp_sp"] == ("dp_shard", "cp", "ulysses") + assert mesh.flat_views["sp"] == ("cp", "ulysses") diff --git a/veomni/distributed/parallel_state.py b/veomni/distributed/parallel_state.py index ebe9b647b2..23a082f501 100644 --- a/veomni/distributed/parallel_state.py +++ b/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") From 482cc3690a7b6b07969ad8879c90cf9aecf0f58d Mon Sep 17 00:00:00 2001 From: "chenzeren.1" Date: Tue, 29 Sep 2026 18:14:57 +0800 Subject: [PATCH 2/2] [ci] test real CPU mesh rank layout for CP and Ulysses --- .../parallel/test_parallel_state_registry.py | 70 +++++++++---------- 1 file changed, 34 insertions(+), 36 deletions(-) diff --git a/tests/parallel/test_parallel_state_registry.py b/tests/parallel/test_parallel_state_registry.py index b72847247b..dd717c5e3c 100644 --- a/tests/parallel/test_parallel_state_registry.py +++ b/tests/parallel/test_parallel_state_registry.py @@ -12,14 +12,13 @@ # 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 ( @@ -70,39 +69,38 @@ def test_a_later_named_init_adopts_the_cached_anonymous_state(single_rank_group) assert get_parallel_state_by_name("base") is anonymous -def test_mesh_places_ulysses_inside_cp_without_enabling_combined_parallelism(monkeypatch): - class RecordingMesh: - def __init__(self): - self.flat_views = {} - self.view = None +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 __getitem__(self, dims): - self.view = dims - return self - - def _flatten(self, mesh_dim_name): - self.flat_views[mesh_dim_name] = self.view - return self + def capture_real_mesh(*args, **kwargs): + mesh = original_init_device_mesh(*args, **kwargs) + captured_mesh["mesh"] = mesh + return mesh - mesh = RecordingMesh() - mesh_args = {} + 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() - def record_mesh(*, device_type, mesh_shape, mesh_dim_names): - mesh_args.update(device_type=device_type, mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names) - return mesh - monkeypatch.setattr(parallel_state_module, "init_device_mesh", record_mesh) - - # The layout can be checked before ParallelState's current CP+Ulysses gate. - # Keep that gate intact: this test only pins the future USP mesh topology. - 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) - - assert mesh_args == { - "device_type": "cpu", - "mesh_shape": (1, 2, 2), - "mesh_dim_names": ("dp_shard", "cp", "ulysses"), - } - assert mesh.flat_views["dp_shard_sp"] == ("dp_shard", "cp", "ulysses") - assert mesh.flat_views["dp_sp"] == ("dp_shard", "cp", "ulysses") - assert mesh.flat_views["sp"] == ("cp", "ulysses") +@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)