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

Integrate packed data context parallel in dump comparator (#19464)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:28:36 +08:00
committed by GitHub
parent 1a32c0db4d
commit 23cbbd6d41
5 changed files with 364 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py
+32 -29
View File
@@ -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]
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
+10 -1
View File
@@ -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
python/sglang/srt/debug_utils/comparator/entrypoint.py
+8 -4
View File
@@ -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,
)
test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py
+68 -2
View File
@@ -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__]))
test/registered/debug_utils/comparator/test_entrypoint.py
+246 -2
View File
@@ -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__]))