config: the derived parallel widths are computed from the leaves (#36790)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
co-authored by
Claude Opus 5
Claude Opus 5
parent
c2928e86d7
commit
43c63a22ff
@@ -7,11 +7,13 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
import dataclasses
|
||||
import json
|
||||
import os
|
||||
import pathlib as _pathlib
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import sglang as _sglang
|
||||
import sglang.srt.server_args as server_args_module
|
||||
from sglang.srt.arg_groups.arg_utils import NS, A, Arg
|
||||
from sglang.srt.runtime_context import (
|
||||
@@ -20,6 +22,7 @@ from sglang.srt.runtime_context import (
|
||||
RuntimeContext,
|
||||
_FlagGroupBase,
|
||||
assert_published,
|
||||
derive_parallel_widths,
|
||||
get_context,
|
||||
get_exec,
|
||||
get_flags,
|
||||
@@ -33,6 +36,7 @@ from sglang.srt.runtime_context import (
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
_SRT = _pathlib.Path(next(iter(_sglang.__path__))).resolve() / "srt"
|
||||
_PS = "sglang.srt.distributed.parallel_state"
|
||||
_DP = "sglang.srt.layers.dp_attention"
|
||||
|
||||
@@ -1385,5 +1389,168 @@ class TestParallelLeafReads(_IsolatedServerArgs):
|
||||
getattr(ParallelContext(), "not_a_leaf")
|
||||
|
||||
|
||||
class TestDerivedWidths(_IsolatedOverrides):
|
||||
"""The widths no flag sets are computed from the leaves and stamped.
|
||||
|
||||
`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
|
||||
distributed init and, after an elastic scale, disagree with the leaves.
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
parallel = get_parallel()
|
||||
self._saved_derived = dict(parallel._derived)
|
||||
parallel.clear_derived_widths()
|
||||
self.addCleanup(
|
||||
lambda: (
|
||||
parallel.clear_derived_widths(),
|
||||
parallel.stamp_derived_widths(**self._saved_derived),
|
||||
)
|
||||
)
|
||||
|
||||
def test_the_quotients_come_from_the_leaves(self):
|
||||
widths = derive_parallel_widths(
|
||||
tp_size=8,
|
||||
attn_cp_size=1,
|
||||
attn_dp_size=2,
|
||||
moe_ep_size=4,
|
||||
moe_dp_size=2,
|
||||
dcp_size=1,
|
||||
dcp_enabled=False,
|
||||
)
|
||||
self.assertEqual(widths["attn_tp_size"], 8 // 2 // 1)
|
||||
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):
|
||||
"""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
|
||||
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`."""
|
||||
widths = derive_parallel_widths(
|
||||
tp_size=4,
|
||||
attn_cp_size=1,
|
||||
attn_dp_size=1,
|
||||
moe_ep_size=1,
|
||||
moe_dp_size=1,
|
||||
dcp_size=1,
|
||||
dcp_enabled=False,
|
||||
)
|
||||
self.assertNotIn("world_size", widths)
|
||||
parallel = get_parallel()
|
||||
parallel.stamp_derived_widths(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):
|
||||
parallel = get_parallel()
|
||||
parallel.stamp_derived_widths(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):
|
||||
parallel = get_parallel()
|
||||
parallel.stamp_derived_widths(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_without_a_stamp_the_live_group_still_answers(self):
|
||||
"""A process that installed groups by hand keeps working."""
|
||||
with patch(f"{_PS}.get_attn_tensor_model_parallel_world_size", return_value=2):
|
||||
self.assertEqual(get_parallel().attn_tp_size, 2)
|
||||
|
||||
def test_with_neither_the_failure_names_the_cause(self):
|
||||
with patch(
|
||||
f"{_PS}.get_attn_tensor_model_parallel_world_size",
|
||||
side_effect=AssertionError("attention tp group is not initialized"),
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, r"derived parallel width"):
|
||||
get_parallel().attn_tp_size
|
||||
|
||||
def test_a_temporary_disable_beats_the_stamp(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."""
|
||||
from sglang.srt.layers import dp_attention
|
||||
|
||||
parallel = get_parallel()
|
||||
parallel.stamp_derived_widths(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):
|
||||
parallel = get_parallel()
|
||||
parallel.stamp_derived_widths(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)
|
||||
self.assertEqual(parallel.attn_dp_size, 4)
|
||||
parallel.clear_derived_widths()
|
||||
with patch(f"{_DP}.get_attention_dp_size", return_value=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.
|
||||
|
||||
`_derived_width` prefers the stamp over the live group, so a stamp that
|
||||
outlived `reset_context()` would let the next test read the previous
|
||||
topology.
|
||||
"""
|
||||
from sglang.srt.runtime_context import reset_context
|
||||
|
||||
parallel = get_parallel()
|
||||
parallel.stamp_derived_widths(attn_tp_size=4)
|
||||
self.assertEqual(parallel.attn_tp_size, 4)
|
||||
reset_context()
|
||||
with patch(f"{_PS}.get_attn_tensor_model_parallel_world_size", return_value=1):
|
||||
self.assertEqual(get_parallel().attn_tp_size, 1)
|
||||
|
||||
def test_the_arithmetic_has_one_home(self):
|
||||
"""`parallel_state` builds its groups from the same dict it stamps, and
|
||||
`dp_attention` derives the pair it needs for the ranks, so a second copy
|
||||
of a quotient would let two answers to one width drift apart."""
|
||||
for rel, spelling in (
|
||||
("distributed/parallel_state.py", "derive_parallel_widths("),
|
||||
("layers/dp_attention.py", "derive_attention_widths("),
|
||||
):
|
||||
source = (_SRT / rel).read_text(encoding="utf-8-sig")
|
||||
self.assertNotIn("// attn_dp_size // attn_cp_size", source, rel)
|
||||
self.assertNotIn("// attn_cp_size // attn_dp_size", source, rel)
|
||||
self.assertNotIn("// moe_ep_size // moe_dp_size", source, rel)
|
||||
self.assertNotIn("if enable_dp_attention else 1", source, rel)
|
||||
self.assertIn(spelling, source, rel)
|
||||
|
||||
def test_the_rank_helper_agrees_with_the_stamp(self):
|
||||
"""`compute_dp_attention_world_info` keeps the ranks and takes the
|
||||
widths from the same derivation the stamp uses."""
|
||||
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)):
|
||||
_, attn_tp_size, _, attn_dp_size = compute_dp_attention_world_info(
|
||||
True, 0, tp_size, dp_size, attn_cp_size
|
||||
)
|
||||
widths = derive_parallel_widths(
|
||||
tp_size=tp_size,
|
||||
attn_cp_size=attn_cp_size,
|
||||
attn_dp_size=attn_dp_size,
|
||||
moe_ep_size=1,
|
||||
moe_dp_size=1,
|
||||
dcp_size=1,
|
||||
dcp_enabled=False,
|
||||
)
|
||||
self.assertEqual(attn_tp_size, widths["attn_tp_size"])
|
||||
self.assertEqual(attn_dp_size, widths["attn_dp_size"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user