config: delete the redundant full stamp in initialize_model_parallel (#39202)
This commit is contained in:
@@ -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"
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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.
|
||||||
|
|||||||
Reference in New Issue
Block a user