minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
+5 -3
View File
@@ -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,
)
),
)
python/sglang/srt/disaggregation/encoder/server.py
+4 -1
View File
@@ -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(
python/sglang/srt/distributed/parallel_state.py
-2
View File
@@ -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"
python/sglang/srt/layers/dp_attention.py
+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
_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
python/sglang/srt/runtime_context.py
+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()`
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