ViperEkura
45cc048fe9
- shard the Muon Newton-Schulz orthogonalization over the FSDP mesh instead of partial local slices - import HF checkpoints faithfully: per-head RoPE permutation for q/k projections and qk-norm, qwen3, shared experts, and qk-norm before RoPE (changes numerics for existing use_qk_norm checkpoints) - make preprocessing and resume self-contained: backfill realigned bucket keys by semantics (masks ones, rest zeros) and snapshot tokenizer files into every checkpoint - keep RL consistent: sync the offline GRPO old_model each optimizer step and validate online strategies through a public one-off-rollout hook that leaves the replay cache untouched - fix streaming serving: withhold partial tool-call prefixes with a stream-end flush, stream tool-call arguments from the raw source span, and terminate SSE frames with a blank line - fix sampling semantics: capture logprobs before top-k/top-p mutate logits in place and detect greedy pipelines polymorphically instead of isinstance bookkeeping
341 lines
11 KiB
Python
341 lines
11 KiB
Python
"""Unit tests for inference sampling strategies."""
|
|
|
|
import torch
|
|
|
|
from astrai.inference.runtime.sample import (
|
|
BaseSamplingStrategy,
|
|
FrequencyPenaltyStrategy,
|
|
SamplingPipeline,
|
|
TemperatureStrategy,
|
|
TopKStrategy,
|
|
TopPStrategy,
|
|
sample,
|
|
)
|
|
|
|
|
|
def test_temperature_scalar():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TemperatureStrategy(0.5)
|
|
result = s.apply(logits.clone())
|
|
assert torch.allclose(result, logits / 0.5)
|
|
|
|
|
|
def test_temperature_skip_when_one():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TemperatureStrategy(1.0)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_temperature_per_sample_tensor():
|
|
logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
s = TemperatureStrategy(torch.tensor([0.5, 0.5]))
|
|
result = s.apply(logits.clone())
|
|
assert torch.allclose(result, logits / 0.5)
|
|
|
|
|
|
def test_top_k_keeps_top():
|
|
logits = torch.tensor([[0.1, 0.5, 0.3, 0.9, 0.2]])
|
|
s = TopKStrategy(top_k=2)
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert kept == 2
|
|
|
|
|
|
def test_top_k_skip_when_zero():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TopKStrategy(top_k=0)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_top_k_batch_tensor():
|
|
"""Each row respects its own top_k."""
|
|
logits = torch.tensor([[0.1, 0.5, 0.3], [0.9, 0.2, 0.1]])
|
|
s = TopKStrategy(top_k=torch.tensor([2, 1]))
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
assert (result[0] > -1e9).sum() == 2
|
|
assert (result[1] > -1e9).sum() == 1
|
|
|
|
|
|
def test_top_p_nucleus_filtering():
|
|
logits = torch.tensor([[10.0, 1.0, 1.0, 1.0, 1.0]])
|
|
s = TopPStrategy(top_p=0.5)
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert kept >= 1
|
|
|
|
|
|
def test_top_p_skip_when_one():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TopPStrategy(top_p=1.0)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_top_p_filter_all_except_max_when_zero():
|
|
logits = torch.tensor([[0.1, 0.5, 0.3, 0.9, 0.2]])
|
|
s = TopPStrategy(top_p=0.0)
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert kept == 1
|
|
|
|
|
|
def test_sampling_pipeline_composes_strategies():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
|
|
pipeline = SamplingPipeline(
|
|
[
|
|
TemperatureStrategy(0.8),
|
|
TopKStrategy(3),
|
|
TopPStrategy(0.95),
|
|
]
|
|
)
|
|
result = pipeline.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert 1 <= kept <= 3
|
|
|
|
|
|
def test_sampling_pipeline_sample_returns_valid_token():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
|
|
pipeline = SamplingPipeline(
|
|
[
|
|
TemperatureStrategy(0.8),
|
|
TopKStrategy(3),
|
|
TopPStrategy(0.95),
|
|
]
|
|
)
|
|
tokens = pipeline.sample(logits)
|
|
assert tokens.shape == (1,)
|
|
assert 0 <= tokens[0] < logits.size(-1)
|
|
|
|
|
|
def test_module_sample_shortcut():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
|
|
tokens = sample(logits, temperature=0.8, top_k=3, top_p=0.95)
|
|
assert tokens.shape == (1,)
|
|
assert 0 <= tokens[0] < logits.size(-1)
|
|
|
|
|
|
def test_module_sample_batch():
|
|
logits = torch.tensor(
|
|
[
|
|
[1.0, 2.0, 3.0, 4.0, 5.0],
|
|
[5.0, 4.0, 3.0, 2.0, 1.0],
|
|
]
|
|
)
|
|
tokens = sample(logits, temperature=0.8, top_k=3, top_p=0.95)
|
|
assert tokens.shape == (2,)
|
|
for t in tokens:
|
|
assert 0 <= t < logits.size(-1)
|
|
|
|
|
|
def test_frequency_penalty_noop_when_zero():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
input_ids = torch.tensor([[0, 2]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.0)
|
|
result = s.apply(logits.clone(), input_ids=input_ids)
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_frequency_penalty_noop_when_no_input_ids():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_frequency_penalty_single_occurrence():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 2]])
|
|
input_mask = torch.tensor([[True, True]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.5
|
|
assert result[0, 1] == 1.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_frequency_penalty_multiple_occurrences():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, True]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.0
|
|
assert result[0, 1] == 1.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_frequency_penalty_respects_padding_mask():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, False]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.5
|
|
assert result[0, 1] == 1.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_frequency_penalty_batch_tensor():
|
|
logits = torch.tensor(
|
|
[
|
|
[4.0, 1.0, 2.0],
|
|
[3.0, 5.0, 1.0],
|
|
]
|
|
)
|
|
input_ids = torch.tensor([[0, 2, 0], [1, 1, 0]])
|
|
input_mask = torch.tensor([[True, True, True], [True, True, False]])
|
|
s = FrequencyPenaltyStrategy(penalty=torch.tensor([0.5, 1.0]))
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.0
|
|
assert result[0, 2] == 1.5
|
|
assert result[1, 1] == 3.0
|
|
|
|
|
|
def test_frequency_penalty_negative_penalty_boosts_repeats():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 0]])
|
|
input_mask = torch.tensor([[True, True]])
|
|
s = FrequencyPenaltyStrategy(penalty=-0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 5.0
|
|
|
|
|
|
def test_frequency_penalty_in_pipeline():
|
|
logits = torch.tensor([[5.0, 1.0, 2.0, 3.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, True]])
|
|
pipeline = SamplingPipeline(
|
|
[
|
|
TemperatureStrategy(1.0),
|
|
FrequencyPenaltyStrategy(0.5),
|
|
]
|
|
)
|
|
result = pipeline.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 4.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_sample_with_frequency_penalty():
|
|
logits = torch.tensor([[5.0, 1.0, 2.0, 3.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, True]])
|
|
tokens = sample(
|
|
logits,
|
|
temperature=1.0,
|
|
top_k=0,
|
|
top_p=1.0,
|
|
frequency_penalty=0.5,
|
|
input_ids=input_ids,
|
|
input_mask=input_mask,
|
|
)
|
|
assert tokens.shape == (1,)
|
|
assert 0 <= tokens[0] < logits.size(-1)
|
|
|
|
|
|
def test_sample_return_logprobs_shape():
|
|
"""``return_logprobs=True`` returns ``[batch]`` logprobs aligned to tokens."""
|
|
logits = torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]])
|
|
out = sample(logits, temperature=1.0, return_logprobs=True)
|
|
tokens, logprobs = out
|
|
assert tokens.shape == (2,)
|
|
assert logprobs.shape == (2,)
|
|
|
|
|
|
def test_sample_return_logprobs_nonpositive():
|
|
"""Probabilities never exceed 1, so logprobs are always ≤ 0."""
|
|
torch.manual_seed(0)
|
|
logits = torch.randn(4, 50)
|
|
_, logprobs = sample(
|
|
logits, temperature=0.8, top_k=20, top_p=0.9, return_logprobs=True
|
|
)
|
|
assert torch.all(logprobs <= 1e-5)
|
|
|
|
|
|
def test_sample_return_logprobs_greedy_path():
|
|
"""Greedy decode (temperature 0) also returns logprobs."""
|
|
logits = torch.tensor([[1.0, 5.0, 2.0]])
|
|
tokens, logprobs = sample(logits, temperature=0.0, return_logprobs=True)
|
|
assert tokens[0].item() == 1
|
|
# log p(token=1) should equal log_softmax(logits)[1]
|
|
expected = torch.log_softmax(logits.float(), dim=-1)[0, 1]
|
|
assert torch.allclose(logprobs[0], expected, atol=1e-5)
|
|
|
|
|
|
def test_sample_return_logprobs_matches_manual_computation():
|
|
"""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)
|
|
expected = torch.gather(
|
|
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]
|
|
|
|
|
|
class _ArgmaxMovingStrategy(BaseSamplingStrategy):
|
|
"""Custom strategy that can move the argmax — must disable greedy."""
|
|
|
|
def apply(
|
|
self, logits, filter_value=-float("inf"), input_ids=None, input_mask=None
|
|
):
|
|
return torch.roll(logits, shifts=1, dims=-1)
|
|
|
|
|
|
def test_greedy_detection_is_polymorphic():
|
|
"""Greedy detection asks strategies polymorphically, no isinstance."""
|
|
base = [TemperatureStrategy(0.0), TopKStrategy(50), TopPStrategy(0.9)]
|
|
assert SamplingPipeline(list(base)).is_greedy is True
|
|
assert SamplingPipeline(base + [FrequencyPenaltyStrategy(0.5)]).is_greedy is False
|
|
|
|
# A custom argmax-moving strategy disables greedy even though the
|
|
# pipeline contains a greedy temperature — this is what isinstance
|
|
# bookkeeping in the old implementation could not see.
|
|
assert SamplingPipeline(base + [_ArgmaxMovingStrategy()]).is_greedy is False
|
|
|
|
|
|
def test_greedy_detection_position_independent():
|
|
"""Greedy temperature anywhere in the pipeline is detected."""
|
|
pipeline = SamplingPipeline([TopKStrategy(50), TemperatureStrategy(0.0)])
|
|
assert pipeline.is_greedy is True
|
|
|
|
|
|
def test_greedy_detection_composes_across_nested_pipelines():
|
|
"""A nested pipeline participates through the same interface."""
|
|
inner = SamplingPipeline([TemperatureStrategy(0.0), TopKStrategy(20)])
|
|
assert inner.is_greedy is True
|
|
assert SamplingPipeline([TopPStrategy(0.9), inner]).is_greedy is True
|
|
assert SamplingPipeline([inner, FrequencyPenaltyStrategy(0.5)]).is_greedy is False
|
|
|
|
|
|
def test_nongreedy_temperature_is_not_greedy():
|
|
pipeline = SamplingPipeline(
|
|
[TemperatureStrategy(0.7), TopKStrategy(0), TopPStrategy(1.0)]
|
|
)
|
|
assert pipeline.is_greedy is False
|