diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index acc4bafc6..cc596e817 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -171,7 +171,9 @@ def _sync_srt_tp_group() -> None: published `srt` config cannot answer: `gpu_worker.py` publishes a dummy carrying *this* package's `tp_size`, which a sequence-parallel launch sets to 1 while the group lent here is as wide as the world. So the widths are - stamped alongside the group, as `srt.initialize_model_parallel` does. + permanently overridden alongside the group -- this runs with no `srt` + config published at all, which is exactly why it cannot go through + `RuntimeContext.override` (it requires one). Only tensor parallelism folds this way, so every other dimension is one. """ @@ -183,7 +185,7 @@ def _sync_srt_tp_group() -> None: if srt_parallel_state._ATTN_TP is None: srt_parallel_state._ATTN_TP = _TP if srt_parallel_state._ATTN_TP is _TP: - get_parallel().stamp_derived_widths( + get_parallel().override_permanently( **derive_parallel_widths( tp_size=_TP.world_size, attn_cp_size=1, @@ -192,7 +194,7 @@ def _sync_srt_tp_group() -> None: moe_dp_size=1, dcp_size=1, dcp_enabled=False, - ) + ), ) diff --git a/python/sglang/srt/disaggregation/encoder/server.py b/python/sglang/srt/disaggregation/encoder/server.py index f31b466ba..68877a0ed 100644 --- a/python/sglang/srt/disaggregation/encoder/server.py +++ b/python/sglang/srt/disaggregation/encoder/server.py @@ -590,7 +590,10 @@ class MMEncoder: distributed_init_method=dist_init_method, local_rank=rank, ) - initialize_model_parallel(tensor_model_parallel_size=get_parallel().tp_size) + initialize_model_parallel( + tensor_model_parallel_size=get_parallel().tp_size, + attention_context_model_parallel_size=get_parallel().attn_cp_size, + ) initialize_dp_attention(server_args, self.model_config) self.model = load_model( diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 3a8a74022..64e0ec26b 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -2832,8 +2832,6 @@ def initialize_model_parallel( group_name="self_pp", ) - get_parallel().stamp_derived_widths(**derived_widths) - def create_custom_parallel_group( group_ranks: List[int], backend: str = "gloo" diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 62a981f3c..71a937bd0 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -71,7 +71,7 @@ def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int): global _ATTN_DP_SIZE, _ATTN_DP_RANK _ATTN_DP_SIZE = new_dp_size _ATTN_DP_RANK = new_dp_rank - get_parallel().stamp_derived_widths(attn_dp_size=new_dp_size) + get_parallel().override_permanently(attn_dp_size=new_dp_size) get_flags().dp.use_world_group_for_gather = True logger.debug( "[Elastic EP] dp_attention switched to WORLD: dp_size=%d dp_rank=%d", @@ -350,7 +350,8 @@ def compute_dp_attention_world_info( """This rank's place in the attention topology, plus the widths it sits in. The widths come from `derive_attention_widths`; what this adds is the two - ranks, which are per-process and so are not part of the stamped set. + ranks, which are per-process and so are not among the widths + `override_permanently` records. """ attn_dp_size, attn_tp_size = derive_attention_widths( tp_size=tp_size, @@ -391,7 +392,7 @@ def initialize_dp_attention( _, _, _ATTN_DP_RANK, _ATTN_DP_SIZE = compute_dp_attention_world_info( enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size ) - get_parallel().stamp_derived_widths(attn_dp_size=_ATTN_DP_SIZE) + get_parallel().override_permanently(attn_dp_size=_ATTN_DP_SIZE) if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size: _ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index b611988e2..559989ee8 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -164,7 +164,7 @@ def derive_parallel_widths( `world_size` is not among them: it is not a quotient, and `get_world_size()` answers with the live WORLD group, which stays right through an elastic - scale-up that a stamp taken at group build would not survive. + scale-up that a value fixed at group build would not survive. """ return { "attn_dp_size": attn_dp_size, @@ -268,7 +268,7 @@ class ParallelContext: def __init__(self): self._overrides = {} self._config = None # parallel config bag, wired at publish - self._derived = {} # widths stamped when the groups are built + self._derived = {} # widths overridden permanently, as the groups are built def __getattr__(self, name): if name.startswith("_"): @@ -291,14 +291,17 @@ class ParallelContext: overrides = self._overrides return overrides[name] if name in overrides else getter() - def stamp_derived_widths(self, **widths) -> None: - """Record the widths derived from the leaves, as the groups are built. + def override_permanently(self, **widths) -> None: + """Permanently correct a derived width the published bag can't answer + or no longer answers correctly -- not `RuntimeContext.override`, + because a derived width is not a resolved config leaf and this must + work with no config published at all (`multimodal_gen` lends a TP + group to `srt` layers with no `srt` config to publish against). - `initialize_model_parallel` computes the set through - `derive_parallel_widths` and hands it here; `initialize_dp_attention` - stamps `attn_dp_size` again once it knows the effective width, and - elastic EP restamps it where it already updates the live one. A stamped - width is what the readers answer with. + Lives beside, not inside, the `@contextmanager` `override` above -- a + name it cannot also have on this class -- because these are permanent + for the process, not scoped to a `with` block: none of the real + callers ever restore the value they set here. """ self._derived.update(widths) @@ -306,13 +309,13 @@ class ParallelContext: self._derived.clear() def _derived_width(self, name): - """A width the configuration implies: override, else stamp, else the - published leaf. + """A width the configuration implies: scoped override, else permanent + override, else the published leaf. - The leaf is computed at publish by `parallel_widths_of`; the stamp sits - above it because an elastic scale-up restamps `attn_dp_size` after - publish, and a scope that swaps in another TP group states the quotients - through `override`. + The leaf is computed at publish by `parallel_widths_of`; the permanent + override sits above it because an elastic scale-up corrects + `attn_dp_size` after publish, and a scope that swaps in another TP + group states the quotients through the scoped `override` above that. Nothing is recomputed on read, so overriding `tp_size` does not move `attn_tp_size`: name the width, or publish a config. @@ -328,10 +331,10 @@ class ParallelContext: return getattr(config, name) raise RuntimeError( f"derived parallel width {name!r} is not available: it is computed " - "from the configured leaves at publish, and restamped when the " - "process groups are built. Nothing is published and nothing has " - "been stamped -- publish a parallel config, or state the width " - f"with get_parallel().override({name}=...)" + "from the configured leaves at publish, and permanently corrected " + "when the process groups are built. Nothing is published and " + "nothing has been set with override_permanently -- publish a " + f"parallel config, or state the width with get_parallel().override({name}=...)" ) @contextmanager @@ -1754,9 +1757,10 @@ def reset_context() -> None: """Clear the context-owned store (unit-test teardown): drop the published ``server_args`` and install fresh ``Flags`` and ``Resources``. - ``parallel`` holds the stamped derived widths, which go with the lifecycle - that stamped them: `_derived_width` prefers the stamp over the leaves, so - leaving one behind lets the next test read the previous topology. + ``parallel`` holds the permanently-overridden derived widths, which go + with the lifecycle that set them: `_derived_width` prefers them over the + published leaves, so leaving one behind lets the next test read the + previous topology. """ _CONTEXT._server_args = None _CONTEXT._config_bags = None diff --git a/test/manual/ep/test_flashinfer_dispatcher.py b/test/manual/ep/test_flashinfer_dispatcher.py index 4847b5c81..707dd5449 100644 --- a/test/manual/ep/test_flashinfer_dispatcher.py +++ b/test/manual/ep/test_flashinfer_dispatcher.py @@ -20,14 +20,9 @@ from sglang.test.test_utils import CustomTestCase class TestFlashinferDispatcher(CustomTestCase): @classmethod def setUpClass(cls): - server_args = ServerArgs(model_path="dummy") - server_args.moe_runner_backend = "flashinfer_cutlass" - server_args.moe_a2a_backend = "flashinfer" - cls.server_args = server_args - set_global_server_args_for_scheduler(server_args) - publish(server_args, role="scheduler") - initialize_moe_config() - + # Dist-init first: world_size (and so the tp/ep width ServerArgs must + # carry) is only known after it, and init_distributed_environment + # itself reads no published config. init_distributed_environment( world_size=-1, # Auto-detect from environment rank=-1, # Auto-detect from environment @@ -38,6 +33,17 @@ class TestFlashinferDispatcher(CustomTestCase): rank = torch.distributed.get_rank() device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") torch.cuda.set_device(device) + + server_args = ServerArgs( + model_path="dummy", tp_size=world_size, ep_size=world_size + ) + server_args.moe_runner_backend = "flashinfer_cutlass" + server_args.moe_a2a_backend = "flashinfer" + cls.server_args = server_args + set_global_server_args_for_scheduler(server_args) + publish(server_args, role="scheduler") + initialize_moe_config() + initialize_model_parallel( tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size ) diff --git a/test/registered/eplb/test_lplb_distributed.py b/test/registered/eplb/test_lplb_distributed.py index 06037297c..59edf9300 100644 --- a/test/registered/eplb/test_lplb_distributed.py +++ b/test/registered/eplb/test_lplb_distributed.py @@ -212,6 +212,13 @@ def _worker_main(local_rank: int, world_size: int): set_global_server_args_for_scheduler( ServerArgs( model_path="dummy", + # Match the tp/ep width initialize_model_parallel is about to + # build below -- get_parallel()'s derived widths (attn_tp_size, + # moe_ep_size, ...) are projected from this at publish time, and + # nothing here should leave that projection reflecting a width + # this process never actually runs at. + tp_size=world_size, + ep_size=world_size, ) ) diff --git a/test/registered/unit/distributed/test_mmencoder_forwards_attn_cp_size.py b/test/registered/unit/distributed/test_mmencoder_forwards_attn_cp_size.py new file mode 100644 index 000000000..c94f74b68 --- /dev/null +++ b/test/registered/unit/distributed/test_mmencoder_forwards_attn_cp_size.py @@ -0,0 +1,67 @@ +"""MMEncoder must forward attn_cp_size to initialize_model_parallel, not just +tp_size. + +Real 2-GPU hardware confirmed a live mismatch this test guards against +statically: with `tp_size=2, attn_cp_size=2` published, calling +`initialize_model_parallel(tensor_model_parallel_size=2)` alone builds the +live attention-TP group at width 2, while `get_parallel().attn_tp_size` +(derived from the published config) answers 1 -- `VisionAttention` +(`layers/attention/vision.py`) reads that derived value as its own +weight-sharding width, so the mismatch is a real, silent wrong-sharding bug, +not just a reporting discrepancy. Forwarding +`attention_context_model_parallel_size=get_parallel().attn_cp_size` too +makes the two agree (confirmed on the same hardware). + +A full `MMEncoder` instantiation needs real weights and a live process +group, so this checks the one line that matters statically: the call passes +`attention_context_model_parallel_size` as well as +`tensor_model_parallel_size`. Not a substitute for testing `MMEncoder` +end-to-end under `--attn-cp-size > 1` on real hardware, but cheap enough to +run everywhere and catches the specific regression class (a future edit +that reverts to the tp-only call). +""" + +import ast +import os + +import sglang.srt.disaggregation.encoder.server as encoder_server_module +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestMMEncoderForwardsAttnCpSize(CustomTestCase): + def test_initialize_model_parallel_call_forwards_attn_cp_size(self): + path = encoder_server_module.__file__ + tree = ast.parse(open(path).read()) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "initialize_model_parallel" + ] + self.assertEqual( + len(calls), + 1, + f"expected exactly one initialize_model_parallel(...) call in " + f"{os.path.basename(path)}, found {len(calls)} -- update this " + "test if that's now intentional", + ) + kwarg_names = {kw.arg for kw in calls[0].keywords} + self.assertIn( + "attention_context_model_parallel_size", + kwarg_names, + "MMEncoder's initialize_model_parallel(...) call must forward " + "attention_context_model_parallel_size (not just " + "tensor_model_parallel_size), or get_parallel().attn_tp_size " + "silently disagrees with the group actually built whenever " + "--attn-cp-size > 1 -- confirmed on real 2-GPU hardware", + ) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 717f9d838..49e319c86 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -1375,7 +1375,8 @@ class TestParallelLeafReads(_IsolatedServerArgs): class TestDerivedWidths(_IsolatedOverrides): - """The widths no flag sets are computed from the leaves and stamped. + """The widths no flag sets are computed from the leaves and permanently + overridable. `attn_tp_size` and its siblings used to be read back off the group coordinator that was built from them, which made the answer depend on @@ -1390,7 +1391,7 @@ class TestDerivedWidths(_IsolatedOverrides): self.addCleanup( lambda: ( parallel.clear_derived_widths(), - parallel.stamp_derived_widths(**self._saved_derived), + parallel.override_permanently(**self._saved_derived), ) ) @@ -1439,13 +1440,14 @@ class TestDerivedWidths(_IsolatedOverrides): get_parallel().attn_tp_size self.assertIn("not available", str(caught.exception)) - def test_a_stamp_and_a_live_group_both_win_over_the_leaves(self): - """Order is stamp, then live group, then the leaves. Where a group - exists it is the truth -- elastic scale-up moves the group without - restamping -- so the leaf derivation only answers where there is none. + def test_a_permanent_override_and_a_live_group_both_win_over_the_leaves(self): + """Order is permanent override, then live group, then the leaves. + Where a group exists it is the truth -- elastic scale-up moves the + group without a fresh override -- so the leaf derivation only + answers where there is none. """ parallel = get_parallel() - parallel.stamp_derived_widths(attn_tp_size=7) + parallel.override_permanently(attn_tp_size=7) self.addCleanup(parallel.clear_derived_widths) with parallel.override(tp_size=8, attn_dp_size=2): self.assertEqual(parallel.attn_tp_size, 7) @@ -1464,9 +1466,9 @@ class TestDerivedWidths(_IsolatedOverrides): self.assertEqual(widths["moe_tp_size"], 8 // 4 // 2) self.assertEqual(widths["attn_dcp_size"], 1) - def test_the_world_size_is_not_stamped(self): + def test_the_world_size_is_not_permanently_overridden(self): """It is not a quotient, and the live getter is right at every moment. - A stamp taken when the groups are built would answer with the launch + A value fixed when the groups are built would answer with the launch count after `try_admit_scale_ranks` expands WORLD, and with the joining cohort's own width on a scale-joiner, which lays its groups out at `tp * pp` while WORLD spans `ep_join_rank_offset + tp * pp`.""" @@ -1481,31 +1483,30 @@ class TestDerivedWidths(_IsolatedOverrides): ) self.assertNotIn("world_size", widths) parallel = get_parallel() - parallel.stamp_derived_widths(attn_tp_size=4) + parallel.override_permanently(attn_tp_size=4) with patch(f"{_PS}.get_world_size", return_value=9): self.assertEqual(parallel.world_size, 9) - def test_a_stamped_width_is_what_the_reader_answers_with(self): + def test_a_permanently_overridden_width_is_what_the_reader_answers_with(self): parallel = get_parallel() - parallel.stamp_derived_widths(attn_tp_size=4, moe_tp_size=1) + parallel.override_permanently(attn_tp_size=4, moe_tp_size=1) with patch( f"{_PS}.get_attn_tensor_model_parallel_world_size", side_effect=AssertionError("the group must not be asked"), ): self.assertEqual(parallel.attn_tp_size, 4) - def test_an_override_still_wins_over_the_stamp(self): + def test_a_scoped_override_still_wins_over_the_permanent_one(self): parallel = get_parallel() - parallel.stamp_derived_widths(attn_tp_size=4) + parallel.override_permanently(attn_tp_size=4) with parallel.override(attn_tp_size=1): self.assertEqual(parallel.attn_tp_size, 1) self.assertEqual(parallel.attn_tp_size, 4) def test_the_group_is_never_consulted(self): - """There is no third source. A quotient comes from an override, a stamp - or the published leaf -- never from a group coordinator, which could - only ever agree, since `initialize_model_parallel` stamps as its last - statement.""" + """There is no third source. A quotient comes from a scoped override, a + permanent override, or the published leaf -- never from a group + coordinator.""" reset_context() self.addCleanup(reset_context) with patch( @@ -1528,50 +1529,51 @@ class TestDerivedWidths(_IsolatedOverrides): with self.assertRaisesRegex(RuntimeError, r"derived parallel width"): get_parallel().attn_tp_size - def test_a_temporary_disable_beats_the_stamp(self): + def test_a_temporary_disable_beats_the_permanent_override(self): """`disable_dp_size()` runs a draft scope without DP attention. It moves the module global the legacy getter reads, so it has to move the derived - width too -- the stamp wins over the live group, and a scope that left - it alone would answer with the target model's width for its duration.""" + width too -- the scoped override wins over the permanent one, and a + scope that left it alone would answer with the target model's width + for its duration.""" from sglang.srt.layers import dp_attention parallel = get_parallel() - parallel.stamp_derived_widths(attn_dp_size=4) + parallel.override_permanently(attn_dp_size=4) with patch.object(dp_attention, "_ATTN_DP_SIZE", 4): with dp_attention.disable_dp_size(): self.assertEqual(dp_attention.get_attention_dp_size(), 1) self.assertEqual(parallel.attn_dp_size, 1) self.assertEqual(parallel.attn_dp_size, 4) - def test_the_stamp_is_cleared_and_restamped(self): + def test_the_permanent_override_is_cleared_and_reset(self): parallel = get_parallel() - parallel.stamp_derived_widths(attn_dp_size=2) + parallel.override_permanently(attn_dp_size=2) self.assertEqual(parallel.attn_dp_size, 2) - # Elastic scaling restamps where it updates the live width. - parallel.stamp_derived_widths(attn_dp_size=4) + # Elastic scaling overrides again where it updates the live width. + parallel.override_permanently(attn_dp_size=4) self.assertEqual(parallel.attn_dp_size, 4) parallel.clear_derived_widths() with parallel.override(tp_size=8, attn_dp_size=1): self.assertEqual(parallel.attn_dp_size, 1) - def test_reset_context_drops_the_stamp(self): - """The stamp belongs to the lifecycle that made it. + def test_reset_context_drops_the_permanent_override(self): + """The permanent override belongs to the lifecycle that made it. - `_derived_width` prefers the stamp over the published leaf, so a stamp - that outlived `reset_context()` would let the next test read the - previous topology. + `_derived_width` prefers it over the published leaf, so one that + outlived `reset_context()` would let the next test read the previous + topology. """ parallel = get_parallel() - parallel.stamp_derived_widths(attn_tp_size=4) + parallel.override_permanently(attn_tp_size=4) self.assertEqual(parallel.attn_tp_size, 4) reset_context() self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy", tp_size=1), role="test") self.assertEqual(get_parallel().attn_tp_size, 1) - def test_the_rank_helper_agrees_with_the_stamp(self): + def test_the_rank_helper_agrees_with_the_override(self): """`compute_dp_attention_world_info` keeps the ranks and takes the - widths from the same derivation the stamp uses.""" + widths from the same derivation `override_permanently`'s callers use.""" from sglang.srt.layers.dp_attention import compute_dp_attention_world_info for tp_size, dp_size, attn_cp_size in ((8, 2, 1), (8, 2, 2), (16, 4, 2)): @@ -1590,6 +1592,126 @@ class TestDerivedWidths(_IsolatedOverrides): self.assertEqual(attn_tp_size, widths["attn_tp_size"]) self.assertEqual(attn_dp_size, widths["attn_dp_size"]) + def test_recomputing_from_published_leaves_matches_the_publish_bag(self): + """`initialize_model_parallel` no longer overrides anything -- see + `test_initialize_model_parallel_no_longer_touches_the_bag` below -- + which makes this the load-bearing half of 16-field-registry-design.md + §6e: every real caller must forward leaves that already match its own + published config, because nothing corrects a mismatch anymore. + `scheduler.py`'s `ps.attn_dp_size`/`ps.moe_ep_size`/etc, and the + weight-cache daemon's own already-published config, both do -- this + pins that the formula they'd recompute from those leaves + (`derive_attention_widths`, `derive_parallel_widths`, the same ones + `publish` itself used) agrees with what's already in the bag, across + the widths `test_the_rank_helper_agrees_with_the_override` does not + vary -- moe_ep_size, moe_dp_size, and dcp_size -- using real + `publish()`. + + A caller that does NOT keep the two in sync is a bug in that caller, + not something this framework silently corrects: two real ones existed + (`test/registered/eplb/test_lplb_distributed.py` and + `test/manual/ep/test_flashinfer_dispatcher.py`, both publishing a + placeholder config and then building real groups at a width it never + reflected) and were fixed by publishing the actual width instead of + relying on a correction to paper over the mismatch. + """ + shapes = ( + dict(tp_size=8), + dict(tp_size=8, dp_size=2, enable_dp_attention=True), + dict(tp_size=8, ep_size=4, moe_dp_size=2), + dict(tp_size=8, dcp_size=8), + ) + for shape in shapes: + with self.subTest(shape=shape): + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy", **shape), role="test") + parallel = get_parallel() + published = { + "attn_tp_size": parallel.attn_tp_size, + "attn_dp_size": parallel.attn_dp_size, + "moe_ep_size": parallel.moe_ep_size, + "moe_tp_size": parallel.moe_tp_size, + "dcp_enabled": parallel.dcp_enabled, + "attn_dcp_size": parallel.attn_dcp_size, + } + # What every real `initialize_model_parallel` caller forwards: + # its own already-published leaves, through the same two + # functions the bag was projected with. + recomputed = derive_parallel_widths( + tp_size=parallel.tp_size, + attn_cp_size=parallel.attn_cp_size, + attn_dp_size=( + parallel.dp_size if parallel.enable_dp_attention else 1 + ), + moe_ep_size=parallel.ep_size, + moe_dp_size=parallel.moe_dp_size, + dcp_size=parallel.dcp_size, + dcp_enabled=parallel.dcp_size > 1, + ) + self.assertEqual(published, recomputed) + + def test_initialize_model_parallel_no_longer_touches_the_bag(self): + """§6e, landed: `initialize_model_parallel` used to recompute and + permanently override the six derived widths on `get_parallel()` + after building its groups; that call is gone. Publish a placeholder + config (tp_size defaults to 1), then build real groups at a + different width -- the published leaf must now stay exactly what it + was, because nothing corrects it. This is the behavior a caller + relies on being told about, loudly, the first time it publishes and + builds inconsistently -- see + `test_recomputing_from_published_leaves_matches_the_publish_bag` + for why every real caller must not do that. + """ + from unittest.mock import Mock + + from sglang.srt.distributed import parallel_state + + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="test") + self.assertEqual(get_parallel().attn_tp_size, 1) + self.assertEqual(get_parallel().moe_ep_size, 1) + + world_size = 8 + with ( + patch.object(parallel_state, "_WORLD", None), + patch.object(parallel_state, "_TP", None), + patch.object(parallel_state, "_DCP", None), + patch.object(parallel_state, "_ATTN_CP", None), + patch.object(parallel_state, "_ATTN_TP", None), + patch.object(parallel_state, "_MOE_DP", None), + patch.object(parallel_state, "_MOE_EP", None), + patch.object(parallel_state, "_MOE_TP", None), + patch.object(parallel_state, "_PP", None), + patch.object(parallel_state, "_SELF_PP", None), + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_world_size", return_value=world_size), + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.get_backend", return_value="nccl"), + patch.object( + parallel_state, + "init_model_parallel_group", + return_value=Mock(device_group=Mock()), + ), + patch.object(parallel_state, "get_world_group") as mock_world_group, + ): + mock_world_group.return_value = Mock(device_group=Mock(), local_rank=0) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=world_size, + expert_model_parallel_size=world_size, + ) + self.addCleanup(parallel_state.destroy_model_parallel) + + self.assertEqual( + get_parallel().attn_tp_size, + 1, + "initialize_model_parallel must not touch the published leaf -- " + "a caller that needs it corrected must publish a config that " + "already matches the width it is about to build", + ) + self.assertEqual(get_parallel().moe_ep_size, 1) + class TestTheDerivedHalfIsDeclared(CustomTestCase): """The quotients are declared beside the leaves, in the same class.