diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 45606dd53..13d1530c6 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -625,6 +625,10 @@ jobs: "num_gpus": 0, "test_file": "test_sample.py" }, + { + "num_gpus": 0, + "test_file": "test_spec_metrics.py" + }, { "num_gpus": 0, "test_file": "test_rollout_validation.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index f58b06483..ec63e818e 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -81,6 +81,7 @@ {'test_file': 'test_rm_math_dapo.py', 'num_gpus': 0}, {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0}, {'test_file': 'test_sample.py', 'num_gpus': 0}, + {'test_file': 'test_spec_metrics.py', 'num_gpus': 0}, {'test_file': 'test_rollout_validation.py', 'num_gpus': 0}, {'test_file': 'test_reloadable_process_group_world.py', 'num_gpus': 0}, {'test_file': 'test_placement_group.py', 'num_gpus': 0}, diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 740ff18e3..c8ee49d25 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -1458,10 +1458,19 @@ def _compute_top_p_kept_vocab_metrics(args, all_samples: list[Sample]): def _compute_spec_metrics(args, all_samples: list[Sample]): if getattr(args, "sglang_speculative_algorithm", None) is None: return {} - num_samples = len(all_samples) + # Pool the raw counters, like _compute_prefix_cache_metrics below. Averaging + # the per-sample ratios weights a 10-token response the same as a 4k-token + # one, and counts samples whose spec counters were never populated (aborted / + # partial rollout) as hard zeros — both only ever bias the report downward. + accept_tokens = sum(sample.spec_info.spec_accept_token_num for sample in all_samples) + draft_tokens = sum(sample.spec_info.spec_draft_token_num for sample in all_samples) + completion_tokens = sum(sample.spec_info.completion_token_num for sample in all_samples) + verify_ct = sum(sample.spec_info.spec_verify_ct for sample in all_samples) metrics = {} - metrics["spec_accept_rate"] = sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples - metrics["spec_accept_length"] = sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples + if draft_tokens > 0: + metrics["spec_accept_rate"] = accept_tokens / draft_tokens + if verify_ct > 0: + metrics["spec_accept_length"] = completion_tokens / verify_ct return metrics diff --git a/tests/test_spec_metrics.py b/tests/test_spec_metrics.py new file mode 100644 index 000000000..b9f998955 --- /dev/null +++ b/tests/test_spec_metrics.py @@ -0,0 +1,117 @@ +"""CPU unit tests for ``slime.ray.rollout._compute_spec_metrics``. + +``Sample.SpecInfo`` deliberately accumulates the four raw counters across +partial-rollout turns ("cannot directly use spec info from sglang because of +partial rollout") precisely so a correct batch statistic can be formed. The +batch metric must therefore pool those counters — ``Σaccept / Σdraft`` and +``Σcompletion / Σverify`` — the same way ``_compute_prefix_cache_metrics`` +pools ``Σcached / Σprompt`` five lines below it. + +The old implementation averaged the *per-sample ratios* instead, which + + * weights a 10-token response the same as a 4k-token one, and + * counts samples whose counters were never populated (aborted / partial + rollout: ``SpecInfo.add`` only runs on terminal meta info) as hard 0.0s, + +both of which only ever bias the report downward — anyone tuning +``--sglang-speculative-*`` off these numbers optimizes against a deflated +signal. +""" + +from __future__ import annotations + +import sys +import types + + +# ``slime.ray.rollout`` imports sglang / sglang_router / wandb / transformers at +# module level; none are installed in the CPU CI env and none are touched by +# the pure functions under test. Stub just enough for the import to resolve +# (same approach as tests/test_agent/test_agent_rollout_cpu.py). +def _stub(name: str, **attrs) -> None: + if name in sys.modules: + return + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + sys.modules[name] = module + + +_stub("sglang_router", __version__="0.2.3") +_stub("sglang_router.launch_router", RouterArgs=object, launch_router=lambda *a, **k: None) +_stub("sglang") +_stub("sglang.srt") +_stub( + "sglang.srt.constants", + GPU_MEMORY_TYPE_CUDA_GRAPH="cuda_graph", + GPU_MEMORY_TYPE_KV_CACHE="kv_cache", + GPU_MEMORY_TYPE_WEIGHTS="weights", +) +_stub("sglang.srt.server_args", ServerArgs=object) +_stub("sglang.srt.utils", kill_process_tree=lambda *a, **k: None) +_stub("wandb") +_stub( + "transformers", + **{n: type(n, (), {}) for n in ("AutoProcessor", "AutoTokenizer", "PreTrainedTokenizerBase", "ProcessorMixin")}, +) + +import pytest + +from slime.ray.rollout import _compute_spec_metrics +from slime.utils.types import Sample + + +NUM_GPUS = 0 + + +class _SpecArgs: + sglang_speculative_algorithm = "EAGLE" + + +def _sample(accept: int, draft: int, verify: int, completion: int) -> Sample: + sample = Sample(index=0, prompt="p") + sample.spec_info.spec_accept_token_num = accept + sample.spec_info.spec_draft_token_num = draft + sample.spec_info.spec_verify_ct = verify + sample.spec_info.completion_token_num = completion + return sample + + +@pytest.mark.unit +def test_disabled_without_speculative_algorithm(): + args = types.SimpleNamespace(sglang_speculative_algorithm=None) + assert _compute_spec_metrics(args, [_sample(1, 2, 1, 2)]) == {} + + +@pytest.mark.unit +def test_pools_counters_across_samples(): + # 7 tiny responses + 1 long one: an unweighted mean of per-sample ratios + # would report accept_rate (7*0.125 + 0.5)/8 = 0.172, accept_length + # (7*1.0 + 3.0)/8 = 1.25 — both far below the pooled truth. + samples = [_sample(1, 8, 1, 1) for _ in range(7)] + [_sample(2000, 4000, 1000, 3000)] + + metrics = _compute_spec_metrics(_SpecArgs(), samples) + + assert metrics["spec_accept_rate"] == pytest.approx((7 * 1 + 2000) / (7 * 8 + 4000)) + assert metrics["spec_accept_length"] == pytest.approx((7 * 1 + 3000) / (7 * 1 + 1000)) + + +@pytest.mark.unit +def test_samples_without_counters_do_not_dilute(): + # SpecInfo.add only runs on terminal meta info, so aborted / partial + # samples legitimately carry all-zero counters. They must not drag the + # batch statistic toward zero. + scored = [_sample(30, 50, 10, 22) for _ in range(6)] + unscored = [_sample(0, 0, 0, 0) for _ in range(2)] + + metrics = _compute_spec_metrics(_SpecArgs(), scored + unscored) + + assert metrics["spec_accept_rate"] == pytest.approx(30 / 50) + assert metrics["spec_accept_length"] == pytest.approx(22 / 10) + + +@pytest.mark.unit +def test_no_counters_at_all_reports_nothing(): + # All-aborted batch (or empty): no fake zeros, and no ZeroDivisionError. + assert _compute_spec_metrics(_SpecArgs(), [_sample(0, 0, 0, 0)]) == {} + assert _compute_spec_metrics(_SpecArgs(), []) == {}