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
This commit is contained in:
2026-09-02 21:25:01 +08:00
parent 92e3cdf044
commit 88c06db096
8 changed files with 181 additions and 40 deletions
+24 -11
View File
@@ -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]