Support multi sharding group on the same dimension in dump comparator (#19601)

This commit is contained in:
fzyzcjy
2026-03-01 10:36:48 +08:00
committed by GitHub
parent 46960e65cf
commit ea6ff7b01f
14 changed files with 469 additions and 151 deletions
@@ -85,8 +85,8 @@ class TestComputePerStepSubPlans:
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),
_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] = [
@@ -177,33 +177,33 @@ class TestComputeAlignerPlan:
class TestComputePerStepSubPlansThd:
def test_thd_zigzag_returns_thd_plans(self) -> None:
"""t(cp,zigzag) h(tp) generates THD-typed unsharder + reorderer plans."""
"""t(cp:zigzag) h(tp) generates THD-typed unsharder + reorderer plans."""
thd_global_seq_lens: list[int] = [100, 64, 92]
result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(
metas=[
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=0,
cp_size=2,
tp_rank=0,
tp_size=2,
),
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=0,
cp_size=2,
tp_rank=1,
tp_size=2,
),
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=1,
cp_size=2,
tp_rank=0,
tp_size=2,
),
_make_meta(
dims="t(cp,zigzag) h(tp)",
dims="t(cp:zigzag) h(tp)",
cp_rank=1,
cp_size=2,
tp_rank=1,