From ea6ff7b01fa34071562dfe7bbd527708b3472371 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 1 Mar 2026 10:36:48 +0800 Subject: [PATCH] Support multi sharding group on the same dimension in dump comparator (#19601) --- .../comparator/aligner/reorderer/planner.py | 15 ++- .../token_aligner/smart/aux_plugins.py | 6 +- .../comparator/aligner/unsharder/planner.py | 46 ++++---- .../sglang/srt/debug_utils/comparator/dims.py | 103 ++++++++++++----- .../srt/debug_utils/comparator/entrypoint.py | 9 -- .../aligner/entrypoint/test_planner.py | 14 +-- .../aligner/reorderer/test_planner.py | 104 ++++++++++++++--- .../aligner/token_aligner/test_aux_loader.py | 18 +-- .../aligner/token_aligner/test_aux_plugins.py | 12 +- .../aligner/unsharder/test_executor.py | 10 +- .../aligner/unsharder/test_planner.py | 83 +++++++++++++- .../debug_utils/comparator/test_dims.py | 89 ++++++++++----- .../debug_utils/comparator/test_entrypoint.py | 105 ++++++++++++++++-- test/registered/debug_utils/test_dumper.py | 6 +- 14 files changed, 469 insertions(+), 151 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py index 11c81cd0e..c8f03362c 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py @@ -26,11 +26,10 @@ def compute_reorderer_plans( plans: list[ReordererPlan] = [] for spec in dim_specs: - if ( - spec.ordering is not None - and spec.ordering != Ordering.NATURAL - and spec.parallel is not None - ): + for modifier in spec.parallel_modifiers: + if modifier.ordering is None or modifier.ordering == Ordering.NATURAL: + continue + if spec.name not in _ALLOWED_ZIGZAG_DIM_NAMES: raise ValueError( f"Zigzag ordering is only supported on sequence dims " @@ -39,11 +38,11 @@ def compute_reorderer_plans( f"but got dim name {spec.name!r} in {spec}" ) - if spec.ordering != Ordering.ZIGZAG: + if modifier.ordering != Ordering.ZIGZAG: raise ValueError( - f"Unsupported ordering {spec.ordering!r} for dim {spec.name!r}" + f"Unsupported ordering {modifier.ordering!r} for dim {spec.name!r}" ) - axis_size: int = parallel_infos[0][spec.parallel].axis_size + axis_size: int = parallel_infos[0][modifier.axis].axis_size if spec.name == TOKEN_DIM_NAME: if thd_global_seq_lens is None: diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py index 58c41eda6..157aa4ba6 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_plugins.py @@ -122,7 +122,7 @@ class _SGLangPlugin(_AuxFrameworkPlugin): will be mishandled. Callers should set dims explicitly for non-zigzag CP. """ if ndim == 1: - return "t(cp,zigzag)" + return "t(cp:zigzag)" raise ValueError( f"SGLang: cannot infer dims for CP-sharded '{name}' with ndim={ndim}" ) @@ -208,9 +208,9 @@ class _MegatronPlugin(_AuxFrameworkPlugin): will be mishandled. Callers should set dims explicitly for non-zigzag CP. """ if ndim == 1: - return "t(cp,zigzag)" + return "t(cp:zigzag)" if ndim == 2: - return "b s(cp,zigzag)" + return "b s(cp:zigzag)" raise ValueError( f"Megatron: cannot infer dims for CP-sharded '{name}' with ndim={ndim}" ) 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 c88a620de..4083e7ada 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py @@ -14,6 +14,7 @@ from sglang.srt.debug_utils.comparator.dims import ( TOKEN_DIM_NAME, DimSpec, ParallelAxis, + ParallelModifier, ) # _CoordsList[tensor_index][axis] = @@ -36,18 +37,21 @@ def compute_unsharder_plan( if not parallel_infos: raise ValueError("parallel_infos must not be empty") - sharded_axis_infos: dict[ParallelAxis, DimSpec] = { - spec.parallel: spec for spec in dim_specs if spec.parallel is not None - } - sharded_axes_raw: set[ParallelAxis] = set(sharded_axis_infos) + # Within each dim spec, reverse modifier order: innermost shard (rightmost) unshards first. + reversed_sharded_modifiers: list[tuple[str, ParallelModifier]] = [ + (spec.name, m) for spec in dim_specs for m in reversed(spec.parallel_modifiers) + ] + sharded_axes_raw: set[ParallelAxis] = { + m.axis for _, m in reversed_sharded_modifiers + } all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info} # axis annotated in dims but absent from all parallel_infos -> axis_size=1, skip sharded_axes: set[ParallelAxis] = sharded_axes_raw & all_axes - sharded_axis_infos = { - k: v for k, v in sharded_axis_infos.items() if k in sharded_axes - } + reversed_sharded_modifiers = [ + (name, m) for name, m in reversed_sharded_modifiers if m.axis in sharded_axes + ] replicated_axes: set[ParallelAxis] = all_axes - sharded_axes if not sharded_axes and not replicated_axes: @@ -67,14 +71,15 @@ def compute_unsharder_plan( (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value) ] + [ ( - axis, + modifier.axis, _resolve_unshard_params( - spec=spec, + modifier=modifier, + dim_name=dim_name, parallel_infos=parallel_infos, thd_global_seq_lens=thd_global_seq_lens, ), ) - for axis, spec in sharded_axis_infos.items() + for dim_name, modifier in reversed_sharded_modifiers ] plans: list[UnsharderPlan] = [] @@ -151,23 +156,20 @@ def _group_and_project( def _resolve_unshard_params( *, - spec: DimSpec, + modifier: ParallelModifier, + dim_name: str, parallel_infos: list[dict[ParallelAxis, AxisInfo]], thd_global_seq_lens: Optional[list[int]] = None, ) -> UnsharderParams: - if spec.reduction is not None: + if modifier.reduction is not None: return ReduceSumParams() if ( - spec.name == TOKEN_DIM_NAME - and spec.parallel == ParallelAxis.CP + dim_name == TOKEN_DIM_NAME + and modifier.axis == ParallelAxis.CP and thd_global_seq_lens is not None ): - if spec.parallel is None: - raise ValueError( - f"THD unshard requires a parallel axis on dim '{spec.name}', but got None" - ) - axis_size: int = parallel_infos[0][spec.parallel].axis_size + axis_size: int = parallel_infos[0][modifier.axis].axis_size for s in thd_global_seq_lens: if s % axis_size != 0: raise ValueError( @@ -175,8 +177,6 @@ def _resolve_unshard_params( f"Sequences must be padded to a multiple of cp_size for CP zigzag." ) seq_lens_per_rank: list[int] = [s // axis_size for s in thd_global_seq_lens] - return CpThdConcatParams( - dim_name=spec.name, seq_lens_per_rank=seq_lens_per_rank - ) + return CpThdConcatParams(dim_name=dim_name, seq_lens_per_rank=seq_lens_per_rank) - return ConcatParams(dim_name=spec.name) + return ConcatParams(dim_name=dim_name) diff --git a/python/sglang/srt/debug_utils/comparator/dims.py b/python/sglang/srt/debug_utils/comparator/dims.py index 0e74e30d3..11c26d959 100644 --- a/python/sglang/srt/debug_utils/comparator/dims.py +++ b/python/sglang/srt/debug_utils/comparator/dims.py @@ -1,10 +1,11 @@ import re -from dataclasses import dataclass from enum import Enum from typing import Optional import torch +from sglang.srt.debug_utils.comparator.utils import _FrozenBase + TOKEN_DIM_NAME: str = "t" BATCH_DIM_NAME: str = "b" SEQ_DIM_NAME: str = "s" @@ -33,14 +34,17 @@ class Reduction(Enum): PARTIAL = "partial" -@dataclass(frozen=True) -class DimSpec: - name: str - parallel: Optional[ParallelAxis] = None +class ParallelModifier(_FrozenBase): + axis: ParallelAxis ordering: Optional[Ordering] = None reduction: Optional[Reduction] = None +class DimSpec(_FrozenBase): + name: str + parallel_modifiers: list[ParallelModifier] = [] + + class _SingletonDimUtil: """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping.""" @@ -83,16 +87,60 @@ class _SingletonDimUtil: _DIM_PATTERN = re.compile(r"^(?P[a-zA-Z_]\w*)(?:\((?P[^)]+)\))?$") -_MODIFIER_FIELDS: list[tuple[type[Enum], str]] = [ - (ParallelAxis, "parallel"), - (Ordering, "ordering"), - (Reduction, "reduction"), -] +_AXIS_LOOKUP: dict[str, ParallelAxis] = {m.value: m for m in ParallelAxis} +_QUALIFIER_LOOKUP: dict[str, Ordering | Reduction] = { + **{m.value: m for m in Ordering}, + **{m.value: m for m in Reduction}, +} -_MODIFIER_LOOKUP: dict[str, tuple[str, Enum]] = {} -for _enum_cls, _field in _MODIFIER_FIELDS: - for _member in _enum_cls: - _MODIFIER_LOOKUP[_member.value] = (_field, _member) + +def _parse_modifier_token(modifier_token: str, dim_token: str) -> ParallelModifier: + """Parse 'sp', 'cp:zigzag', 'tp:partial', or 'cp:zigzag+partial' → ParallelModifier. + + Format: ``axis`` or ``axis:qual`` or ``axis:qual+qual``. + Colon separates axis from qualifiers; ``+`` separates multiple qualifiers. + """ + axis_str: str + qualifiers_str: str + if ":" in modifier_token: + axis_str, qualifiers_str = modifier_token.split(":", maxsplit=1) + else: + axis_str, qualifiers_str = modifier_token, "" + + axis_str = axis_str.strip() + axis: Optional[ParallelAxis] = _AXIS_LOOKUP.get(axis_str) + if axis is None: + raise ValueError( + f"Unknown axis {axis_str!r} in modifier {modifier_token!r} " + f"of dim spec: {dim_token!r}" + ) + + ordering: Optional[Ordering] = None + reduction: Optional[Reduction] = None + + for q_str in (q.strip() for q in qualifiers_str.split("+") if q.strip()): + qualifier: Optional[Ordering | Reduction] = _QUALIFIER_LOOKUP.get(q_str) + if qualifier is None: + raise ValueError( + f"Unknown qualifier {q_str!r} in modifier " + f"{modifier_token!r} of dim spec: {dim_token!r}" + ) + if isinstance(qualifier, Ordering): + if ordering is not None: + raise ValueError( + f"Multiple ordering values in modifier " + f"{modifier_token!r} of dim spec: {dim_token!r}" + ) + ordering = qualifier + else: + if reduction is not None: + raise ValueError( + f"Multiple reduction values in modifier " + f"{modifier_token!r} of dim spec: {dim_token!r}" + ) + reduction = qualifier + + return ParallelModifier(axis=axis, ordering=ordering, reduction=reduction) def parse_dim(token: str) -> DimSpec: @@ -103,26 +151,29 @@ def parse_dim(token: str) -> DimSpec: if match is None: raise ValueError(f"Invalid dim token: {token!r}") - name = match.group("name") - modifiers_str = match.group("modifiers") + name: str = match.group("name") + modifiers_str: Optional[str] = match.group("modifiers") if modifiers_str is None: return DimSpec(name=name) - fields: dict[str, Enum] = {} - for part in (p.strip() for p in modifiers_str.split(",")): - if part not in _MODIFIER_LOOKUP: - raise ValueError(f"Unknown modifier {part!r} in dim spec: {token!r}") - field_name, enum_value = _MODIFIER_LOOKUP[part] - if field_name in fields: - raise ValueError(f"Multiple {field_name} values in dim token: {token!r}") - fields[field_name] = enum_value + modifiers: list[ParallelModifier] = [] + seen_axes: set[ParallelAxis] = set() - return DimSpec(name=name, **fields) + for modifier_token in (p.strip() for p in modifiers_str.split(",")): + modifier: ParallelModifier = _parse_modifier_token(modifier_token, token) + if modifier.axis in seen_axes: + raise ValueError( + f"Duplicate axis {modifier.axis.value!r} in dim spec: {token!r}" + ) + seen_axes.add(modifier.axis) + modifiers.append(modifier) + + return DimSpec(name=name, parallel_modifiers=modifiers) def parse_dims(dims_str: str) -> list[DimSpec]: - """Parse 'b s(cp,zigzag) h(tp) d' -> list[DimSpec].""" + """Parse 'b s(cp:zigzag) h(tp) d' -> list[DimSpec].""" if not dims_str.strip(): raise ValueError("dims string must not be empty") diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index ce52aa57a..d1ce7ce61 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -36,9 +36,6 @@ from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, ) -from sglang.srt.debug_utils.comparator.per_token_visualizer import ( - generate_per_token_heatmap, -) from sglang.srt.debug_utils.comparator.utils import Pair from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path @@ -258,12 +255,6 @@ def _consume_comparison_records( return summary, skipped_names - if visualize_per_token is not None and collected_comparisons: - generate_per_token_heatmap( - records=collected_comparisons, - output_path=visualize_per_token, - ) - def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() 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 36788f7c7..b6bec9b27 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py @@ -85,8 +85,8 @@ class TestComputePerStepSubPlans: def test_zigzag_returns_both_plans(self) -> None: result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( metas=[ - _make_meta(dims="b s(cp,zigzag) h", cp_rank=0, cp_size=2), - _make_meta(dims="b s(cp,zigzag) h", cp_rank=1, cp_size=2), + _make_meta(dims="b s(cp:zigzag) h", cp_rank=0, cp_size=2), + _make_meta(dims="b s(cp:zigzag) h", cp_rank=1, cp_size=2), ] ) unsharder_plans: list[UnsharderPlan] = [ @@ -177,33 +177,33 @@ class TestComputeAlignerPlan: class TestComputePerStepSubPlansThd: def test_thd_zigzag_returns_thd_plans(self) -> None: - """t(cp,zigzag) h(tp) generates THD-typed unsharder + reorderer plans.""" + """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)", + 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)", + 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)", + 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)", + dims="t(cp:zigzag) h(tp)", cp_rank=1, cp_size=2, tp_rank=1, diff --git a/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py b/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py index 1bc96ac97..4c3040e7d 100644 --- a/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py @@ -25,8 +25,8 @@ register_cpu_ci(est_time=10, suite="default", nightly=True) class TestComputeReordererPlans: def test_compute_reorderer_plans_zigzag(self) -> None: - """s(cp,zigzag) produces a ReordererPlan.""" - dim_specs = parse_dims("b s(cp,zigzag) h(tp)") + """s(cp:zigzag) produces a ReordererPlan.""" + dim_specs = parse_dims("b s(cp:zigzag) h(tp)") parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ { ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2), @@ -43,8 +43,8 @@ class TestComputeReordererPlans: assert plans[0].params.cp_size == 2 def test_compute_reorderer_plans_thd_zigzag(self) -> None: - """t(cp,zigzag) produces a ZigzagToNaturalThdParams plan.""" - dim_specs = parse_dims("t(cp,zigzag) h(tp)") + """t(cp:zigzag) produces a ZigzagToNaturalThdParams plan.""" + dim_specs = parse_dims("t(cp:zigzag) h(tp)") parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ { ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2), @@ -64,8 +64,8 @@ class TestComputeReordererPlans: assert plans[0].params.seq_lens == [100, 64, 92] def test_non_seq_dim_still_raises(self) -> None: - """Zigzag on non-sequence/non-token dim (e.g. h(cp,zigzag)) raises ValueError.""" - dim_specs = parse_dims("h(cp,zigzag) d") + """Zigzag on non-sequence/non-token dim (e.g. h(cp:zigzag)) raises ValueError.""" + dim_specs = parse_dims("h(cp:zigzag) d") parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ {ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}, ] @@ -73,8 +73,8 @@ class TestComputeReordererPlans: compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos) def test_thd_zigzag_without_seq_lens_raises(self) -> None: - """t(cp,zigzag) without thd_global_seq_lens raises ValueError.""" - dim_specs = parse_dims("t(cp,zigzag) h(tp)") + """t(cp:zigzag) without thd_global_seq_lens raises ValueError.""" + dim_specs = parse_dims("t(cp:zigzag) h(tp)") parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ { ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2), @@ -85,8 +85,8 @@ class TestComputeReordererPlans: compute_reorderer_plans(dim_specs=dim_specs, parallel_infos=parallel_infos) def test_thd_natural_no_reorder(self) -> None: - """t(cp,natural) and t(cp) produce no reorder plans.""" - for dims_str in ["t(cp,natural) h(tp)", "t(cp) h(tp)"]: + """t(cp:natural) and t(cp) produce no reorder plans.""" + for dims_str in ["t(cp:natural) h(tp)", "t(cp) h(tp)"]: dim_specs = parse_dims(dims_str) parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ { @@ -100,8 +100,8 @@ class TestComputeReordererPlans: assert plans == [] def test_compute_reorderer_plans_natural(self) -> None: - """s(cp) and s(cp,natural) produce no reorder plans.""" - for dims_str in ["b s(cp) h(tp)", "b s(cp,natural) h(tp)"]: + """s(cp) and s(cp:natural) produce no reorder plans.""" + for dims_str in ["b s(cp) h(tp)", "b s(cp:natural) h(tp)"]: dim_specs = parse_dims(dims_str) parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ { @@ -141,7 +141,7 @@ class TestCpZigzagTpE2E: } ) - dim_specs: list[DimSpec] = parse_dims("b s(cp,zigzag) h(tp)") + dim_specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp)") dim_names: list[str] = [s.name for s in dim_specs] unsharder_plans = compute_unsharder_plan( @@ -166,5 +166,83 @@ class TestCpZigzagTpE2E: assert torch.allclose(current[0].rename(None), full_tensor) +class TestCpZigzagSpSameDimE2E: + """E2E test for t(cp:zigzag,sp) — two axes sharding the same token dim.""" + + def test_cp2_sp2_zigzag_e2e(self) -> None: + """CP=2 zigzag + SP=2 on same token dim: full unshard + reorder round-trip. + + Shard order (outer to inner, matching left-to-right in dims annotation): + 1. CP zigzag splits token dim into 2 CP chunks (zigzag order) + 2. SP splits each CP chunk into 2 SP chunks + + Unshard order (inner to outer, right-to-left): + 1. SP concat (inner): merge SP chunks back + 2. CP concat (outer): merge CP chunks back + 3. Zigzag reorder: restore natural token order + """ + torch.manual_seed(42) + total_tokens: int = 16 + hidden: int = 8 + full_tensor: torch.Tensor = torch.randn(total_tokens, hidden) + + # Step 1: CP zigzag split — split into 2*cp_size=4 natural chunks, reorder by zigzag + cp_size: int = 2 + sp_size: int = 2 + n_natural_chunks: int = cp_size * 2 + natural_chunks: list[torch.Tensor] = list( + full_tensor.chunk(n_natural_chunks, dim=0) + ) + zigzag_order: list[int] = [0, 3, 1, 2] + zigzagged: torch.Tensor = torch.cat( + [natural_chunks[i] for i in zigzag_order], dim=0 + ) + cp_chunks: list[torch.Tensor] = list(zigzagged.chunk(cp_size, dim=0)) + + # Step 2: SP split within each CP chunk + tensors: list[torch.Tensor] = [] + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(cp_size): + sp_chunks: list[torch.Tensor] = list( + cp_chunks[cp_rank].chunk(sp_size, dim=0) + ) + for sp_rank in range(sp_size): + tensors.append(sp_chunks[sp_rank]) + parallel_infos.append( + { + ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=cp_size), + ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=sp_size), + } + ) + + dim_specs: list[DimSpec] = parse_dims("t(cp:zigzag,sp) h") + dim_names: list[str] = [s.name for s in dim_specs] + + unsharder_plans = compute_unsharder_plan( + dim_specs=dim_specs, parallel_infos=parallel_infos + ) + reorderer_plans = compute_reorderer_plans( + dim_specs=dim_specs, + parallel_infos=parallel_infos, + thd_global_seq_lens=[total_tokens], + ) + all_plans = [*unsharder_plans, *reorderer_plans] + + assert len(unsharder_plans) == 2 # SP concat, CP concat + assert unsharder_plans[0].axis == ParallelAxis.SP + assert unsharder_plans[1].axis == ParallelAxis.CP + assert len(reorderer_plans) == 1 # zigzag reorder + + current: list[torch.Tensor] = [t.refine_names(*dim_names) for t in tensors] + for plan in all_plans: + if isinstance(plan, ReordererPlan): + current = execute_reorderer_plan(plan, current) + else: + current = execute_unsharder_plan(plan, current).tensors + + assert len(current) == 1 + assert torch.allclose(current[0].rename(None), full_tensor) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py index 7f5c6de11..0b36871e5 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_loader.py @@ -90,7 +90,7 @@ class TestEnsureDimsInMetas: assert result is metas def test_cp_sharded_sglang_input_ids_infers_dims(self): - """CP + input_ids in sglang infers dims 't(cp,zigzag)'.""" + """CP + input_ids in sglang infers dims 't(cp:zigzag)'.""" metas: list[dict] = [ self._make_meta(cp_size=2, cp_rank=0), self._make_meta(cp_size=2, cp_rank=1), @@ -99,11 +99,11 @@ class TestEnsureDimsInMetas: name="input_ids", plugin=_sglang_plugin, metas=metas, ndim=1 ) assert result is not metas - assert result[0]["dims"] == "t(cp,zigzag)" - assert result[1]["dims"] == "t(cp,zigzag)" + assert result[0]["dims"] == "t(cp:zigzag)" + assert result[1]["dims"] == "t(cp:zigzag)" def test_cp_sharded_sglang_positions_infers_dims(self): - """CP + positions in sglang infers dims 't(cp,zigzag)'.""" + """CP + positions in sglang infers dims 't(cp:zigzag)'.""" metas: list[dict] = [ self._make_meta(cp_size=2, cp_rank=0), self._make_meta(cp_size=2, cp_rank=1), @@ -111,10 +111,10 @@ class TestEnsureDimsInMetas: result = _ensure_dims_in_metas( name="positions", plugin=_sglang_plugin, metas=metas, ndim=1 ) - assert result[0]["dims"] == "t(cp,zigzag)" + assert result[0]["dims"] == "t(cp:zigzag)" def test_cp_sharded_megatron_input_ids_infers_dims_1d(self): - """CP + input_ids in megatron (1D) infers dims 't(cp,zigzag)'.""" + """CP + input_ids in megatron (1D) infers dims 't(cp:zigzag)'.""" metas: list[dict] = [ {"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}}, {"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}}, @@ -122,10 +122,10 @@ class TestEnsureDimsInMetas: result = _ensure_dims_in_metas( name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=1 ) - assert result[0]["dims"] == "t(cp,zigzag)" + assert result[0]["dims"] == "t(cp:zigzag)" def test_cp_sharded_megatron_input_ids_infers_dims_2d(self): - """CP + input_ids in megatron (2D) infers dims 'b s(cp,zigzag)'.""" + """CP + input_ids in megatron (2D) infers dims 'b s(cp:zigzag)'.""" metas: list[dict] = [ {"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}}, {"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}}, @@ -133,7 +133,7 @@ class TestEnsureDimsInMetas: result = _ensure_dims_in_metas( name="input_ids", plugin=_megatron_plugin, metas=metas, ndim=2 ) - assert result[0]["dims"] == "b s(cp,zigzag)" + assert result[0]["dims"] == "b s(cp:zigzag)" def test_cp_non_sharded_name_returns_metas_unchanged(self): """CP + non-sharded tensor name (seq_lens) returns metas as-is.""" diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py index 2da3d78e1..4dcf797b8 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py @@ -218,19 +218,19 @@ class TestInferCpShardedDims: """Tests for infer_cp_sharded_dims on each plugin.""" def test_megatron_infer_1d(self) -> None: - """Megatron 1D → 't(cp,zigzag)'.""" + """Megatron 1D → 't(cp:zigzag)'.""" result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1) - assert result == "t(cp,zigzag)" + assert result == "t(cp:zigzag)" def test_megatron_infer_2d(self) -> None: - """Megatron 2D → 'b s(cp,zigzag)'.""" + """Megatron 2D → 'b s(cp:zigzag)'.""" result: str = _megatron_plugin.infer_cp_sharded_dims(name="input_ids", ndim=2) - assert result == "b s(cp,zigzag)" + assert result == "b s(cp:zigzag)" def test_sglang_infer_1d(self) -> None: - """SGLang 1D → 't(cp,zigzag)'.""" + """SGLang 1D → 't(cp:zigzag)'.""" result: str = _sglang_plugin.infer_cp_sharded_dims(name="input_ids", ndim=1) - assert result == "t(cp,zigzag)" + assert result == "t(cp:zigzag)" def test_megatron_infer_3d_raises(self) -> None: """Megatron 3D raises ValueError.""" 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 57283e1ce..f41cbc66f 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py @@ -654,7 +654,7 @@ class TestReduceSum: part_a = full_tensor * 0.6 part_b = full_tensor * 0.4 - dim_specs = parse_dims("h(tp,partial) d") + dim_specs = parse_dims("h(tp:partial) d") parallel_infos = [ {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2) ] @@ -676,7 +676,7 @@ class TestReduceSum: full_tensor = torch.randn(4, 8) parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)] - dim_specs = parse_dims("h(tp,partial) d") + dim_specs = parse_dims("h(tp:partial) d") parallel_infos = [ {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4) ] @@ -710,7 +710,7 @@ class TestReduceSum: } ) - dim_specs = parse_dims("b s(cp) h(tp,partial)") + dim_specs = parse_dims("b s(cp) h(tp:partial)") plans = compute_unsharder_plan(dim_specs, parallel_infos) assert len(plans) == 2 @@ -739,7 +739,7 @@ class TestReduceSum: {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") + 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) @@ -752,7 +752,7 @@ class TestReduceSum: def test_reduce_preserves_named_dims(self) -> None: """Named tensor dimensions are preserved through reduce_sum.""" - dim_specs = parse_dims("h(tp,partial) d") + 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") 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 55e012e74..cd359811c 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_planner.py @@ -175,7 +175,7 @@ class TestComputeUnsharderPlan: compute_unsharder_plan(dim_specs, parallel_infos) def test_reduction_partial_returns_reduce_sum(self) -> None: - dim_specs = parse_dims("h(tp,partial)") + dim_specs = parse_dims("h(tp:partial)") parallel_infos = [ {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2) ] @@ -188,7 +188,7 @@ class TestComputeUnsharderPlan: def test_reduction_partial_tp4(self) -> None: """TP=4 with partial reduction produces a single ReduceSumParams step.""" - dim_specs = parse_dims("h(tp,partial)") + dim_specs = parse_dims("h(tp:partial)") parallel_infos = [ {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4) ] @@ -200,7 +200,7 @@ class TestComputeUnsharderPlan: 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)") + 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): @@ -221,7 +221,7 @@ class TestComputeUnsharderPlan: def test_reduction_scrambled_ranks(self) -> None: """Scrambled world_rank order with partial reduction.""" - dim_specs = parse_dims("h(tp,partial)") + 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)}, @@ -235,7 +235,7 @@ class TestComputeUnsharderPlan: assert plans[0].groups == [[1, 3, 0, 2]] def test_ordering_zigzag_accepted(self) -> None: - dim_specs = parse_dims("s(cp,zigzag)") + dim_specs = parse_dims("s(cp:zigzag)") parallel_infos = [ {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2) ] @@ -244,7 +244,7 @@ class TestComputeUnsharderPlan: assert plans[0].axis == ParallelAxis.CP def test_ordering_natural_accepted(self) -> None: - dim_specs = parse_dims("s(cp,natural)") + dim_specs = parse_dims("s(cp:natural)") parallel_infos = [ {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2) ] @@ -288,6 +288,77 @@ class TestComputeUnsharderPlan: assert len(plans[2].groups) == 1 assert len(plans[2].groups[0]) == 2 + def test_same_dim_cp_sp_plan(self) -> None: + """t(cp:zigzag,sp) with CP=2 SP=2: SP unshards first (inner), then CP.""" + dim_specs = parse_dims("t(cp:zigzag,sp) 1 h") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(2): + for sp_rank in range(2): + parallel_infos.append( + { + ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2), + ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=2), + } + ) + + plans = compute_unsharder_plan(dim_specs, parallel_infos) + + assert len(plans) == 2 + + # SP unshards first (rightmost modifier = innermost shard) + sp_plan = plans[0] + assert sp_plan.axis == ParallelAxis.SP + assert isinstance(sp_plan.params, ConcatParams) + assert sp_plan.params.dim_name == "t" + assert len(sp_plan.groups) == 2 + for group in sp_plan.groups: + assert len(group) == 2 + + # CP unshards second (leftmost modifier = outermost shard) + cp_plan = plans[1] + assert cp_plan.axis == ParallelAxis.CP + assert isinstance(cp_plan.params, ConcatParams) + assert cp_plan.params.dim_name == "t" + assert len(cp_plan.groups) == 1 + assert len(cp_plan.groups[0]) == 2 + + def test_same_dim_cp_sp_with_thd(self) -> None: + """t(cp:zigzag,sp) with THD: SP → ConcatParams, CP → CpThdConcatParams.""" + from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( + CpThdConcatParams, + ) + + dim_specs = parse_dims("t(cp:zigzag,sp) h") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(2): + for sp_rank in range(2): + parallel_infos.append( + { + ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2), + ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=2), + } + ) + + thd_global_seq_lens: list[int] = [100, 64] + plans = compute_unsharder_plan( + dim_specs, parallel_infos, thd_global_seq_lens=thd_global_seq_lens + ) + + assert len(plans) == 2 + + # SP unshards first: plain concat (SP is not CP, no THD special handling) + sp_plan = plans[0] + assert sp_plan.axis == ParallelAxis.SP + assert isinstance(sp_plan.params, ConcatParams) + assert sp_plan.params.dim_name == "t" + + # CP unshards second: THD concat because dim is 't' + axis is CP + thd_global_seq_lens provided + cp_plan = plans[1] + assert cp_plan.axis == ParallelAxis.CP + assert isinstance(cp_plan.params, CpThdConcatParams) + assert cp_plan.params.dim_name == "t" + assert cp_plan.params.seq_lens_per_rank == [50, 32] + def test_sp_in_dims_but_not_in_parallel_info(self) -> None: """s(sp) in dims but SP absent from parallel_info (SP disabled), should auto-skip.""" dim_specs = parse_dims("s(sp) b h(tp)") diff --git a/test/registered/debug_utils/comparator/test_dims.py b/test/registered/debug_utils/comparator/test_dims.py index 0717e6800..408d8f9bc 100644 --- a/test/registered/debug_utils/comparator/test_dims.py +++ b/test/registered/debug_utils/comparator/test_dims.py @@ -11,6 +11,7 @@ from sglang.srt.debug_utils.comparator.dims import ( DimSpec, Ordering, ParallelAxis, + ParallelModifier, Reduction, _SingletonDimUtil, apply_dim_names, @@ -32,48 +33,78 @@ class TestParseDim: assert parse_dim("b") == DimSpec(name="b") def test_parallel_axis(self) -> None: - assert parse_dim("h(tp)") == DimSpec(name="h", parallel=ParallelAxis.TP) + assert parse_dim("h(tp)") == DimSpec( + name="h", + parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)], + ) def test_all_parallel_axes(self) -> None: - assert parse_dim("a(tp)").parallel == ParallelAxis.TP - assert parse_dim("a(cp)").parallel == ParallelAxis.CP - assert parse_dim("a(ep)").parallel == ParallelAxis.EP - assert parse_dim("a(sp)").parallel == ParallelAxis.SP + assert parse_dim("a(tp)").parallel_modifiers[0].axis == ParallelAxis.TP + assert parse_dim("a(cp)").parallel_modifiers[0].axis == ParallelAxis.CP + assert parse_dim("a(ep)").parallel_modifiers[0].axis == ParallelAxis.EP + assert parse_dim("a(sp)").parallel_modifiers[0].axis == ParallelAxis.SP def test_ordering(self) -> None: - assert parse_dim("s(cp,zigzag)").ordering == Ordering.ZIGZAG - assert parse_dim("s(cp,natural)").ordering == Ordering.NATURAL + assert ( + parse_dim("s(cp:zigzag)").parallel_modifiers[0].ordering == Ordering.ZIGZAG + ) + assert ( + parse_dim("s(cp:natural)").parallel_modifiers[0].ordering + == Ordering.NATURAL + ) def test_reduction(self) -> None: - assert parse_dim("h(tp,partial)").reduction == Reduction.PARTIAL - - def test_all_modifiers(self) -> None: - assert parse_dim("s(cp,zigzag,partial)") == DimSpec( - name="s", - parallel=ParallelAxis.CP, - ordering=Ordering.ZIGZAG, - reduction=Reduction.PARTIAL, + assert ( + parse_dim("h(tp:partial)").parallel_modifiers[0].reduction + == Reduction.PARTIAL ) + def test_all_qualifiers(self) -> None: + assert parse_dim("s(cp:zigzag+partial)") == DimSpec( + name="s", + parallel_modifiers=[ + ParallelModifier( + axis=ParallelAxis.CP, + ordering=Ordering.ZIGZAG, + reduction=Reduction.PARTIAL, + ), + ], + ) + + def test_multi_axis(self) -> None: + result: DimSpec = parse_dim("t(cp:zigzag,sp)") + assert result.name == "t" + assert len(result.parallel_modifiers) == 2 + assert result.parallel_modifiers[0] == ParallelModifier( + axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG + ) + assert result.parallel_modifiers[1] == ParallelModifier(axis=ParallelAxis.SP) + def test_invalid_token_raises(self) -> None: with pytest.raises(ValueError, match="Invalid dim token"): parse_dim("h()") with pytest.raises(ValueError, match="Invalid dim token"): parse_dim("h(tp(x))") - def test_unknown_modifier_raises(self) -> None: - with pytest.raises(ValueError, match="Unknown modifier"): + def test_unknown_axis_raises(self) -> None: + with pytest.raises(ValueError, match="Unknown axis"): parse_dim("h(xyz)") - with pytest.raises(ValueError, match="Unknown modifier"): - parse_dim("h(tp,foobar)") + + def test_unknown_qualifier_raises(self) -> None: + with pytest.raises(ValueError, match="Unknown qualifier"): + parse_dim("h(tp:foobar)") def test_multiple_ordering_raises(self) -> None: with pytest.raises(ValueError, match="Multiple ordering"): - parse_dim("s(cp,zigzag,natural)") + parse_dim("s(cp:zigzag+natural)") def test_multiple_reduction_raises(self) -> None: with pytest.raises(ValueError, match="Multiple reduction"): - parse_dim("h(tp,partial,partial)") + parse_dim("h(tp:partial+partial)") + + def test_duplicate_axis_raises(self) -> None: + with pytest.raises(ValueError, match="Duplicate axis"): + parse_dim("h(tp,tp)") def test_squeeze_dim(self) -> None: assert parse_dim("1") == DimSpec(name="1") @@ -96,10 +127,18 @@ class TestParseDims: assert parse_dims("t") == [DimSpec(name="t")] def test_mixed_annotated(self) -> None: - assert parse_dims("b s(cp,zigzag) h(tp) d") == [ + assert parse_dims("b s(cp:zigzag) h(tp) d") == [ DimSpec(name="b"), - DimSpec(name="s", parallel=ParallelAxis.CP, ordering=Ordering.ZIGZAG), - DimSpec(name="h", parallel=ParallelAxis.TP), + DimSpec( + name="s", + parallel_modifiers=[ + ParallelModifier(axis=ParallelAxis.CP, ordering=Ordering.ZIGZAG), + ], + ), + DimSpec( + name="h", + parallel_modifiers=[ParallelModifier(axis=ParallelAxis.TP)], + ), DimSpec(name="d"), ] @@ -158,7 +197,7 @@ class TestFindDimIndex: assert find_dim_index(specs, "d") == 3 def test_with_modifiers(self) -> None: - specs: list[DimSpec] = parse_dims("b s(cp,zigzag) h(tp) d") + specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp) d") assert find_dim_index(specs, "h") == 2 def test_empty_list(self) -> None: diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index add958808..78c0dd79e 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -837,7 +837,7 @@ class TestEntrypointGroupingLogical: tp_size=1, seq_dim=1, head_dim=2, - dims_str="b s(cp,zigzag) h", + dims_str="b s(cp:zigzag) h", ) args = _make_args( @@ -871,7 +871,7 @@ class TestEntrypointGroupingLogical: tp_size=2, seq_dim=1, head_dim=2, - dims_str="b s(cp,zigzag) h(tp)", + dims_str="b s(cp:zigzag) h(tp)", ) args = _make_args( @@ -959,14 +959,14 @@ class TestEntrypointGroupingLogical: full_tensor=full_baseline, name="attn_out", tp_size=2, - dims_str="b h(tp,partial)", + 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)", + dims_str="b h(tp:partial)", ) args = _make_args(baseline_path, target_path, diff_threshold=0.01) @@ -997,7 +997,7 @@ class TestEntrypointGroupingLogical: full_tensor=target_full, name="attn_out", tp_size=2, - dims_str="b h(tp,partial)", + dims_str="b h(tp:partial)", ) args = _make_args(baseline_path, target_path, diff_threshold=0.01) @@ -1026,7 +1026,7 @@ class TestEntrypointGroupingLogical: rank=rank, name="hidden", tensor=cp_chunks[cp_rank] / 2, - dims="b s(cp) h(tp,partial)", + dims="b s(cp) h(tp:partial)", parallel_info={ "cp_rank": cp_rank, "cp_size": 2, @@ -1046,6 +1046,38 @@ class TestEntrypointGroupingLogical: comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" + def test_cp_zigzag_sp_same_dim_unshard(self, tmp_path, capsys): + """CP=2 zigzag + SP=2 on same seq dim: multi-axis unshard + reorder.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8, 6) + full_target = full_baseline + torch.randn(4, 8, 6) * 0.001 + + 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_cp_zigzag_sp_sharded_dumps( + side_dir, + full_tensor=full_tensor, + name="hidden", + cp_size=2, + sp_size=2, + dims_str="b s(cp:zigzag,sp) h", + ) + + 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.name == "hidden" + class TestEntrypointConcatMode: """Test concat token-aligner mode through the full entrypoint pipeline.""" @@ -2578,6 +2610,63 @@ def _create_cp_zigzag_tp_sharded_dumps( return directory / _FIXED_EXP_NAME +def _create_cp_zigzag_sp_sharded_dumps( + directory: Path, + *, + full_tensor: torch.Tensor, + name: str, + cp_size: int, + sp_size: int, + dims_str: str, + seq_dim: int = 1, + num_steps: int = 1, +) -> Path: + """Create CP-zigzag + SP sharded dump files for a seq dim (b s h format). + + Shard order (outer to inner, matching left-to-right in dims annotation): + 1. CP zigzag splits seq dim into cp_size chunks (zigzag order) + 2. SP splits each CP chunk into sp_size chunks + """ + num_chunks: int = cp_size * 2 + natural_chunks: list[torch.Tensor] = list( + full_tensor.chunk(num_chunks, dim=seq_dim) + ) + + zigzag_order: list[int] = [] + for i in range(cp_size): + zigzag_order.append(i) + zigzag_order.append(num_chunks - 1 - i) + + zigzagged: torch.Tensor = torch.cat( + [natural_chunks[idx] for idx in zigzag_order], dim=seq_dim + ) + cp_chunks: list[torch.Tensor] = list(zigzagged.chunk(cp_size, dim=seq_dim)) + + rank: int = 0 + for cp_rank in range(cp_size): + sp_chunks: list[torch.Tensor] = list( + cp_chunks[cp_rank].chunk(sp_size, dim=seq_dim) + ) + for sp_rank in range(sp_size): + _create_rank_dump( + directory, + rank=rank, + name=name, + tensor=sp_chunks[sp_rank], + dims=dims_str, + parallel_info={ + "cp_rank": cp_rank, + "cp_size": cp_size, + "sp_rank": sp_rank, + "sp_size": sp_size, + }, + num_steps=num_steps, + ) + rank += 1 + + return directory / _FIXED_EXP_NAME + + def _create_replicated_tp_sharded_cp_dumps( directory: Path, *, @@ -2772,7 +2861,7 @@ def _create_thd_cp_zigzag_dumps( seq_lens: list[int], cp_size: int, total_per_rank: int, - dims_str: str = "t(cp,zigzag)", + dims_str: str = "t(cp:zigzag)", num_steps: int = 1, ) -> Path: """Create THD CP-zigzag sharded dump files simulating Megatron forward. @@ -2981,7 +3070,7 @@ class TestEntrypointThdCpZigzag: rank=cp_rank, name="hidden_states", tensor=rank_hidden, - dims="t(cp,zigzag) h", + dims="t(cp:zigzag) h", parallel_info={"cp_rank": cp_rank, "cp_size": cp_size}, framework="megatron", extra_dumps=[ diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 643250012..394948dac 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -2348,7 +2348,7 @@ class TestDumperDims: ) tensor = torch.randn(4, 8, requires_grad=True) - dumper.dump("hidden", tensor, dims="b h(tp)", dims_grad="b h(tp,partial)") + dumper.dump("hidden", tensor, dims="b h(tp)", dims_grad="b h(tp:partial)") dumper.step() tensor.backward(torch.ones_like(tensor)) @@ -2362,10 +2362,10 @@ class TestDumperDims: value_data = torch.load(value_file, weights_only=False) assert value_data["meta"]["dims"] == "b h(tp)" - assert value_data["meta"]["dims_grad"] == "b h(tp,partial)" + assert value_data["meta"]["dims_grad"] == "b h(tp:partial)" grad_data = torch.load(grad_file, weights_only=False) - assert grad_data["meta"]["dims"] == "b h(tp,partial)" + assert grad_data["meta"]["dims"] == "b h(tp:partial)" def test_dims_grad_inherits(self, tmp_path) -> None: dumper = _Dumper(