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/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,

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")

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."""