[PD] Skip singleton transfer-status all-reduces (#40003)
Co-authored-by: Pranjal Shankhdhar <pranjalssh@users.noreply.github.com> Co-authored-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
This commit is contained in:
metamergebot
co-authored by
Pranjal Shankhdhar
Jialin Ouyang
Pranjal Shankhdhar
Jialin Ouyang
parent
5b42d10edf
commit
8189e3896b
@@ -2,12 +2,13 @@ import struct
|
||||
import threading
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import ANY, Mock, call, patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.disaggregation.base.conn import KVArgs, StateType
|
||||
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType
|
||||
from sglang.srt.disaggregation.common.conn import CommonKVManager
|
||||
from sglang.srt.disaggregation.common.staging_buffer import (
|
||||
StagingAllocator,
|
||||
@@ -36,6 +37,9 @@ from sglang.srt.disaggregation.utils import (
|
||||
build_transfer_entry_pairs,
|
||||
compute_mamba_state_slice_byte_blocks,
|
||||
get_qsa_pending_state_indices,
|
||||
poll_and_all_reduce,
|
||||
poll_and_all_reduce_attn_cp_tp_group,
|
||||
poll_and_all_reduce_with_staging,
|
||||
setup_state_kv_args,
|
||||
should_send_replicated_state,
|
||||
)
|
||||
@@ -158,6 +162,117 @@ class TestDisaggregationWire(unittest.TestCase):
|
||||
self.assertEqual(unpack_list_of_buffers(pack_list_of_buffers(bufs)), bufs)
|
||||
|
||||
|
||||
class TestPollCollectives(CustomTestCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
failure_prob = patch.object(
|
||||
envs.SGLANG_TEST_DISAGG_FAILURE_PROB, "get", return_value=0
|
||||
)
|
||||
failure_prob.start()
|
||||
self.addCleanup(failure_prob.stop)
|
||||
|
||||
def test_tp_cp_consensus_skips_only_singleton_groups(self):
|
||||
"""Poll states retain TP/CP consensus without singleton collectives."""
|
||||
local = [KVPoll.Success, KVPoll.Success, KVPoll.Success]
|
||||
tp_peers = [KVPoll.Success, KVPoll.Transferring, KVPoll.Success]
|
||||
cp_peers = [KVPoll.Bootstrapping, KVPoll.Success, KVPoll.Failed]
|
||||
for tp_size, cp_size in [(4, 1), (2, 2), (1, 4), (1, 1)]:
|
||||
with self.subTest(tp_size=tp_size, cp_size=cp_size):
|
||||
tp, cp = object(), object()
|
||||
sizes = {tp: tp_size, cp: cp_size}
|
||||
peers = {tp: tp_peers, cp: cp_peers}
|
||||
pollers = [Mock(poll=Mock(return_value=state)) for state in local]
|
||||
|
||||
def reduce(tensor, op, group):
|
||||
self.assertEqual(op, dist.ReduceOp.MIN)
|
||||
if sizes[group] > 1:
|
||||
tensor.copy_(
|
||||
torch.minimum(
|
||||
tensor, torch.tensor(peers[group], dtype=tensor.dtype)
|
||||
)
|
||||
)
|
||||
|
||||
expected = local.copy()
|
||||
expected_calls = []
|
||||
for group in [tp, cp]:
|
||||
if sizes[group] > 1:
|
||||
expected = [min(a, b) for a, b in zip(expected, peers[group])]
|
||||
expected_calls.append(
|
||||
call(ANY, op=dist.ReduceOp.MIN, group=group)
|
||||
)
|
||||
with (
|
||||
patch.object(dist, "get_world_size", side_effect=sizes.__getitem__),
|
||||
patch.object(dist, "all_reduce", side_effect=reduce) as all_reduce,
|
||||
):
|
||||
result = poll_and_all_reduce_attn_cp_tp_group(pollers, cp, tp)
|
||||
|
||||
self.assertEqual(result, expected)
|
||||
self.assertEqual(all_reduce.call_args_list, expected_calls)
|
||||
for poller in pollers:
|
||||
poller.poll.assert_called_once_with()
|
||||
|
||||
def test_singleton_polling_needs_no_tensor_or_collective(self):
|
||||
"""A single rank can poll transfers without allocating a reduction tensor."""
|
||||
states = [KVPoll.Failed, KVPoll.Transferring, KVPoll.Success]
|
||||
pollers = [Mock(poll=Mock(return_value=state)) for state in states]
|
||||
with (
|
||||
patch.object(dist, "get_world_size", return_value=1),
|
||||
patch.object(dist, "all_reduce") as all_reduce,
|
||||
patch.object(
|
||||
torch, "tensor", side_effect=AssertionError("No tensor needed")
|
||||
),
|
||||
):
|
||||
self.assertEqual(poll_and_all_reduce(pollers, object()), states)
|
||||
all_reduce.assert_not_called()
|
||||
|
||||
def test_singleton_still_waits_for_decode_metadata(self):
|
||||
pollers = [Mock(poll=Mock(return_value=KVPoll.Success)) for _ in range(2)]
|
||||
requests = [
|
||||
SimpleNamespace(
|
||||
req=SimpleNamespace(bootstrap_host="127.0.0.1"),
|
||||
metadata_buffer_index=index,
|
||||
)
|
||||
for index in range(2)
|
||||
]
|
||||
metadata = SimpleNamespace(bootstrap_room=torch.tensor([[0], [42]]))
|
||||
with (
|
||||
patch.object(dist, "get_world_size", return_value=1),
|
||||
patch.object(dist, "all_reduce") as all_reduce,
|
||||
):
|
||||
result = poll_and_all_reduce(pollers, object(), requests, metadata)
|
||||
self.assertEqual(result, [KVPoll.Transferring, KVPoll.Success])
|
||||
all_reduce.assert_not_called()
|
||||
|
||||
def test_singleton_still_advances_and_waits_for_staging(self):
|
||||
request = SimpleNamespace(
|
||||
kv_receiver=Mock(
|
||||
require_staging=True, poll=Mock(return_value=KVPoll.Success)
|
||||
)
|
||||
)
|
||||
staging = Mock(
|
||||
is_done=Mock(return_value=False), is_failed=Mock(return_value=False)
|
||||
)
|
||||
with (
|
||||
patch.object(dist, "get_world_size", return_value=1),
|
||||
patch.object(dist, "all_reduce") as all_reduce,
|
||||
):
|
||||
result = poll_and_all_reduce_with_staging([request], staging, object())
|
||||
self.assertEqual(result, [KVPoll.Transferring])
|
||||
staging.advance_scatter.assert_called_once_with(request)
|
||||
all_reduce.assert_not_called()
|
||||
|
||||
def test_singleton_preserves_failure_injection(self):
|
||||
poller = Mock(poll=Mock(return_value=KVPoll.Success))
|
||||
with (
|
||||
patch.object(envs.SGLANG_TEST_DISAGG_FAILURE_PROB, "get", return_value=1),
|
||||
patch.object(dist, "get_world_size", return_value=1),
|
||||
patch.object(dist, "all_reduce") as all_reduce,
|
||||
):
|
||||
result = poll_and_all_reduce([poller], object())
|
||||
self.assertEqual(result, [KVPoll.Failed])
|
||||
all_reduce.assert_not_called()
|
||||
|
||||
|
||||
class TestCPReplicatedStateTransfer(unittest.TestCase):
|
||||
def test_only_nonzero_cp_ranks_without_layer_split_skip_state(self):
|
||||
cases = [
|
||||
|
||||
Reference in New Issue
Block a user