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:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user