config: delete the redundant full stamp in initialize_model_parallel (#39202)

This commit is contained in:
Cheng Wan
2026-09-12 17:27:21 -07:00
committed by GitHub
parent fa260f26da
commit fa663e7297
9 changed files with 283 additions and 73 deletions
@@ -171,7 +171,9 @@ def _sync_srt_tp_group() -> None:
published `srt` config cannot answer: `gpu_worker.py` publishes a dummy published `srt` config cannot answer: `gpu_worker.py` publishes a dummy
carrying *this* package's `tp_size`, which a sequence-parallel launch sets 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 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. 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: if srt_parallel_state._ATTN_TP is None:
srt_parallel_state._ATTN_TP = _TP srt_parallel_state._ATTN_TP = _TP
if srt_parallel_state._ATTN_TP is _TP: if srt_parallel_state._ATTN_TP is _TP:
get_parallel().stamp_derived_widths( get_parallel().override_permanently(
**derive_parallel_widths( **derive_parallel_widths(
tp_size=_TP.world_size, tp_size=_TP.world_size,
attn_cp_size=1, attn_cp_size=1,
@@ -192,7 +194,7 @@ def _sync_srt_tp_group() -> None:
moe_dp_size=1, moe_dp_size=1,
dcp_size=1, dcp_size=1,
dcp_enabled=False, dcp_enabled=False,
) ),
) )
@@ -590,7 +590,10 @@ class MMEncoder:
distributed_init_method=dist_init_method, distributed_init_method=dist_init_method,
local_rank=rank, 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) initialize_dp_attention(server_args, self.model_config)
self.model = load_model( self.model = load_model(
@@ -2832,8 +2832,6 @@ def initialize_model_parallel(
group_name="self_pp", group_name="self_pp",
) )
get_parallel().stamp_derived_widths(**derived_widths)
def create_custom_parallel_group( def create_custom_parallel_group(
group_ranks: List[int], backend: str = "gloo" group_ranks: List[int], backend: str = "gloo"
+4 -3
View File
@@ -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 global _ATTN_DP_SIZE, _ATTN_DP_RANK
_ATTN_DP_SIZE = new_dp_size _ATTN_DP_SIZE = new_dp_size
_ATTN_DP_RANK = new_dp_rank _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 get_flags().dp.use_world_group_for_gather = True
logger.debug( logger.debug(
"[Elastic EP] dp_attention switched to WORLD: dp_size=%d dp_rank=%d", "[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. """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 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( attn_dp_size, attn_tp_size = derive_attention_widths(
tp_size=tp_size, tp_size=tp_size,
@@ -391,7 +392,7 @@ def initialize_dp_attention(
_, _, _ATTN_DP_RANK, _ATTN_DP_SIZE = compute_dp_attention_world_info( _, _, _ATTN_DP_RANK, _ATTN_DP_SIZE = compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size 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: 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 _ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset
+26 -22
View File
@@ -164,7 +164,7 @@ def derive_parallel_widths(
`world_size` is not among them: it is not a quotient, and `get_world_size()` `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 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 { return {
"attn_dp_size": attn_dp_size, "attn_dp_size": attn_dp_size,
@@ -268,7 +268,7 @@ class ParallelContext:
def __init__(self): def __init__(self):
self._overrides = {} self._overrides = {}
self._config = None # parallel config bag, wired at publish 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): def __getattr__(self, name):
if name.startswith("_"): if name.startswith("_"):
@@ -291,14 +291,17 @@ class ParallelContext:
overrides = self._overrides overrides = self._overrides
return overrides[name] if name in overrides else getter() return overrides[name] if name in overrides else getter()
def stamp_derived_widths(self, **widths) -> None: def override_permanently(self, **widths) -> None:
"""Record the widths derived from the leaves, as the groups are built. """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 Lives beside, not inside, the `@contextmanager` `override` above -- a
`derive_parallel_widths` and hands it here; `initialize_dp_attention` name it cannot also have on this class -- because these are permanent
stamps `attn_dp_size` again once it knows the effective width, and for the process, not scoped to a `with` block: none of the real
elastic EP restamps it where it already updates the live one. A stamped callers ever restore the value they set here.
width is what the readers answer with.
""" """
self._derived.update(widths) self._derived.update(widths)
@@ -306,13 +309,13 @@ class ParallelContext:
self._derived.clear() self._derived.clear()
def _derived_width(self, name): def _derived_width(self, name):
"""A width the configuration implies: override, else stamp, else the """A width the configuration implies: scoped override, else permanent
published leaf. override, else the published leaf.
The leaf is computed at publish by `parallel_widths_of`; the stamp sits The leaf is computed at publish by `parallel_widths_of`; the permanent
above it because an elastic scale-up restamps `attn_dp_size` after override sits above it because an elastic scale-up corrects
publish, and a scope that swaps in another TP group states the quotients `attn_dp_size` after publish, and a scope that swaps in another TP
through `override`. group states the quotients through the scoped `override` above that.
Nothing is recomputed on read, so overriding `tp_size` does not move Nothing is recomputed on read, so overriding `tp_size` does not move
`attn_tp_size`: name the width, or publish a config. `attn_tp_size`: name the width, or publish a config.
@@ -328,10 +331,10 @@ class ParallelContext:
return getattr(config, name) return getattr(config, name)
raise RuntimeError( raise RuntimeError(
f"derived parallel width {name!r} is not available: it is computed " f"derived parallel width {name!r} is not available: it is computed "
"from the configured leaves at publish, and restamped when the " "from the configured leaves at publish, and permanently corrected "
"process groups are built. Nothing is published and nothing has " "when the process groups are built. Nothing is published and "
"been stamped -- publish a parallel config, or state the width " "nothing has been set with override_permanently -- publish a "
f"with get_parallel().override({name}=...)" f"parallel config, or state the width with get_parallel().override({name}=...)"
) )
@contextmanager @contextmanager
@@ -1754,9 +1757,10 @@ def reset_context() -> None:
"""Clear the context-owned store (unit-test teardown): drop the published """Clear the context-owned store (unit-test teardown): drop the published
``server_args`` and install fresh ``Flags`` and ``Resources``. ``server_args`` and install fresh ``Flags`` and ``Resources``.
``parallel`` holds the stamped derived widths, which go with the lifecycle ``parallel`` holds the permanently-overridden derived widths, which go
that stamped them: `_derived_width` prefers the stamp over the leaves, so with the lifecycle that set them: `_derived_width` prefers them over the
leaving one behind lets the next test read the previous topology. published leaves, so leaving one behind lets the next test read the
previous topology.
""" """
_CONTEXT._server_args = None _CONTEXT._server_args = None
_CONTEXT._config_bags = None _CONTEXT._config_bags = None
+14 -8
View File
@@ -20,14 +20,9 @@ from sglang.test.test_utils import CustomTestCase
class TestFlashinferDispatcher(CustomTestCase): class TestFlashinferDispatcher(CustomTestCase):
@classmethod @classmethod
def setUpClass(cls): def setUpClass(cls):
server_args = ServerArgs(model_path="dummy") # Dist-init first: world_size (and so the tp/ep width ServerArgs must
server_args.moe_runner_backend = "flashinfer_cutlass" # carry) is only known after it, and init_distributed_environment
server_args.moe_a2a_backend = "flashinfer" # itself reads no published config.
cls.server_args = server_args
set_global_server_args_for_scheduler(server_args)
publish(server_args, role="scheduler")
initialize_moe_config()
init_distributed_environment( init_distributed_environment(
world_size=-1, # Auto-detect from environment world_size=-1, # Auto-detect from environment
rank=-1, # Auto-detect from environment rank=-1, # Auto-detect from environment
@@ -38,6 +33,17 @@ class TestFlashinferDispatcher(CustomTestCase):
rank = torch.distributed.get_rank() rank = torch.distributed.get_rank()
device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
torch.cuda.set_device(device) 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( initialize_model_parallel(
tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
) )
@@ -212,6 +212,13 @@ def _worker_main(local_rank: int, world_size: int):
set_global_server_args_for_scheduler( set_global_server_args_for_scheduler(
ServerArgs( ServerArgs(
model_path="dummy", 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,
) )
) )
@@ -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()
+156 -34
View File
@@ -1375,7 +1375,8 @@ class TestParallelLeafReads(_IsolatedServerArgs):
class TestDerivedWidths(_IsolatedOverrides): 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 `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 coordinator that was built from them, which made the answer depend on
@@ -1390,7 +1391,7 @@ class TestDerivedWidths(_IsolatedOverrides):
self.addCleanup( self.addCleanup(
lambda: ( lambda: (
parallel.clear_derived_widths(), 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 get_parallel().attn_tp_size
self.assertIn("not available", str(caught.exception)) self.assertIn("not available", str(caught.exception))
def test_a_stamp_and_a_live_group_both_win_over_the_leaves(self): def test_a_permanent_override_and_a_live_group_both_win_over_the_leaves(self):
"""Order is stamp, then live group, then the leaves. Where a group """Order is permanent override, then live group, then the leaves.
exists it is the truth -- elastic scale-up moves the group without Where a group exists it is the truth -- elastic scale-up moves the
restamping -- so the leaf derivation only answers where there is none. group without a fresh override -- so the leaf derivation only
answers where there is none.
""" """
parallel = get_parallel() parallel = get_parallel()
parallel.stamp_derived_widths(attn_tp_size=7) parallel.override_permanently(attn_tp_size=7)
self.addCleanup(parallel.clear_derived_widths) self.addCleanup(parallel.clear_derived_widths)
with parallel.override(tp_size=8, attn_dp_size=2): with parallel.override(tp_size=8, attn_dp_size=2):
self.assertEqual(parallel.attn_tp_size, 7) 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["moe_tp_size"], 8 // 4 // 2)
self.assertEqual(widths["attn_dcp_size"], 1) 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. """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 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 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`.""" `tp * pp` while WORLD spans `ep_join_rank_offset + tp * pp`."""
@@ -1481,31 +1483,30 @@ class TestDerivedWidths(_IsolatedOverrides):
) )
self.assertNotIn("world_size", widths) self.assertNotIn("world_size", widths)
parallel = get_parallel() 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): with patch(f"{_PS}.get_world_size", return_value=9):
self.assertEqual(parallel.world_size, 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 = 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( with patch(
f"{_PS}.get_attn_tensor_model_parallel_world_size", f"{_PS}.get_attn_tensor_model_parallel_world_size",
side_effect=AssertionError("the group must not be asked"), side_effect=AssertionError("the group must not be asked"),
): ):
self.assertEqual(parallel.attn_tp_size, 4) 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 = get_parallel()
parallel.stamp_derived_widths(attn_tp_size=4) parallel.override_permanently(attn_tp_size=4)
with parallel.override(attn_tp_size=1): with parallel.override(attn_tp_size=1):
self.assertEqual(parallel.attn_tp_size, 1) self.assertEqual(parallel.attn_tp_size, 1)
self.assertEqual(parallel.attn_tp_size, 4) self.assertEqual(parallel.attn_tp_size, 4)
def test_the_group_is_never_consulted(self): def test_the_group_is_never_consulted(self):
"""There is no third source. A quotient comes from an override, a stamp """There is no third source. A quotient comes from a scoped override, a
or the published leaf -- never from a group coordinator, which could permanent override, or the published leaf -- never from a group
only ever agree, since `initialize_model_parallel` stamps as its last coordinator."""
statement."""
reset_context() reset_context()
self.addCleanup(reset_context) self.addCleanup(reset_context)
with patch( with patch(
@@ -1528,50 +1529,51 @@ class TestDerivedWidths(_IsolatedOverrides):
with self.assertRaisesRegex(RuntimeError, r"derived parallel width"): with self.assertRaisesRegex(RuntimeError, r"derived parallel width"):
get_parallel().attn_tp_size 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 """`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 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 width too -- the scoped override wins over the permanent one, and a
it alone would answer with the target model's width for its duration.""" scope that left it alone would answer with the target model's width
for its duration."""
from sglang.srt.layers import dp_attention from sglang.srt.layers import dp_attention
parallel = get_parallel() 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 patch.object(dp_attention, "_ATTN_DP_SIZE", 4):
with dp_attention.disable_dp_size(): with dp_attention.disable_dp_size():
self.assertEqual(dp_attention.get_attention_dp_size(), 1) self.assertEqual(dp_attention.get_attention_dp_size(), 1)
self.assertEqual(parallel.attn_dp_size, 1) self.assertEqual(parallel.attn_dp_size, 1)
self.assertEqual(parallel.attn_dp_size, 4) 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 = get_parallel()
parallel.stamp_derived_widths(attn_dp_size=2) parallel.override_permanently(attn_dp_size=2)
self.assertEqual(parallel.attn_dp_size, 2) self.assertEqual(parallel.attn_dp_size, 2)
# Elastic scaling restamps where it updates the live width. # Elastic scaling overrides again where it updates the live width.
parallel.stamp_derived_widths(attn_dp_size=4) parallel.override_permanently(attn_dp_size=4)
self.assertEqual(parallel.attn_dp_size, 4) self.assertEqual(parallel.attn_dp_size, 4)
parallel.clear_derived_widths() parallel.clear_derived_widths()
with parallel.override(tp_size=8, attn_dp_size=1): with parallel.override(tp_size=8, attn_dp_size=1):
self.assertEqual(parallel.attn_dp_size, 1) self.assertEqual(parallel.attn_dp_size, 1)
def test_reset_context_drops_the_stamp(self): def test_reset_context_drops_the_permanent_override(self):
"""The stamp belongs to the lifecycle that made it. """The permanent override belongs to the lifecycle that made it.
`_derived_width` prefers the stamp over the published leaf, so a stamp `_derived_width` prefers it over the published leaf, so one that
that outlived `reset_context()` would let the next test read the outlived `reset_context()` would let the next test read the previous
previous topology. topology.
""" """
parallel = get_parallel() 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) self.assertEqual(parallel.attn_tp_size, 4)
reset_context() reset_context()
self.addCleanup(reset_context) self.addCleanup(reset_context)
publish(ServerArgs(model_path="dummy", tp_size=1), role="test") publish(ServerArgs(model_path="dummy", tp_size=1), role="test")
self.assertEqual(get_parallel().attn_tp_size, 1) 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 """`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 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)): 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_tp_size, widths["attn_tp_size"])
self.assertEqual(attn_dp_size, widths["attn_dp_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): class TestTheDerivedHalfIsDeclared(CustomTestCase):
"""The quotients are declared beside the leaves, in the same class. """The quotients are declared beside the leaves, in the same class.