From 88c06db096f197acac2a66953bde445c3d720121 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 2 Sep 2026 21:25:01 +0800 Subject: [PATCH] fix: resolve audited training and inference bugs - reject prompts that encode to zero tokens in add_task instead of admitting a task whose prefill can never run, and surface empty-id run_batch calls as prompt_empty errors - deliver the STOP stream callback when cancelling a live task so clients observe termination instead of hanging until socket timeout - strip the torch.compile _orig_mod. prefix at every unwrap_model site and when loading checkpoints so FSDP state dicts and saved weights no longer leak the wrapper name into downstream keys - reject online_* train strategies with nprocs > 1 at config validation time, explaining the NCCL all-gather deadlock they would otherwise hit mid-run - apply the frequency penalty before temperature scaling (OpenAI semantics) so the penalty survives temperature=0 instead of being annihilated by the 1e8 logit blowup, and exclude penalty pipelines from the greedy fast path - return logprobs from the raw pre-strategy distribution so they match training-side policy logprobs for PPO/GRPO importance ratios --- astrai/config/train_config.py | 17 +++++++--- astrai/inference/runtime/sample.py | 50 +++++++++++++++++++++--------- astrai/inference/scheduler.py | 4 +++ astrai/inference/task.py | 26 +++++++++++++--- astrai/parallel/executor.py | 29 +++++++++++++++-- astrai/trainer/train_context.py | 19 ++++++++++-- tests/inference/test_sample.py | 35 ++++++++++++++------- tests/inference/test_task.py | 41 +++++++++++++++++++++++- 8 files changed, 181 insertions(+), 40 deletions(-) diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 31330d3..4b2dc54 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -213,8 +213,17 @@ class TrainConfig(BaseConfig): @model_validator(mode="after") def _validate_online_strategy(self) -> "TrainConfig": - if self.strategy.startswith("online_") and self.reward_model_fn is None: - raise ValueError( - f"reward_model_fn is required for online RL strategy {self.strategy!r}" - ) + if self.strategy.startswith("online_"): + if self.reward_model_fn is None: + raise ValueError( + f"reward_model_fn is required for online RL strategy " + f"{self.strategy!r}" + ) + if self.nprocs > 1: + raise ValueError( + f"online RL strategy {self.strategy!r} requires single-process " + f"training (nprocs=1): per-rank rollouts issue different " + f"numbers of forward passes and desynchronize the " + f"ddp/fsdp collectives, deadlocking NCCL" + ) return self diff --git a/astrai/inference/runtime/sample.py b/astrai/inference/runtime/sample.py index b921961..38f7aa9 100644 --- a/astrai/inference/runtime/sample.py +++ b/astrai/inference/runtime/sample.py @@ -289,8 +289,8 @@ class SamplingPipeline(BaseSamplingStrategy): input_mask: Boolean mask for ``input_ids`` padding. return_logprobs: If ``True``, return ``(tokens, logprobs)`` where ``logprobs[i]`` is the log-probability of - ``tokens[i]`` under the (post-strategy) sampling - distribution. + ``tokens[i]`` under the raw (pre-strategy) model + distribution, matching training-side policy logprobs. Returns: Sampled token IDs ``[batch]``, or — when ``return_logprobs`` @@ -310,18 +310,32 @@ class SamplingPipeline(BaseSamplingStrategy): ).squeeze(-1) if not return_logprobs: return tokens - log_probs = torch.log_softmax(transformed.float(), dim=-1) + # Log-probabilities of the raw (pre-strategy) model distribution, + # matching the training-side policy logprobs exactly: the behaviour + # logprobs recorded for online RL must live in the same + # distribution the trainer differentiates, not the + # temperature/top-p filtered one tokens were drawn from. + log_probs = torch.log_softmax(logits.float(), dim=-1) chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1) return tokens, chosen def _is_greedy_pipeline(self) -> bool: - """True if the first strategy is greedy temperature (temp=0).""" + """True if sampling reduces to argmax over the raw logits. + + A greedy temperature with only top-k/top-p strategies does: the + filters always keep the argmax token. A frequency penalty can + change the argmax, so those pipelines must run the full + transformation even at ``temperature=0``. + """ if not self.strategies: return False first = self.strategies[0] - return isinstance(first, TemperatureStrategy) and self._is_greedy( - first.temperature - ) + if not ( + isinstance(first, TemperatureStrategy) + and self._is_greedy(first.temperature) + ): + return False + return not any(isinstance(s, FrequencyPenaltyStrategy) for s in self.strategies) @torch.inference_mode() @@ -354,9 +368,9 @@ def sample( input_ids: Previously generated token IDs ``[batch, seq_len]``. input_mask: Boolean mask for ``input_ids`` padding. return_logprobs: If ``True``, also return the log-probability - of each sampled token under the (post-strategy) sampling - distribution — useful for RL rollout (PPO/GRPO importance - ratios). + of each sampled token under the raw (pre-strategy) model + distribution — usable directly for RL rollout (PPO/GRPO + importance ratios against the training-side policy logprobs). Returns: Sampled token IDs ``[batch]``, or — when ``return_logprobs`` is @@ -369,13 +383,19 @@ def sample( else frequency_penalty != 0 ) - strategies: List[BaseSamplingStrategy] = [ - TemperatureStrategy(temperature), - TopKStrategy(top_k), - TopPStrategy(top_p), - ] + strategies: List[BaseSamplingStrategy] = [] if has_freq: + # Penalty first, on the raw logits (OpenAI semantics): applying it + # after a temperature scaling would shrink it by the temperature + # and annihilate it entirely at temperature=0. strategies.append(FrequencyPenaltyStrategy(frequency_penalty)) + strategies.extend( + [ + TemperatureStrategy(temperature), + TopKStrategy(top_k), + TopPStrategy(top_p), + ] + ) return SamplingPipeline(strategies).sample( logits, diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index ab1f22d..3121ae3 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -444,6 +444,10 @@ class InferenceScheduler: tasks: List[Optional[Task]] = [] error_reasons: List[Optional[str]] = [] for ids in prompt_ids_list: + if not ids: + tasks.append(None) + error_reasons.append("prompt_empty") + continue if len(ids) >= seq_cap: tasks.append(None) error_reasons.append("prompt_too_long") diff --git a/astrai/inference/task.py b/astrai/inference/task.py index 53eeff3..55ce442 100644 --- a/astrai/inference/task.py +++ b/astrai/inference/task.py @@ -185,6 +185,11 @@ class TaskManager: ) -> str: task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}" prompt_ids = self.tokenizer.encode(prompt) + if not prompt_ids: + # An empty prompt never completes prefill (``prefill_done`` stays + # False) and would crash the decode path on ``prompt_ids[-1]``; + # rejecting it here keeps the scheduling loop alive. + raise ValueError("prompt encoded to zero tokens; refusing to schedule") if len(prompt_ids) > self.max_seq_len: prompt_ids = prompt_ids[-self.max_seq_len :] @@ -219,10 +224,19 @@ class TaskManager: return task_id def cancel_task(self, task_id: str) -> Tuple[List[Task], bool]: - """Mark a task cancelled and return tasks safe to clean immediately.""" + """Mark a task cancelled and return tasks safe to clean immediately. + + Registered stream callbacks receive the terminal ``STOP`` sentinel + for every live cancellation: the scheduling loop drains ABORTED + tasks without invoking callbacks, so skipping it here would leave + consumers (e.g. ``GenerateResult.wait_completion``) waiting forever. + """ + callback = None + cancelled = False + immediate: List[Task] = [] with self._lock: task = self._tasks.get(task_id) - self._callbacks.pop(task_id, None) + callback = self._callbacks.pop(task_id, None) if task is None or task.status in ( TaskStatus.FINISHED, TaskStatus.ABORTED, @@ -231,13 +245,17 @@ class TaskManager: task.status = TaskStatus.ABORTED self._cancelled_total += 1 + cancelled = True if task in self.waiting_queue: self.waiting_queue = deque( waiting for waiting in self.waiting_queue if waiting is not task ) self._tasks.pop(task_id, None) - return [task], True - return [], True + immediate = [task] + + if cancelled and callback is not None: + callback(STOP) + return immediate, cancelled def remove_task(self, task_id: str) -> List[Task]: """Backward-compatible alias for cancellation.""" diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index d0c12c8..8040587 100644 --- a/astrai/parallel/executor.py +++ b/astrai/parallel/executor.py @@ -23,6 +23,27 @@ from astrai.parallel.setup import get_rank, get_world_size logger = logging.getLogger(__name__) +_COMPILE_PREFIX = "_orig_mod." + + +def strip_compile_prefix( + state_dict: Dict[str, torch.Tensor], +) -> Dict[str, torch.Tensor]: + """Drop the ``_orig_mod.`` key prefix ``torch.compile`` adds. + + ``OptimizedModule.state_dict()`` prefixes every key, so checkpoints or + reference-model copies taken from a compiled model fail to load into a + plain module (strict) or silently load nothing (non-strict). Stripping + here, at the single source every consumer reads from, keeps saved keys + canonical regardless of compile mode. + """ + if any(key.startswith(_COMPILE_PREFIX) for key in state_dict): + state_dict = { + key.removeprefix(_COMPILE_PREFIX): value + for key, value in state_dict.items() + } + return state_dict + def broadcast_state_dict( state_dict: Optional[Dict[str, torch.Tensor]], @@ -91,6 +112,7 @@ def create_ref_model( if state_dict is None: return None + state_dict = strip_compile_prefix(state_dict) ref_model = model_fn() ref_model.load_state_dict(state_dict) ref_model.requires_grad_(False) @@ -206,7 +228,7 @@ class BaseExecutor: loss.backward() def unwrap_model(self, model: nn.Module): - return model.state_dict() + return strip_compile_prefix(model.state_dict()) @contextmanager def checkpoint_context(self, model: nn.Module): @@ -308,8 +330,8 @@ class DDPExecutor(BaseExecutor): def unwrap_model(self, model: nn.Module): if isinstance(model, DDP): - return model.module.state_dict() - return model.state_dict() + return strip_compile_prefix(model.module.state_dict()) + return strip_compile_prefix(model.state_dict()) @ExecutorFactory.register("fsdp") @@ -411,6 +433,7 @@ class FSDPExecutor(BaseExecutor): state_dict = model.state_dict() result = {} for k, v in state_dict.items(): + k = k.removeprefix(_COMPILE_PREFIX) if isinstance(v, DTensor): full = v.full_tensor() if get_rank() == 0: diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 5753c92..ff6ac9a 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -13,7 +13,12 @@ from astrai.config.train_config import TrainConfig from astrai.dataset import RDSampler from astrai.inference.scheduler import InferenceScheduler from astrai.model.components.lora import inject_lora -from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model +from astrai.parallel.executor import ( + BaseExecutor, + ExecutorFactory, + create_ref_model, + strip_compile_prefix, +) from astrai.parallel.setup import get_current_device, get_rank, get_world_size from astrai.protocols import OptimizerProtocol, SchedulerProtocol from astrai.serialization import ( @@ -145,6 +150,7 @@ class TrainContextBuilder: checkpoint.state_dict, ConfigFactory.load(checkpoint.config or state.model_config), ) + checkpoint.state_dict = strip_compile_prefix(checkpoint.state_dict) state.state_dict = checkpoint.state_dict state.model_config = checkpoint.config or state.model_config if self._resume: @@ -192,7 +198,16 @@ class TrainContextBuilder: target_modules=set(cfg.lora.target_modules), ) if state.state_dict is not None: - model.load_state_dict(state.state_dict, strict=False) + result = model.load_state_dict(state.state_dict, strict=False) + if result.missing_keys or result.unexpected_keys: + logger.warning( + "preloaded state dict mismatch: %d missing, %d unexpected " + "(first missing: %s, first unexpected: %s)", + len(result.missing_keys), + len(result.unexpected_keys), + result.missing_keys[:3], + result.unexpected_keys[:3], + ) return model def after_wrap(model): diff --git a/tests/inference/test_sample.py b/tests/inference/test_sample.py index 63489dd..991fa11 100644 --- a/tests/inference/test_sample.py +++ b/tests/inference/test_sample.py @@ -263,22 +263,35 @@ def test_sample_return_logprobs_greedy_path(): def test_sample_return_logprobs_matches_manual_computation(): - """Returned logprob equals log_softmax(transformed_logits)[token].""" + """Returned logprob equals log_softmax(raw_logits)[token]. + + Logprobs live in the raw (pre-strategy) model distribution so they + line up with training-side policy logprobs for RL importance ratios. + """ torch.manual_seed(1) logits = torch.randn(2, 30) tokens, logprobs = sample(logits, temperature=0.7, top_p=0.95, return_logprobs=True) - # Recompute with the same pipeline - from astrai.inference.runtime.sample import ( - SamplingPipeline, - TemperatureStrategy, - TopPStrategy, - ) - - pipeline = SamplingPipeline([TemperatureStrategy(0.7), TopPStrategy(0.95)]) - transformed = pipeline.apply(logits.clone()) expected = torch.gather( - torch.log_softmax(transformed.float(), dim=-1), + torch.log_softmax(logits.float(), dim=-1), -1, tokens.unsqueeze(-1), ).squeeze(-1) assert torch.allclose(logprobs, expected, atol=1e-5) + + +def test_greedy_respects_frequency_penalty(): + """temperature=0 must not silently skip the frequency penalty.""" + torch.manual_seed(0) + logits = torch.tensor([[5.0, 4.0, 3.0]]) + + plain = sample(logits.clone(), temperature=0.0) + assert plain.tolist() == [0] + + penalized = sample( + logits.clone(), + temperature=0.0, + frequency_penalty=2.0, + input_ids=torch.tensor([[0, 0, 0, 0]]), + ) + # Token 0 saw four occurrences: 5 - 2*4 < 4, so the argmax flips. + assert penalized.tolist() == [1] diff --git a/tests/inference/test_task.py b/tests/inference/test_task.py index c65eaa1..97ab2a4 100644 --- a/tests/inference/test_task.py +++ b/tests/inference/test_task.py @@ -2,7 +2,9 @@ from unittest.mock import MagicMock -from astrai.inference import Task, TaskManager, TaskStatus +import pytest + +from astrai.inference import STOP, Task, TaskManager, TaskStatus def _make_mock_tokenizer(): @@ -178,3 +180,40 @@ def test_task_manager_get_stats(): assert stats["total_tasks"] == 1 assert stats["waiting_queue"] == 1 assert stats["active_tasks"] == 0 + + +def test_task_manager_add_task_rejects_empty_prompt(): + tm = TaskManager(tokenizer=_make_mock_tokenizer()) + tm.tokenizer.encode.return_value = [] + + with pytest.raises(ValueError, match="zero tokens"): + tm.add_task("") + + +def test_task_manager_cancel_delivers_stop_callback(): + tm = TaskManager(tokenizer=_make_mock_tokenizer()) + received = [] + tm.add_task("test", stream_callback=received.append) + + immediate, cancelled = tm.cancel_task("does-not-exist") + assert not cancelled and immediate == [] and received == [] + + task_id = next(iter(tm._tasks)) + immediate, cancelled = tm.cancel_task(task_id) + assert cancelled + assert len(immediate) == 1 + assert received == [STOP] + + +def test_task_manager_cancel_active_task_delivers_stop_callback(): + tm = TaskManager(tokenizer=_make_mock_tokenizer()) + received = [] + task_id = tm.add_task("test", stream_callback=received.append) + task = tm._tasks[task_id] + tm.waiting_queue.clear() + tm.active_tasks.append(task) + task.status = TaskStatus.RUNNING + + immediate, cancelled = tm.cancel_task(task_id) + assert cancelled and immediate == [] + assert received == [STOP]