From 003ad6daaad7c808bcc53dbf40a82d6f8652a54e Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 1 Mar 2026 10:33:39 +0800 Subject: [PATCH] Support partial tensors waiting for reduction and pipeline parallel in dump comparator (#19595) --- .../comparator/aligner/unsharder/executor.py | 9 + .../comparator/aligner/unsharder/planner.py | 5 +- .../comparator/aligner/unsharder/types.py | 6 +- .../source_patcher/code_patcher.py | 2 +- .../aligner/unsharder/test_executor.py | 200 +++++++++--------- .../aligner/unsharder/test_planner.py | 59 +++++- .../comparator/test_bundle_matcher.py | 183 +++++++++++++++- .../debug_utils/comparator/test_entrypoint.py | 128 +++++++++++ 8 files changed, 484 insertions(+), 108 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py index a033730bc..c52a99153 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py @@ -6,6 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( ConcatParams, CpThdConcatParams, PickParams, + ReduceSumParams, UnsharderParams, UnsharderPlan, ) @@ -63,6 +64,14 @@ def _apply_unshard( seq_lens_per_rank=params.seq_lens_per_rank, ) + if isinstance(params, ReduceSumParams): + stripped: list[torch.Tensor] = [t.rename(None) for t in ordered_tensors] + result: torch.Tensor = torch.stack(stripped).sum(dim=0) + names: tuple[Optional[str], ...] = ordered_tensors[0].names + if names[0] is not None: + result = result.refine_names(*names) + return result + raise ValueError(f"Unsupported unshard operation: {type(params).__name__}") diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py index 061365705..c88a620de 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py @@ -6,6 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( ConcatParams, CpThdConcatParams, PickParams, + ReduceSumParams, UnsharderParams, UnsharderPlan, ) @@ -155,9 +156,7 @@ def _resolve_unshard_params( thd_global_seq_lens: Optional[list[int]] = None, ) -> UnsharderParams: if spec.reduction is not None: - raise NotImplementedError( - f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)" - ) + return ReduceSumParams() if ( spec.name == TOKEN_DIM_NAME diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py index 2ad4b4cf9..19c752d3a 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py @@ -38,8 +38,12 @@ class PickParams(_FrozenBase): op: Literal["pick"] = "pick" +class ReduceSumParams(_FrozenBase): + op: Literal["reduce_sum"] = "reduce_sum" + + UnsharderParams = Annotated[ - Union[ConcatParams, CpThdConcatParams, PickParams], + Union[ConcatParams, CpThdConcatParams, PickParams, ReduceSumParams], Field(discriminator="op"), ] diff --git a/python/sglang/srt/debug_utils/source_patcher/code_patcher.py b/python/sglang/srt/debug_utils/source_patcher/code_patcher.py index 421f2ccd2..c1f2de063 100644 --- a/python/sglang/srt/debug_utils/source_patcher/code_patcher.py +++ b/python/sglang/srt/debug_utils/source_patcher/code_patcher.py @@ -118,7 +118,7 @@ def _inject_preamble(*, config: PatchConfig, extra_imports: list[str]) -> PatchC for spec in config.patches: existing: str = spec.preamble combined: str = ( - existing + "\n" + import_block if existing.strip() else import_block + import_block + "\n" + existing if existing.strip() else import_block ) new_patches.append( PatchSpec(target=spec.target, edits=spec.edits, preamble=combined) diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py index eca1fe27a..e72a25964 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py @@ -15,6 +15,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( AxisInfo, CpThdConcatParams, PickParams, + ReduceSumParams, UnsharderPlan, ) from sglang.srt.debug_utils.comparator.dims import ( @@ -639,118 +640,125 @@ class TestThdCpConcat: ) -class TestThdCpConcat: - def test_single_seq(self) -> None: - """Single seq THD unshard: 2 ranks → per-seq concat.""" - rank0 = torch.tensor([1, 2, 3]).refine_names("t") - rank1 = torch.tensor([4, 5, 6]).refine_names("t") - - plan = UnsharderPlan( - axis=ParallelAxis.CP, - params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]), - groups=[[0, 1]], - ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) - - assert len(result) == 1 - expected = torch.tensor([1, 2, 3, 4, 5, 6]) - assert torch.equal(result[0].rename(None), expected) - - def test_multi_seq(self) -> None: - """Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46].""" - # rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)] - # rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)] - seq_a_r0 = torch.arange(0, 50) - seq_b_r0 = torch.arange(100, 132) - pad_r0 = torch.full((46,), -1) - rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t") - - seq_a_r1 = torch.arange(50, 100) - seq_b_r1 = torch.arange(132, 164) - pad_r1 = torch.full((46,), -2) - rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t") - - plan = UnsharderPlan( - axis=ParallelAxis.CP, - params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]), - groups=[[0, 1]], - ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) - - assert len(result) == 1 - unsharded: torch.Tensor = result[0].rename(None) - - # seqA: r0(50) + r1(50) = 100 tokens, values 0..99 - assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1])) - # seqB: r0(32) + r1(32) = 64 tokens - assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1])) - # pad: r0(46) + r1(46) = 92 tokens - assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1])) - - def test_with_hidden_dim(self) -> None: - """THD unshard with trailing hidden dim: shape [T, H].""" +class TestReduceSum: + def test_basic_tp2_reduce(self) -> None: + """2 partial tensors sum to full tensor.""" torch.manual_seed(42) - hidden: int = 4 - # rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)] - # rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)] - seq_a_r0 = torch.randn(3, hidden) - seq_b_r0 = torch.randn(2, hidden) - rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h") + full_tensor = torch.randn(4, 8) + part_a = full_tensor * 0.6 + part_b = full_tensor * 0.4 - seq_a_r1 = torch.randn(3, hidden) - seq_b_r1 = torch.randn(2, hidden) - rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h") + dim_specs = parse_dims("h(tp,partial) d") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2) + ] + plans = compute_unsharder_plan(dim_specs, parallel_infos) + assert len(plans) == 1 + assert isinstance(plans[0].params, ReduceSumParams) - plan = UnsharderPlan( - axis=ParallelAxis.CP, - params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]), - groups=[[0, 1]], - ) + named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs) with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) + result = execute_unsharder_plan(plans[0], named_parts) assert len(result) == 1 - unsharded: torch.Tensor = result[0].rename(None) + assert torch.allclose(result[0].rename(None), full_tensor) - assert unsharded.shape == (10, hidden) - assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1])) - assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1])) - - def test_with_leading_batch_dim(self) -> None: - """THD unshard with leading batch dim: shape [B, T, H], t is dim=1.""" + def test_tp4_reduce(self) -> None: + """4 partial tensors sum to full tensor.""" torch.manual_seed(42) - batch: int = 2 - hidden: int = 4 - # rank0: [seqA_r0(3) | seqB_r0(2)] per batch item - # rank1: [seqA_r1(3) | seqB_r1(2)] per batch item - seq_a_r0 = torch.randn(batch, 3, hidden) - seq_b_r0 = torch.randn(batch, 2, hidden) - rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h") + full_tensor = torch.randn(4, 8) + parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)] - seq_a_r1 = torch.randn(batch, 3, hidden) - seq_b_r1 = torch.randn(batch, 2, hidden) - rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h") + dim_specs = parse_dims("h(tp,partial) d") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4) + ] + plans = compute_unsharder_plan(dim_specs, parallel_infos) + assert len(plans) == 1 + + named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs) + with warning_sink.context(): + result = execute_unsharder_plan(plans[0], named_parts) + + assert len(result) == 1 + assert torch.allclose(result[0].rename(None), full_tensor) + + def test_multi_axis_concat_then_reduce(self) -> None: + """CP concat + TP reduce end-to-end.""" + torch.manual_seed(42) + full_tensor = torch.randn(4, 8, 16) + + cp_chunks = list(full_tensor.chunk(2, dim=1)) + # Each CP chunk is held as partial sums across TP ranks + 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] * 0.5) + 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) h(tp,partial)") + plans = compute_unsharder_plan(dim_specs, parallel_infos) + assert len(plans) == 2 + + current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) + with warning_sink.context(): + for plan in plans: + current = execute_unsharder_plan(plan, current) + + assert len(current) == 1 + assert torch.allclose(current[0].rename(None), full_tensor) + + def test_reduce_scrambled_ranks(self) -> None: + """Scrambled rank order — sum is commutative so result is the same.""" + torch.manual_seed(42) + full_tensor = torch.randn(4, 8) + parts: list[torch.Tensor] = [ + full_tensor * 0.1, + full_tensor * 0.2, + full_tensor * 0.3, + full_tensor * 0.4, + ] + + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)}, + {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)}, + {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)}, + {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)}, + ] + dim_specs = parse_dims("h(tp,partial) d") + plans = compute_unsharder_plan(dim_specs, parallel_infos) + + named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs) + with warning_sink.context(): + result = execute_unsharder_plan(plans[0], named_parts) + + assert len(result) == 1 + assert torch.allclose(result[0].rename(None), full_tensor) + + def test_reduce_preserves_named_dims(self) -> None: + """Named tensor dimensions are preserved through reduce_sum.""" + dim_specs = parse_dims("h(tp,partial) d") + part_a = torch.randn(4, 8).refine_names("h", "d") + part_b = torch.randn(4, 8).refine_names("h", "d") plan = UnsharderPlan( - axis=ParallelAxis.CP, - params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]), + axis=ParallelAxis.TP, + params=ReduceSumParams(), groups=[[0, 1]], ) with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) + result = execute_unsharder_plan(plan, [part_a, part_b]) assert len(result) == 1 - unsharded: torch.Tensor = result[0].rename(None) - - assert unsharded.shape == (batch, 10, hidden) - # seqA: r0(3) + r1(3) = 6 tokens per batch - assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1)) - # seqB: r0(2) + r1(2) = 4 tokens per batch - assert torch.equal( - unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1) - ) + assert result[0].names == ("h", "d") + expected = (part_a.rename(None) + part_b.rename(None)).refine_names("h", "d") + assert torch.allclose(result[0].rename(None), expected.rename(None)) if __name__ == "__main__": diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py index 97d70585c..55e012e74 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py @@ -9,6 +9,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( AxisInfo, ConcatParams, PickParams, + ReduceSumParams, ) from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims from sglang.test.ci.ci_register import register_cpu_ci @@ -173,13 +174,65 @@ class TestComputeUnsharderPlan: with pytest.raises(ValueError, match="axis_rank coverage.*incomplete"): compute_unsharder_plan(dim_specs, parallel_infos) - def test_reduction_not_implemented_raises(self) -> None: + def test_reduction_partial_returns_reduce_sum(self) -> None: dim_specs = parse_dims("h(tp,partial)") parallel_infos = [ {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2) ] - with pytest.raises(NotImplementedError, match="reduction"): - compute_unsharder_plan(dim_specs, parallel_infos) + plans = compute_unsharder_plan(dim_specs, parallel_infos) + + assert len(plans) == 1 + assert plans[0].axis == ParallelAxis.TP + assert isinstance(plans[0].params, ReduceSumParams) + assert plans[0].groups == [[0, 1]] + + def test_reduction_partial_tp4(self) -> None: + """TP=4 with partial reduction produces a single ReduceSumParams step.""" + dim_specs = parse_dims("h(tp,partial)") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4) + ] + plans = compute_unsharder_plan(dim_specs, parallel_infos) + + assert len(plans) == 1 + assert isinstance(plans[0].params, ReduceSumParams) + assert plans[0].groups == [[0, 1, 2, 3]] + + def test_multi_axis_with_reduction_on_one(self) -> None: + """CP concat + TP reduce produces a 2-step plan.""" + dim_specs = parse_dims("s(cp) h(tp,partial)") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(2): + for tp_rank in range(2): + parallel_infos.append( + { + ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2), + ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2), + } + ) + + plans = compute_unsharder_plan(dim_specs, parallel_infos) + + assert len(plans) == 2 + assert plans[0].axis == ParallelAxis.CP + assert isinstance(plans[0].params, ConcatParams) + assert plans[1].axis == ParallelAxis.TP + assert isinstance(plans[1].params, ReduceSumParams) + + def test_reduction_scrambled_ranks(self) -> None: + """Scrambled world_rank order with partial reduction.""" + dim_specs = parse_dims("h(tp,partial)") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)}, + {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)}, + {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)}, + {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)}, + ] + plans = compute_unsharder_plan(dim_specs, parallel_infos) + + assert len(plans) == 1 + assert isinstance(plans[0].params, ReduceSumParams) + assert plans[0].groups == [[1, 3, 0, 2]] def test_ordering_zigzag_accepted(self) -> None: dim_specs = parse_dims("s(cp,zigzag)") diff --git a/test/registered/debug_utils/comparator/test_bundle_matcher.py b/test/registered/debug_utils/comparator/test_bundle_matcher.py index ec2702fd3..cab011061 100644 --- a/test/registered/debug_utils/comparator/test_bundle_matcher.py +++ b/test/registered/debug_utils/comparator/test_bundle_matcher.py @@ -17,15 +17,38 @@ 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 + *, + name: str, + step: int = 0, + rank: int = 0, + layer_id: int | None = None, + 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} + layer_part: str = f"___layer_id={layer_id}" if layer_id is not None else "" + filename = f"name={name}___step={step}___rank={rank}{layer_part}.pt" + row: dict[str, Any] = { + "name": name, + "step": step, + "rank": rank, + "filename": filename, + } + if layer_id is not None: + row["layer_id"] = layer_id + return row def _make_df(rows: list[dict[str, Any]]) -> pl.DataFrame: - return pl.DataFrame(rows) + if not rows: + return pl.DataFrame(rows) + + all_keys: set[str] = set() + for row in rows: + all_keys.update(row.keys()) + normalized: list[dict[str, Any]] = [ + {k: row.get(k, None) for k in all_keys} for row in rows + ] + return pl.DataFrame(normalized) class TestMatchBundles: @@ -147,6 +170,158 @@ class TestMatchBundles: assert len(results[0].y) == 2 +class TestMatchBundlesPipelineParallel: + """Tests verifying that PP works correctly with the existing matching logic.""" + + LOGICAL_SKIP_KEYS: set[str] = {"filename", "rank", "dump_index", "recompute_status"} + + def test_same_layer_id_different_ranks_match(self) -> None: + """SGLang PP=2 rank 0 (layers 0-31) vs Megatron PP=4 rank 2 (layers 16-31): + layer_id=20 should match regardless of world rank.""" + target_df: pl.DataFrame = _make_df( + [_make_row(name="hidden", rank=0, layer_id=20)] + ) + baseline_df: pl.DataFrame = _make_df( + [_make_row(name="hidden", rank=2, layer_id=20)] + ) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys=self.LOGICAL_SKIP_KEYS, + ) + + assert len(results) == 1 + assert len(results[0].x) == 1 + assert len(results[0].y) == 1 + + def test_layer_id_none_non_layer_tensors_match(self) -> None: + """Non-layer tensors (embedding, lm_head) have no layer_id. + They should match across different PP ranks.""" + target_df: pl.DataFrame = _make_df([_make_row(name="embed_tokens", rank=0)]) + baseline_df: pl.DataFrame = _make_df([_make_row(name="embed_tokens", rank=0)]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys=self.LOGICAL_SKIP_KEYS, + ) + + assert len(results) == 1 + assert len(results[0].x) == 1 + assert len(results[0].y) == 1 + + def test_different_pp_sizes_layer_and_non_layer_bundles(self) -> None: + """SGLang PP=2 TP=2 (4 ranks) vs Megatron PP=4 TP=2 (8 ranks). + Layer tensors match by (name, layer_id); non-layer tensors match by name. + All ranks are grouped into the same bundle when rank is skipped.""" + target_df: pl.DataFrame = _make_df( + [ + # SGLang: pp_stage=0 has ranks 0,1 (TP=2) + _make_row(name="hidden", rank=0, layer_id=20), + _make_row(name="hidden", rank=1, layer_id=20), + # SGLang: embedding on pp_stage=0 + _make_row(name="embed_tokens", rank=0), + _make_row(name="embed_tokens", rank=1), + # SGLang: lm_head on pp_stage=1, ranks 2,3 + _make_row(name="lm_head", rank=2), + _make_row(name="lm_head", rank=3), + ] + ) + baseline_df: pl.DataFrame = _make_df( + [ + # Megatron: pp_stage=1 has ranks 2,3 for layer 20 + _make_row(name="hidden", rank=2, layer_id=20), + _make_row(name="hidden", rank=3, layer_id=20), + # Megatron: embedding on pp_stage=0 + _make_row(name="embed_tokens", rank=0), + _make_row(name="embed_tokens", rank=1), + # Megatron: lm_head on pp_stage=3, ranks 6,7 + _make_row(name="lm_head", rank=6), + _make_row(name="lm_head", rank=7), + ] + ) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys=self.LOGICAL_SKIP_KEYS, + ) + + assert len(results) == 3 + names_to_pairs: dict[str, Pair[TensorBundleInfo]] = {} + for pair in results: + key: str = pair.y[0].name + layer_suffix: str = "" + if "layer_id" in target_df.columns: + row_match = [ + r + for r in target_df.to_dicts() + if r["filename"] == pair.y[0].filename + ] + if row_match and row_match[0].get("layer_id") is not None: + layer_suffix = f"_{row_match[0]['layer_id']}" + names_to_pairs[key + layer_suffix] = pair + + assert len(names_to_pairs["hidden_20"].x) == 2 + assert len(names_to_pairs["hidden_20"].y) == 2 + assert len(names_to_pairs["embed_tokens"].x) == 2 + assert len(names_to_pairs["embed_tokens"].y) == 2 + assert len(names_to_pairs["lm_head"].x) == 2 + assert len(names_to_pairs["lm_head"].y) == 2 + + def test_unmatched_layer_id_creates_empty_baseline(self) -> None: + """If target has a layer_id that baseline doesn't, the baseline side + should be empty (not incorrectly matched to a different layer).""" + target_df: pl.DataFrame = _make_df( + [ + _make_row(name="hidden", rank=0, layer_id=10), + _make_row(name="hidden", rank=0, layer_id=20), + ] + ) + baseline_df: pl.DataFrame = _make_df( + [ + _make_row(name="hidden", rank=0, layer_id=10), + ] + ) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys=self.LOGICAL_SKIP_KEYS, + ) + + assert len(results) == 2 + matched: list[Pair[TensorBundleInfo]] = [r for r in results if r.x] + unmatched: list[Pair[TensorBundleInfo]] = [r for r in results if not r.x] + assert len(matched) == 1 + assert len(unmatched) == 1 + + def test_pp1_vs_pp_gt1_matches_by_layer_id(self) -> None: + """PP=1 (all layers on 1 rank) vs PP>1 (layers split across ranks). + Should match correctly by layer_id regardless of rank.""" + target_df: pl.DataFrame = _make_df( + [ + # PP=1: all on rank 0 + _make_row(name="hidden", rank=0, layer_id=0), + _make_row(name="hidden", rank=0, layer_id=1), + ] + ) + baseline_df: pl.DataFrame = _make_df( + [ + # PP=2: layer 0 on rank 0, layer 1 on rank 1 + _make_row(name="hidden", rank=0, layer_id=0), + _make_row(name="hidden", rank=1, layer_id=1), + ] + ) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys=self.LOGICAL_SKIP_KEYS, + ) + + assert len(results) == 2 + for pair in results: + assert len(pair.x) == 1 + assert len(pair.y) == 1 + + class TestRowsToTensorInfos: def test_filters_extra_columns(self) -> None: rows: list[dict[str, Any]] = [ diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 028bb5a40..55aef9ed6 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -944,6 +944,107 @@ class TestEntrypointGroupingLogical: ] assert len(recompute_warnings) > 0 + def test_tp_partial_reduction_unshard(self, tmp_path, capsys): + """TP=2 with partial reduction: element-wise sum reconstructs full tensor.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8) + full_target = full_baseline + torch.randn(4, 8) * 0.001 + + baseline_dir = tmp_path / "baseline" + target_dir = tmp_path / "target" + + baseline_path = _create_tp_partial_dumps( + baseline_dir, + full_tensor=full_baseline, + name="attn_out", + tp_size=2, + dims_str="b h(tp,partial)", + ) + target_path = _create_tp_partial_dumps( + target_dir, + full_tensor=full_target, + name="attn_out", + tp_size=2, + dims_str="b h(tp,partial)", + ) + + args = _make_args(baseline_path, target_path, diff_threshold=0.01) + + records = _run_and_parse(args, capsys) + comp = _assert_single_comparison_passed(records) + assert comp.name == "attn_out" + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.total == 1 + assert summary.passed == 1 + + def test_tp_partial_vs_single_rank(self, tmp_path, capsys): + """Baseline single rank vs target TP=2 partial: unshard target then compare.""" + torch.manual_seed(42) + full_tensor = torch.randn(4, 8) + target_full = full_tensor + torch.randn(4, 8) * 0.001 + + baseline_dir = tmp_path / "baseline" + target_dir = tmp_path / "target" + + baseline_path = _create_rank_dump( + baseline_dir, rank=0, name="attn_out", tensor=full_tensor + ) + target_path = _create_tp_partial_dumps( + target_dir, + full_tensor=target_full, + name="attn_out", + tp_size=2, + dims_str="b h(tp,partial)", + ) + + args = _make_args(baseline_path, target_path, diff_threshold=0.01) + + records = _run_and_parse(args, capsys) + comp = _assert_single_comparison_passed(records) + assert comp.name == "attn_out" + + def test_cp_concat_tp_partial_reduction(self, tmp_path, capsys): + """CP=2 concat + TP=2 partial reduction: multi-axis unshard.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8, 16) + full_target = full_baseline + torch.randn(4, 8, 16) * 0.001 + + for side_dir, full_tensor in [ + (tmp_path / "baseline", full_baseline), + (tmp_path / "target", full_target), + ]: + side_dir.mkdir() + cp_chunks = list(full_tensor.chunk(2, dim=1)) + rank = 0 + for cp_rank in range(2): + for tp_rank in range(2): + _create_rank_dump( + side_dir, + rank=rank, + name="hidden", + tensor=cp_chunks[cp_rank] / 2, + dims="b s(cp) h(tp,partial)", + parallel_info={ + "cp_rank": cp_rank, + "cp_size": 2, + "tp_rank": tp_rank, + "tp_size": 2, + }, + ) + rank += 1 + + args = _make_args( + tmp_path / "baseline" / _FIXED_EXP_NAME, + tmp_path / "target" / _FIXED_EXP_NAME, + diff_threshold=0.01, + ) + + records = _run_and_parse(args, capsys) + comp = _assert_single_comparison_passed(records) + assert comp.name == "hidden" + class TestEntrypointAxisAligner: """Test cross-framework dim reordering through the full entrypoint pipeline.""" @@ -2047,6 +2148,33 @@ def _create_tp_sharded_dumps( return directory / _FIXED_EXP_NAME +def _create_tp_partial_dumps( + directory: Path, + *, + full_tensor: torch.Tensor, + name: str, + tp_size: int, + dims_str: str, + num_steps: int = 1, +) -> Path: + """Create TP-partial dump files where each rank holds full_tensor / tp_size. + + Each rank stores an equal fraction of the full tensor so that + element-wise summation across ranks reconstructs the original. + """ + for tp_rank in range(tp_size): + _create_rank_dump( + directory, + rank=tp_rank, + name=name, + tensor=full_tensor / tp_size, + dims=dims_str, + parallel_info={"tp_rank": tp_rank, "tp_size": tp_size}, + num_steps=num_steps, + ) + return directory / _FIXED_EXP_NAME + + def _create_recompute_rank_dump( directory: Path, *,