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

61
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py
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]

11
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
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

12
python/sglang/srt/debug_utils/comparator/entrypoint.py
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,
)