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

Support partial tensors waiting for reduction and pipeline parallel in dump comparator (#19595)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:33:39 +08:00
committed by GitHub
parent 67810828cf
commit 003ad6daaa
8 changed files with 484 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py
View File

@@ -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__}")

5
python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py
View File

@@ -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

6
python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py
View File

@@ -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"),
]

2
python/sglang/srt/debug_utils/source_patcher/code_patcher.py
View File

@@ -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)

200
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
View File

@@ -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__":

59
test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py
View File

@@ -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)")

183
test/registered/debug_utils/comparator/test_bundle_matcher.py
View File

@@ -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]] = [

128
test/registered/debug_utils/comparator/test_entrypoint.py
View File

@@ -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,
*,