dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

under test, not sure no errors

Browse Source
This commit is contained in:
DP Migration
2026-09-01 10:24:14 +00:00
parent 8c9d913f3f
commit 94d77cf0b4
15 changed files with 1859 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

482
tests/python/test_qwen3_5_parallel.py Normal file
View File

@@ -0,0 +1,482 @@
"""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"])