Pipeline parallelism x speculative decoding (EAGLE/MTP) compatibility (#30775)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: YAMY1234 <74099316+YAMY1234@users.noreply.github.com> Co-authored-by: Yangmin Li <yangminl@nvidia.com>
This commit is contained in:
eeecho
co-authored by
Claude Fable 5
YAMY1234
Yangmin Li
Claude Fable 5
YAMY1234
Yangmin Li
parent
2733afe54e
commit
25ce8063f7
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_pp_spec.TestPPSpecConsistency.test_pp_matches_non_pp
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_DRAFT_MODEL_EAGLE,
|
||||
DEFAULT_TARGET_MODEL_EAGLE,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="extra-b", runner_config="4-gpu-h100")
|
||||
|
||||
# topk=1 chains and a topk=2 tree: the relayed topology is constant for the
|
||||
# former and data-dependent for the latter, so both shapes are covered.
|
||||
SPEC_SHAPES = {
|
||||
"chain": ("2", "1", "3"), # num_steps, eagle_topk, num_draft_tokens
|
||||
"tree": ("2", "2", "4"),
|
||||
}
|
||||
|
||||
|
||||
class TestPPSpecConsistency(CustomTestCase):
|
||||
"""PP x speculative decoding must match plain speculative decoding.
|
||||
|
||||
Every PP stage rebuilds the verify tree from state relayed by the last
|
||||
stage, so a mis-sized proxy buffer or a mis-rebuilt tree shows up as an
|
||||
accuracy drop or as speculation that never gets accepted -- not as a
|
||||
crash. Pin both.
|
||||
"""
|
||||
|
||||
def _run(self, pp_size: int, shape: str):
|
||||
num_steps, topk, num_draft_tokens = SPEC_SHAPES[shape]
|
||||
other_args = [
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_DRAFT_MODEL_EAGLE,
|
||||
"--speculative-num-steps",
|
||||
num_steps,
|
||||
"--speculative-eagle-topk",
|
||||
topk,
|
||||
"--speculative-num-draft-tokens",
|
||||
num_draft_tokens,
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
]
|
||||
if pp_size > 1:
|
||||
other_args += ["--pp-size", str(pp_size), "--disable-overlap-schedule"]
|
||||
|
||||
process = popen_launch_server(
|
||||
DEFAULT_TARGET_MODEL_EAGLE,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
env={**os.environ, "SGLANG_ENABLE_PP_SPEC": "1"},
|
||||
)
|
||||
try:
|
||||
metrics = run_eval(
|
||||
SimpleNamespace(
|
||||
base_url=DEFAULT_URL_FOR_TEST,
|
||||
model=DEFAULT_TARGET_MODEL_EAGLE,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=256,
|
||||
num_threads=32,
|
||||
)
|
||||
)
|
||||
server_info = requests.get(f"{DEFAULT_URL_FOR_TEST}/get_server_info")
|
||||
accept_length = server_info.json()["internal_states"][0][
|
||||
"avg_spec_accept_length"
|
||||
]
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
return metrics["score"], accept_length
|
||||
|
||||
def _assert_matches(self, shape: str):
|
||||
base_score, base_accept = self._run(pp_size=1, shape=shape)
|
||||
pp_score, pp_accept = self._run(pp_size=2, shape=shape)
|
||||
print(
|
||||
f"[PP spec {shape}] no-PP: score={base_score:.4f} accept={base_accept:.2f}"
|
||||
f" | PP2: score={pp_score:.4f} accept={pp_accept:.2f}"
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(
|
||||
pp_score,
|
||||
base_score - 0.02,
|
||||
msg=(
|
||||
f"PP+spec accuracy dropped more than 2% against plain spec. "
|
||||
f"no-PP: {base_score:.2%}, PP2: {pp_score:.2%}"
|
||||
),
|
||||
)
|
||||
# A relay that loses the drafted tree still produces correct output --
|
||||
# the bonus token is force-accepted and the rest is rejected -- so
|
||||
# accuracy alone cannot tell drafting from a no-op. Acceptance can.
|
||||
self.assertGreaterEqual(
|
||||
pp_accept,
|
||||
base_accept - 0.2,
|
||||
msg=(
|
||||
f"PP+spec accept length collapsed against plain spec. "
|
||||
f"no-PP: {base_accept:.2f}, PP2: {pp_accept:.2f}"
|
||||
),
|
||||
)
|
||||
|
||||
def test_pp_matches_non_pp(self):
|
||||
self._assert_matches("chain")
|
||||
|
||||
def test_pp_matches_non_pp_tree(self):
|
||||
self._assert_matches("tree")
|
||||
|
||||
|
||||
class TestPPSpecGate(CustomTestCase):
|
||||
"""The gate is off by default, and the combinations the relay cannot
|
||||
reproduce identically on every stage are rejected rather than silently
|
||||
mis-rebuilt."""
|
||||
|
||||
def _server_args(self, **overrides):
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.mock_model.utils import MOCK_MODEL_PATH
|
||||
|
||||
# Argument validation only reads the config, so the small mock model
|
||||
# keeps these cases off the GPU entirely.
|
||||
args = dict(
|
||||
model_path=MOCK_MODEL_PATH,
|
||||
pp_size=2,
|
||||
disable_overlap_schedule=True,
|
||||
speculative_algorithm="EAGLE",
|
||||
speculative_draft_model_path=MOCK_MODEL_PATH,
|
||||
speculative_num_steps=2,
|
||||
speculative_eagle_topk=1,
|
||||
speculative_num_draft_tokens=3,
|
||||
)
|
||||
args.update(overrides)
|
||||
server_args = ServerArgs(**args)
|
||||
# check_server_args reads resolution-filled fields (served_model_name,
|
||||
# chunked_prefill_size); the bare constructor leaves them None.
|
||||
server_args.resolve_once()
|
||||
return server_args
|
||||
|
||||
def test_gate_off_keeps_the_ban(self):
|
||||
os.environ.pop("SGLANG_ENABLE_PP_SPEC", None)
|
||||
with self.assertRaises(AssertionError):
|
||||
self._server_args().check_server_args()
|
||||
|
||||
def test_gate_on_rejects_unsupported_combinations(self):
|
||||
os.environ["SGLANG_ENABLE_PP_SPEC"] = "1"
|
||||
try:
|
||||
self._server_args().check_server_args()
|
||||
# DP attention partitions the batch per DP rank, so the stages
|
||||
# would no longer rebuild the same verify tree.
|
||||
with self.assertRaises(AssertionError):
|
||||
self._server_args(
|
||||
tp_size=2, dp_size=2, enable_dp_attention=True
|
||||
).check_server_args()
|
||||
# Adaptive spec changes num_draft_tokens at runtime, which the
|
||||
# relay slices results with.
|
||||
with self.assertRaises(AssertionError):
|
||||
self._server_args(
|
||||
speculative_adaptive=True, speculative_num_steps=3
|
||||
).check_server_args()
|
||||
# The relay carries an EAGLE-shaped tree; other algorithms would
|
||||
# be mis-rebuilt on the non-last stages.
|
||||
with self.assertRaises(AssertionError):
|
||||
self._server_args(
|
||||
speculative_algorithm="NGRAM", speculative_draft_model_path=None
|
||||
).check_server_args()
|
||||
# PD prefill needs the RelayPayload draft fields the gated flow
|
||||
# does not carry.
|
||||
with self.assertRaises(AssertionError):
|
||||
self._server_args(disaggregation_mode="prefill").check_server_args()
|
||||
finally:
|
||||
os.environ.pop("SGLANG_ENABLE_PP_SPEC", None)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,120 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_pp_spec_embed_scan.TestDraftEmbedScan
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from transformers import MistralConfig, PretrainedConfig
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
# Shrunk from inclusionAI/Ling-mini-2.0; the field names are the real
|
||||
# checkpoint's, only the sizes are cut down for test speed.
|
||||
_BAILING_CONFIG = {
|
||||
"architectures": ["BailingMoeV2ForCausalLM"],
|
||||
"model_type": "bailing_moe",
|
||||
"num_hidden_layers": 1,
|
||||
"num_nextn_predict_layers": 1,
|
||||
"hidden_size": 256,
|
||||
"intermediate_size": 512,
|
||||
"moe_intermediate_size": 128,
|
||||
"moe_shared_expert_intermediate_size": 128,
|
||||
"num_shared_experts": 1,
|
||||
"num_experts": 16,
|
||||
"num_experts_per_tok": 4,
|
||||
"n_group": 4,
|
||||
"topk_group": 2,
|
||||
"norm_topk_prob": True,
|
||||
"moe_router_enable_expert_bias": True,
|
||||
"routed_scaling_factor": 2.5,
|
||||
"score_function": "sigmoid",
|
||||
"router_dtype": "fp32",
|
||||
"first_k_dense_replace": 0,
|
||||
"num_attention_heads": 4,
|
||||
"num_key_value_heads": 2,
|
||||
"head_dim": 64,
|
||||
"partial_rotary_factor": 0.5,
|
||||
"use_qk_norm": True,
|
||||
"use_qkv_bias": False,
|
||||
"use_bias": False,
|
||||
"hidden_act": "silu",
|
||||
"rms_norm_eps": 1e-6,
|
||||
"max_position_embeddings": 4096,
|
||||
"rope_parameters": {"rope_theta": 600000, "rope_type": "default"},
|
||||
"vocab_size": 1024,
|
||||
"tie_word_embeddings": False,
|
||||
}
|
||||
|
||||
|
||||
class TestDraftEmbedScan(CustomTestCase):
|
||||
"""PP+spec loads the draft input embedding via a type scan plus a
|
||||
checkpoint-name table; a draft family whose embedding hangs under a new
|
||||
attribute name (or grows a second VocabParallelEmbedding) would misload
|
||||
silently. Pin the scan on the families the runtime GLM/DSv4 tests miss.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
||||
server_args.resolve_once()
|
||||
get_context().set_server_args(server_args)
|
||||
|
||||
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
||||
os.environ.setdefault("MASTER_PORT", "29611")
|
||||
init_distributed_environment(
|
||||
world_size=1, rank=0, local_rank=0, distributed_init_method="env://"
|
||||
)
|
||||
initialize_model_parallel(tensor_model_parallel_size=1)
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.cuda.set_device(0)
|
||||
|
||||
def test_bailing_nextn_embedding_is_found(self):
|
||||
from sglang.srt.models.bailing_moe_nextn import BailingMoeForCausalLMNextN
|
||||
from sglang.srt.speculative.eagle_worker_v2 import _find_draft_input_embedding
|
||||
|
||||
config = PretrainedConfig.from_dict(dict(_BAILING_CONFIG))
|
||||
with torch.device("cuda"):
|
||||
model = BailingMoeForCausalLMNextN(config)
|
||||
self.assertIs(_find_draft_input_embedding(model), model.model.word_embeddings)
|
||||
|
||||
def test_mistral_eagle_embedding_is_found(self):
|
||||
from sglang.srt.models.mistral_eagle import MistralForCausalLMEagle
|
||||
from sglang.srt.speculative.eagle_worker_v2 import _find_draft_input_embedding
|
||||
|
||||
config = MistralConfig(
|
||||
vocab_size=1024,
|
||||
hidden_size=256,
|
||||
intermediate_size=512,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
)
|
||||
with torch.device("cuda"):
|
||||
model = MistralForCausalLMEagle(config)
|
||||
self.assertIs(_find_draft_input_embedding(model), model.model.embed_tokens)
|
||||
|
||||
def test_name_table_matches_published_checkpoints(self):
|
||||
from sglang.srt.speculative.eagle_worker_v2 import _EMBED_TENSOR_NAMES
|
||||
|
||||
# External-source literals: the embedding tensor's spelling in each
|
||||
# family's published target checkpoint (inclusionAI/Ling-*-2.0
|
||||
# model.safetensors.index.json; Mistral-Large-3 consolidated index).
|
||||
self.assertIn("model.word_embeddings.weight", _EMBED_TENSOR_NAMES)
|
||||
self.assertIn("tok_embeddings.weight", _EMBED_TENSOR_NAMES)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -167,7 +167,7 @@ def test_pipeline_sampling_mask_round_trip_without_logprobs():
|
||||
next_token_ids=torch.tensor([3]),
|
||||
)
|
||||
payload = Scheduler._pp_prepare_tensor_dict(
|
||||
SimpleNamespace(), result, SimpleNamespace(return_logprob=False)
|
||||
object.__new__(Scheduler), result, SimpleNamespace(return_logprob=False)
|
||||
)
|
||||
output, _, _ = get_logprob_from_pp_outputs(PPProxyTensors(payload))
|
||||
for name in ("token_ids", "lengths", "selected_logprobs", "statuses"):
|
||||
|
||||
Reference in New Issue
Block a user