ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: resolve audited training and inference bugs

Browse Source 
- 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:
ViperEkura
2026-09-02 21:25:01 +08:00
parent 92e3cdf044
commit 88c06db096
8 changed files with 181 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/config/train_config.py
+13 -4
View File
@@ -213,8 +213,17 @@ class TrainConfig(BaseConfig):
@model_validator(mode="after")
def _validate_online_strategy(self) -> "TrainConfig":
if self.strategy.startswith("online_") and self.reward_model_fn is None:
raise ValueError(
f"reward_model_fn is required for online RL strategy {self.strategy!r}"
)
if self.strategy.startswith("online_"):
if self.reward_model_fn is None:
raise ValueError(
f"reward_model_fn is required for online RL strategy "
f"{self.strategy!r}"
)
if self.nprocs > 1:
raise ValueError(
f"online RL strategy {self.strategy!r} requires single-process "
f"training (nprocs=1): per-rank rollouts issue different "
f"numbers of forward passes and desynchronize the "
f"ddp/fsdp collectives, deadlocking NCCL"
)
return self
astrai/inference/runtime/sample.py
+35 -15
View File
@@ -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
+4
View File
@@ -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
+22 -4
View File
@@ -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."""
astrai/parallel/executor.py
+26 -3
View File
@@ -23,6 +23,27 @@ from astrai.parallel.setup import get_rank, get_world_size
logger = logging.getLogger(__name__)
_COMPILE_PREFIX = "_orig_mod."
def strip_compile_prefix(
state_dict: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
"""Drop the ``_orig_mod.`` key prefix ``torch.compile`` adds.
``OptimizedModule.state_dict()`` prefixes every key, so checkpoints or
reference-model copies taken from a compiled model fail to load into a
plain module (strict) or silently load nothing (non-strict). Stripping
here, at the single source every consumer reads from, keeps saved keys
canonical regardless of compile mode.
"""
if any(key.startswith(_COMPILE_PREFIX) for key in state_dict):
state_dict = {
key.removeprefix(_COMPILE_PREFIX): value
for key, value in state_dict.items()
}
return state_dict
def broadcast_state_dict(
state_dict: Optional[Dict[str, torch.Tensor]],
@@ -91,6 +112,7 @@ def create_ref_model(
if state_dict is None:
return None
state_dict = strip_compile_prefix(state_dict)
ref_model = model_fn()
ref_model.load_state_dict(state_dict)
ref_model.requires_grad_(False)
@@ -206,7 +228,7 @@ class BaseExecutor:
loss.backward()
def unwrap_model(self, model: nn.Module):
return model.state_dict()
return strip_compile_prefix(model.state_dict())
@contextmanager
def checkpoint_context(self, model: nn.Module):
@@ -308,8 +330,8 @@ class DDPExecutor(BaseExecutor):
def unwrap_model(self, model: nn.Module):
if isinstance(model, DDP):
return model.module.state_dict()
return model.state_dict()
return strip_compile_prefix(model.module.state_dict())
return strip_compile_prefix(model.state_dict())
@ExecutorFactory.register("fsdp")
@@ -411,6 +433,7 @@ class FSDPExecutor(BaseExecutor):
state_dict = model.state_dict()
result = {}
for k, v in state_dict.items():
k = k.removeprefix(_COMPILE_PREFIX)
if isinstance(v, DTensor):
full = v.full_tensor()
if get_rank() == 0:
astrai/trainer/train_context.py
+17 -2
View File
@@ -13,7 +13,12 @@ from astrai.config.train_config import TrainConfig
from astrai.dataset import RDSampler
from astrai.inference.scheduler import InferenceScheduler
from astrai.model.components.lora import inject_lora
from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
from astrai.parallel.executor import (
BaseExecutor,
ExecutorFactory,
create_ref_model,
strip_compile_prefix,
)
from astrai.parallel.setup import get_current_device, get_rank, get_world_size
from astrai.protocols import OptimizerProtocol, SchedulerProtocol
from astrai.serialization import (
@@ -145,6 +150,7 @@ class TrainContextBuilder:
checkpoint.state_dict,
ConfigFactory.load(checkpoint.config or state.model_config),
)
checkpoint.state_dict = strip_compile_prefix(checkpoint.state_dict)
state.state_dict = checkpoint.state_dict
state.model_config = checkpoint.config or state.model_config
if self._resume:
@@ -192,7 +198,16 @@ class TrainContextBuilder:
target_modules=set(cfg.lora.target_modules),
)
if state.state_dict is not None:
model.load_state_dict(state.state_dict, strict=False)
result = model.load_state_dict(state.state_dict, strict=False)
if result.missing_keys or result.unexpected_keys:
logger.warning(
"preloaded state dict mismatch: %d missing, %d unexpected "
"(first missing: %s, first unexpected: %s)",
len(result.missing_keys),
len(result.unexpected_keys),
result.missing_keys[:3],
result.unexpected_keys[:3],
)
return model
def after_wrap(model):
tests/inference/test_sample.py
+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]
tests/inference/test_task.py
+40 -1
View File
@@ -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]