Add aligner entrypoint and bundle handler in dump comparator (#19375)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,199 @@
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
|
||||
AlignerResult,
|
||||
_execute_step_plans,
|
||||
execute_aligner_plan,
|
||||
execute_sub_plan,
|
||||
execute_sub_plans,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
|
||||
AlignerPerStepPlan,
|
||||
AlignerPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
|
||||
ConcatParams,
|
||||
UnsharderPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestExecuteSubPlans:
|
||||
def test_empty_tensors_returns_none(self) -> None:
|
||||
result: Optional[torch.Tensor] = execute_sub_plans(tensors=[], plans=[])
|
||||
assert result is None
|
||||
|
||||
def test_no_plans_single_tensor_passthrough(self) -> None:
|
||||
tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
|
||||
result: Optional[torch.Tensor] = execute_sub_plans(tensors=[tensor], plans=[])
|
||||
assert result is not None
|
||||
assert torch.equal(result, tensor)
|
||||
|
||||
def test_no_plans_multiple_tensors_returns_none(self) -> None:
|
||||
tensors: list[torch.Tensor] = [
|
||||
torch.tensor([1.0]),
|
||||
torch.tensor([2.0]),
|
||||
]
|
||||
result: Optional[torch.Tensor] = execute_sub_plans(tensors=tensors, plans=[])
|
||||
assert result is None
|
||||
|
||||
def test_with_unsharder_plan(self) -> None:
|
||||
t0: torch.Tensor = torch.tensor([[1.0, 2.0]])
|
||||
t1: torch.Tensor = torch.tensor([[3.0, 4.0]])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.TP,
|
||||
params=ConcatParams(dim=1),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
|
||||
result: Optional[torch.Tensor] = execute_sub_plans(
|
||||
tensors=[t0, t1], plans=[plan]
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
|
||||
assert torch.equal(result, expected)
|
||||
|
||||
|
||||
class TestExecuteSubPlan:
|
||||
def test_unknown_plan_type_raises(self) -> None:
|
||||
class _FakePlan:
|
||||
pass
|
||||
|
||||
with pytest.raises(NotImplementedError, match="Unknown"):
|
||||
execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan()) # type: ignore[arg-type]
|
||||
|
||||
|
||||
class TestExecuteStepPlans:
|
||||
def test_step_with_none_result_omitted(self) -> None:
|
||||
tensors: list[torch.Tensor] = [
|
||||
torch.tensor([1.0]),
|
||||
torch.tensor([2.0]),
|
||||
]
|
||||
|
||||
step_plan = AlignerPerStepPlan(
|
||||
step=0,
|
||||
input_object_indices=[0, 1],
|
||||
sub_plans=[],
|
||||
)
|
||||
|
||||
result: dict[int, torch.Tensor] = _execute_step_plans(
|
||||
tensors=tensors, step_plans=[step_plan]
|
||||
)
|
||||
|
||||
assert result == {}
|
||||
|
||||
def test_single_step_passthrough(self) -> None:
|
||||
tensor: torch.Tensor = torch.tensor([1.0, 2.0])
|
||||
|
||||
step_plan = AlignerPerStepPlan(
|
||||
step=5,
|
||||
input_object_indices=[0],
|
||||
sub_plans=[],
|
||||
)
|
||||
|
||||
result: dict[int, torch.Tensor] = _execute_step_plans(
|
||||
tensors=[tensor], step_plans=[step_plan]
|
||||
)
|
||||
|
||||
assert 5 in result
|
||||
assert torch.equal(result[5], tensor)
|
||||
|
||||
|
||||
class TestExecuteAlignerPlan:
|
||||
def _make_step_plan(self, *, step: int, indices: list[int]) -> AlignerPerStepPlan:
|
||||
return AlignerPerStepPlan(step=step, input_object_indices=indices, sub_plans=[])
|
||||
|
||||
def test_x_side_empty_returns_failed_x(self) -> None:
|
||||
plan = AlignerPlan(
|
||||
per_step_plans=Pair(
|
||||
x=[self._make_step_plan(step=0, indices=[0, 1])],
|
||||
y=[self._make_step_plan(step=0, indices=[0])],
|
||||
),
|
||||
)
|
||||
|
||||
tensors_pair: Pair[list[torch.Tensor]] = Pair(
|
||||
x=[torch.tensor([1.0]), torch.tensor([2.0])],
|
||||
y=[torch.tensor([3.0])],
|
||||
)
|
||||
|
||||
result: AlignerResult = execute_aligner_plan(
|
||||
tensors_pair=tensors_pair, plan=plan
|
||||
)
|
||||
|
||||
assert result.tensors is None
|
||||
assert result.failed_side_xy == "x"
|
||||
|
||||
def test_y_side_empty_returns_failed_y(self) -> None:
|
||||
plan = AlignerPlan(
|
||||
per_step_plans=Pair(
|
||||
x=[self._make_step_plan(step=0, indices=[0])],
|
||||
y=[self._make_step_plan(step=0, indices=[0, 1])],
|
||||
),
|
||||
)
|
||||
|
||||
tensors_pair: Pair[list[torch.Tensor]] = Pair(
|
||||
x=[torch.tensor([1.0])],
|
||||
y=[torch.tensor([2.0]), torch.tensor([3.0])],
|
||||
)
|
||||
|
||||
result: AlignerResult = execute_aligner_plan(
|
||||
tensors_pair=tensors_pair, plan=plan
|
||||
)
|
||||
|
||||
assert result.tensors is None
|
||||
assert result.failed_side_xy == "y"
|
||||
|
||||
def test_single_step(self) -> None:
|
||||
plan = AlignerPlan(
|
||||
per_step_plans=Pair(
|
||||
x=[self._make_step_plan(step=0, indices=[0])],
|
||||
y=[self._make_step_plan(step=0, indices=[0])],
|
||||
),
|
||||
)
|
||||
|
||||
t_x: torch.Tensor = torch.tensor([1.0, 2.0])
|
||||
t_y: torch.Tensor = torch.tensor([3.0, 4.0])
|
||||
tensors_pair: Pair[list[torch.Tensor]] = Pair(x=[t_x], y=[t_y])
|
||||
|
||||
result: AlignerResult = execute_aligner_plan(
|
||||
tensors_pair=tensors_pair, plan=plan
|
||||
)
|
||||
|
||||
assert result.tensors is not None
|
||||
assert result.failed_side_xy is None
|
||||
assert torch.equal(result.tensors.x, t_x)
|
||||
assert torch.equal(result.tensors.y, t_y)
|
||||
|
||||
def test_success_returns_none_failed_side(self) -> None:
|
||||
plan = AlignerPlan(
|
||||
per_step_plans=Pair(
|
||||
x=[self._make_step_plan(step=0, indices=[0])],
|
||||
y=[self._make_step_plan(step=0, indices=[0])],
|
||||
),
|
||||
)
|
||||
|
||||
tensors_pair: Pair[list[torch.Tensor]] = Pair(
|
||||
x=[torch.tensor([10.0])],
|
||||
y=[torch.tensor([20.0])],
|
||||
)
|
||||
|
||||
result: AlignerResult = execute_aligner_plan(
|
||||
tensors_pair=tensors_pair, plan=plan
|
||||
)
|
||||
|
||||
assert result.failed_side_xy is None
|
||||
assert result.tensors is not None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,146 @@
|
||||
import sys
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
|
||||
_compute_per_step_plans,
|
||||
compute_aligner_plan,
|
||||
compute_per_step_sub_plans,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
|
||||
AlignerPerStepPlan,
|
||||
AlignerPerStepSubPlan,
|
||||
AlignerPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
|
||||
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
|
||||
def _make_meta(
|
||||
*,
|
||||
step: int = 0,
|
||||
dims: Optional[str] = None,
|
||||
tp_rank: int = 0,
|
||||
tp_size: int = 1,
|
||||
cp_rank: int = 0,
|
||||
cp_size: int = 1,
|
||||
) -> dict[str, Any]:
|
||||
meta: dict[str, Any] = {"step": step}
|
||||
if dims is not None:
|
||||
meta["dims"] = dims
|
||||
meta["sglang_parallel_info"] = {
|
||||
"tp_rank": tp_rank,
|
||||
"tp_size": tp_size,
|
||||
"cp_rank": cp_rank,
|
||||
"cp_size": cp_size,
|
||||
}
|
||||
return meta
|
||||
|
||||
|
||||
class TestComputePerStepSubPlans:
|
||||
def test_empty_metas(self) -> None:
|
||||
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(metas=[])
|
||||
assert result == []
|
||||
|
||||
def test_single_meta(self) -> None:
|
||||
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
|
||||
metas=[_make_meta(dims="b h(tp)", tp_size=2)]
|
||||
)
|
||||
assert result == []
|
||||
|
||||
def test_dims_none(self) -> None:
|
||||
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
|
||||
metas=[
|
||||
_make_meta(tp_rank=0, tp_size=2),
|
||||
_make_meta(tp_rank=1, tp_size=2),
|
||||
]
|
||||
)
|
||||
assert result == []
|
||||
|
||||
def test_tp_sharded_returns_unsharder_plan(self) -> None:
|
||||
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
|
||||
metas=[
|
||||
_make_meta(dims="b h(tp)", tp_rank=0, tp_size=2),
|
||||
_make_meta(dims="b h(tp)", tp_rank=1, tp_size=2),
|
||||
]
|
||||
)
|
||||
assert len(result) >= 1
|
||||
unsharder_plans: list[UnsharderPlan] = [
|
||||
p for p in result if isinstance(p, UnsharderPlan)
|
||||
]
|
||||
assert len(unsharder_plans) >= 1
|
||||
|
||||
def test_zigzag_returns_both_plans(self) -> None:
|
||||
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
|
||||
metas=[
|
||||
_make_meta(dims="b s(cp,zigzag) h", cp_rank=0, cp_size=2),
|
||||
_make_meta(dims="b s(cp,zigzag) h", cp_rank=1, cp_size=2),
|
||||
]
|
||||
)
|
||||
unsharder_plans: list[UnsharderPlan] = [
|
||||
p for p in result if isinstance(p, UnsharderPlan)
|
||||
]
|
||||
reorderer_plans: list[ReordererPlan] = [
|
||||
p for p in result if isinstance(p, ReordererPlan)
|
||||
]
|
||||
assert len(unsharder_plans) >= 1
|
||||
assert len(reorderer_plans) >= 1
|
||||
|
||||
|
||||
class TestComputePerStepPlans:
|
||||
def test_groups_by_step(self) -> None:
|
||||
metas: list[dict[str, Any]] = [
|
||||
_make_meta(step=0, tp_rank=0, tp_size=2),
|
||||
_make_meta(step=0, tp_rank=1, tp_size=2),
|
||||
_make_meta(step=1, tp_rank=0, tp_size=1),
|
||||
]
|
||||
result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0].step == 0
|
||||
assert result[0].input_object_indices == [0, 1]
|
||||
assert result[1].step == 1
|
||||
assert result[1].input_object_indices == [2]
|
||||
|
||||
def test_sorted_by_step(self) -> None:
|
||||
metas: list[dict[str, Any]] = [
|
||||
_make_meta(step=2),
|
||||
_make_meta(step=0),
|
||||
_make_meta(step=1),
|
||||
]
|
||||
result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas)
|
||||
|
||||
steps: list[int] = [p.step for p in result]
|
||||
assert steps == [0, 1, 2]
|
||||
|
||||
def test_single_meta_per_step_empty_sub_plans(self) -> None:
|
||||
metas: list[dict[str, Any]] = [
|
||||
_make_meta(step=0),
|
||||
_make_meta(step=1),
|
||||
]
|
||||
result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas)
|
||||
|
||||
assert len(result) == 2
|
||||
assert all(plan.sub_plans == [] for plan in result)
|
||||
|
||||
|
||||
class TestComputeAlignerPlan:
|
||||
def test_wraps_both_sides(self) -> None:
|
||||
metas_x: list[dict[str, Any]] = [_make_meta(step=0)]
|
||||
metas_y: list[dict[str, Any]] = [_make_meta(step=0)]
|
||||
|
||||
plan: AlignerPlan = compute_aligner_plan(
|
||||
metas_pair=Pair(x=metas_x, y=metas_y),
|
||||
)
|
||||
|
||||
assert len(plan.per_step_plans.x) == 1
|
||||
assert len(plan.per_step_plans.y) == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user