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__]))
|
||||
150
test/registered/debug_utils/comparator/test_bundle_matcher.py
Normal file
150
test/registered/debug_utils/comparator/test_bundle_matcher.py
Normal file
@@ -0,0 +1,150 @@
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import polars as pl
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.bundle_matcher import (
|
||||
TensorBundleInfo,
|
||||
TensorFileInfo,
|
||||
_rows_to_tensor_infos,
|
||||
match_bundles,
|
||||
)
|
||||
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_row(
|
||||
*, name: str, step: int = 0, rank: int = 0, filename: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
if filename is None:
|
||||
filename = f"name={name}___step={step}___rank={rank}.pt"
|
||||
return {"name": name, "step": step, "rank": rank, "filename": filename}
|
||||
|
||||
|
||||
def _make_df(rows: list[dict[str, Any]]) -> pl.DataFrame:
|
||||
return pl.DataFrame(rows)
|
||||
|
||||
|
||||
class TestMatchBundles:
|
||||
def test_single_tensor_single_step(self) -> None:
|
||||
target_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
|
||||
baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
|
||||
|
||||
results: list[Pair[TensorBundleInfo]] = match_bundles(
|
||||
dfs=Pair(x=baseline_df, y=target_df),
|
||||
skip_keys={"filename"},
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0].x) == 1
|
||||
assert len(results[0].y) == 1
|
||||
assert results[0].y[0].name == "t_a"
|
||||
|
||||
def test_multiple_names_separate_bundles(self) -> None:
|
||||
target_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a"),
|
||||
_make_row(name="t_b"),
|
||||
])
|
||||
baseline_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a"),
|
||||
_make_row(name="t_b"),
|
||||
])
|
||||
|
||||
results: list[Pair[TensorBundleInfo]] = match_bundles(
|
||||
dfs=Pair(x=baseline_df, y=target_df),
|
||||
skip_keys={"filename"},
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
result_names: list[str] = [r.y[0].name for r in results]
|
||||
assert "t_a" in result_names
|
||||
assert "t_b" in result_names
|
||||
|
||||
def test_skip_rank_groups_across_ranks(self) -> None:
|
||||
target_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a", rank=0),
|
||||
_make_row(name="t_a", rank=1),
|
||||
])
|
||||
baseline_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a", rank=0),
|
||||
_make_row(name="t_a", rank=1),
|
||||
])
|
||||
|
||||
results: list[Pair[TensorBundleInfo]] = match_bundles(
|
||||
dfs=Pair(x=baseline_df, y=target_df),
|
||||
skip_keys={"filename", "rank"},
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0].y) == 2
|
||||
|
||||
def test_baseline_missing_tensor(self) -> None:
|
||||
target_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a"),
|
||||
_make_row(name="t_extra"),
|
||||
])
|
||||
baseline_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a"),
|
||||
])
|
||||
|
||||
results: list[Pair[TensorBundleInfo]] = match_bundles(
|
||||
dfs=Pair(x=baseline_df, y=target_df),
|
||||
skip_keys={"filename"},
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
extra_pair: Pair[TensorBundleInfo] = [
|
||||
r for r in results if r.y[0].name == "t_extra"
|
||||
][0]
|
||||
assert extra_pair.x == []
|
||||
|
||||
def test_empty_target_returns_empty(self) -> None:
|
||||
target_df: pl.DataFrame = _make_df([])
|
||||
baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
|
||||
|
||||
results: list[Pair[TensorBundleInfo]] = match_bundles(
|
||||
dfs=Pair(x=baseline_df, y=target_df),
|
||||
skip_keys={"filename"},
|
||||
)
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_skip_step_groups_across_steps(self) -> None:
|
||||
target_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a", step=0),
|
||||
_make_row(name="t_a", step=1),
|
||||
])
|
||||
baseline_df: pl.DataFrame = _make_df([
|
||||
_make_row(name="t_a", step=0),
|
||||
_make_row(name="t_a", step=1),
|
||||
])
|
||||
|
||||
results: list[Pair[TensorBundleInfo]] = match_bundles(
|
||||
dfs=Pair(x=baseline_df, y=target_df),
|
||||
skip_keys={"filename", "step"},
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0].y) == 2
|
||||
|
||||
|
||||
class TestRowsToTensorInfos:
|
||||
def test_filters_extra_columns(self) -> None:
|
||||
rows: list[dict[str, Any]] = [
|
||||
{"filename": "a.pt", "name": "t_a", "step": 0, "rank": 7}
|
||||
]
|
||||
infos: list[TensorFileInfo] = _rows_to_tensor_infos(rows)
|
||||
|
||||
assert len(infos) == 1
|
||||
assert infos[0] == TensorFileInfo(filename="a.pt", name="t_a", step=0)
|
||||
|
||||
def test_empty_rows(self) -> None:
|
||||
infos: list[TensorFileInfo] = _rows_to_tensor_infos([])
|
||||
assert infos == []
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user