Support s≡t dimension name equivalence in dump comparator (#21027)

2026-03-20 22:03:34 +08:00
parent cc22601d28
commit 154395ab7d
3 changed files with 292 additions and 3 deletions
@@ -432,5 +432,115 @@ class TestEndToEndThreeWayFused:
        assert torch.equal(y_aligned, x_tensor)


+class TestSeqTokenEquivalencePlan:
+    """Tests for s≡t dimension name equivalence in compute_axis_aligner_plan."""
+
+    def test_s_t_equivalence_squeeze(self) -> None:
+        """sglang 't h' vs megatron 's 1 h': plan squeezes y-side singleton, x-side no-op."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h", y="s 1 h")
+        )
+        assert result is not None
+        assert result.pattern.x is None
+        assert result.pattern.y == "s 1 h -> s h"
+
+    def test_s_t_equivalence_same_shape(self) -> None:
+        """'t h d' vs 's h d': same order after normalization → no plan needed."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h d", y="s h d")
+        )
+        assert result is None
+
+    def test_s_t_equivalence_with_swap(self) -> None:
+        """'t d h' vs 's h d': plan not None, x-pattern reorders."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t d h", y="s h d")
+        )
+        assert result is not None
+        assert result.pattern.x is not None
+
+    def test_s_t_equivalence_with_fused(self) -> None:
+        """'t (a*b)' vs 's a b': plan not None, y-pattern flattens."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t (a*b)", y="s a b")
+        )
+        assert result is not None
+        assert result.pattern.y is not None
+
+    def test_s_t_equivalence_with_squeeze_and_fused(self) -> None:
+        """'t (num_heads*head_dim)' vs 's 1 num_heads head_dim': plan squeezes + flattens."""
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t (num_heads*head_dim)", y="s 1 num_heads head_dim")
+        )
+        assert result is not None
+
+
+class TestSeqTokenEquivalenceExecute:
+    """Tests for s≡t dimension name equivalence in execute_axis_aligner_plan."""
+
+    def test_execute_s_t_squeeze(self) -> None:
+        """Tensor [4,1,8] with pattern 's 1 h -> s h' → shape [4,8]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 1, 8)
+        plan = AxisAlignerPlan(pattern=Pair(x=None, y="s 1 h -> s h"))
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="y"
+        )
+
+        assert result.shape == (4, 8)
+        assert torch.equal(result, tensor.squeeze(1))
+
+
+class TestEndToEndSeqTokenEquivalence:
+    """End-to-end tests for s≡t equivalence through compute + execute pipeline."""
+
+    def test_s_t_squeeze_full_pipeline(self) -> None:
+        """x=tensor(4,8) dims='t h', y=tensor(4,1,8) dims='s 1 h' → both aligned to (4,8)."""
+        torch.manual_seed(42)
+        data: torch.Tensor = torch.randn(4, 8)
+        x_tensor: torch.Tensor = data.clone()
+        y_tensor: torch.Tensor = data.unsqueeze(1)
+
+        plan: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h", y="s 1 h")
+        )
+        assert plan is not None
+
+        x_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=x_tensor, plan=plan, side="x"
+        )
+        y_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=y_tensor, plan=plan, side="y"
+        )
+
+        assert x_aligned.shape == y_aligned.shape == (4, 8)
+        assert torch.equal(x_aligned, y_aligned)
+
+    def test_s_t_fused_full_pipeline(self) -> None:
+        """x=tensor(4,128) dims='t (nh*hd)', y=tensor(4,1,8,16) dims='s 1 nh hd' → both (4,128)."""
+        torch.manual_seed(42)
+        num_heads: int = 8
+        head_dim: int = 16
+        data: torch.Tensor = torch.randn(4, num_heads * head_dim)
+        x_tensor: torch.Tensor = data.clone()
+        y_tensor: torch.Tensor = data.reshape(4, num_heads, head_dim).unsqueeze(1)
+
+        plan: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t (nh*hd)", y="s 1 nh hd")
+        )
+        assert plan is not None
+
+        x_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=x_tensor, plan=plan, side="x"
+        )
+        y_aligned: torch.Tensor = execute_axis_aligner_plan(
+            tensor=y_tensor, plan=plan, side="y"
+        )
+
+        assert x_aligned.shape == y_aligned.shape == (4, num_heads * head_dim)
+        assert torch.equal(x_aligned, y_aligned)
+
+
 if __name__ == "__main__":
    sys.exit(pytest.main([__file__]))
@@ -1718,6 +1718,157 @@ class TestEntrypointAxisAligner:
        assert comp.target.shape == [4, 8]


+class TestEntrypointSeqTokenEquivalence:
+    """Test s≡t dim name equivalence through the full entrypoint pipeline."""
+
+    def test_s_t_squeeze_single_rank(self, tmp_path, capsys):
+        """Baseline dims='t h' (2D [4,8]), target dims='s 1 h' (3D [4,1,8]) → comparator passes."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor,
+            dims="t h",
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor.unsqueeze(1),
+            dims="s 1 h",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+        assert comp.baseline.shape == [4, 8]
+        assert comp.target.shape == [4, 8]
+
+    def test_s_t_squeeze_with_tp_unshard(self, tmp_path, capsys):
+        """Baseline TP=2 dims='t h[tp]', target TP=2 dims='s 1 h[tp]' → unshard + squeeze + s≡t."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_tp_sharded_dumps(
+            baseline_dir,
+            full_tensor=full_tensor,
+            name="hidden",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="t h[tp]",
+        )
+        _create_tp_sharded_dumps(
+            target_dir,
+            full_tensor=full_tensor.unsqueeze(1),
+            name="hidden",
+            tp_size=2,
+            shard_dim=2,
+            dims_str="s 1 h[tp]",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+
+    def test_s_t_fused_with_squeeze(self, tmp_path, capsys):
+        """Baseline dims='t (num_heads*head_dim)[tp]' (2D), target dims='s 1 num_heads[tp] head_dim' (4D)."""
+        torch.manual_seed(42)
+        num_heads = 8
+        head_dim = 16
+        full_tensor_2d = torch.randn(4, num_heads * head_dim)
+        full_tensor_4d = full_tensor_2d.reshape(4, num_heads, head_dim).unsqueeze(1)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_tp_sharded_dumps(
+            baseline_dir,
+            full_tensor=full_tensor_2d,
+            name="attn_pre_o_proj",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="t (num_heads*head_dim)[tp]",
+        )
+        _create_tp_sharded_dumps(
+            target_dir,
+            full_tensor=full_tensor_4d,
+            name="attn_pre_o_proj",
+            tp_size=2,
+            shard_dim=2,
+            dims_str="s 1 num_heads[tp] head_dim",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "attn_pre_o_proj"
+
+    def test_s_t_mismatch_with_named_batch_fails(self, tmp_path, capsys):
+        """Baseline dims='t h', target dims='s b h' (named b, not constant 1) → dim mismatch → skip/error."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor,
+            dims="t h",
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor.unsqueeze(1).expand(4, 1, 8).contiguous(),
+            dims="s b h",
+        )
+
+        argv = _make_argv(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records, _ = _run_and_parse(argv, capsys)
+        comparisons = [r for r in records if isinstance(r, ComparisonTensorRecord)]
+        assert len(comparisons) == 1
+        comp = comparisons[0]
+        assert (
+            comp.shape_mismatch
+            or (comp.diff is not None and not comp.diff.passed)
+            or len(comp.errors) > 0
+        )
+
+
 class TestEntrypointReplicatedAxis:
    """Test replicated-axis scenarios through the full entrypoint pipeline."""