chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Reorganize modules and pipeline in dump comparator (#19374)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-26 10:00:13 +08:00
committed by GitHub
parent 508b8e3387
commit 2739d7df62
35 changed files with 459 additions and 333 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/conftest.py
View File
test/registered/debug_utils/comparator/aligner/reorderer/conftest.py
View File
test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py
+50
View File
@@ -0,0 +1,50 @@
import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
_reorder_zigzag_to_natural,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestZigzagToNatural:
def test_zigzag_to_natural_cp2(self) -> None:
"""cp_size=2: zigzag order [0,3,1,2] -> natural [0,1,2,3]."""
natural = torch.arange(24).reshape(4, 6)
chunks = list(natural.chunk(4, dim=0))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=2)
assert torch.equal(result, natural)
def test_zigzag_to_natural_cp3(self) -> None:
"""cp_size=3: zigzag 162534 -> natural 123456 (1-indexed)."""
natural = torch.arange(60).reshape(6, 10)
chunks = list(natural.chunk(6, dim=0))
zigzag_order: list[int] = [0, 5, 1, 4, 2, 3]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=3)
assert torch.equal(result, natural)
def test_zigzag_to_natural_arbitrary_dim(self) -> None:
"""Reorder along dim=1 instead of dim=0."""
natural = torch.arange(48).reshape(3, 4, 4)
chunks = list(natural.chunk(4, dim=1))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=1)
result = _reorder_zigzag_to_natural(zigzagged, dim=1, cp_size=2)
assert torch.equal(result, natural)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/aligner/test_reorder.py → test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py
+31 -63
View File
@@ -3,63 +3,30 @@ import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.aligner.reorder import (
ReorderPlan,
_reorder_zigzag_to_natural,
compute_reorder_plans,
execute_reorder_plan,
from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
execute_reorderer_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
execute_unshard_plan,
from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import (
compute_reorderer_plans,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
execute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
compute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestZigzagToNatural:
def test_zigzag_to_natural_cp2(self) -> None:
"""cp_size=2: zigzag order [0,3,1,2] -> natural [0,1,2,3]."""
natural = torch.arange(24).reshape(4, 6)
chunks = list(natural.chunk(4, dim=0))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=2)
assert torch.equal(result, natural)
def test_zigzag_to_natural_cp3(self) -> None:
"""cp_size=3: zigzag 162534 -> natural 123456 (1-indexed)."""
natural = torch.arange(60).reshape(6, 10)
chunks = list(natural.chunk(6, dim=0))
zigzag_order: list[int] = [0, 5, 1, 4, 2, 3]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=3)
assert torch.equal(result, natural)
def test_zigzag_to_natural_arbitrary_dim(self) -> None:
"""Reorder along dim=1 instead of dim=0."""
natural = torch.arange(48).reshape(3, 4, 4)
chunks = list(natural.chunk(4, dim=1))
zigzag_order: list[int] = [0, 3, 1, 2]
zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=1)
result = _reorder_zigzag_to_natural(zigzagged, dim=1, cp_size=2)
assert torch.equal(result, natural)
class TestComputeReorderPlans:
def test_compute_reorder_plans_zigzag(self) -> None:
"""s(cp,zigzag) produces a ReorderPlan."""
class TestComputeReordererPlans:
def test_compute_reorderer_plans_zigzag(self) -> None:
"""s(cp,zigzag) produces a ReordererPlan."""
dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
{
@@ -67,7 +34,7 @@ class TestComputeReorderPlans:
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
},
]
plans = compute_reorder_plans(
plans = compute_reorderer_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
@@ -76,7 +43,7 @@ class TestComputeReorderPlans:
assert plans[0].params.dim == 1
assert plans[0].params.cp_size == 2
def test_compute_reorder_plans_non_seq_dim_raises(self) -> None:
def test_compute_reorderer_plans_non_seq_dim_raises(self) -> None:
"""Zigzag on non-sequence dim (e.g. t(cp,zigzag)) raises ValueError."""
dim_specs = parse_dims("t(cp,zigzag) h(tp)")
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
@@ -86,9 +53,9 @@ class TestComputeReorderPlans:
},
]
with pytest.raises(ValueError, match="only supported on sequence dims"):
compute_reorder_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
def test_compute_reorder_plans_natural(self) -> None:
def test_compute_reorderer_plans_natural(self) -> None:
"""s(cp) and s(cp,natural) produce no reorder plans."""
for dims_str in ["b s(cp) h(tp)", "b s(cp,natural) h(tp)"]:
dim_specs = parse_dims(dims_str)
@@ -98,7 +65,7 @@ class TestComputeReorderPlans:
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
},
]
plans = compute_reorder_plans(
plans = compute_reorderer_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
assert plans == []
@@ -132,23 +99,24 @@ class TestCpZigzagTpE2E:
dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
unshard_plans = compute_unshard_plan(
unsharder_plans = compute_unsharder_plan(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
reorder_plans = compute_reorder_plans(
reorderer_plans = compute_reorderer_plans(
dim_specs=dim_specs, parallel_infos=parallel_infos
)
all_plans = [*unshard_plans, *reorder_plans]
all_plans = [*unsharder_plans, *reorderer_plans]
assert len(unshard_plans) == 2
assert len(reorder_plans) == 1
assert len(unsharder_plans) == 2
assert len(reorderer_plans) == 1
current: list[torch.Tensor] = tensors
for plan in all_plans:
if isinstance(plan, ReorderPlan):
current = execute_reorder_plan(plan, current)
else:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context():
for plan in all_plans:
if isinstance(plan, ReordererPlan):
current = execute_reorderer_plan(plan, current)
else:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
test/registered/debug_utils/comparator/aligner/unsharder/conftest.py
View File
test/registered/debug_utils/comparator/aligner/unshard/test_execute.py → test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
+79 -60
View File
@@ -3,25 +3,26 @@ import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
_apply_unshard,
_verify_replicated_group,
execute_unshard_plan,
execute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
compute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
AxisInfo,
PickParams,
)
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestExecuteUnshardPlan:
class TestExecuteUnsharderPlan:
def test_tp4_concat(self) -> None:
full_tensor = torch.randn(2, 8, 16)
shards = list(full_tensor.chunk(4, dim=1))
@@ -30,10 +31,11 @@ class TestExecuteUnshardPlan:
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
result, warnings = execute_unshard_plan(plans[0], shards)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], shards)
assert len(result) == 1
assert torch.allclose(result[0], full_tensor)
assert warnings == []
@@ -49,7 +51,7 @@ class TestExecuteUnshardPlan:
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
]
dim_specs = parse_dims("h(tp) d")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
tensors_ordered_by_world_rank = [
@@ -59,7 +61,8 @@ class TestExecuteUnshardPlan:
shards[1], # world_rank=3, axis_rank=1
]
result, warnings = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], tensors_ordered_by_world_rank)
assert len(result) == 1
assert torch.allclose(result[0], full_tensor)
assert warnings == []
@@ -82,7 +85,7 @@ class TestExecuteUnshardPlan:
}
)
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
tensors: list[torch.Tensor] = []
@@ -91,10 +94,12 @@ class TestExecuteUnshardPlan:
for tp_rank in range(4):
tensors.append(source[tp_rank])
intermediate, _ = execute_unshard_plan(plans[0], tensors)
with warning_sink.context() as _warnings:
intermediate = execute_unsharder_plan(plans[0], tensors)
assert len(intermediate) == 4
final, _ = execute_unshard_plan(plans[1], intermediate)
with warning_sink.context() as _warnings:
final = execute_unsharder_plan(plans[1], intermediate)
assert len(final) == 1
def test_cp_tp_concat(self) -> None:
@@ -117,12 +122,13 @@ class TestExecuteUnshardPlan:
)
dim_specs = parse_dims("b s(cp) h(tp)")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
current = tensors
for plan in plans:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context() as _warnings:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
@@ -158,12 +164,13 @@ class TestExecuteUnshardPlan:
)
dim_specs = parse_dims("b s(cp) h(tp)")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
current = tensors
for plan in plans:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context() as _warnings:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
@@ -211,12 +218,13 @@ class TestExecuteUnshardPlan:
)
dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 3
current = tensors
for plan in plans:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context() as _warnings:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
@@ -259,12 +267,13 @@ class TestExecuteUnshardPlan:
)
dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 3
current = tensors
for plan in plans:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context() as _warnings:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
@@ -280,11 +289,12 @@ class TestPickOperation:
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert isinstance(plans[0].params, PickParams)
result, warnings = execute_unshard_plan(plans[0], [tensor, tensor.clone()])
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], [tensor, tensor.clone()])
assert len(result) == 1
assert torch.allclose(result[0], tensor)
assert warnings == []
@@ -311,7 +321,7 @@ class TestPickOperation:
},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
assert len(pick_plans) == 1
assert pick_plans[0].axis == ParallelAxis.CP
@@ -319,7 +329,8 @@ class TestPickOperation:
tensor = torch.randn(4)
tensors = [tensor.clone() for _ in range(4)]
result, warnings = execute_unshard_plan(pick_plans[0], tensors)
with warning_sink.context() as warnings:
result = execute_unsharder_plan(pick_plans[0], tensors)
assert len(result) == 2
assert warnings == []
@@ -342,18 +353,19 @@ class TestPickOperation:
)
dim_specs = parse_dims("b s(cp) d")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
current = tensors
for plan in plans:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context() as _warnings:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
def test_fully_replicated_e2e(self) -> None:
"""CP2 TP2, dims='b h d': fully replicated → 2 pick steps → 1 tensor."""
"""CP2 TP2, dims='b h d': fully replicated -> 2 pick steps -> 1 tensor."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
@@ -370,13 +382,14 @@ class TestPickOperation:
)
dim_specs = parse_dims("b h d")
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
assert all(isinstance(p.params, PickParams) for p in plans)
current = tensors
for plan in plans:
current, _ = execute_unshard_plan(plan, current)
with warning_sink.context() as _warnings:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
assert torch.allclose(current[0], full_tensor)
@@ -388,11 +401,12 @@ class TestVerifyReplicatedGroup:
tensor_a = torch.ones(4)
tensor_b = torch.ones(4) + 0.1
warnings = _verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.TP,
group_index=0,
)
with warning_sink.context() as warnings:
_verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(warnings) == 1
assert warnings[0].axis == "tp"
assert warnings[0].group_index == 0
@@ -404,11 +418,12 @@ class TestVerifyReplicatedGroup:
"""_verify_replicated_group produces no warning for identical replicas."""
tensor = torch.randn(4, 8)
warnings = _verify_replicated_group(
[tensor, tensor.clone()],
axis=ParallelAxis.TP,
group_index=0,
)
with warning_sink.context() as warnings:
_verify_replicated_group(
[tensor, tensor.clone()],
axis=ParallelAxis.TP,
group_index=0,
)
assert warnings == []
def test_multiple_mismatches(self) -> None:
@@ -417,55 +432,59 @@ class TestVerifyReplicatedGroup:
other_a = torch.ones(4)
other_b = torch.ones(4) * 2
warnings = _verify_replicated_group(
[baseline, other_a, other_b],
axis=ParallelAxis.CP,
group_index=1,
)
with warning_sink.context() as warnings:
_verify_replicated_group(
[baseline, other_a, other_b],
axis=ParallelAxis.CP,
group_index=1,
)
assert len(warnings) == 2
assert warnings[0].differing_index == 1
assert warnings[1].differing_index == 2
assert warnings[1].max_abs_diff == pytest.approx(2.0, abs=1e-5)
def test_execute_returns_warnings(self) -> None:
"""execute_unshard_plan returns warnings for replicated mismatch."""
"""execute_unsharder_plan emits warnings for replicated mismatch."""
dim_specs = parse_dims("h d")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
tensor_a = torch.zeros(4)
tensor_b = torch.ones(4)
result, warnings = execute_unshard_plan(plans[0], [tensor_a, tensor_b])
with warning_sink.context() as warnings:
result = execute_unsharder_plan(plans[0], [tensor_a, tensor_b])
assert len(result) == 1
assert len(warnings) == 1
assert torch.allclose(result[0], tensor_a)
def test_atol_boundary_within(self) -> None:
"""Difference exactly at atol (1e-6) → torch.allclose passes → no warning."""
"""Difference exactly at atol (1e-6) -> torch.allclose passes -> no warning."""
baseline = torch.zeros(4)
other = torch.full((4,), 1e-6)
warnings = _verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
with warning_sink.context() as warnings:
_verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
assert warnings == []
def test_atol_boundary_exceeded(self) -> None:
"""Difference just above atol (1e-6 + 1e-9) → torch.allclose fails → warning."""
"""Difference just above atol (1e-6 + 1e-9) -> torch.allclose fails -> warning."""
baseline = torch.zeros(4)
other = torch.full((4,), 1e-6 + 1e-9)
warnings = _verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
with warning_sink.context() as warnings:
_verify_replicated_group(
[baseline, other],
axis=ParallelAxis.TP,
group_index=0,
)
assert len(warnings) == 1
assert warnings[0].differing_index == 1
test/registered/debug_utils/comparator/aligner/unshard/test_parallel_info.py → test/registered/debug_utils/comparator/aligner/unsharder/test_parallel_info.py
+2 -2
View File
@@ -2,10 +2,10 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.aligner.unshard.parallel_info import (
from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
normalize_parallel_info,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
from sglang.test.ci.ci_register import register_cpu_ci
test/registered/debug_utils/comparator/aligner/unshard/test_plan.py → test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py
+25 -25
View File
@@ -2,10 +2,10 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
compute_unshard_plan,
from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
compute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
AxisInfo,
ConcatParams,
PickParams,
@@ -16,13 +16,13 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
class TestComputeUnshardPlan:
class TestComputeUnsharderPlan:
def test_tp4_plan(self) -> None:
dim_specs = parse_dims("b s h(tp) d")
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.TP
@@ -36,18 +36,18 @@ class TestComputeUnshardPlan:
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
]
with pytest.raises(ValueError, match="Inconsistent axis_size"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
def test_missing_axis_in_parallel_info_raises(self) -> None:
dim_specs = parse_dims("h(tp)")
parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}]
with pytest.raises(ValueError, match="missing parallel_info"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
def test_empty_parallel_infos_raises(self) -> None:
dim_specs = parse_dims("h(tp)")
with pytest.raises(ValueError, match="must not be empty"):
compute_unshard_plan(dim_specs, [])
compute_unsharder_plan(dim_specs, [])
def test_scrambled_world_ranks(self) -> None:
"""world_rank order != axis_rank order."""
@@ -58,14 +58,14 @@ class TestComputeUnshardPlan:
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].groups == [[1, 3, 0, 2]]
def test_no_sharded_axes_returns_empty(self) -> None:
dim_specs = parse_dims("b s d")
parallel_infos = [{}]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert plans == []
def test_multi_axis_plan(self) -> None:
@@ -89,7 +89,7 @@ class TestComputeUnshardPlan:
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
assert plans[0].axis == ParallelAxis.CP
@@ -108,7 +108,7 @@ class TestComputeUnshardPlan:
}
)
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -144,7 +144,7 @@ class TestComputeUnshardPlan:
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
@@ -168,7 +168,7 @@ class TestComputeUnshardPlan:
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
]
with pytest.raises(ValueError, match="axis_rank coverage.*incomplete"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
def test_reduction_not_implemented_raises(self) -> None:
dim_specs = parse_dims("h(tp,partial)")
@@ -176,14 +176,14 @@ class TestComputeUnshardPlan:
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
with pytest.raises(NotImplementedError, match="reduction"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
def test_ordering_zigzag_accepted(self) -> None:
dim_specs = parse_dims("s(cp,zigzag)")
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.CP
@@ -192,7 +192,7 @@ class TestComputeUnshardPlan:
parallel_infos = [
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert plans[0].axis == ParallelAxis.CP
@@ -211,7 +211,7 @@ class TestComputeUnshardPlan:
}
)
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 3
assert plans[0].axis == ParallelAxis.EP
@@ -246,7 +246,7 @@ class TestComputeUnshardPlan:
},
]
with pytest.raises(ValueError, match="missing parallel_info"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
class TestReplicatedAxes:
@@ -271,7 +271,7 @@ class TestReplicatedAxes:
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
assert plans[0].axis == ParallelAxis.TP
@@ -305,7 +305,7 @@ class TestReplicatedAxes:
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
assert all(isinstance(p.params, PickParams) for p in plans)
@@ -327,7 +327,7 @@ class TestReplicatedAxes:
}
)
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 3
pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
@@ -360,7 +360,7 @@ class TestReplicatedAxes:
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
},
]
plans = compute_unshard_plan(dim_specs, parallel_infos)
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
assert plans[0].axis == ParallelAxis.CP
@@ -382,7 +382,7 @@ class TestReplicatedAxes:
},
]
with pytest.raises(ValueError, match="Inconsistent axis_size"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
def test_replicated_axis_missing_from_rank_raises(self) -> None:
"""A rank missing a replicated axis that other ranks have raises ValueError."""
@@ -398,7 +398,7 @@ class TestReplicatedAxes:
},
]
with pytest.raises(ValueError, match="missing parallel_info"):
compute_unshard_plan(dim_specs, parallel_infos)
compute_unsharder_plan(dim_specs, parallel_infos)
if __name__ == "__main__":