Support replication axis in dump comparator (#19282)
This commit is contained in:
fzyzcjy
@@ -2,30 +2,80 @@ import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
|
||||
ConcatParams,
|
||||
PickParams,
|
||||
UnshardParams,
|
||||
UnshardPlan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
AlignWarning,
|
||||
ReplicatedMismatchWarning,
|
||||
)
|
||||
|
||||
|
||||
def execute_unshard_plan(
|
||||
plan: UnshardPlan,
|
||||
tensors: list[torch.Tensor],
|
||||
) -> list[torch.Tensor]:
|
||||
) -> tuple[list[torch.Tensor], list[AlignWarning]]:
|
||||
all_warnings: list[AlignWarning] = []
|
||||
result: list[torch.Tensor] = []
|
||||
for group in plan.groups:
|
||||
|
||||
for group_idx, group in enumerate(plan.groups):
|
||||
group_tensors = [tensors[i] for i in group]
|
||||
result.append(_apply_unshard(plan.params, group_tensors))
|
||||
return result
|
||||
tensor, warnings = _apply_unshard(
|
||||
plan.params,
|
||||
group_tensors,
|
||||
axis=plan.axis,
|
||||
group_index=group_idx,
|
||||
)
|
||||
result.append(tensor)
|
||||
all_warnings.extend(warnings)
|
||||
|
||||
return result, all_warnings
|
||||
|
||||
|
||||
def _apply_unshard(
|
||||
params: UnshardParams, ordered_tensors: list[torch.Tensor]
|
||||
) -> torch.Tensor:
|
||||
params: UnshardParams,
|
||||
ordered_tensors: list[torch.Tensor],
|
||||
*,
|
||||
axis: ParallelAxis,
|
||||
group_index: int,
|
||||
) -> tuple[torch.Tensor, list[AlignWarning]]:
|
||||
if isinstance(params, PickParams):
|
||||
warnings = _verify_replicated_group(
|
||||
ordered_tensors,
|
||||
axis=axis,
|
||||
group_index=group_index,
|
||||
)
|
||||
return ordered_tensors[0], warnings
|
||||
|
||||
if isinstance(params, ConcatParams):
|
||||
return _unshard_concat(ordered_tensors, dim=params.dim)
|
||||
return torch.cat(ordered_tensors, dim=params.dim), []
|
||||
|
||||
# Phase 2: ReduceSumParams, CpZigzagParams
|
||||
raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
|
||||
|
||||
|
||||
def _unshard_concat(tensors: list[torch.Tensor], dim: int) -> torch.Tensor:
|
||||
return torch.cat(tensors, dim=dim)
|
||||
def _verify_replicated_group(
|
||||
ordered_tensors: list[torch.Tensor],
|
||||
*,
|
||||
axis: ParallelAxis,
|
||||
group_index: int,
|
||||
) -> list[ReplicatedMismatchWarning]:
|
||||
warnings: list[ReplicatedMismatchWarning] = []
|
||||
baseline = ordered_tensors[0]
|
||||
|
||||
for i in range(1, len(ordered_tensors)):
|
||||
other = ordered_tensors[i]
|
||||
if not torch.allclose(baseline, other, atol=1e-6):
|
||||
warnings.append(
|
||||
ReplicatedMismatchWarning(
|
||||
axis=axis.value,
|
||||
group_index=group_index,
|
||||
differing_index=i,
|
||||
baseline_index=0,
|
||||
max_abs_diff=(baseline - other).abs().max().item(),
|
||||
)
|
||||
)
|
||||
|
||||
return warnings
|
||||
|
||||
@@ -4,6 +4,7 @@ from typing import NamedTuple
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
|
||||
AxisInfo,
|
||||
ConcatParams,
|
||||
PickParams,
|
||||
UnshardParams,
|
||||
UnshardPlan,
|
||||
)
|
||||
@@ -32,31 +33,38 @@ def compute_unshard_plan(
|
||||
for dim_idx, spec in enumerate(dim_specs)
|
||||
if spec.parallel is not None
|
||||
}
|
||||
if not sharded_axis_infos:
|
||||
sharded_axes: set[ParallelAxis] = set(sharded_axis_infos)
|
||||
|
||||
all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info}
|
||||
replicated_axes: set[ParallelAxis] = all_axes - sharded_axes
|
||||
|
||||
if not sharded_axes and not replicated_axes:
|
||||
return []
|
||||
|
||||
_validate(sharded_axes=set(sharded_axis_infos), parallel_infos=parallel_infos)
|
||||
_validate(
|
||||
axes_to_validate=sharded_axes | replicated_axes,
|
||||
parallel_infos=parallel_infos,
|
||||
)
|
||||
|
||||
current_coords: _CoordsList = [
|
||||
{axis: info[axis].axis_rank for axis in sharded_axis_infos}
|
||||
{axis: info[axis].axis_rank for axis in sharded_axes | replicated_axes}
|
||||
for info in parallel_infos
|
||||
]
|
||||
|
||||
axis_and_params: list[tuple[ParallelAxis, UnshardParams]] = [
|
||||
(axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
|
||||
] + [
|
||||
(axis, _resolve_unshard_params(spec=spec, dim_index=dim_index))
|
||||
for axis, (dim_index, spec) in sharded_axis_infos.items()
|
||||
]
|
||||
|
||||
plans: list[UnshardPlan] = []
|
||||
for axis, (dim_index, spec) in sharded_axis_infos.items():
|
||||
for axis, params in axis_and_params:
|
||||
result = _group_and_project(
|
||||
current_coords=current_coords,
|
||||
target_axis=axis,
|
||||
)
|
||||
|
||||
plans.append(
|
||||
UnshardPlan(
|
||||
axis=axis,
|
||||
params=_resolve_unshard_params(spec=spec, dim_index=dim_index),
|
||||
groups=result.groups,
|
||||
)
|
||||
)
|
||||
|
||||
plans.append(UnshardPlan(axis=axis, params=params, groups=result.groups))
|
||||
current_coords = result.projected_coords
|
||||
|
||||
return plans
|
||||
@@ -64,18 +72,18 @@ def compute_unshard_plan(
|
||||
|
||||
def _validate(
|
||||
*,
|
||||
sharded_axes: set[ParallelAxis],
|
||||
axes_to_validate: set[ParallelAxis],
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]],
|
||||
) -> None:
|
||||
"""Check that every rank has all sharded axes, sizes are consistent, and ranks are complete."""
|
||||
"""Check that every rank has all axes, sizes are consistent, and ranks are complete."""
|
||||
axis_sizes: dict[ParallelAxis, int] = {}
|
||||
|
||||
for world_rank, parallel_info in enumerate(parallel_infos):
|
||||
for axis in sharded_axes:
|
||||
for axis in axes_to_validate:
|
||||
if axis not in parallel_info:
|
||||
raise ValueError(
|
||||
f"world_rank={world_rank} missing parallel_info for "
|
||||
f"sharded axis {axis.value!r}"
|
||||
f"axis {axis.value!r}"
|
||||
)
|
||||
|
||||
axis_info = parallel_info[axis]
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
from typing import Annotated, Literal, Union
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.utils import _FrozenBase
|
||||
@@ -16,7 +18,14 @@ class ConcatParams(_FrozenBase):
|
||||
dim: int
|
||||
|
||||
|
||||
UnshardParams = ConcatParams
|
||||
class PickParams(_FrozenBase):
|
||||
op: Literal["pick"] = "pick"
|
||||
|
||||
|
||||
UnshardParams = Annotated[
|
||||
Union[ConcatParams, PickParams],
|
||||
Field(discriminator="op"),
|
||||
]
|
||||
|
||||
|
||||
class UnshardPlan(_FrozenBase):
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from abc import abstractmethod
|
||||
from typing import Annotated, Literal, Union
|
||||
|
||||
from pydantic import Discriminator, TypeAdapter
|
||||
from pydantic import Discriminator, Field, TypeAdapter
|
||||
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
|
||||
format_comparison,
|
||||
@@ -12,9 +12,38 @@ from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
|
||||
from sglang.srt.debug_utils.comparator.utils import _StrictBase
|
||||
|
||||
|
||||
class ReplicatedMismatchWarning(_StrictBase):
|
||||
kind: Literal["replicated_mismatch"] = "replicated_mismatch"
|
||||
axis: str
|
||||
group_index: int
|
||||
differing_index: int
|
||||
baseline_index: int
|
||||
max_abs_diff: float
|
||||
|
||||
def to_text(self) -> str:
|
||||
return (
|
||||
f"Replicated along {self.axis}: group {self.group_index}, "
|
||||
f"index {self.differing_index} differs from {self.baseline_index} "
|
||||
f"(max_abs_diff={self.max_abs_diff:.6e})"
|
||||
)
|
||||
|
||||
|
||||
AlignWarning = (
|
||||
ReplicatedMismatchWarning # future: Annotated[Union[...], Discriminator("kind")]
|
||||
)
|
||||
|
||||
|
||||
class _OutputRecord(_StrictBase):
|
||||
align_warnings: list[AlignWarning] = Field(default_factory=list)
|
||||
|
||||
@abstractmethod
|
||||
def to_text(self) -> str: ...
|
||||
def _format_body(self) -> str: ...
|
||||
|
||||
def to_text(self) -> str:
|
||||
body = self._format_body()
|
||||
if self.align_warnings:
|
||||
body += "\n" + "\n".join(f" ⚠ {w.to_text()}" for w in self.align_warnings)
|
||||
return body
|
||||
|
||||
|
||||
class ConfigRecord(_OutputRecord):
|
||||
@@ -25,7 +54,7 @@ class ConfigRecord(_OutputRecord):
|
||||
start_step: int
|
||||
end_step: int
|
||||
|
||||
def to_text(self) -> str:
|
||||
def _format_body(self) -> str:
|
||||
return (
|
||||
f"Config: baseline={self.baseline_path} target={self.target_path}\n"
|
||||
f"diff_threshold={self.diff_threshold} "
|
||||
@@ -39,10 +68,12 @@ class SkipRecord(_OutputRecord):
|
||||
reason: str
|
||||
|
||||
@property
|
||||
def category(self):
|
||||
def category(self) -> str:
|
||||
if self.align_warnings:
|
||||
return "failed"
|
||||
return "skipped"
|
||||
|
||||
def to_text(self) -> str:
|
||||
def _format_body(self) -> str:
|
||||
return f"Skip: {self.name} ({self.reason})"
|
||||
|
||||
|
||||
@@ -50,10 +81,12 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
|
||||
type: Literal["comparison"] = "comparison"
|
||||
|
||||
@property
|
||||
def category(self):
|
||||
def category(self) -> str:
|
||||
if self.align_warnings:
|
||||
return "failed"
|
||||
return "passed" if self.diff is not None and self.diff.passed else "failed"
|
||||
|
||||
def to_text(self) -> str:
|
||||
def _format_body(self) -> str:
|
||||
return format_comparison(self)
|
||||
|
||||
|
||||
@@ -64,7 +97,7 @@ class SummaryRecord(_OutputRecord):
|
||||
failed: int
|
||||
skipped: int
|
||||
|
||||
def to_text(self) -> str:
|
||||
def _format_body(self) -> str:
|
||||
return (
|
||||
f"Summary: {self.passed} passed, {self.failed} failed, "
|
||||
f"{self.skipped} skipped (total {self.total})"
|
||||
|
||||
@@ -20,6 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
|
||||
from sglang.srt.debug_utils.comparator.dims import parse_dims
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
AlignWarning,
|
||||
ComparisonRecord,
|
||||
SkipRecord,
|
||||
)
|
||||
@@ -50,12 +51,13 @@ def process_tensor_group(
|
||||
t_extracted = _extract_tensors(t_tensors)
|
||||
del b_tensors, t_tensors
|
||||
|
||||
b_tensor = _execute_plans(b_extracted, b_plans)
|
||||
t_tensor = _execute_plans(t_extracted, t_plans)
|
||||
b_tensor, b_warns = _execute_plans(b_extracted, b_plans)
|
||||
t_tensor, t_warns = _execute_plans(t_extracted, t_plans)
|
||||
all_warnings: list[AlignWarning] = b_warns + t_warns
|
||||
|
||||
if b_tensor is None or t_tensor is None:
|
||||
reason = "baseline_load_failed" if b_tensor is None else "target_load_failed"
|
||||
return SkipRecord(name=name, reason=reason)
|
||||
return SkipRecord(name=name, reason=reason, align_warnings=all_warnings)
|
||||
|
||||
info = compare_tensors(
|
||||
x_baseline=b_tensor,
|
||||
@@ -64,7 +66,7 @@ def process_tensor_group(
|
||||
diff_threshold=diff_threshold,
|
||||
)
|
||||
|
||||
return ComparisonRecord(**info.model_dump())
|
||||
return ComparisonRecord(**info.model_dump(), align_warnings=all_warnings)
|
||||
|
||||
|
||||
def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
|
||||
@@ -112,27 +114,32 @@ def _extract_tensors(
|
||||
def _execute_plans(
|
||||
tensors: list[torch.Tensor],
|
||||
plans: list[Plan],
|
||||
) -> Optional[torch.Tensor]:
|
||||
) -> tuple[Optional[torch.Tensor], list[AlignWarning]]:
|
||||
if not tensors:
|
||||
return None
|
||||
return None, []
|
||||
|
||||
if not plans:
|
||||
if len(tensors) != 1:
|
||||
return None
|
||||
return tensors[0]
|
||||
return None, []
|
||||
return tensors[0], []
|
||||
|
||||
warnings: list[AlignWarning] = []
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = _execute_plan(current, plan)
|
||||
current, new_warnings = _execute_plan(current, plan)
|
||||
warnings.extend(new_warnings)
|
||||
|
||||
assert len(current) == 1
|
||||
return current[0]
|
||||
return current[0], warnings
|
||||
|
||||
|
||||
def _execute_plan(tensors, plan):
|
||||
def _execute_plan(
|
||||
tensors: list[torch.Tensor],
|
||||
plan: Plan,
|
||||
) -> tuple[list[torch.Tensor], list[AlignWarning]]:
|
||||
if isinstance(plan, UnshardPlan):
|
||||
return execute_unshard_plan(plan, tensors)
|
||||
elif isinstance(plan, ReorderPlan):
|
||||
return execute_reorder_plan(plan, tensors)
|
||||
return execute_reorder_plan(plan, tensors), []
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown {plan=}")
|
||||
|
||||
@@ -148,7 +148,7 @@ class TestCpZigzagTpE2E:
|
||||
if isinstance(plan, ReorderPlan):
|
||||
current = execute_reorder_plan(plan, current)
|
||||
else:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
current, _ = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
@@ -5,12 +5,16 @@ import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
|
||||
_apply_unshard,
|
||||
_verify_replicated_group,
|
||||
execute_unshard_plan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
|
||||
compute_unshard_plan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
|
||||
AxisInfo,
|
||||
PickParams,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -29,9 +33,10 @@ class TestExecuteUnshardPlan:
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
|
||||
result = execute_unshard_plan(plans[0], shards)
|
||||
result, warnings = execute_unshard_plan(plans[0], shards)
|
||||
assert len(result) == 1
|
||||
assert torch.allclose(result[0], full_tensor)
|
||||
assert warnings == []
|
||||
|
||||
def test_scrambled_world_ranks_correct_result(self) -> None:
|
||||
full_tensor = torch.randn(4, 8)
|
||||
@@ -54,9 +59,10 @@ class TestExecuteUnshardPlan:
|
||||
shards[1], # world_rank=3, axis_rank=1
|
||||
]
|
||||
|
||||
result = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
|
||||
result, warnings = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
|
||||
assert len(result) == 1
|
||||
assert torch.allclose(result[0], full_tensor)
|
||||
assert warnings == []
|
||||
|
||||
def test_single_step_reduces_tensor_count(self) -> None:
|
||||
"""8 tensors with 2 groups of 4 produce 2 output tensors."""
|
||||
@@ -85,10 +91,10 @@ class TestExecuteUnshardPlan:
|
||||
for tp_rank in range(4):
|
||||
tensors.append(source[tp_rank])
|
||||
|
||||
intermediate = execute_unshard_plan(plans[0], tensors)
|
||||
intermediate, _ = execute_unshard_plan(plans[0], tensors)
|
||||
assert len(intermediate) == 4
|
||||
|
||||
final = execute_unshard_plan(plans[1], intermediate)
|
||||
final, _ = execute_unshard_plan(plans[1], intermediate)
|
||||
assert len(final) == 1
|
||||
|
||||
def test_cp_tp_concat(self) -> None:
|
||||
@@ -116,7 +122,7 @@ class TestExecuteUnshardPlan:
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
current, _ = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
@@ -157,7 +163,7 @@ class TestExecuteUnshardPlan:
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
current, _ = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
@@ -169,7 +175,12 @@ class TestExecuteUnshardPlan:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported unshard"):
|
||||
_apply_unshard(_FakeParams(), [torch.randn(2, 2)])
|
||||
_apply_unshard(
|
||||
_FakeParams(),
|
||||
[torch.randn(2, 2)],
|
||||
axis=ParallelAxis.TP,
|
||||
group_index=0,
|
||||
)
|
||||
|
||||
def test_cp_tp_ep_three_axis_concat(self) -> None:
|
||||
"""CP=2 + TP=2 + EP=2: three-step unshard reconstructs original tensor."""
|
||||
@@ -205,7 +216,7 @@ class TestExecuteUnshardPlan:
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
current, _ = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
@@ -253,11 +264,211 @@ class TestExecuteUnshardPlan:
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
current, _ = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
|
||||
class TestPickOperation:
|
||||
def test_pick_single_group(self) -> None:
|
||||
"""PickParams picks the first tensor from a single group."""
|
||||
tensor = torch.randn(4, 8)
|
||||
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)
|
||||
assert len(plans) == 1
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
|
||||
result, warnings = execute_unshard_plan(plans[0], [tensor, tensor.clone()])
|
||||
assert len(result) == 1
|
||||
assert torch.allclose(result[0], tensor)
|
||||
assert warnings == []
|
||||
|
||||
def test_pick_multiple_groups(self) -> None:
|
||||
"""PickParams with multiple groups picks one from each."""
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
|
||||
plans = compute_unshard_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
|
||||
|
||||
tensor = torch.randn(4)
|
||||
tensors = [tensor.clone() for _ in range(4)]
|
||||
|
||||
result, warnings = execute_unshard_plan(pick_plans[0], tensors)
|
||||
assert len(result) == 2
|
||||
assert warnings == []
|
||||
|
||||
def test_replicated_tp_sharded_cp_e2e(self) -> None:
|
||||
"""CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
cp_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
tensors.append(cp_chunks[cp_rank].clone())
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b s(cp) d")
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 2
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current, _ = execute_unshard_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."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
tensors.append(full_tensor.clone())
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b h d")
|
||||
plans = compute_unshard_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)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
|
||||
class TestVerifyReplicatedGroup:
|
||||
def test_warns_on_mismatch(self) -> None:
|
||||
"""_verify_replicated_group produces warning when replicas differ."""
|
||||
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,
|
||||
)
|
||||
assert len(warnings) == 1
|
||||
assert warnings[0].axis == "tp"
|
||||
assert warnings[0].group_index == 0
|
||||
assert warnings[0].differing_index == 1
|
||||
assert warnings[0].baseline_index == 0
|
||||
assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
|
||||
|
||||
def test_no_warn_when_identical(self) -> None:
|
||||
"""_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,
|
||||
)
|
||||
assert warnings == []
|
||||
|
||||
def test_multiple_mismatches(self) -> None:
|
||||
"""_verify_replicated_group reports each differing replica."""
|
||||
baseline = torch.zeros(4)
|
||||
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,
|
||||
)
|
||||
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."""
|
||||
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)
|
||||
|
||||
tensor_a = torch.zeros(4)
|
||||
tensor_b = torch.ones(4)
|
||||
|
||||
result, warnings = execute_unshard_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."""
|
||||
baseline = torch.zeros(4)
|
||||
other = torch.full((4,), 1e-6)
|
||||
|
||||
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."""
|
||||
baseline = torch.zeros(4)
|
||||
other = torch.full((4,), 1e-6 + 1e-9)
|
||||
|
||||
warnings = _verify_replicated_group(
|
||||
[baseline, other],
|
||||
axis=ParallelAxis.TP,
|
||||
group_index=0,
|
||||
)
|
||||
assert len(warnings) == 1
|
||||
assert warnings[0].differing_index == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
@@ -5,7 +5,11 @@ import pytest
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
|
||||
compute_unshard_plan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
|
||||
from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
|
||||
AxisInfo,
|
||||
ConcatParams,
|
||||
PickParams,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -228,7 +232,7 @@ class TestComputeUnshardPlan:
|
||||
assert len(plans[2].groups) == 1
|
||||
assert len(plans[2].groups[0]) == 2
|
||||
|
||||
def test_replicated_axis_raises(self) -> None:
|
||||
def test_sharded_axis_missing_from_rank_raises(self) -> None:
|
||||
"""A world_rank missing a sharded axis raises ValueError."""
|
||||
dim_specs = parse_dims("s(cp) h(tp)")
|
||||
parallel_infos = [
|
||||
@@ -238,7 +242,159 @@ class TestComputeUnshardPlan:
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
# missing TP — replicated
|
||||
# missing TP — sharded axis absent from rank
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="missing parallel_info"):
|
||||
compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
|
||||
class TestReplicatedAxes:
|
||||
def test_replicated_tp_with_sharded_cp(self) -> None:
|
||||
"""CP2 TP2, dims='b s(cp) d' → PickPlan(TP) + ConcatPlan(CP)."""
|
||||
dim_specs = parse_dims("b s(cp) d")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.TP
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert len(plans[0].groups) == 2
|
||||
for group in plans[0].groups:
|
||||
assert len(group) == 2
|
||||
|
||||
assert plans[1].axis == ParallelAxis.CP
|
||||
assert isinstance(plans[1].params, ConcatParams)
|
||||
assert plans[1].params.dim == 1
|
||||
|
||||
def test_fully_replicated(self) -> None:
|
||||
"""CP2 TP2, dims='b h d' → PickPlan(CP) + PickPlan(TP)."""
|
||||
dim_specs = parse_dims("b h d")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert all(isinstance(p.params, PickParams) for p in plans)
|
||||
axes = {p.axis for p in plans}
|
||||
assert axes == {ParallelAxis.CP, ParallelAxis.TP}
|
||||
|
||||
def test_multiple_replicated_one_sharded(self) -> None:
|
||||
"""CP2 TP2 EP2, dims='h(tp)' → PickPlan(CP) + PickPlan(EP) + ConcatPlan(TP)."""
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for ep_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 3
|
||||
pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
|
||||
concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)]
|
||||
assert len(pick_plans) == 2
|
||||
assert len(concat_plans) == 1
|
||||
assert concat_plans[0].axis == ParallelAxis.TP
|
||||
|
||||
replicated_axes = {p.axis for p in pick_plans}
|
||||
assert replicated_axes == {ParallelAxis.CP, ParallelAxis.EP}
|
||||
|
||||
def test_replicated_scrambled_ranks(self) -> None:
|
||||
"""Scrambled world_rank order with replicated axis."""
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.CP
|
||||
assert isinstance(plans[0].params, PickParams)
|
||||
assert plans[1].axis == ParallelAxis.TP
|
||||
assert isinstance(plans[1].params, ConcatParams)
|
||||
|
||||
def test_replicated_axis_inconsistent_size_raises(self) -> None:
|
||||
"""Replicated axis with inconsistent sizes raises ValueError."""
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=4),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="Inconsistent axis_size"):
|
||||
compute_unshard_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."""
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
# missing CP — replicated axis absent from this rank
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="missing parallel_info"):
|
||||
|
||||
@@ -6,6 +6,7 @@ import pytest
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
ComparisonRecord,
|
||||
ConfigRecord,
|
||||
ReplicatedMismatchWarning,
|
||||
SkipRecord,
|
||||
SummaryRecord,
|
||||
parse_record_json,
|
||||
@@ -118,5 +119,71 @@ class TestRecordTypes:
|
||||
assert restored == record
|
||||
|
||||
|
||||
def _make_warning(**overrides) -> ReplicatedMismatchWarning:
|
||||
defaults: dict = dict(
|
||||
axis="tp",
|
||||
group_index=0,
|
||||
differing_index=1,
|
||||
baseline_index=0,
|
||||
max_abs_diff=0.1,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return ReplicatedMismatchWarning(**defaults)
|
||||
|
||||
|
||||
class TestAlignWarnings:
|
||||
def test_comparison_record_failed_when_diff_passed_but_warnings(self):
|
||||
"""ComparisonRecord with diff.passed=True but align_warnings → category=='failed'."""
|
||||
record = ComparisonRecord(
|
||||
name="hidden",
|
||||
baseline=_make_tensor_info(),
|
||||
target=_make_tensor_info(),
|
||||
unified_shape=[4, 8],
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(passed=True),
|
||||
align_warnings=[_make_warning()],
|
||||
)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_skip_record_failed_when_warnings(self):
|
||||
"""SkipRecord with align_warnings → category=='failed' instead of 'skipped'."""
|
||||
record = SkipRecord(
|
||||
name="x",
|
||||
reason="no_baseline",
|
||||
align_warnings=[_make_warning()],
|
||||
)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_align_warnings_json_round_trip(self):
|
||||
"""align_warnings survive model_dump_json → parse_record_json round-trip."""
|
||||
warning = _make_warning(
|
||||
axis="cp",
|
||||
group_index=2,
|
||||
differing_index=3,
|
||||
baseline_index=0,
|
||||
max_abs_diff=0.42,
|
||||
)
|
||||
record = ComparisonRecord(
|
||||
name="mlp",
|
||||
baseline=_make_tensor_info(),
|
||||
target=_make_tensor_info(),
|
||||
unified_shape=[4, 8],
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(),
|
||||
align_warnings=[warning],
|
||||
)
|
||||
|
||||
restored = parse_record_json(record.model_dump_json())
|
||||
assert isinstance(restored, ComparisonRecord)
|
||||
assert len(restored.align_warnings) == 1
|
||||
|
||||
restored_warning = restored.align_warnings[0]
|
||||
assert restored_warning.axis == "cp"
|
||||
assert restored_warning.group_index == 2
|
||||
assert restored_warning.differing_index == 3
|
||||
assert restored_warning.baseline_index == 0
|
||||
assert restored_warning.max_abs_diff == pytest.approx(0.42)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
@@ -878,6 +878,138 @@ class TestEntrypointGroupingLogical:
|
||||
assert comp.name == "hidden"
|
||||
|
||||
|
||||
class TestEntrypointReplicatedAxis:
|
||||
"""Test replicated-axis scenarios through the full entrypoint pipeline."""
|
||||
|
||||
def test_replicated_axis_identical_replicas_passed(self, tmp_path, capsys):
|
||||
"""CP2 TP2, TP replicated and identical → passed, no align_warnings."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 6)
|
||||
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
for side_dir, full_tensor in [
|
||||
(baseline_dir, full_baseline),
|
||||
(target_dir, full_target),
|
||||
]:
|
||||
_create_replicated_tp_sharded_cp_dumps(
|
||||
side_dir,
|
||||
full_tensor=full_tensor,
|
||||
name="attn_out",
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s(cp) d",
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=0.01,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.align_warnings == []
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.passed == 1
|
||||
|
||||
def test_replicated_mismatch_fails(self, tmp_path, capsys):
|
||||
"""CP2 TP2, TP replicas differ (> atol) → failed with align_warnings."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 6)
|
||||
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
for side_dir, full_tensor in [
|
||||
(baseline_dir, full_baseline),
|
||||
(target_dir, full_target),
|
||||
]:
|
||||
_create_replicated_tp_sharded_cp_dumps(
|
||||
side_dir,
|
||||
full_tensor=full_tensor,
|
||||
name="attn_out",
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s(cp) d",
|
||||
tp_noise=0.5,
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=0.01,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comparisons = _get_comparisons(records)
|
||||
assert len(comparisons) == 1
|
||||
assert comparisons[0].category == "failed"
|
||||
assert len(comparisons[0].align_warnings) > 0
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.failed == 1
|
||||
|
||||
def test_summary_counts_failed_from_align_warnings_only(self, tmp_path, capsys):
|
||||
"""Diff itself passes but TP replicas differ → summary.failed=1 from align_warnings."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 6)
|
||||
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
_create_replicated_tp_sharded_cp_dumps(
|
||||
baseline_dir,
|
||||
full_tensor=full_baseline,
|
||||
name="attn_out",
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s(cp) d",
|
||||
tp_noise=0.5,
|
||||
)
|
||||
_create_replicated_tp_sharded_cp_dumps(
|
||||
target_dir,
|
||||
full_tensor=full_target,
|
||||
name="attn_out",
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
dims_str="b s(cp) d",
|
||||
tp_noise=0.5,
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=0.5,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comparisons = _get_comparisons(records)
|
||||
assert len(comparisons) == 1
|
||||
|
||||
comp = comparisons[0]
|
||||
assert comp.diff is not None
|
||||
assert comp.diff.passed
|
||||
assert len(comp.align_warnings) > 0
|
||||
assert comp.category == "failed"
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.failed == 1
|
||||
assert summary.passed == 0
|
||||
|
||||
|
||||
# --------------------------- Assertion helpers -------------------
|
||||
|
||||
|
||||
@@ -1138,6 +1270,49 @@ def _create_cp_zigzag_tp_sharded_dumps(
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
def _create_replicated_tp_sharded_cp_dumps(
|
||||
directory: Path,
|
||||
*,
|
||||
full_tensor: torch.Tensor,
|
||||
name: str,
|
||||
cp_size: int,
|
||||
tp_size: int,
|
||||
seq_dim: int,
|
||||
dims_str: str,
|
||||
tp_noise: float = 0.0,
|
||||
) -> Path:
|
||||
"""Create CP-sharded + TP-replicated dump files from a full tensor.
|
||||
|
||||
CP direction: chunks along seq_dim (sharded).
|
||||
TP direction: clones (replicated), with optional noise to simulate mismatch.
|
||||
"""
|
||||
cp_chunks: list[torch.Tensor] = list(full_tensor.chunk(cp_size, dim=seq_dim))
|
||||
|
||||
rank: int = 0
|
||||
for cp_rank in range(cp_size):
|
||||
for tp_rank in range(tp_size):
|
||||
shard = cp_chunks[cp_rank].clone()
|
||||
if tp_noise > 0 and tp_rank > 0:
|
||||
shard = shard + torch.randn_like(shard) * tp_noise
|
||||
|
||||
_create_rank_dump(
|
||||
directory,
|
||||
rank=rank,
|
||||
name=name,
|
||||
tensor=shard,
|
||||
dims=dims_str,
|
||||
parallel_info={
|
||||
"cp_rank": cp_rank,
|
||||
"cp_size": cp_size,
|
||||
"tp_rank": tp_rank,
|
||||
"tp_size": tp_size,
|
||||
},
|
||||
)
|
||||
rank += 1
|
||||
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
def _create_tp_sharded_dumps(
|
||||
directory: Path,
|
||||
*,
|
||||
|
||||
Reference in New Issue
Block a user