483 lines
17 KiB
Python
483 lines
17 KiB
Python
"""Parallel-layout tests for the Qwen3.5 Python model (DP/EP).
|
|
|
|
Ported from xLLM upstream commit 78aa2a85 (PR #2258)
|
|
Original: tests/python/test_deepseek_v32_parallel.py
|
|
Adapted: Qwen3.5 MoeSparseBlock with 256 routed experts + shared expert
|
|
|
|
Test coverage maps to the issue's test cases:
|
|
TC-01 MoE expert routing DP isolation
|
|
TC-02 CUDA/ACL graph capture per DP rank
|
|
TC-03 DP broadcast/gather correctness
|
|
TC-04 Executor DP initialization
|
|
TC-05 End-to-end DP parallel test suite
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Insert repo root so "python.*" resolves to project_6/python/*, not stdlib
|
|
_repo_root = str(Path(__file__).resolve().parents[2])
|
|
if _repo_root not in sys.path:
|
|
sys.path.insert(0, _repo_root)
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
# Block the repo-root __init__.py from pulling in vllm
|
|
sys.modules.setdefault("project_6", MagicMock())
|
|
|
|
from python.layers.fused_moe import DPAwareMoEMixin # noqa: E402
|
|
from python.models.qwen3_5 import ( # noqa: E402
|
|
_dp_all_gather,
|
|
_dp_all_gather_variable,
|
|
configure_dp,
|
|
dp_forward_moe_wrapper,
|
|
)
|
|
from python.model_executor.runners.decode_cuda_graph import ( # noqa: E402
|
|
DecodeCudaGraphRunner,
|
|
_decode_bucket,
|
|
_decode_graph_buckets,
|
|
)
|
|
from python.model_executor.runners.decode_acl_graph import ( # noqa: E402
|
|
DecodeAclGraphRunner,
|
|
)
|
|
from python.model_executor.executor import ModelExecutor # noqa: E402
|
|
from python.attention.backend import ( # noqa: E402
|
|
DPBackendConfig, create_attention_backend, register_backend,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mock MoeSparseBlock for testing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class MockMoeSparseBlock(nn.Module):
|
|
"""Minimal mock of Qwen3_5MoeSparseBlock for DP testing."""
|
|
|
|
def __init__(self, hidden_size: int = 64, num_experts: int = 256):
|
|
super().__init__()
|
|
self.hidden_size = hidden_size
|
|
self.num_experts = num_experts
|
|
self.dp_size = 1
|
|
self.dp_rank = 0
|
|
# Dummy weight so next(model.parameters()) works
|
|
self.dummy = nn.Parameter(torch.zeros(1))
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
"""Identity forward — just returns input (MoE logic mocked)."""
|
|
return hidden_states
|
|
|
|
|
|
def _make_mock_moe(dp_size: int = 1, dp_rank: int = 0) -> MockMoeSparseBlock:
|
|
moe = MockMoeSparseBlock()
|
|
configure_dp(moe, dp_size, dp_rank)
|
|
return moe
|
|
|
|
|
|
def _mock_metadata(
|
|
dp_token_counts=(4,),
|
|
is_prefill=False,
|
|
is_chunked_prefill=False,
|
|
dp_is_decode=None,
|
|
execution_state=None,
|
|
):
|
|
metadata = SimpleNamespace(
|
|
dp_token_counts=dp_token_counts,
|
|
is_prefill=is_prefill,
|
|
is_chunked_prefill=is_chunked_prefill,
|
|
execution_state=execution_state,
|
|
)
|
|
if dp_is_decode is not None:
|
|
metadata.dp_is_decode = dp_is_decode
|
|
return metadata
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TC-01: MoE expert routing DP isolation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMoEDPIsolation:
|
|
"""Verify that DP replicas route experts independently."""
|
|
|
|
def test_dp1_no_gather(self):
|
|
"""With dp_size=1, no all-gather should occur."""
|
|
moe = _make_mock_moe(dp_size=1)
|
|
hidden = torch.randn(4, 64)
|
|
metadata = _mock_metadata(dp_token_counts=(4,))
|
|
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
assert result.shape == hidden.shape
|
|
|
|
def test_dp2_calls_gather(self):
|
|
"""With dp_size=2, all-gather should be invoked."""
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=0)
|
|
hidden = torch.randn(3, 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), execution_state="graph"
|
|
)
|
|
|
|
with patch(
|
|
"python.models.qwen3_5._dp_all_gather"
|
|
) as mock_gather:
|
|
mock_gather.side_effect = lambda x, **kw: x.repeat(
|
|
kw.get("world_size", 1), *([1] * (x.dim() - 1))
|
|
)
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
mock_gather.assert_called_once()
|
|
call_kwargs = mock_gather.call_args[1]
|
|
assert call_kwargs["dim"] == 0
|
|
assert call_kwargs["world_size"] == 2
|
|
assert call_kwargs["group_name"] == "dp"
|
|
|
|
def test_dp2_different_inputs_route_independently(self):
|
|
"""Different DP replicas with different inputs produce different
|
|
routing decisions (verified by output shapes being correct)."""
|
|
moe_r0 = _make_mock_moe(dp_size=2, dp_rank=0)
|
|
moe_r1 = _make_mock_moe(dp_size=2, dp_rank=1)
|
|
|
|
hidden_r0 = torch.randn(3, 64)
|
|
hidden_r1 = torch.randn(4, 64)
|
|
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
|
|
)
|
|
|
|
# Rank 0
|
|
result_r0 = dp_forward_moe_wrapper(
|
|
moe_r0, hidden_r0, moe_r0.forward, metadata
|
|
)
|
|
assert result_r0.shape[0] == 3
|
|
|
|
# Rank 1
|
|
result_r1 = dp_forward_moe_wrapper(
|
|
moe_r1, hidden_r1, moe_r1.forward, metadata
|
|
)
|
|
assert result_r1.shape[0] == 4
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TC-02: CUDA/ACL graph capture per DP rank
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCudaGraphDPCapture:
|
|
"""CUDA graph bucket computation and graph key with DP."""
|
|
|
|
def test_bucket_dp1(self):
|
|
buckets = _decode_graph_buckets(32, dp_size=1)
|
|
assert 1 in buckets
|
|
assert buckets[-1] == 32
|
|
|
|
def test_bucket_dp2_halves_capacity(self):
|
|
buckets = _decode_graph_buckets(32, dp_size=2)
|
|
# max_local_batch = ceil(32/2) = 16
|
|
assert buckets[-1] <= 16
|
|
|
|
def test_graph_key_dp1(self):
|
|
runner = DecodeCudaGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=16
|
|
)
|
|
input_ids = torch.zeros(4, dtype=torch.int32)
|
|
key = runner.graph_key(input_ids)
|
|
assert key is not None
|
|
padded, counts = key
|
|
assert padded == _decode_bucket(4)
|
|
assert counts == (padded,)
|
|
|
|
def test_graph_key_dp2(self):
|
|
runner = DecodeCudaGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
|
|
dp_size=2, dp_rank=0
|
|
)
|
|
input_ids = torch.zeros(4, dtype=torch.int32)
|
|
key = runner.graph_key(
|
|
input_ids,
|
|
dp_token_counts=(4, 3),
|
|
dp_is_decode=(1, 1),
|
|
)
|
|
assert key is not None
|
|
padded, counts = key
|
|
assert len(counts) == 2
|
|
assert counts[0] == counts[1] == padded
|
|
|
|
def test_graph_key_dp2_exceeds_capacity_returns_none(self):
|
|
runner = DecodeCudaGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=8,
|
|
dp_size=2, dp_rank=0
|
|
)
|
|
input_ids = torch.zeros(100, dtype=torch.int32)
|
|
key = runner.graph_key(
|
|
input_ids,
|
|
dp_token_counts=(100, 100),
|
|
dp_is_decode=(1, 1),
|
|
)
|
|
assert key is None
|
|
|
|
def test_graph_key_dp_mismatched_counts_raises(self):
|
|
runner = DecodeCudaGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
|
|
dp_size=2, dp_rank=0
|
|
)
|
|
input_ids = torch.zeros(4, dtype=torch.int32)
|
|
with pytest.raises(RuntimeError, match="dp_token_counts"):
|
|
runner.graph_key(input_ids, dp_token_counts=(4,))
|
|
|
|
|
|
class TestAclGraphDPCapture:
|
|
"""ACL graph runner DP validation."""
|
|
|
|
def test_acl_dp2_init(self):
|
|
runner = DecodeAclGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
|
|
dp_size=2, dp_rank=0
|
|
)
|
|
assert runner.dp_size == 2
|
|
assert runner.dp_rank == 0
|
|
assert runner.max_batch == 16 # ceil(32/2)
|
|
|
|
def test_acl_dp_validate_wrong_counts_raises(self):
|
|
runner = DecodeAclGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
|
|
dp_size=2, dp_rank=0
|
|
)
|
|
with pytest.raises(RuntimeError, match="dp_token_counts"):
|
|
runner._validate_dp_token_counts((4,))
|
|
|
|
def test_acl_dp1_validate_accepts_none(self):
|
|
runner = DecodeAclGraphRunner(
|
|
nn.Linear(1, 1), torch.device("cpu"), max_batch=32,
|
|
dp_size=1, dp_rank=0
|
|
)
|
|
# dp_size=1: validation is a no-op
|
|
runner._validate_dp_token_counts(None)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TC-03: DP broadcast/gather correctness
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDPBroadcastGather:
|
|
"""Verify that gather + scatter preserves local token identity."""
|
|
|
|
def test_padded_gather_rank0_recovers_local(self):
|
|
"""dp_rank=0: after padded gather → slice, output shape matches input."""
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=0)
|
|
hidden = torch.randn(3, 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), execution_state="graph"
|
|
)
|
|
|
|
with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
|
|
# Simulate all_gather: pad to 4, gather 2 replicas → [8, 64]
|
|
mock_g.side_effect = lambda x, **kw: x.repeat(
|
|
kw.get("world_size", 1), *([1] * (x.dim() - 1))
|
|
)
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
# dp_rank=0, padded=4: narrow(0, 0, 3) → [3, 64]
|
|
assert result.shape[0] == 3
|
|
|
|
def test_padded_gather_rank1_recovers_local(self):
|
|
"""dp_rank=1: output is sliced from the second replica's region."""
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=1)
|
|
hidden = torch.randn(4, 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), execution_state="graph"
|
|
)
|
|
|
|
with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
|
|
mock_g.side_effect = lambda x, **kw: x.repeat(
|
|
kw.get("world_size", 1), *([1] * (x.dim() - 1))
|
|
)
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
# dp_rank=1, padded=4: narrow(0, 4, 4) → [4, 64]
|
|
assert result.shape[0] == 4
|
|
|
|
def test_compact_gather_rank0(self):
|
|
"""Eager decode path uses compact gather; rank 0 gets first slice."""
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=0)
|
|
hidden = torch.randn(3, 64)
|
|
compact = torch.randn(7, 64) # 3 + 4 tokens
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
|
|
)
|
|
|
|
with patch("python.models.qwen3_5._dp_all_gather_variable") as mock_gv:
|
|
mock_gv.return_value = compact
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
mock_gv.assert_called_once()
|
|
# dp_rank=0: offset=0, narrow(0, 0, 3)
|
|
assert result.shape[0] == 3
|
|
|
|
def test_compact_gather_rank1(self):
|
|
"""Eager decode path: rank 1 gets second slice."""
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=1)
|
|
hidden = torch.randn(4, 64)
|
|
compact = torch.randn(7, 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
|
|
)
|
|
|
|
with patch("python.models.qwen3_5._dp_all_gather_variable") as mock_gv:
|
|
mock_gv.return_value = compact
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
# dp_rank=1: offset=sum([3])=3, narrow(0, 3, 4)
|
|
assert result.shape[0] == 4
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TC-04: Executor DP initialization
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExecutorDPInit:
|
|
"""Verify ModelExecutor reads and validates DP config."""
|
|
|
|
def _make_model(self):
|
|
return nn.Linear(10, 10)
|
|
|
|
def test_dp_size_stored(self):
|
|
model = self._make_model()
|
|
executor = ModelExecutor(model, {"dp_size": 2, "dp_rank": 0}, 32)
|
|
assert executor.dp_size == 2
|
|
|
|
def test_dp1_default(self):
|
|
model = self._make_model()
|
|
executor = ModelExecutor(model, {}, 32)
|
|
assert executor.dp_size == 1
|
|
|
|
def test_dp_with_unsupported_backend_raises(self):
|
|
model = self._make_model()
|
|
with pytest.raises(NotImplementedError, match="data parallel"):
|
|
ModelExecutor(
|
|
model,
|
|
{"dp_size": 2, "dp_rank": 0, "python_graph_backend": "inductor"},
|
|
32,
|
|
)
|
|
|
|
def test_dp_with_cudagraphs_accepted(self):
|
|
model = self._make_model()
|
|
executor = ModelExecutor(
|
|
model,
|
|
{"dp_size": 2, "dp_rank": 0, "python_graph_backend": "cudagraphs"},
|
|
32,
|
|
)
|
|
assert executor.decode_graph_runner is not None
|
|
|
|
def test_dp_with_aclgraph_accepted(self):
|
|
model = self._make_model()
|
|
executor = ModelExecutor(
|
|
model,
|
|
{"dp_size": 2, "dp_rank": 0, "python_graph_backend": "aclgraph"},
|
|
32,
|
|
)
|
|
assert executor.decode_graph_runner is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TC-05: End-to-end DP parallel test suite
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestEndToEndDP:
|
|
"""Integration tests combining executor + MoE + graph runner."""
|
|
|
|
def test_2way_dp_moe_shapes(self):
|
|
"""2-way DP: both ranks produce correct output shapes."""
|
|
for rank in (0, 1):
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=rank)
|
|
local_tokens = 3 if rank == 0 else 4
|
|
hidden = torch.randn(local_tokens, 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(3, 4), dp_is_decode=(1, 1)
|
|
)
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
assert result.shape[0] == local_tokens
|
|
|
|
def test_4way_dp_moe_shapes(self):
|
|
"""4-way DP: all ranks produce correct output shapes."""
|
|
counts = (2, 3, 4, 5)
|
|
for rank in range(4):
|
|
moe = _make_mock_moe(dp_size=4, dp_rank=rank)
|
|
hidden = torch.randn(counts[rank], 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=counts, dp_is_decode=(1, 1, 1, 1)
|
|
)
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
assert result.shape[0] == counts[rank]
|
|
|
|
def test_dp_plus_tp_cuda_graph_buckets(self):
|
|
"""DP=2, TP=4 on 8 devices: graph buckets respect DP-reduced capacity."""
|
|
# max_batch=64 across 2 DP replicas → 32 per replica
|
|
buckets = _decode_graph_buckets(64, dp_size=2)
|
|
assert buckets[-1] <= 32
|
|
|
|
def test_varying_batch_sizes_padded_path(self):
|
|
"""Different batch sizes per DP rank use padded gather."""
|
|
moe = _make_mock_moe(dp_size=2, dp_rank=0)
|
|
for local_size in (1, 5, 8, 16):
|
|
other_size = local_size + 2
|
|
hidden = torch.randn(local_size, 64)
|
|
metadata = _mock_metadata(
|
|
dp_token_counts=(local_size, other_size),
|
|
execution_state="graph",
|
|
)
|
|
with patch("python.models.qwen3_5._dp_all_gather") as mock_g:
|
|
mock_g.side_effect = lambda x, **kw: x.repeat(
|
|
kw.get("world_size", 1), *([1] * (x.dim() - 1))
|
|
)
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
assert result.shape[0] == local_size
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TC-07: Existing tests unbroken (sanity)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBackwardCompat:
|
|
"""Verify dp_size=1 (default) doesn't change existing behaviour."""
|
|
|
|
def test_dp1_identity(self):
|
|
moe = _make_mock_moe(dp_size=1)
|
|
hidden = torch.randn(8, 64)
|
|
metadata = _mock_metadata(dp_token_counts=(8,))
|
|
result = dp_forward_moe_wrapper(
|
|
moe, hidden, moe.forward, metadata
|
|
)
|
|
assert torch.equal(result, hidden)
|
|
|
|
def test_executor_dp1_no_graph_runner(self):
|
|
model = nn.Linear(10, 10)
|
|
executor = ModelExecutor(model, {}, 32)
|
|
assert executor.decode_graph_runner is None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|