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

fix: resolve audited training, import, and serving bugs

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-09-03 20:27:41 +08:00
parent 7e98a419a7
commit 45cc048fe9
21 changed files with 834 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/network/openai.py
+22 -1
View File
@@ -204,10 +204,31 @@ class OpenAIResponseBuilder(ResponseBuilder):
)
def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
events: List[str] = []
if self._parser is not None:
for d in self._parser.finalize(stop.body):
if "content" in d:
events.append(
sse_event(
{
"id": self._resp_id,
"object": "chat.completion.chunk",
"created": ctx.created,
"model": self._model,
"choices": [
{
"index": 0,
"delta": {"content": d["content"]},
"finish_reason": None,
}
],
}
)
)
finish_reason = "stop"
if self._parser is not None and self._parser.has_tool_calls:
finish_reason = "tool_calls"
return [
return events + [
sse_event(
{
"id": self._resp_id,
astrai/inference/network/protocol.py
+4 -2
View File
@@ -20,8 +20,10 @@ def sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
if event:
lines.append(f"event: {event}")
lines.append(f"data: {json.dumps(data, ensure_ascii=False)}")
lines.append("")
return "\n".join(lines)
# The SSE spec dispatches an event only at a blank line, so the frame
# must end with "\n\n" (a single trailing newline keeps clients waiting
# and concatenates consecutive events into one corrupt payload).
return "\n".join(lines) + "\n\n"
def sse_done() -> str:
astrai/inference/network/tool_parser.py
+85 -6
View File
@@ -69,6 +69,10 @@ class BaseToolParser(ABC):
def has_tool_calls(self) -> bool:
"""True if the parser detected at least one tool call in the stream."""
def finalize(self, body: str) -> List[Dict]:
"""Flush parser state once generation ended. Default: nothing."""
return []
class ToolParserFactory(BaseFactory["BaseToolParser"]):
pass
@@ -108,6 +112,30 @@ def _scan_json(text: str, start: int = 0):
return len(text), False
def _raw_arguments_span(json_str: str) -> Optional[str]:
"""Extract the raw text span of the ``arguments`` value, if possible.
Streaming diffs emit the *source* text of ``arguments``; using the
same span here keeps the completed arguments a prefix-extension of
what was already streamed (``json.dumps`` would re-quote and
re-space the value and corrupt the concatenated result).
"""
m = re.search(r'"arguments"\s*:\s*', json_str)
if not m:
return None
rest = json_str[m.end() :]
if rest[:1] in ("{", "["):
end, ok = _scan_json(rest, 0)
if ok:
return rest[1 : end - 1]
return None
str_match = re.match(r'"(?:[^"\\]|\\.)*"', rest)
if str_match:
return str_match.group(0)
bare = re.match(r"[^,}\s][^,}]*", rest)
return bare.group(0).rstrip() if bare else None
def _parse_tool_call_json(json_str: str, complete: bool):
"""Extract *name* and *arguments* from a tool-call JSON string.
@@ -122,14 +150,23 @@ def _parse_tool_call_json(json_str: str, complete: bool):
if not isinstance(name, str) or not name:
return None, "", False
args = obj.get("arguments")
raw = _raw_arguments_span(json_str)
if isinstance(args, dict):
if not args:
args = ""
else:
args = json.dumps(args, ensure_ascii=False)
args = args[1:-1].rstrip()
# Prefer the source span: it matches what streaming
# already emitted (prefix-consistent completion).
args = (
raw
if raw is not None
else json.dumps(args, ensure_ascii=False)[1:-1].rstrip()
)
elif isinstance(args, list):
args = json.dumps(args, ensure_ascii=False) if args else ""
if not args:
args = ""
else:
args = raw if raw is not None else json.dumps(args, ensure_ascii=False)
elif isinstance(args, str):
pass
else:
@@ -306,10 +343,52 @@ class SimpleJsonToolParser(BaseToolParser):
return deltas
def _emit_plain_content(self, body: str, deltas: List[Dict]) -> List[Dict]:
new_content = body[self._emitted_content_len :]
if new_content:
safe_end = self._safe_content_end(body)
if safe_end > self._emitted_content_len:
deltas.append({"content": body[self._emitted_content_len : safe_end]})
self._emitted_content_len = safe_end
return deltas
_POSSIBLE_NAME_PREFIX_RE = re.compile(r'^\s*"?(?:n|na|nam|name)?"?\s*:?\s*$')
@classmethod
def _safe_content_end(cls, body: str) -> int:
"""End of *body* that is safe to emit as plain content.
A trailing unclosed ``{`` whose remainder could still grow into
``{"name": ...`` (the name-prefix itself, or any further ``{``
opened after it) is withheld, otherwise a partial ``{"na`` prefix
would leak into user-visible content and never be retracted.
Anything withheld is flushed by :meth:`finalize` when generation
ends without a tool call, so plain text is never lost.
"""
pos = 0
while True:
brace = body.find("{", pos)
if brace == -1:
return len(body)
end, complete = _scan_json(body, brace)
if complete:
pos = end
continue
tail = body[brace + 1 :]
if cls._POSSIBLE_NAME_PREFIX_RE.match(tail) or "{" in tail:
return brace
return len(body)
def finalize(self, body: str) -> List[Dict]:
"""Flush content withheld as a possible tool-call prefix.
Called once generation completed: if no tool call materialised,
emit the withheld remainder so streamed content matches the
non-streaming response.
"""
if self._has_tool_calls:
return []
deltas: List[Dict] = []
if len(body) > self._emitted_content_len:
deltas.append({"content": body[self._emitted_content_len :]})
self._emitted_content_len = len(body)
deltas.append({"content": new_content})
return deltas
# -------------------------------------------------------- complete
astrai/inference/runtime/sample.py
+67 -27
View File
@@ -42,6 +42,22 @@ class BaseSamplingStrategy(ABC):
"""
raise NotImplementedError
@property
def preserves_argmax(self) -> bool:
"""Whether ``apply`` never moves the argmax token.
Conservative default: strategies must opt in. The greedy
short-circuit in :class:`SamplingPipeline` asks this
polymorphically, so a new strategy that can move the argmax
automatically disables it — no isinstance bookkeeping.
"""
return False
@property
def is_greedy(self) -> bool:
"""Whether this strategy collapses sampling onto the argmax token."""
return False
class TemperatureStrategy(BaseSamplingStrategy):
"""Divides logits by temperature to control randomness.
@@ -53,6 +69,22 @@ class TemperatureStrategy(BaseSamplingStrategy):
def __init__(self, temperature: Union[float, Tensor] = 1.0):
self.temperature = temperature
@staticmethod
def is_greedy_temperature(temperature: Union[float, Tensor]) -> bool:
if isinstance(temperature, Tensor):
return bool((temperature == 0).all())
return temperature == 0
@property
def is_greedy(self) -> bool:
return self.is_greedy_temperature(self.temperature)
@property
def preserves_argmax(self) -> bool:
# Scaling by a positive constant (1/t, clamped away from zero)
# preserves logit order; t=0 degenerates onto the argmax itself.
return True
def apply(
self,
logits: Tensor,
@@ -78,6 +110,11 @@ class TopKStrategy(BaseSamplingStrategy):
top_k: Scalar or ``[batch]`` tensor (0 disables).
"""
@property
def preserves_argmax(self) -> bool:
# The argmax token always ranks first, so any k >= 1 keeps it.
return True
def __init__(self, top_k: Union[int, Tensor] = 0):
self.top_k = top_k
@@ -121,6 +158,11 @@ class TopPStrategy(BaseSamplingStrategy):
top_p: Scalar or ``[batch]`` tensor (1.0 disables).
"""
@property
def preserves_argmax(self) -> bool:
# Nucleus filtering always keeps the highest-probability token.
return True
def __init__(self, top_p: Union[float, Tensor] = 1.0):
self.top_p = top_p
@@ -250,6 +292,22 @@ class SamplingPipeline(BaseSamplingStrategy):
def __init__(self, strategies: List[BaseSamplingStrategy]):
self.strategies = strategies
@property
def preserves_argmax(self) -> bool:
# A composite preserves the argmax iff every stage does.
return all(s.preserves_argmax for s in self.strategies)
@property
def is_greedy(self) -> bool:
"""Whether sampling always yields the argmax of the raw logits.
True iff some stage forces greedy and no stage can move the
argmax before or after it. Both facts are declared
polymorphically by each strategy, so composing in a new strategy
type (or a nested pipeline) updates this automatically.
"""
return any(s.is_greedy for s in self.strategies) and self.preserves_argmax
def apply(
self,
logits: Tensor,
@@ -261,12 +319,6 @@ class SamplingPipeline(BaseSamplingStrategy):
logits = strategy.apply(logits, filter_value, input_ids, input_mask)
return logits
@staticmethod
def _is_greedy(temperature: Union[float, Tensor]) -> bool:
if isinstance(temperature, Tensor):
return bool((temperature == 0).all())
return temperature == 0
@torch.inference_mode()
def sample(
self,
@@ -294,7 +346,7 @@ class SamplingPipeline(BaseSamplingStrategy):
Sampled token IDs ``[batch]``, or — when ``return_logprobs``
is ``True`` — a ``(token_ids, chosen_logprobs)`` tuple.
"""
if self._is_greedy_pipeline():
if self.is_greedy:
tokens = logits.argmax(dim=-1)
if not return_logprobs:
return tokens
@@ -302,6 +354,13 @@ class SamplingPipeline(BaseSamplingStrategy):
chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
return tokens, chosen
# Capture the raw distribution before the strategy pipeline runs:
# top-k/top-p mutate the logits tensor in place, so computing this
# after ``apply`` would read the filtered distribution instead of
# the raw model distribution the caller documented.
if return_logprobs:
raw_log_probs = torch.log_softmax(logits.float(), dim=-1)
transformed = self.apply(logits, filter_value, input_ids, input_mask)
tokens = torch.multinomial(
torch.softmax(transformed, dim=-1), num_samples=1
@@ -313,28 +372,9 @@ class SamplingPipeline(BaseSamplingStrategy):
# 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)
chosen = torch.gather(raw_log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
return tokens, chosen
def _is_greedy_pipeline(self) -> bool:
"""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]
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()
def sample(