Support loading token aligner data in dump comparator (#19376)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,278 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
|
||||
_detect_plugin,
|
||||
_ensure_dims_in_metas,
|
||||
_load_and_align_aux_tensor,
|
||||
_load_non_tensor_aux,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
|
||||
_MegatronPlugin,
|
||||
_SGLangPlugin,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
|
||||
from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
_sglang_plugin = _SGLangPlugin()
|
||||
_megatron_plugin = _MegatronPlugin()
|
||||
|
||||
|
||||
def _save_pt(
|
||||
dump_path: Path,
|
||||
*,
|
||||
name: str,
|
||||
step: int,
|
||||
rank: int,
|
||||
value: object,
|
||||
meta: dict | None = None,
|
||||
) -> str:
|
||||
filename: str = f"name={name}___step={step}___rank={rank}.pt"
|
||||
payload: dict = {"value": value, "meta": meta or {}}
|
||||
torch.save(payload, dump_path / filename)
|
||||
return filename
|
||||
|
||||
|
||||
def _make_df_from_filenames(filenames: list[str]) -> pl.DataFrame:
|
||||
rows: list[dict] = []
|
||||
for fn in filenames:
|
||||
parts: dict = {}
|
||||
stem: str = fn.removesuffix(".pt")
|
||||
for kv in stem.split("___"):
|
||||
if "=" in kv:
|
||||
k, v = kv.split("=", 1)
|
||||
parts[k] = v
|
||||
rows.append(
|
||||
{
|
||||
"filename": fn,
|
||||
"name": parts["name"],
|
||||
"step": int(parts["step"]),
|
||||
"rank": int(parts["rank"]),
|
||||
}
|
||||
)
|
||||
return pl.DataFrame(rows)
|
||||
|
||||
|
||||
class TestEnsureDimsInMetas:
|
||||
"""Tests for _ensure_dims_in_metas."""
|
||||
|
||||
def _make_meta(self, *, cp_size: int = 1, cp_rank: int = 0) -> dict:
|
||||
return {
|
||||
"sglang_parallel_info": {
|
||||
"tp_rank": 0,
|
||||
"tp_size": 1,
|
||||
"cp_rank": cp_rank,
|
||||
"cp_size": cp_size,
|
||||
}
|
||||
}
|
||||
|
||||
def test_no_cp_returns_metas_unchanged(self):
|
||||
"""Without CP parallelism, metas are returned as-is."""
|
||||
metas: list[dict] = [self._make_meta(cp_size=1)]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_sglang_plugin, metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
def test_dims_already_present_returns_metas_unchanged(self):
|
||||
"""If dims is already in meta, metas are returned as-is."""
|
||||
metas: list[dict] = [{**self._make_meta(cp_size=2, cp_rank=0), "dims": "t"}]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_sglang_plugin, metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
def test_cp_sharded_sglang_input_ids_raises(self):
|
||||
"""CP + input_ids in sglang raises NotImplementedError."""
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="CP-sharded"):
|
||||
_ensure_dims_in_metas(name="input_ids", plugin=_sglang_plugin, metas=metas)
|
||||
|
||||
def test_cp_sharded_sglang_positions_raises(self):
|
||||
"""CP + positions in sglang raises NotImplementedError."""
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="CP-sharded"):
|
||||
_ensure_dims_in_metas(name="positions", plugin=_sglang_plugin, metas=metas)
|
||||
|
||||
def test_cp_sharded_megatron_input_ids_raises(self):
|
||||
"""CP + input_ids in megatron raises NotImplementedError."""
|
||||
metas: list[dict] = [
|
||||
{"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
|
||||
{"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="CP-sharded"):
|
||||
_ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_megatron_plugin, metas=metas
|
||||
)
|
||||
|
||||
def test_cp_non_sharded_name_returns_metas_unchanged(self):
|
||||
"""CP + non-sharded tensor name (seq_lens) returns metas as-is."""
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="seq_lens", plugin=_sglang_plugin, metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
def test_unknown_plugin_returns_metas_unchanged(self):
|
||||
"""CP + plugin with empty cp_sharded_names returns metas as-is."""
|
||||
|
||||
class _DummyPlugin(_SGLangPlugin):
|
||||
@property
|
||||
def cp_sharded_names(self) -> frozenset[str]:
|
||||
return frozenset()
|
||||
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_DummyPlugin(), metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
|
||||
class TestDetectPlugin:
|
||||
def test_discriminating_names_sglang(self, tmp_path: Path) -> None:
|
||||
fn: str = _save_pt(
|
||||
tmp_path, name="seq_lens", step=0, rank=0, value=torch.tensor([3])
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn])
|
||||
|
||||
result = _detect_plugin(df, dump_path=tmp_path)
|
||||
|
||||
assert result is not None
|
||||
assert result.name == "sglang"
|
||||
|
||||
def test_fallback_to_meta_based_detection(self, tmp_path: Path) -> None:
|
||||
fn: str = _save_pt(
|
||||
tmp_path,
|
||||
name="input_ids",
|
||||
step=0,
|
||||
rank=0,
|
||||
value=torch.tensor([1, 2, 3]),
|
||||
meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn])
|
||||
|
||||
result = _detect_plugin(df, dump_path=tmp_path)
|
||||
|
||||
assert result is not None
|
||||
assert result.name == "sglang"
|
||||
|
||||
def test_returns_none_no_match(self, tmp_path: Path) -> None:
|
||||
fn: str = _save_pt(
|
||||
tmp_path, name="unrelated_tensor", step=0, rank=0, value=torch.tensor([1])
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn])
|
||||
|
||||
result = _detect_plugin(df, dump_path=tmp_path)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestLoadNonTensorAux:
|
||||
def test_multi_rank_mismatch_warning(self, tmp_path: Path) -> None:
|
||||
fn0: str = _save_pt(tmp_path, name="rids", step=0, rank=0, value=["req_A"])
|
||||
fn1: str = _save_pt(tmp_path, name="rids", step=0, rank=1, value=["req_B"])
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
|
||||
|
||||
sink = WarningSink()
|
||||
with sink.context() as warnings:
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch(
|
||||
"sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
|
||||
sink,
|
||||
):
|
||||
result = _load_non_tensor_aux(
|
||||
name="rids", step=0, df=df, dump_path=tmp_path
|
||||
)
|
||||
|
||||
assert result == ["req_A"]
|
||||
assert len(warnings) == 1
|
||||
assert isinstance(warnings[0], GeneralWarning)
|
||||
assert "rids_mismatch" in warnings[0].category
|
||||
|
||||
def test_no_rows_returns_none(self, tmp_path: Path) -> None:
|
||||
df: pl.DataFrame = _make_df_from_filenames([])
|
||||
|
||||
result = _load_non_tensor_aux(name="rids", step=0, df=df, dump_path=tmp_path)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestLoadAndAlignAuxTensor:
|
||||
def test_multi_rank_no_dims_emits_warning(self, tmp_path: Path) -> None:
|
||||
fn0: str = _save_pt(
|
||||
tmp_path,
|
||||
name="input_ids",
|
||||
step=0,
|
||||
rank=0,
|
||||
value=torch.tensor([1, 2, 3]),
|
||||
meta={
|
||||
"sglang_parallel_info": {
|
||||
"tp_rank": 0,
|
||||
"tp_size": 2,
|
||||
"cp_rank": 0,
|
||||
"cp_size": 1,
|
||||
}
|
||||
},
|
||||
)
|
||||
fn1: str = _save_pt(
|
||||
tmp_path,
|
||||
name="input_ids",
|
||||
step=0,
|
||||
rank=1,
|
||||
value=torch.tensor([4, 5, 6]),
|
||||
meta={
|
||||
"sglang_parallel_info": {
|
||||
"tp_rank": 1,
|
||||
"tp_size": 2,
|
||||
"cp_rank": 0,
|
||||
"cp_size": 1,
|
||||
}
|
||||
},
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
|
||||
|
||||
sink = WarningSink()
|
||||
with sink.context() as warnings:
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch(
|
||||
"sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
|
||||
sink,
|
||||
):
|
||||
result = _load_and_align_aux_tensor(
|
||||
name="input_ids",
|
||||
step=0,
|
||||
df=df,
|
||||
dump_path=tmp_path,
|
||||
plugin=_sglang_plugin,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert torch.equal(result, torch.tensor([1, 2, 3]))
|
||||
assert len(warnings) == 1
|
||||
assert isinstance(warnings[0], GeneralWarning)
|
||||
assert "aux_no_dims" in warnings[0].category
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,144 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
|
||||
_infer_positions,
|
||||
_MegatronPlugin,
|
||||
_SGLangPlugin,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
|
||||
PositionalSeqId,
|
||||
SGLangSeqId,
|
||||
TokenAlignerStepAux,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
_sglang_plugin = _SGLangPlugin()
|
||||
_megatron_plugin = _MegatronPlugin()
|
||||
|
||||
|
||||
class TestNormalizeSGLang:
|
||||
"""Tests for SGLang aux tensor normalization."""
|
||||
|
||||
def test_with_rids(self):
|
||||
"""SGLang tensors with rids produce string seq_ids."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30]),
|
||||
"positions": torch.tensor([0, 1, 2]),
|
||||
"seq_lens": torch.tensor([3]),
|
||||
"rids": ["A"],
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.input_ids == [10, 20, 30]
|
||||
assert result.positions == [0, 1, 2]
|
||||
assert result.seq_lens == [3]
|
||||
assert result.seq_ids == [SGLangSeqId(rid="A")]
|
||||
|
||||
def test_rids_none_fallback(self):
|
||||
"""Missing rids results in (step, index) fallback seq_ids."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20]),
|
||||
"positions": torch.tensor([0, 1]),
|
||||
"seq_lens": torch.tensor([2]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=3
|
||||
)
|
||||
assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
|
||||
|
||||
def test_multiple_seqs_with_rids(self):
|
||||
"""Multiple sequences with rids."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"positions": torch.tensor([0, 1, 2, 0, 1]),
|
||||
"seq_lens": torch.tensor([3, 2]),
|
||||
"rids": ["A", "B"],
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
|
||||
|
||||
|
||||
class TestNormalizeMegatron:
|
||||
"""Tests for Megatron aux tensor normalization."""
|
||||
|
||||
def test_cu_seqlens_to_seq_lens(self):
|
||||
"""cu_seqlens_q is converted to seq_lens via differencing."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"cu_seqlens_q": torch.tensor([0, 3, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.seq_lens == [3, 2]
|
||||
|
||||
def test_positions_inferred_thd(self):
|
||||
"""Positions inferred from seq_lens in thd layout."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"cu_seqlens_q": torch.tensor([0, 3, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.positions == [0, 1, 2, 0, 1]
|
||||
|
||||
def test_position_ids_passthrough(self):
|
||||
"""Explicit position_ids used directly instead of inference."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"position_ids": torch.tensor([5, 6, 7, 8, 9]),
|
||||
"cu_seqlens_q": torch.tensor([0, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.positions == [5, 6, 7, 8, 9]
|
||||
|
||||
def test_seq_ids_are_step_index_tuples(self):
|
||||
"""Megatron seq_ids are (step, seq_index) tuples."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"cu_seqlens_q": torch.tensor([0, 3, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=5
|
||||
)
|
||||
assert result.seq_ids == [
|
||||
PositionalSeqId(step=5, seq_index=0),
|
||||
PositionalSeqId(step=5, seq_index=1),
|
||||
]
|
||||
|
||||
|
||||
class TestInferPositions:
|
||||
"""Tests for position inference helper."""
|
||||
|
||||
def test_thd_multiple_sequences(self):
|
||||
"""thd: positions reset to 0 for each sequence."""
|
||||
result = _infer_positions(
|
||||
seq_lens=torch.tensor([2, 3]),
|
||||
)
|
||||
assert torch.equal(result, torch.tensor([0, 1, 0, 1, 2]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user