diff --git a/miles/backends/sglang_utils/arguments.py b/miles/backends/sglang_utils/arguments.py index 71e3f48fb4..81e50adeb5 100644 --- a/miles/backends/sglang_utils/arguments.py +++ b/miles/backends/sglang_utils/arguments.py @@ -107,9 +107,15 @@ def new_add_argument_wrapper(*name_or_flags, **kwargs): # Avoid double prefixing if dest somehow already starts with sglang_ if not original_dest.startswith("sglang_"): final_kwargs["dest"] = f"sglang_{original_dest}" - # If 'dest' is not explicitly provided (or is None/not a string), - # argparse will derive 'dest' from the (now prefixed) flag names. - # E.g., if the first flag is "--sglang-foo-bar", argparse sets dest to "sglang_foo_bar". + elif "dest" not in final_kwargs: + # argparse derives dest from the first alias, so store parallel sizes under SGLang's short field names. + for item_flag in name_or_flags: + if not isinstance(item_flag, str) or not item_flag.startswith("--"): + continue + canonical_dest = item_flag[2:].replace("-", "_") + if canonical_dest in ("tp_size", "dp_size", "pp_size", "ep_size"): + final_kwargs["dest"] = f"sglang_{canonical_dest}" + break old_add_argument(*new_name_or_flags_list, **final_kwargs) diff --git a/miles/utils/chat_template_utils/tito_tokenizer.py b/miles/utils/chat_template_utils/tito_tokenizer.py index cbd8802302..7bdeb58026 100644 --- a/miles/utils/chat_template_utils/tito_tokenizer.py +++ b/miles/utils/chat_template_utils/tito_tokenizer.py @@ -768,6 +768,10 @@ def __init__( }, allowed_append_roles=allowed_append_roles, ) + self.chat_template_kwargs = { + **self.chat_template_kwargs, + "thinking": deepseek.V32.render_thinking_enabled(self.chat_template_kwargs), + } # --------------------------------------------------------------------------- diff --git a/miles/utils/test_utils/session_verify_runner.py b/miles/utils/test_utils/session_verify_runner.py index 0ec4a70672..23468157f4 100644 --- a/miles/utils/test_utils/session_verify_runner.py +++ b/miles/utils/test_utils/session_verify_runner.py @@ -70,7 +70,8 @@ "ci_test": True, "colocate": True, "train_backend": "fsdp", - "sglang_expert_parallel_size": 1, + "sglang_ep_size": 1, + "enable_spec": False, } @@ -166,8 +167,17 @@ def namespace_to_train_args(ns: argparse.Namespace) -> str: # DeepSeek V3.2 (and other NSA/MoE archs) requires expert-parallel > 1 in # sglang; the default is 1, which is fatal at engine init. Only emit the # flag when the caller asks for ep>1 so single-expert models stay untouched. - if ns.sglang_expert_parallel_size > 1: - parts.append(f"--sglang-expert-parallel-size {ns.sglang_expert_parallel_size}") + if ns.sglang_ep_size > 1: + parts.append(f"--sglang-expert-parallel-size {ns.sglang_ep_size}") + if ns.enable_spec: + parts.extend( + [ + "--sglang-speculative-algorithm EAGLE", + "--sglang-speculative-num-steps 2", + "--sglang-speculative-eagle-topk 1", + "--sglang-speculative-num-draft-tokens 3", + ] + ) if ns.use_session_server: parts.append("--use-session-server") if ns.debug_rollout_only: diff --git a/tests/e2e/sglang/test_session_server_multi_role/_common.py b/tests/e2e/sglang/test_session_server_multi_role/_common.py index 772c6e8b4c..25f5b99fb7 100644 --- a/tests/e2e/sglang/test_session_server_multi_role/_common.py +++ b/tests/e2e/sglang/test_session_server_multi_role/_common.py @@ -3,7 +3,7 @@ Each test file in this directory owns a single ``ModelConfig`` and drives it through ``run_one(cfg)``. The runner is a thin wrapper around ``miles.utils.test_utils.session_verify_runner.run_session_verify`` with the -4-GPU H200 ``num_gpus`` override applied centrally. +model-specific GPU topology applied centrally. """ import argparse @@ -28,6 +28,7 @@ class ModelConfig: # sglang expert-parallel size. MoE archs like DeepSeek V4 hit a fused-moe # shape assert at ep=1; mirror the family's serving recipe (usually =tp). ep_size: int = 1 + enable_spec: bool = False cycles: int = 3 n_samples_per_prompt: int = 4 # Soft-threshold override for assistant_text mismatch ratio. Default @@ -44,7 +45,8 @@ class ModelConfig: def run_one(cfg: ModelConfig) -> None: invariants = dict(SESSION_VERIFY_INVARIANT_ARGS) - invariants["sglang_expert_parallel_size"] = cfg.ep_size + invariants["sglang_ep_size"] = cfg.ep_size + invariants["enable_spec"] = cfg.enable_spec args = argparse.Namespace( hf_checkpoint=cfg.model_name, tito_model=cfg.tito_model, diff --git a/tests/e2e/sglang/test_session_server_multi_role/test_deepseekv4.py b/tests/e2e/sglang/test_session_server_multi_role/test_deepseekv4.py index 59e38bebf7..42313f5b49 100644 --- a/tests/e2e/sglang/test_session_server_multi_role/test_deepseekv4.py +++ b/tests/e2e/sglang/test_session_server_multi_role/test_deepseekv4.py @@ -16,6 +16,7 @@ tp_size=4, # V4-Flash serving recipe (scripts/run_deepseek_v4.py): tp=4, ep=4. ep_size=4, + enable_spec=True, cycles=2, # V4 sorts tool_result blocks by the preceding assistant's tool_calls # order, so a sentinel tool_call_id would not roundtrip; use the diff --git a/tests/fast/utils/chat_template_utils/test_tito_tokenizer.py b/tests/fast/utils/chat_template_utils/test_tito_tokenizer.py index 31b2aff54c..c74637d350 100644 --- a/tests/fast/utils/chat_template_utils/test_tito_tokenizer.py +++ b/tests/fast/utils/chat_template_utils/test_tito_tokenizer.py @@ -46,12 +46,14 @@ from __future__ import annotations from pathlib import Path +from unittest.mock import MagicMock import pytest from transformers import AutoTokenizer from miles.utils.chat_template_utils import MismatchType, apply_chat_template, resolve_fixed_chat_template from miles.utils.chat_template_utils.tito_tokenizer import ( + DeepSeekV32TITOTokenizer, GLM47TITOTokenizer, Qwen3TITOTokenizer, Qwen35TITOTokenizer, @@ -221,6 +223,33 @@ def test_default(self, default_tito: TITOTokenizer): assert default_tito._assistant_start_str is None assert default_tito.trailing_token_ids == frozenset() + @pytest.mark.parametrize( + "chat_template_kwargs, expected", + [ + pytest.param({}, True, id="default-thinking"), + pytest.param({"enable_thinking": False}, False, id="disable-via-miles-kwarg"), + pytest.param({"thinking": False}, False, id="disable-via-sglang-kwarg"), + pytest.param( + {"enable_thinking": False, "thinking": True}, + False, + id="miles-kwarg-precedes-sglang-kwarg", + ), + pytest.param( + {"thinking_mode": "thinking", "thinking": False}, + True, + id="explicit-mode-precedes-sglang-kwarg", + ), + pytest.param({"thinking_mode": "chat"}, False, id="explicit-chat-mode"), + ], + ) + def test_deepseek_v32_forwards_effective_thinking_mode(self, chat_template_kwargs, expected): + tokenizer = MagicMock() + tokenizer.convert_tokens_to_ids.side_effect = [1, 2] + + tito = DeepSeekV32TITOTokenizer(tokenizer, chat_template_kwargs=chat_template_kwargs) + + assert tito.chat_template_kwargs["thinking"] is expected + def test_comparator_inherits_trailing_ids(self, qwen3_tito: Qwen3TITOTokenizer): """create_comparator propagates trailing_token_ids to the comparator's trim set.""" comp = qwen3_tito.create_comparator() diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 20a7bfb758..108db29c9e 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -6,6 +6,8 @@ import pytest +from miles.backends.sglang_utils.arguments import add_sglang_arguments +from miles.backends.sglang_utils.arguments import validate_args as validate_sglang_args from miles.utils.arguments import ( _maybe_apply_dumper_overrides, _resolve_ft_components, @@ -208,3 +210,68 @@ def test_enabled_with_components_returns_distinct_copy(self) -> None: assert result == ["train", "rollout"] assert result is not components + + +@pytest.mark.parametrize( + ("parallel_args", "expected"), + [ + ([], (1, 1, 1, 1)), + ( + [ + "--sglang-tensor-parallel-size", + "2", + "--sglang-data-parallel-size", + "3", + "--sglang-pipeline-parallel-size", + "4", + "--sglang-expert-parallel-size", + "5", + "--sglang-enable-dp-attention", + ], + (2, 3, 4, 5), + ), + ( + [ + "--sglang-tp-size", + "2", + "--sglang-dp-size", + "3", + "--sglang-pp-size", + "4", + "--sglang-ep-size", + "5", + "--sglang-enable-dp-attention", + ], + (2, 3, 4, 5), + ), + ], +) +def test_sglang_parallel_sizes_use_short_namespace_fields(parallel_args, expected): + parser = argparse.ArgumentParser() + add_sglang_arguments(parser) + args = parser.parse_args(parallel_args) + + assert (args.sglang_tp_size, args.sglang_dp_size, args.sglang_pp_size, args.sglang_ep_size) == expected + assert not hasattr(args, "sglang_tensor_parallel_size") + assert not hasattr(args, "sglang_data_parallel_size") + assert not hasattr(args, "sglang_pipeline_parallel_size") + assert not hasattr(args, "sglang_expert_parallel_size") + + args.rollout_num_gpus_per_engine = 8 + args.true_on_policy_mode = False + args.recompute_logprobs_via_prefill = False + args.sglang_router_policy = None + + validate_sglang_args(args) + + assert args.sglang_tp_size == 8 + assert (args.sglang_dp_size, args.sglang_pp_size, args.sglang_ep_size) == expected[1:] + + +def test_sglang_parallel_size_aliases_keep_last_value(): + parser = argparse.ArgumentParser() + add_sglang_arguments(parser) + + args = parser.parse_args(["--sglang-data-parallel-size", "2", "--sglang-dp-size", "3"]) + + assert args.sglang_dp_size == 3 diff --git a/tests/fast/utils/test_utils/test_session_verify_runner.py b/tests/fast/utils/test_utils/test_session_verify_runner.py index 770f7f5401..b2b5b8637d 100644 --- a/tests/fast/utils/test_utils/test_session_verify_runner.py +++ b/tests/fast/utils/test_utils/test_session_verify_runner.py @@ -55,11 +55,26 @@ def test_namespace_to_train_args_omits_expert_parallel_for_single_expert(): def test_namespace_to_train_args_emits_expert_parallel_for_moe(): - train_args = _build_args(sglang_expert_parallel_size=8) + train_args = _build_args(sglang_ep_size=8) assert "--sglang-expert-parallel-size 8" in train_args +def test_namespace_to_train_args_omits_speculative_decoding_by_default(): + train_args = _build_args() + + assert "--sglang-speculative-" not in train_args + + +def test_namespace_to_train_args_enables_eagle_speculative_decoding(): + train_args = _build_args(enable_spec=True) + + assert "--sglang-speculative-algorithm EAGLE" in train_args + assert "--sglang-speculative-num-steps 2" in train_args + assert "--sglang-speculative-eagle-topk 1" in train_args + assert "--sglang-speculative-num-draft-tokens 3" in train_args + + def _write_metrics(path, entries: list[dict]) -> None: path.write_text("\n".join(json.dumps(entry) for entry in entries) + "\n") diff --git a/tests/manual/tito/__init__.py b/tests/manual/tito/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/manual/tito/deepseek_v32.py b/tests/manual/tito/deepseek_v32.py new file mode 100644 index 0000000000..0b56a7b543 --- /dev/null +++ b/tests/manual/tito/deepseek_v32.py @@ -0,0 +1,23 @@ +import os + +from tests.e2e.sglang.test_session_server_multi_role._common import ModelConfig, run_one + +CONFIG = ModelConfig( + model_name=os.environ.get("DEEPSEEK_V32_MODEL", "deepseek-ai/DeepSeek-V3.2"), + reasoning_parser="deepseek-v3", + tool_call_parser="deepseekv32", + tito_model="deepseekv32", + allowed_append_roles=("tool",), + num_gpus=8, + tp_size=8, + ep_size=8, + enable_spec=True, +) + + +def test_deepseek_v32_session_tito(): + run_one(CONFIG) + + +if __name__ == "__main__": + test_deepseek_v32_session_tito()