From 7fe28afcfcff9fbfb1ef737d4507fdbf00055106 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 19:45:20 +0800 Subject: [PATCH 1/2] perf(eagle3): gate the drafter quality diagnostics on the log level compute_loss collected per-step quality stats unconditionally. Each step ran valid_position.any() and three .item() calls to build quality_step_stats, then the summary line fired at WARNING with six more .item() calls. Every one of them is a device sync, they run on every microbatch, and the numbers are log-only: nothing in the returned loss dict reads them. The EAGLE-1/2 backend already guards the same diagnostics with logger.isEnabledFor(logging.DEBUG). Do the same here, extend it to the non-finite-position and sparse-restricted-CE logs, and move the summary line from WARNING to DEBUG so routine per-step training metrics stop being reported as warnings. Behaviour at DEBUG is unchanged. Signed-off-by: khazic --- ...test_eagle3_diagnostics_gating_contract.py | 136 ++++++++++++++++++ verl_speco/backends/eagle3_trainer_backend.py | 18 ++- 2 files changed, 149 insertions(+), 5 deletions(-) create mode 100644 tests/integration/test_eagle3_diagnostics_gating_contract.py diff --git a/tests/integration/test_eagle3_diagnostics_gating_contract.py b/tests/integration/test_eagle3_diagnostics_gating_contract.py new file mode 100644 index 00000000..41a340b9 --- /dev/null +++ b/tests/integration/test_eagle3_diagnostics_gating_contract.py @@ -0,0 +1,136 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Contract tests for the EAGLE-3 drafter quality diagnostics. + +They are log-only and every one of them forces a device sync, so they must not +run on the hot path unless the log will actually be emitted. +""" + +from __future__ import annotations + +import logging + +import pytest + +QUALITY_LOG_PREFIX = "[drafter logits quality]" +BACKEND_LOGGER = "verl_speco.backends.eagle3_trainer_backend" + + +def _backend_and_inputs(): + import torch + from omegaconf import OmegaConf + + from verl_speco.backends.eagle3_trainer_backend import Eagle3TrainerBackend + + vocab_size, seq_len, ttt_length = 6, 4, 2 + + backend = Eagle3TrainerBackend( + OmegaConf.create( + { + "rollout": { + "drafter": {"training": {"use_logits": False, "ttt_length": ttt_length}} + }, + "model": {"path": "/tmp/none"}, + } + ), + OmegaConf.create({}), + ) + backend.target_model = lambda last_hidden: torch.zeros( + *last_hidden.shape[:-1], vocab_size + ) + + class _FakeDraft: + t2d = torch.ones(vocab_size, dtype=torch.bool) + + def __call__(self, **kwargs): + return { + "logits": [ + torch.randn(1, seq_len, vocab_size) for _ in range(ttt_length) + ], + "position_masks": [ + torch.ones(1, seq_len) for _ in range(ttt_length) + ], + } + + batch = { + "input_ids": torch.zeros(1, seq_len, dtype=torch.long), + "hidden_states": torch.zeros(1, seq_len, 8), + "last_hidden_states": torch.zeros(1, seq_len, 8), + "attention_mask": torch.ones(1, seq_len, dtype=torch.long), + "loss_mask": torch.ones(1, seq_len), + "position_ids": torch.arange(seq_len).unsqueeze(0), + } + return backend, _FakeDraft(), batch + + +def test_quality_diagnostics_are_silent_above_debug(caplog) -> None: + pytest.importorskip("torch") + pytest.importorskip("transformers") + + backend, model, batch = _backend_and_inputs() + logging.getLogger(BACKEND_LOGGER).setLevel(logging.INFO) + + with caplog.at_level(logging.DEBUG, logger=BACKEND_LOGGER): + backend.compute_loss(model, batch, 0) + + assert not [r for r in caplog.records if QUALITY_LOG_PREFIX in r.getMessage()] + + +def test_quality_diagnostics_are_emitted_at_debug(caplog) -> None: + pytest.importorskip("torch") + pytest.importorskip("transformers") + + backend, model, batch = _backend_and_inputs() + logging.getLogger(BACKEND_LOGGER).setLevel(logging.DEBUG) + + try: + with caplog.at_level(logging.DEBUG, logger=BACKEND_LOGGER): + backend.compute_loss(model, batch, 0) + finally: + logging.getLogger(BACKEND_LOGGER).setLevel(logging.INFO) + + quality_records = [ + r for r in caplog.records if QUALITY_LOG_PREFIX in r.getMessage() + ] + assert quality_records + # Routine per-step training metrics belong at DEBUG, not WARNING. + assert all(r.levelno == logging.DEBUG for r in quality_records) + + +def test_quality_diagnostics_do_not_sync_above_debug() -> None: + """Above DEBUG the loss must not pay for the diagnostics' device syncs.""" + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + + original_item = torch.Tensor.item + + def run_at(level: int) -> int: + backend, model, batch = _backend_and_inputs() + logging.getLogger(BACKEND_LOGGER).setLevel(level) + calls = 0 + + def counting_item(self): + nonlocal calls + calls += 1 + return original_item(self) + + torch.Tensor.item = counting_item + try: + backend.compute_loss(model, batch, 0) + finally: + torch.Tensor.item = original_item + logging.getLogger(BACKEND_LOGGER).setLevel(logging.INFO) + return calls + + assert run_at(logging.INFO) < run_at(logging.DEBUG) diff --git a/verl_speco/backends/eagle3_trainer_backend.py b/verl_speco/backends/eagle3_trainer_backend.py index 039d625f..46f5a90a 100644 --- a/verl_speco/backends/eagle3_trainer_backend.py +++ b/verl_speco/backends/eagle3_trainer_backend.py @@ -1178,6 +1178,10 @@ def compute_loss(self, model, batch, _current_pad_size): 0.0, device=input_ids.device, dtype=torch.float32 ) gamma = 0.8 + # The per-step quality stats below are log-only, and every one of them + # forces a device sync. Collect them only when the log will be emitted, + # matching how the EAGLE-1/2 backend gates the same diagnostics. + diagnostics_enabled = logger.isEnabledFor(logging.DEBUG) # Preprocess shifted targets for idx in range(length): @@ -1234,7 +1238,7 @@ def compute_loss(self, model, batch, _current_pad_size): position_mask=position_mask, ) target_top1 = target_p.argmax(dim=-1) - if ( + if diagnostics_enabled and ( base_valid_position.any() and not valid_position[base_valid_position].all() ): @@ -1244,7 +1248,7 @@ def compute_loss(self, model, batch, _current_pad_size): int(dropped_tokens.detach().cpu().item()), ) with torch.no_grad(): - if valid_position.any(): + if diagnostics_enabled and valid_position.any(): draft_top1 = logits.argmax(dim=-1) step_top1_correct = ( (draft_top1[valid_position] == target_top1[valid_position]) @@ -1303,7 +1307,11 @@ def compute_loss(self, model, batch, _current_pad_size): total_local_ploss += (gamma**idx) * step_loss_sum total_local_tokens += valid_position.float().sum() - if use_sparse_restricted_ce and sparse_base_tokens.detach().float().item() > 0: + if ( + diagnostics_enabled + and use_sparse_restricted_ce + and sparse_base_tokens.detach().float().item() > 0 + ): logger.debug( "[drafter sparse restricted ce] base_tokens=%s valid_tokens=%s dropped=%s " "intersection_mean=%.6f hit_mass_mean=%.6f min_intersection=%s min_hit_mass=%s", @@ -1326,8 +1334,8 @@ def compute_loss(self, model, batch, _current_pad_size): logits_sparse_min_mass, ) - if quality_tokens.detach().float().item() > 0: - logger.warning( + if diagnostics_enabled and quality_tokens.detach().float().item() > 0: + logger.debug( "[drafter logits quality] valid_tokens=%s top1_acc=%.6f top%s_acc=%.6f " "local_ploss_sum=%.6f local_tokens=%s per_step=%s", int(quality_tokens.detach().cpu().item()), From 87c08aca3110739ddf54fd3096afc7965c983b60 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 19:48:22 +0800 Subject: [PATCH 2/2] test(eagle3): record the backend logger directly instead of via caplog caplog.at_level raises the logger to DEBUG, which is the exact condition under test, so the above-DEBUG case could never be observed through it. Signed-off-by: khazic --- ...test_eagle3_diagnostics_gating_contract.py | 58 ++++++++++++------- 1 file changed, 37 insertions(+), 21 deletions(-) diff --git a/tests/integration/test_eagle3_diagnostics_gating_contract.py b/tests/integration/test_eagle3_diagnostics_gating_contract.py index 41a340b9..8366903f 100644 --- a/tests/integration/test_eagle3_diagnostics_gating_contract.py +++ b/tests/integration/test_eagle3_diagnostics_gating_contract.py @@ -74,35 +74,49 @@ def __call__(self, **kwargs): return backend, _FakeDraft(), batch -def test_quality_diagnostics_are_silent_above_debug(caplog) -> None: - pytest.importorskip("torch") - pytest.importorskip("transformers") +class _Recorder(logging.Handler): + """Records everything the backend logger emits, whatever its level is. - backend, model, batch = _backend_and_inputs() - logging.getLogger(BACKEND_LOGGER).setLevel(logging.INFO) + ``caplog.at_level`` would raise the logger to DEBUG, which is exactly the + condition under test, so the handler is attached directly instead. + """ - with caplog.at_level(logging.DEBUG, logger=BACKEND_LOGGER): - backend.compute_loss(model, batch, 0) + def __init__(self): + super().__init__(level=logging.DEBUG) + self.records: list[logging.LogRecord] = [] + + def emit(self, record): + self.records.append(record) - assert not [r for r in caplog.records if QUALITY_LOG_PREFIX in r.getMessage()] + +def _compute_loss_at(level: int) -> list[logging.LogRecord]: + backend, model, batch = _backend_and_inputs() + backend_logger = logging.getLogger(BACKEND_LOGGER) + previous_level = backend_logger.level + recorder = _Recorder() + backend_logger.setLevel(level) + backend_logger.addHandler(recorder) + try: + backend.compute_loss(model, batch, 0) + finally: + backend_logger.removeHandler(recorder) + backend_logger.setLevel(previous_level) + return [r for r in recorder.records if QUALITY_LOG_PREFIX in r.getMessage()] -def test_quality_diagnostics_are_emitted_at_debug(caplog) -> None: +def test_quality_diagnostics_are_silent_above_debug() -> None: pytest.importorskip("torch") pytest.importorskip("transformers") - backend, model, batch = _backend_and_inputs() - logging.getLogger(BACKEND_LOGGER).setLevel(logging.DEBUG) + assert _compute_loss_at(logging.INFO) == [] - try: - with caplog.at_level(logging.DEBUG, logger=BACKEND_LOGGER): - backend.compute_loss(model, batch, 0) - finally: - logging.getLogger(BACKEND_LOGGER).setLevel(logging.INFO) - quality_records = [ - r for r in caplog.records if QUALITY_LOG_PREFIX in r.getMessage() - ] +def test_quality_diagnostics_are_emitted_at_debug() -> None: + pytest.importorskip("torch") + pytest.importorskip("transformers") + + quality_records = _compute_loss_at(logging.DEBUG) + assert quality_records # Routine per-step training metrics belong at DEBUG, not WARNING. assert all(r.levelno == logging.DEBUG for r in quality_records) @@ -117,7 +131,9 @@ def test_quality_diagnostics_do_not_sync_above_debug() -> None: def run_at(level: int) -> int: backend, model, batch = _backend_and_inputs() - logging.getLogger(BACKEND_LOGGER).setLevel(level) + backend_logger = logging.getLogger(BACKEND_LOGGER) + previous_level = backend_logger.level + backend_logger.setLevel(level) calls = 0 def counting_item(self): @@ -130,7 +146,7 @@ def counting_item(self): backend.compute_loss(model, batch, 0) finally: torch.Tensor.item = original_item - logging.getLogger(BACKEND_LOGGER).setLevel(logging.INFO) + backend_logger.setLevel(previous_level) return calls assert run_at(logging.INFO) < run_at(logging.DEBUG)