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

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]

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]