From 23cbbd6d4118829979a22c553ed1c533615efd3e Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 27 Feb 2026 08:28:36 +0800 Subject: [PATCH] Integrate packed data context parallel in dump comparator (#19464) --- .../comparator/aligner/entrypoint/planner.py | 61 +++-- .../comparator/bundle_comparator.py | 11 +- .../srt/debug_utils/comparator/entrypoint.py | 12 +- .../aligner/entrypoint/test_planner.py | 70 ++++- .../debug_utils/comparator/test_entrypoint.py | 248 +++++++++++++++++- 5 files changed, 364 insertions(+), 38 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py index fb5675fc3..9859dec2b 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py @@ -23,11 +23,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import ( from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import ( compute_unsharder_plan, ) -from sglang.srt.debug_utils.comparator.dims import ( - TOKEN_DIM_NAME, - find_dim_index, - parse_dims, -) +from sglang.srt.debug_utils.comparator.dims import DimSpec, parse_dims from sglang.srt.debug_utils.comparator.utils import Pair @@ -35,6 +31,9 @@ def compute_aligner_plan( *, metas_pair: Pair[list[dict[str, Any]]], token_aligner_plan: Optional[TokenAlignerPlan], + thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair( + x=None, y=None + ), ) -> AlignerPlan: dims_str_pair: Pair[Optional[str]] = metas_pair.map( lambda metas: metas[0].get("dims") if metas else None @@ -44,32 +43,26 @@ def compute_aligner_plan( ) return AlignerPlan( - per_step_plans=metas_pair.map( - lambda metas: _compute_per_step_plans(metas=metas) + per_step_plans=Pair( + x=_compute_per_step_plans( + metas=metas_pair.x, + thd_seq_lens_by_step=thd_seq_lens_by_step_pair.x, + ), + y=_compute_per_step_plans( + metas=metas_pair.y, + thd_seq_lens_by_step=thd_seq_lens_by_step_pair.y, + ), ), token_aligner_plan=token_aligner_plan, axis_swapper_plan=axis_swapper_plan, ) -def _compute_token_dim(metas: list[dict[str, Any]]) -> int: - fallback_dim = 0 - - if not metas: - return fallback_dim - - dims_str: Optional[str] = metas[0].get("dims") - if dims_str is None: - return fallback_dim - - idx: Optional[int] = find_dim_index(parse_dims(dims_str), TOKEN_DIM_NAME) - if idx is None: - return fallback_dim - - return idx - - -def _compute_per_step_plans(metas: list[dict[str, Any]]) -> list[AlignerPerStepPlan]: +def _compute_per_step_plans( + metas: list[dict[str, Any]], + *, + thd_seq_lens_by_step: Optional[dict[int, list[int]]] = None, +) -> list[AlignerPerStepPlan]: step_to_input_indices: dict[int, list[int]] = {} for i, meta in enumerate(metas): step: int = int(meta["step"]) @@ -79,8 +72,12 @@ def _compute_per_step_plans(metas: list[dict[str, Any]]) -> list[AlignerPerStepP for step in sorted(step_to_input_indices): input_indices: list[int] = step_to_input_indices[step] step_metas: list[dict[str, Any]] = [metas[idx] for idx in input_indices] + step_seq_lens: Optional[list[int]] = ( + thd_seq_lens_by_step.get(step) if thd_seq_lens_by_step is not None else None + ) plans: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( - metas=step_metas + metas=step_metas, + thd_global_seq_lens=step_seq_lens, ) result.append( AlignerPerStepPlan( @@ -93,6 +90,8 @@ def _compute_per_step_plans(metas: list[dict[str, Any]]) -> list[AlignerPerStepP def compute_per_step_sub_plans( metas: list[dict[str, Any]], + *, + thd_global_seq_lens: Optional[list[int]] = None, ) -> list[AlignerPerStepSubPlan]: if not metas or len(metas) == 1: return [] @@ -101,13 +100,17 @@ def compute_per_step_sub_plans( if dims_str is None: return [] - dim_specs = parse_dims(dims_str) + dim_specs: list[DimSpec] = parse_dims(dims_str) parallel_infos = [normalize_parallel_info(meta) for meta in metas] unsharder_plans = compute_unsharder_plan( - dim_specs=dim_specs, parallel_infos=parallel_infos + dim_specs=dim_specs, + parallel_infos=parallel_infos, + thd_global_seq_lens=thd_global_seq_lens, ) reorderer_plans = compute_reorderer_plans( - dim_specs=dim_specs, parallel_infos=parallel_infos + dim_specs=dim_specs, + parallel_infos=parallel_infos, + thd_global_seq_lens=thd_global_seq_lens, ) return [*unsharder_plans, *reorderer_plans] diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index e11eec961..e211bb3c1 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -41,6 +41,9 @@ def compare_bundle_pair( target_path: Path, token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold: float, + thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair( + x=None, y=None + ), ) -> Union[ComparisonRecord, SkipRecord]: with warning_sink.context() as collected_warnings: result = _compare_bundle_pair_raw( @@ -50,6 +53,7 @@ def compare_bundle_pair( target_path=target_path, token_aligner_plan=token_aligner_plan, diff_threshold=diff_threshold, + thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, ) return result.model_copy(update={"warnings": collected_warnings}) @@ -63,6 +67,9 @@ def _compare_bundle_pair_raw( target_path: Path, token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold: float, + thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair( + x=None, y=None + ), ) -> Union[ComparisonRecord, SkipRecord]: # 1. Load (tensor + meta, ungrouped) valid_pair: Pair[list[ValueWithMeta]] = Pair( @@ -79,7 +86,9 @@ def _compare_bundle_pair_raw( lambda items: [it.meta for it in items] ) plan: AlignerPlan = compute_aligner_plan( - metas_pair=metas_pair, token_aligner_plan=token_aligner_plan + metas_pair=metas_pair, + token_aligner_plan=token_aligner_plan, + thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, ) # 3. Apply dim names to tensors, then execute diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 4bddc1982..9920f18fc 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -10,7 +10,8 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import ( AUX_NAMES, ) from sglang.srt.debug_utils.comparator.aligner.token_aligner.entrypoint import ( - compute_maybe_token_aligner_plan, + TokenAlignerResult, + compute_maybe_token_aligner_result, ) from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenAlignerPlan, @@ -46,14 +47,14 @@ def run(args: argparse.Namespace) -> None: warning_sink.set_output_format(args.output_format) dfs: Pair[pl.DataFrame] = _read_df(args) - token_aligner_plan = compute_maybe_token_aligner_plan(args, dfs) + ta_result: TokenAlignerResult = compute_maybe_token_aligner_result(args, dfs) dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES))) bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles( dfs=dfs, skip_keys=_compute_skip_keys( - args, has_token_aligner_plan=token_aligner_plan is not None + args, has_token_aligner_plan=ta_result.plan is not None ), ) @@ -61,8 +62,9 @@ def run(args: argparse.Namespace) -> None: bundle_info_pairs=bundle_info_pairs, baseline_path=Path(args.baseline_path), target_path=Path(args.target_path), - token_aligner_plan=token_aligner_plan, + token_aligner_plan=ta_result.plan, diff_threshold=args.diff_threshold, + thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair, ) _consume_comparison_records( comparison_records=comparison_records, output_format=args.output_format @@ -99,6 +101,7 @@ def _compare_bundle_pairs( target_path: Path, token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold: float, + thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]], ) -> Iterator[Union[ComparisonRecord, SkipRecord]]: for bundle_info_pair in bundle_info_pairs: if not bundle_info_pair.y: @@ -115,6 +118,7 @@ def _compare_bundle_pairs( target_path=target_path, token_aligner_plan=token_aligner_plan, diff_threshold=diff_threshold, + thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, ) diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py index 80789338f..a8f938f35 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py @@ -13,8 +13,14 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( AlignerPerStepSubPlan, AlignerPlan, ) -from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan -from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan +from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ( + ReordererPlan, + ZigzagToNaturalThdParams, +) +from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( + CpThdConcatParams, + UnsharderPlan, +) from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.utils import Pair from sglang.test.ci.ci_register import register_cpu_ci @@ -166,5 +172,65 @@ class TestComputeAlignerPlan: assert plan.token_aligner_plan is ta_plan +class TestComputePerStepSubPlansThd: + def test_thd_zigzag_returns_thd_plans(self) -> None: + """t(cp,zigzag) h(tp) generates THD-typed unsharder + reorderer plans.""" + thd_global_seq_lens: list[int] = [100, 64, 92] + result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( + metas=[ + _make_meta( + dims="t(cp,zigzag) h(tp)", + cp_rank=0, + cp_size=2, + tp_rank=0, + tp_size=2, + ), + _make_meta( + dims="t(cp,zigzag) h(tp)", + cp_rank=0, + cp_size=2, + tp_rank=1, + tp_size=2, + ), + _make_meta( + dims="t(cp,zigzag) h(tp)", + cp_rank=1, + cp_size=2, + tp_rank=0, + tp_size=2, + ), + _make_meta( + dims="t(cp,zigzag) h(tp)", + cp_rank=1, + cp_size=2, + tp_rank=1, + tp_size=2, + ), + ], + thd_global_seq_lens=thd_global_seq_lens, + ) + + unsharder_plans: list[UnsharderPlan] = [ + p for p in result if isinstance(p, UnsharderPlan) + ] + reorderer_plans: list[ReordererPlan] = [ + p for p in result if isinstance(p, ReordererPlan) + ] + + # Should have at least one THD concat plan for CP axis + thd_concat_plans: list[UnsharderPlan] = [ + p for p in unsharder_plans if isinstance(p.params, CpThdConcatParams) + ] + assert len(thd_concat_plans) == 1 + assert thd_concat_plans[0].params.seq_lens_per_rank == [50, 32, 46] + + # Should have exactly one THD reorder plan + assert len(reorderer_plans) == 1 + assert isinstance(reorderer_plans[0].params, ZigzagToNaturalThdParams) + assert reorderer_plans[0].params.cp_size == 2 + # Reorder seq_lens = global seq_lens (reorder happens after unshard) + assert reorderer_plans[0].params.seq_lens == [100, 64, 92] + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 2a32b1b49..16e06b0e2 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -1399,9 +1399,14 @@ def _create_rank_dump( tensor: torch.Tensor, dims: str | None = None, parallel_info: dict | None = None, + framework: str = "sglang", num_steps: int = 1, + extra_dumps: list[tuple[str, object]] | None = None, ) -> Path: - """Create a dump file via the real dumper, as if running on the given rank.""" + """Create a dump file via the real dumper, as if running on the given rank. + + extra_dumps: additional (name, value) pairs to dump alongside the main tensor each step. + """ with pytest.MonkeyPatch.context() as mp: mp.setattr(_dumper_module, "_get_rank", lambda: rank) @@ -1416,11 +1421,13 @@ def _create_rank_dump( static_meta: dict = {"world_rank": rank, "world_size": 1} if parallel_info is not None: - static_meta["sglang_parallel_info"] = parallel_info + static_meta[f"{framework}_parallel_info"] = parallel_info dumper.__dict__["_static_meta"] = static_meta for _ in range(num_steps): dumper.dump(name, tensor, dims=dims) + for extra_name, extra_value in extra_dumps or []: + dumper.dump(extra_name, extra_value) dumper.step() return directory / _FIXED_EXP_NAME @@ -1631,5 +1638,242 @@ def _create_tp_sharded_dumps( return directory / _FIXED_EXP_NAME +def _zigzag_split_seq(seq_natural: torch.Tensor, *, cp_size: int) -> list[torch.Tensor]: + """Split a natural-order seq into per-rank zigzag segments.""" + num_chunks: int = cp_size * 2 + chunks: list[torch.Tensor] = list(seq_natural.chunk(num_chunks, dim=0)) + order: list[int] = [] + for i in range(cp_size): + order.append(i) + order.append(num_chunks - 1 - i) + zigzagged: torch.Tensor = torch.cat([chunks[i] for i in order], dim=0) + return list(zigzagged.chunk(cp_size, dim=0)) + + +def _create_thd_cp_zigzag_dumps( + directory: Path, + *, + full_tensor: torch.Tensor, + name: str, + seq_lens: list[int], + cp_size: int, + total_per_rank: int, + dims_str: str = "t(cp,zigzag)", + num_steps: int = 1, +) -> Path: + """Create THD CP-zigzag sharded dump files simulating Megatron forward. + + Args: + full_tensor: 1D tensor of shape [T] in natural order. + seq_lens: per-seq token counts in natural order (e.g. [100, 64]). + cp_size: context parallelism size. + total_per_rank: total tokens per rank (including padding). + dims_str: dims annotation for the main tensor. + """ + # Build per-rank tensors from natural-order full_tensor + offset: int = 0 + rank_segments: list[list[torch.Tensor]] = [[] for _ in range(cp_size)] + + for seq_len in seq_lens: + seq_natural: torch.Tensor = full_tensor[offset : offset + seq_len] + seq_ranks: list[torch.Tensor] = _zigzag_split_seq(seq_natural, cp_size=cp_size) + for rank_idx in range(cp_size): + rank_segments[rank_idx].append(seq_ranks[rank_idx]) + offset += seq_len + + # Build cu_seqlens from seq_lens (global, replicated across ranks) + cu_seqlens_values: list[int] = [0] + for slen in seq_lens: + cu_seqlens_values.append(cu_seqlens_values[-1] + slen) + + # Pad to total_per_rank per rank (global pad = last cu_seqlens entry to total_per_rank * cp_size) + total_global: int = total_per_rank * cp_size + if cu_seqlens_values[-1] < total_global: + pad_global: int = total_global - cu_seqlens_values[-1] + cu_seqlens_values.append(total_global) + pad_per_rank: int = pad_global // cp_size + for rank_idx in range(cp_size): + rank_segments[rank_idx].append(torch.zeros(pad_per_rank)) + + cu_seqlens_q: torch.Tensor = torch.tensor(cu_seqlens_values, dtype=torch.int64) + + # Dump each rank + for cp_rank in range(cp_size): + rank_tensor: torch.Tensor = torch.cat(rank_segments[cp_rank], dim=0) + assert ( + rank_tensor.shape[0] == total_per_rank + ), f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}" + + _create_rank_dump( + directory, + rank=cp_rank, + name=name, + tensor=rank_tensor, + dims=dims_str, + parallel_info={ + "cp_rank": cp_rank, + "cp_size": cp_size, + }, + framework="megatron", + num_steps=num_steps, + extra_dumps=[ + ("cu_seqlens_q", cu_seqlens_q), + ("input_ids", rank_tensor.to(torch.int64)), + ], + ) + + return directory / _FIXED_EXP_NAME + + +class TestEntrypointThdCpZigzag: + """E2E entrypoint tests for THD CP zigzag format. + + Tests the full pipeline: dump creation → metadata loading → aligner plan → + unshard + reorder → tensor comparison. + """ + + def test_sglang_vs_megatron_zigzag_cp(self, tmp_path: Path, capsys) -> None: + """SGLang single-rank THD baseline vs Megatron CP=2 zigzag target.""" + torch.manual_seed(42) + hidden_dim: int = 8 + cp_size: int = 2 + + # Two sequences: 8 and 4 tokens (divisible by cp_size*2=4 for clean zigzag) + seq_a_ids: list[int] = [10, 20, 30, 40, 50, 60, 70, 80] + seq_b_ids: list[int] = [100, 200, 300, 400] + all_ids: list[int] = seq_a_ids + seq_b_ids + total_tokens: int = len(all_ids) + seq_lens: list[int] = [len(seq_a_ids), len(seq_b_ids)] + + hidden_states: torch.Tensor = torch.randn(total_tokens, hidden_dim) + + # --- SGLang baseline: single rank, 1 step --- + sglang_dir: Path = tmp_path / "baseline" + sglang_dir.mkdir() + sglang_dumper = _Dumper( + config=DumperConfig( + enable=True, + dir=str(sglang_dir), + exp_name=_FIXED_EXP_NAME, + enable_http_server=False, + ) + ) + + positions: list[int] = list(range(seq_lens[0])) + list(range(seq_lens[1])) + sglang_dumper.dump("input_ids", torch.tensor(all_ids)) + sglang_dumper.dump("positions", torch.tensor(positions)) + sglang_dumper.dump("seq_lens", torch.tensor(seq_lens)) + sglang_dumper.dump("rids", ["A", "B"]) + sglang_dumper.dump("hidden_states", hidden_states) + sglang_dumper.step() + + # --- Megatron target: CP=2, zigzag, 1 step --- + megatron_dir: Path = tmp_path / "target" + megatron_dir.mkdir() + + # Zigzag-split input_ids and hidden_states per sequence, then concat + ids_tensor: torch.Tensor = torch.tensor(all_ids, dtype=torch.int64) + offset: int = 0 + rank_id_segments: list[list[torch.Tensor]] = [[] for _ in range(cp_size)] + rank_hidden_segments: list[list[torch.Tensor]] = [[] for _ in range(cp_size)] + for slen in seq_lens: + seq_ids: torch.Tensor = ids_tensor[offset : offset + slen] + seq_hidden: torch.Tensor = hidden_states[offset : offset + slen] + zigzag_ids: list[torch.Tensor] = _zigzag_split_seq(seq_ids, cp_size=cp_size) + zigzag_hidden: list[torch.Tensor] = _zigzag_split_seq( + seq_hidden, cp_size=cp_size + ) + for rank_idx in range(cp_size): + rank_id_segments[rank_idx].append(zigzag_ids[rank_idx]) + rank_hidden_segments[rank_idx].append(zigzag_hidden[rank_idx]) + offset += slen + + cu_seqlens_q: torch.Tensor = torch.tensor( + [0] + [sum(seq_lens[: i + 1]) for i in range(len(seq_lens))], + dtype=torch.int64, + ) + + for cp_rank in range(cp_size): + rank_ids: torch.Tensor = torch.cat(rank_id_segments[cp_rank]) + rank_hidden: torch.Tensor = torch.cat(rank_hidden_segments[cp_rank]) + _create_rank_dump( + megatron_dir, + rank=cp_rank, + name="hidden_states", + tensor=rank_hidden, + dims="t(cp,zigzag) h", + parallel_info={"cp_rank": cp_rank, "cp_size": cp_size}, + framework="megatron", + extra_dumps=[ + ("cu_seqlens_q", cu_seqlens_q), + ("input_ids", rank_ids), + ], + ) + + # --- Run comparison --- + args: Namespace = _make_args( + sglang_dir / _FIXED_EXP_NAME, + megatron_dir / _FIXED_EXP_NAME, + grouping="logical", + diff_threshold=1e-3, + ) + records: list[AnyRecord] = _run_and_parse(args, capsys) + + comparisons: list[ComparisonRecord] = _get_comparisons(records) + hidden_comparisons: list[ComparisonRecord] = [ + c for c in comparisons if c.name == "hidden_states" + ] + assert len(hidden_comparisons) >= 1 + assert all(c.diff is not None and c.diff.passed for c in hidden_comparisons) + + def test_thd_cp_zigzag_unshard(self, tmp_path: Path, capsys) -> None: + """Both sides THD CP=2 zigzag, comparison should pass.""" + torch.manual_seed(42) + cp_size: int = 2 + seq_lens: list[int] = [100, 64] + total_tokens: int = sum(seq_lens) + total_per_rank: int = 128 + + full_tensor: torch.Tensor = torch.randn(total_tokens + 92) + + baseline_dir: Path = tmp_path / "baseline" + target_dir: Path = tmp_path / "target" + baseline_dir.mkdir() + target_dir.mkdir() + + baseline_path: Path = _create_thd_cp_zigzag_dumps( + baseline_dir, + full_tensor=full_tensor, + name="hidden_states", + seq_lens=seq_lens, + cp_size=cp_size, + total_per_rank=total_per_rank, + ) + + # Target: same data with small noise + target_tensor: torch.Tensor = full_tensor + torch.randn_like(full_tensor) * 1e-5 + target_path: Path = _create_thd_cp_zigzag_dumps( + target_dir, + full_tensor=target_tensor, + name="hidden_states", + seq_lens=seq_lens, + cp_size=cp_size, + total_per_rank=total_per_rank, + ) + + args: Namespace = _make_args( + baseline_path, target_path, grouping="logical", diff_threshold=1e-3 + ) + records: list[AnyRecord] = _run_and_parse(args, capsys) + + # hidden_states should pass comparison (after unshard + reorder) + comparisons: list[ComparisonRecord] = _get_comparisons(records) + hidden_comparisons: list[ComparisonRecord] = [ + c for c in comparisons if c.name == "hidden_states" + ] + assert len(hidden_comparisons) >= 1 + assert all(c.diff is not None and c.diff.passed for c in hidden_comparisons) + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))