diff --git a/astrai/inference/network/openai.py b/astrai/inference/network/openai.py index 0688c19..718d495 100644 --- a/astrai/inference/network/openai.py +++ b/astrai/inference/network/openai.py @@ -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, diff --git a/astrai/inference/network/protocol.py b/astrai/inference/network/protocol.py index db6939c..f8f785a 100644 --- a/astrai/inference/network/protocol.py +++ b/astrai/inference/network/protocol.py @@ -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: diff --git a/astrai/inference/network/tool_parser.py b/astrai/inference/network/tool_parser.py index 2e871a4..d8bb6a0 100644 --- a/astrai/inference/network/tool_parser.py +++ b/astrai/inference/network/tool_parser.py @@ -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 diff --git a/astrai/inference/runtime/sample.py b/astrai/inference/runtime/sample.py index de0dc80..431ede2 100644 --- a/astrai/inference/runtime/sample.py +++ b/astrai/inference/runtime/sample.py @@ -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( diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py index f3d300f..88f0dee 100644 --- a/astrai/model/components/attention.py +++ b/astrai/model/components/attention.py @@ -69,11 +69,15 @@ class GQA(nn.Module): q = self._split_heads(self.q_proj(x), self.n_heads) k = self._split_heads(self.k_proj(x), self.n_kv_heads) v = self._split_heads(self.v_proj(x), self.n_kv_heads) - q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb) + # Match the HuggingFace convention (Qwen2/Gemma): normalize Q/K + # before RoPE. RMSNorm's per-channel gain does not commute with + # the rotation, so the order changes numerics. if self.use_qk_norm: q, k = self.q_norm(q), self.k_norm(k) + q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb) + sdqa_out = attention( q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd ).reshape(*x.shape[:-1], self.dim) diff --git a/astrai/optim/muon_adamw.py b/astrai/optim/muon_adamw.py index f27f5f3..a2a1715 100644 --- a/astrai/optim/muon_adamw.py +++ b/astrai/optim/muon_adamw.py @@ -1,9 +1,16 @@ """Legacy Muon + AdamW combined optimizer.""" +from collections.abc import Mapping from typing import Any import torch from torch import Tensor, nn, optim +from torch.distributed.tensor import DTensor, distribute_tensor +from torch.optim._muon import ( + _adjust_lr, + _single_tensor_muon, + _zeropower_via_newtonschulz, +) from astrai.optim.composite import ( OptimizerFactory, @@ -14,6 +21,91 @@ from astrai.optim.composite import ( ) +def _scalar_lr(lr: Any) -> float: + return lr.item() if isinstance(lr, Tensor) else lr + + +def _sharded_orthogonalize(update: Tensor, group: Mapping) -> Tensor: + """Newton-Schulz for a sharded DTensor momentum update. + + NS needs global matmuls, so gather the update to the full matrix, + orthogonalize it, and scatter the result back onto the update's + shard layout. ``full_tensor()`` returns the same gathered matrix on + every rank, so the scatter is a uniform collective. + """ + full = update.full_tensor() + ortho = _zeropower_via_newtonschulz( + full, group["ns_coefficients"], group["ns_steps"], group["eps"] + ) + return distribute_tensor(ortho, update.device_mesh, update.placements) + + +class _ShardedMuon(optim.Muon): + """Muon that materializes sharded DTensor params around Newton-Schulz. + + FSDP2 hands this optimizer dim-0 sharded DTensor parameters. The NS + iteration needs global matmuls: run it on the gathered full matrix, + then scatter the orthogonalized update back onto the parameter's + sharded layout so momentum buffers and weight decay stay sharded. + Without this, ``og @ og.T`` produces ``Partial(sum)`` DTensors that + downstream ``addmm`` calls consume without completing the reduction, + silently corrupting every update (measured 2e-4-9e-4 relative error + per step at world_size=2). + + Plain (non-DTensor) params are routed through torch's own + ``_single_tensor_muon`` so unsharded runs stay bit-for-bit identical + to ``optim.Muon`` and this class carries only the DTensor delta. + Element-wise ops (momentum lerp, weight decay, the final ``add_``) + are DTensor-safe and run directly on the shards. + """ + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params: list[Tensor] = [] + grads: list[Tensor] = [] + bufs: list[Tensor] = [] + self._init_group(group, params, grads, bufs) + + plain, sharded = [], [] + for param, grad, buf in zip(params, grads, bufs): + (sharded if isinstance(param, DTensor) else plain).append( + (param, grad, buf) + ) + + if plain: + pp, gg, bb = (list(t) for t in zip(*plain)) + _single_tensor_muon( + pp, + gg, + bb, + lr=group["lr"], + weight_decay=group["weight_decay"], + momentum=group["momentum"], + nesterov=group["nesterov"], + ns_coefficients=group["ns_coefficients"], + ns_steps=group["ns_steps"], + eps=group["eps"], + adjust_lr_fn=group["adjust_lr_fn"], + has_complex=False, + ) + + lr = _scalar_lr(group["lr"]) + for param, grad, buf in sharded: + buf.lerp_(grad, 1 - group["momentum"]) + update = grad.lerp(buf, group["momentum"]) if group["nesterov"] else buf + + adjusted_lr = _adjust_lr(lr, group["adjust_lr_fn"], param.shape) + param.mul_(1 - lr * group["weight_decay"]) + param.add_(_sharded_orthogonalize(update, group), alpha=-adjusted_lr) + return loss + + @OptimizerFactory.register("muon_adamw") class MuonAdamW(optim.Optimizer): """Combined Muon (matrix) + AdamW (non-matrix) optimizer.""" @@ -57,7 +149,7 @@ class MuonAdamW(optim.Optimizer): else: other_params.append(param) - self.muon = optim.Muon( + self.muon = _ShardedMuon( matrix_params, lr=lr, weight_decay=weight_decay, diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index 8040587..bcd214b 100644 --- a/astrai/parallel/executor.py +++ b/astrai/parallel/executor.py @@ -61,12 +61,13 @@ def broadcast_state_dict( rank = dist.get_rank() - # Broadcast metadata (keys, shapes, dtypes, device) so non-src ranks - # can allocate matching empty tensors on the correct device. + # Broadcast metadata (keys, shapes, dtypes, device kind) so non-src + # ranks can allocate matching empty tensors on their own device. if rank == src: - device = next(iter(state_dict.values())).device + first = next(iter(state_dict.values()), None) + on_cuda = first is not None and first.is_cuda metadata = [ - (k, tuple(v.shape), v.dtype, str(device)) for k, v in state_dict.items() + (k, tuple(v.shape), v.dtype, on_cuda) for k, v in state_dict.items() ] else: metadata = None @@ -75,10 +76,18 @@ def broadcast_state_dict( metadata = metadata_list[0] # Non-src ranks allocate empty tensors with the broadcasted metadata. + # The allocation must sit on *this* rank's local device: process groups + # are pinned per-rank via ``device_id=``, so NCCL rejects a tensor + # allocated on the src rank's device (and the cross-GPU allocation + # would be wrong here anyway). if rank != src: + device = ( + torch.device("cuda", torch.cuda.current_device()) + if any(m[3] for m in metadata) + else torch.device("cpu") + ) state_dict = { - k: torch.empty(s, dtype=d, device=torch.device(dev)) - for k, s, d, dev in metadata + k: torch.empty(s, dtype=d, device=device) for k, s, d, _ in metadata } # Broadcast each tensor in-place. @@ -400,17 +409,16 @@ class FSDPExecutor(BaseExecutor): if not self.use_distributed: return super().clip_grad_norm(model, max_norm) - # FSDP params are DTensors (sharded across ranks). - # torch.nn.utils.clip_grad_norm_ computes LOCAL norm per rank, - # so we must all-reduce to get the global norm before clipping. - local_norm = torch.nn.utils.get_total_norm( + # FSDP params are DTensors (sharded across ranks), and + # get_total_norm reduces them to a globally-reduced replicated + # DTensor — no extra all-reduce is needed. Manually summing the + # local slice again would inflate the norm by sqrt(world_size) + # and over-clip every step. + total_norm = torch.nn.utils.get_total_norm( [p.grad for p in model.parameters() if p.grad is not None], ) - if isinstance(local_norm, DTensor): - local_norm = local_norm.to_local() - total_norm_sq = local_norm**2 - dist.all_reduce(total_norm_sq, op=dist.ReduceOp.SUM) - total_norm = total_norm_sq.sqrt() + if isinstance(total_norm, DTensor): + total_norm = total_norm.full_tensor() clip_coef = max_norm / (total_norm + 1e-6) clip_coef_clamped = torch.clamp(clip_coef, max=1.0) diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index 012e12a..034b00d 100644 --- a/astrai/preprocessing/pipeline.py +++ b/astrai/preprocessing/pipeline.py @@ -146,11 +146,17 @@ class Pipeline: @staticmethod def _align_bucket(bucket: dict, result: dict, ids: list): - """Pad previously-accumulated keys that are missing from *result*.""" + """Pad previously-accumulated keys that are missing from *result*. + + Builders omit all-ones masks (``loss_mask`` / ``*_mask``) to save + space, so an omitted mask means "train on every token" and is + back-filled with ones; every other missing key pads with zeros. + """ for key in list(bucket.keys()): if key in result: continue - bucket[key].append([0] * len(ids)) + fill = 1 if key == "loss_mask" or key.endswith("_mask") else 0 + bucket[key].append([fill] * len(ids)) def _iter_items(self): for path in self.paths: diff --git a/astrai/serialization/hf_adapter.py b/astrai/serialization/hf_adapter.py index 3fc2481..27babec 100644 --- a/astrai/serialization/hf_adapter.py +++ b/astrai/serialization/hf_adapter.py @@ -36,6 +36,7 @@ HF_MODEL_TYPES = frozenset( "mixtral", "qwen2", "qwen2_moe", + "qwen3", "gemma", "gemma2", "phi3", @@ -58,13 +59,29 @@ _MOE_EXPERTS = re.compile( r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(weight|bias)$" ) _MOE_SHARED = re.compile( - r"^model\.layers\.(\d+)\.mlp\.shared_expert(?:s)?\.(\d+)\." + r"^model\.layers\.(\d+)\.mlp\.shared_expert(?:s)?(?:\.(\d+))?\." r"(gate|up|down)_proj\.(weight|bias)$" ) _ASTR_PREFIXES = ("embed_tokens.", "layers.", "norm.", "lm_head.") +def _half_to_interleaved(head_dim: int) -> torch.Tensor: + """Row permutation converting HF half-split RoPE coordinates to + AstrAI interleaved coordinates. + + HF rotate_half pairs channels ``(i, i + head_dim/2)``; AstrAI pairs + ``(2i, 2i + 1)``. Both use frequency ``i`` for the pair, so AstrAI + channel ``2i`` takes the HF value at channel ``i`` and AstrAI + ``2i + 1`` takes HF ``i + head_dim/2``. + """ + half = head_dim // 2 + perm = torch.empty(head_dim, dtype=torch.long) + perm[0::2] = torch.arange(half) + perm[1::2] = torch.arange(half, head_dim) + return perm + + def looks_like_hf_state_dict(state_dict: Mapping[str, Any]) -> bool: """Return True if *state_dict* uses HuggingFace key names.""" return any( @@ -168,8 +185,11 @@ def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]: cfg["n_activated_experts"] = raw["n_activated_experts"] if "n_shared_experts" in raw: cfg["n_shared_experts"] = raw["n_shared_experts"] + elif raw.get("shared_expert_intermediate_size"): + # Qwen2-MoE exposes a single un-indexed shared expert. + cfg["n_shared_experts"] = 1 else: - # Mixtral has no shared experts; AstrAI defaults to one. + # Mixtral has no shared experts. cfg["n_shared_experts"] = 0 if cfg.get("moe_intermediate_size") is None and "intermediate_size" in raw: # MoE configs store the per-expert FFN size in intermediate_size. @@ -201,6 +221,14 @@ def convert_hf_weights( ) ffn_type = getattr(config, "ffn_type", "mlp") + permute_rope = getattr(config, "attn_type", "gqa") != "mla" + head_dim = None + if permute_rope: + head_dim = config.hidden_size // config.num_attention_heads + if head_dim % 2 != 0: + raise ValueError( + f"head_dim={head_dim} is odd; rotary permutation requires even" + ) converted: Dict[str, torch.Tensor] = {} skipped: list[str] = [] for key, tensor in state_dict.items(): @@ -223,8 +251,9 @@ def convert_hf_weights( else: m = _MOE_SHARED.match(key) if m: + shared_idx = m.group(2) if m.group(2) is not None else "0" new_key = ( - f"layers.{m.group(1)}.mlp.shared_experts.{m.group(2)}." + f"layers.{m.group(1)}.mlp.shared_experts.{shared_idx}." f"{m.group(3)}.{m.group(4)}" ) if new_key is None: @@ -242,10 +271,31 @@ def convert_hf_weights( new_key = ( f"layers.{m.group(1)}.attention.{m.group(2)}_proj.{m.group(3)}" ) + if permute_rope and m.group(2) in ("q", "k"): + rows = tensor.shape[0] + if rows % head_dim != 0: + raise ValueError( + f"{key}: {rows} output rows not divisible by " + f"head_dim={head_dim}" + ) + base = _half_to_interleaved(head_dim).to(tensor.device) + blocks = ( + torch.arange(rows // head_dim, device=tensor.device) * head_dim + ) + perm = (blocks[:, None] + base[None, :]).flatten() + tensor = tensor.index_select(0, perm) elif (m := _Q_NORM.match(key)) is not None: new_key = f"layers.{m.group(1)}.attention.q_norm.weight" + if permute_rope and tensor.shape[0] == head_dim: + tensor = tensor.index_select( + 0, _half_to_interleaved(head_dim).to(tensor.device) + ) elif (m := _K_NORM.match(key)) is not None: new_key = f"layers.{m.group(1)}.attention.k_norm.weight" + if permute_rope and tensor.shape[0] == head_dim: + tensor = tensor.index_select( + 0, _half_to_interleaved(head_dim).to(tensor.device) + ) elif (m := _INPUT_NORM.match(key)) is not None: new_key = f"layers.{m.group(1)}.input_norm.weight" elif (m := _POST_NORM.match(key)) is not None: diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index a012353..6686c46 100644 --- a/astrai/trainer/rollout.py +++ b/astrai/trainer/rollout.py @@ -551,3 +551,16 @@ class RolloutRunner: # cache publication. Reward scoring itself intentionally remains # outside the policy lock because it may call an external service. return self.generator.with_policy_snapshot(commit) + + def evaluate(self, batch: Dict) -> RolloutResult: + """One-off rollout + scoring that leaves the replay cache untouched. + + Used by validation on online strategies: the training cache, its + cadence counter, and the cache key stay intact, so evaluation + prompts never disturb the rollout replay schedule. + """ + raw = self.generator.generate(batch) + self._validate_policy_version(raw) + scored = self._score(raw) + self._validate_policy_version(scored) + return scored diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 151c7e6..13274d9 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -1,7 +1,7 @@ """Training strategy implementations with factory pattern.""" from abc import ABC -from typing import Callable, Dict, List, Optional, TypedDict, Union +from typing import Any, Callable, Dict, List, Optional, TypedDict, Union import torch import torch.nn as nn @@ -202,6 +202,21 @@ class BaseStrategy(ABC): def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: return self._normalize_output(self.compute_loss(batch)) + def validate_online(self, batch: Dict[str, Any]) -> Optional[LossOutput]: + """Validate one batch through a one-off rollout. + + Online strategies with an injected rollout runner evaluate a + fresh, throw-away rollout so the training replay cache and its + cadence stay untouched. Returns ``None`` when no runner is + configured (offline mode); callers then fall back to + ``strategy(batch)``. + """ + if self._rollout_runner is None: + return None + result = self._rollout_runner.evaluate(batch) + prepared = self.prepare_from_rollout(result) + return self.compute_loss_output(prepared) + def _loss_output( self, task_loss: Tensor, @@ -595,6 +610,19 @@ class GRPOStrategy(BaseStrategy): if state_dict is not None: self.old_model.load_state_dict(state_dict) + def optimizer_step(self, optimizer: Optimizer): + """Step the optimizer, then refresh the offline behaviour policy. + + Without this sync the frozen ``old_model`` drifts away from the + training policy, so the PPO ratio degenerates and clipping shuts + learning down. Online GRPO passes ``logprobs_old`` instead and + runs with ``old_model=None``, skipping the sync. + """ + result = super().optimizer_step(optimizer) + if self.old_model is not None: + self.sync_old_model() + return result + def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: batch = move_to_device(batch, self.device) prompts = batch["prompts"] diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index a1c8677..831b669 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -1,6 +1,7 @@ import json import logging import os +import shutil import sys import time from functools import partial @@ -29,6 +30,32 @@ from astrai.trainer.train_context import TrainContext logger = logging.getLogger(__name__) +_TOKENIZER_FILES = ( + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", +) + + +def _copy_tokenizer_files(param_path: Optional[str], save_path: str): + """Snapshot tokenizer files into the checkpoint directory. + + ``param_path`` is the launch model directory (or, on resume, a + previous self-contained checkpoint), so the copy makes every + checkpoint independently resumable for online training, which + loads its tokenizer from ``param_path``. + """ + if not param_path: + return + for name in _TOKENIZER_FILES: + src = os.path.join(param_path, name) + dst = os.path.join(save_path, name) + if not os.path.isfile(src) or ( + os.path.isfile(dst) and os.path.samefile(src, dst) + ): + continue + shutil.copy2(src, dst) + @runtime_checkable class TrainCallback(Protocol): @@ -176,6 +203,7 @@ class CheckpointCallback(TrainCallback): meta=meta, ) context.checkpoint.save(save_path) + _copy_tokenizer_files(context.param_path, save_path) self.last_ckpt_step = context.optimizer_step def after_optimizer_step(self, context: TrainContext): @@ -325,7 +353,12 @@ class MetricCallback(TrainCallback): with torch.no_grad(): for batch in context.val_dataloader: - loss_output = context.strategy(batch) + # Online strategies evaluate a one-off rollout (leaving + # the replay cache untouched) via the public hook; None + # means offline — validate the batch directly. + loss_output = context.strategy.validate_online(batch) + if loss_output is None: + loss_output = context.strategy(batch) total_loss += loss_output["loss"].item() num_batches += 1 diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 5bac667..37ff5ca 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -59,6 +59,7 @@ class TrainContext: world_size: int = field(default=1) rank: int = field(default=0) kwargs: Dict[str, Any] = field(default_factory=dict) + param_path: Optional[str] = field(default=None) _stop_event: threading.Event = field(default_factory=threading.Event) @@ -180,6 +181,7 @@ class TrainContextBuilder: epoch=state.epoch, consumed_samples=state.consumed_samples, checkpoint=state.checkpoint, + param_path=self._param_path, ) def _prepare_model( diff --git a/tests/data/test_pipeline_align.py b/tests/data/test_pipeline_align.py new file mode 100644 index 0000000..0a1eb9c --- /dev/null +++ b/tests/data/test_pipeline_align.py @@ -0,0 +1,26 @@ +"""Tests for preprocessing pipeline bucket alignment.""" + +from astrai.preprocessing.pipeline import Pipeline + + +def test_align_bucket_backfills_missing_mask_with_ones(): + bucket = { + "sequence": [[1, 2], [3, 4]], + "loss_mask": [[0, 1]], + "chosen_mask": [[1]], + "position_ids": [[0, 1]], + } + result = {"sequence": [5, 6, 7]} + Pipeline._align_bucket(bucket, result, [5, 6, 7]) + assert bucket["loss_mask"][-1] == [1, 1, 1] + assert bucket["chosen_mask"][-1] == [1, 1, 1] + assert bucket["position_ids"][-1] == [0, 0, 0] + assert bucket["sequence"] == [[1, 2], [3, 4]] + + +def test_align_bucket_keeps_present_keys(): + bucket = {"sequence": [[1, 2]], "loss_mask": [[0, 1]]} + result = {"sequence": [9], "loss_mask": [1]} + Pipeline._align_bucket(bucket, result, [9]) + assert bucket["loss_mask"] == [[0, 1]] + assert bucket["sequence"] == [[1, 2]] diff --git a/tests/inference/test_sample.py b/tests/inference/test_sample.py index 991fa11..f950fd3 100644 --- a/tests/inference/test_sample.py +++ b/tests/inference/test_sample.py @@ -3,6 +3,7 @@ import torch from astrai.inference.runtime.sample import ( + BaseSamplingStrategy, FrequencyPenaltyStrategy, SamplingPipeline, TemperatureStrategy, @@ -295,3 +296,45 @@ def test_greedy_respects_frequency_penalty(): ) # 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 diff --git a/tests/inference/test_tool_parser.py b/tests/inference/test_tool_parser.py index ddfd19f..2b63235 100644 --- a/tests/inference/test_tool_parser.py +++ b/tests/inference/test_tool_parser.py @@ -565,3 +565,40 @@ def test_parser_uses_token_ids_for_detection(): parser = TokenIdParser() parser.feed("hello", current_token_ids=[1, 999, 3]) assert parser.has_tool_calls + + +def test_streaming_partial_name_prefix_never_leaks_into_content(): + parser = SimpleJsonToolParser() + parts = ["Hello ", '{"', '{"n', '{"na', '{"name"'] + emitted = [] + body = "" + for part in parts: + body += part + for d in parser.feed(body): + if "content" in d: + emitted.append(d["content"]) + assert "".join(emitted) == "Hello " + + +def test_finalize_flushes_withheld_plain_json_content(): + parser = SimpleJsonToolParser() + text = 'Answer: {"price": 1}' + deltas = parser.feed(text) + streamed = "".join(d["content"] for d in deltas if "content" in d) + flushed = parser.finalize(text) + joined = streamed + "".join(d["content"] for d in flushed if "content" in d) + assert joined == text + assert not parser.has_tool_calls + assert parser.finalize(text) == [] + + +def test_streaming_args_concat_matches_parse_complete(): + parser = SimpleJsonToolParser() + # Compact spacing: json.dumps would re-space this and desync the + # streamed arguments diff. + text = '{"name": "get_weather","arguments": {"city":"Beijing","unit":"c"}}' + _, args_chunks = _simulate_streaming(parser, text) + streamed = "".join(args_chunks) + completed = parser.parse_complete(text)["tool_calls"][0]["function"]["arguments"] + assert streamed == completed + assert streamed == '"city":"Beijing","unit":"c"' diff --git a/tests/serialization/test_hf_adapter.py b/tests/serialization/test_hf_adapter.py index b79818a..cf5d45a 100644 --- a/tests/serialization/test_hf_adapter.py +++ b/tests/serialization/test_hf_adapter.py @@ -15,6 +15,7 @@ from astrai.serialization import ( looks_like_hf_state_dict, save_model, ) +from astrai.serialization.hf_adapter import _half_to_interleaved from tests.helpers import assert_state_dicts_equal, make_tiny_config LLAMA_RAW = { @@ -47,10 +48,38 @@ MOE_RAW = { } -def to_hf_keys(state_dict): - """Rename AstrAI state dict keys to HuggingFace LLaMA-style names.""" +def to_hf_keys(state_dict, head_dim=None): + """Rename AstrAI state dict keys to HuggingFace LLaMA-style names. + + When *head_dim* is given, q/k projections and q/k norm weights are + also converted from AstrAI interleaved RoPE coordinates to the HF + half-split (rotate_half) convention, so the produced state dict is a + faithful HF-layout checkpoint. + """ out = {} for key, tensor in state_dict.items(): + if head_dim is not None: + name = key.split(".") + is_qk_proj = ( + len(name) >= 4 + and name[2] == "attention" + and name[3] in ("q_proj", "k_proj") + ) + is_qk_norm = ( + len(name) >= 4 + and name[2] == "attention" + and name[3] in ("q_norm", "k_norm") + and name[4] == "weight" + ) + if is_qk_proj or is_qk_norm: + inv = torch.argsort(_half_to_interleaved(head_dim)) + rows = tensor.shape[0] + if rows > head_dim: + blocks = torch.arange(rows // head_dim) * head_dim + idx = (blocks[:, None] + inv[None, :]).flatten() + else: + idx = inv + tensor = tensor.index_select(0, idx) if key == "embed_tokens.weight": out["model.embed_tokens.weight"] = tensor elif key == "norm.weight": @@ -141,7 +170,9 @@ def test_adapt_config_passthrough(): def test_convert_hf_weights_dense_roundtrip(): cfg = make_tiny_config() model = AutoRegressiveLM(cfg) - converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg) + converted = convert_hf_weights( + to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg + ) assert_state_dicts_equal(converted, model.state_dict()) @@ -157,7 +188,10 @@ def test_convert_hf_weights_moe_roundtrip(): model = AutoRegressiveLM(cfg) hf_raw = convert_hf_config(MOE_RAW) hf_cfg = ConfigFactory.load(hf_raw) - converted = convert_hf_weights(to_hf_keys(model.state_dict()), hf_cfg) + converted = convert_hf_weights( + to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), + hf_cfg, + ) assert_state_dicts_equal(converted, model.state_dict()) @@ -216,7 +250,9 @@ def test_convert_hf_weights_moe_with_dense_layers_roundtrip(): decoder_sparse_step=1, ) model = AutoRegressiveLM(cfg) - converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg) + converted = convert_hf_weights( + to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg + ) assert_state_dicts_equal(converted, model.state_dict()) @@ -230,7 +266,7 @@ def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip(): shared_expert_intermediate_size=16, ) model = AutoRegressiveLM(cfg) - hf_sd = to_hf_keys(model.state_dict()) + hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads) hf_sd = { k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items() } @@ -241,7 +277,9 @@ def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip(): def test_convert_hf_weights_gemma_qk_norm_roundtrip(): cfg = make_tiny_config(use_qk_norm=True) model = AutoRegressiveLM(cfg) - converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg) + converted = convert_hf_weights( + to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads), cfg + ) assert_state_dicts_equal(converted, model.state_dict()) @@ -250,7 +288,9 @@ def test_from_pretrained_hf_directory(tmp_path): model = AutoRegressiveLM(cfg).eval() save_model( config=LLAMA_RAW, - state_dict=to_hf_keys(model.state_dict()), + state_dict=to_hf_keys( + model.state_dict(), cfg.hidden_size // cfg.num_attention_heads + ), save_directory=str(tmp_path), ) loaded = AutoModel.from_pretrained(tmp_path).eval() @@ -292,7 +332,9 @@ def test_from_pretrained_weights_format_astrai_rejects_hf(tmp_path): model = AutoRegressiveLM(cfg) save_model( config=LLAMA_RAW, - state_dict=to_hf_keys(model.state_dict()), + state_dict=to_hf_keys( + model.state_dict(), cfg.hidden_size // cfg.num_attention_heads + ), save_directory=str(tmp_path), ) with pytest.raises(ValueError): @@ -313,7 +355,7 @@ def test_from_pretrained_invalid_weights_format(tmp_path): def test_from_pretrained_hf_directory_sharded(tmp_path): cfg = make_tiny_config() model = AutoRegressiveLM(cfg).eval() - hf_sd = to_hf_keys(model.state_dict()) + hf_sd = to_hf_keys(model.state_dict(), cfg.hidden_size // cfg.num_attention_heads) keys = sorted(hf_sd) split = len(keys) // 2 shard_a = {k: hf_sd[k] for k in keys[:split]} @@ -342,6 +384,128 @@ def test_from_pretrained_hf_directory_sharded(tmp_path): ) +def _half_split_rope(q, theta=10000.0): + """HF llama-style rotate_half RoPE on [batch, seq, heads, head_dim].""" + b, s, h, d = q.shape + inv_freq = theta ** (-torch.arange(0, d, 2, dtype=torch.float64) / d) + freqs = torch.outer(torch.arange(s, dtype=torch.float64), inv_freq).float() + cos, sin = freqs.cos()[None, :, None, :], freqs.sin()[None, :, None, :] + q1, q2 = q[..., : d // 2], q[..., d // 2 :] + return torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1) + + +def _rms_norm_hf(t, weight, eps): + t = t.float() + t = t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + eps) + return weight.float() * t + + +def _hf_reference_attn( + x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim, q_norm_w=None, k_norm_w=None, eps=1e-5 +): + """Ground-truth HF attention: per-head RMSNorm BEFORE RoPE (half-split).""" + import torch.nn.functional as F + + b, s, dim = x.shape + q = (x @ Wq.T).reshape(b, s, n_heads, head_dim).float() + k = (x @ Wk.T).reshape(b, s, n_kv, head_dim).float() + v = (x @ Wv.T).reshape(b, s, n_kv, head_dim).float() + if q_norm_w is not None: + q = _rms_norm_hf(q, q_norm_w, eps) + k = _rms_norm_hf(k, k_norm_w, eps) + q, k = _half_split_rope(q), _half_split_rope(k) + rep = n_heads // n_kv + k = k.repeat_interleave(rep, dim=2).transpose(1, 2) + v = v.repeat_interleave(rep, dim=2).transpose(1, 2) + out = F.scaled_dot_product_attention(q.transpose(1, 2), k, v, is_causal=True) + out = out.transpose(1, 2).reshape(b, s, n_heads * head_dim) + return out @ Wo.T + + +def _run_converted_gqa(x, hf_sd, cfg): + from astrai.model.components.attention import GQA + from astrai.model.components.rope import get_rotary_emb + + attn = GQA( + dim=cfg.hidden_size, + n_heads=cfg.num_attention_heads, + n_kv_heads=cfg.num_key_value_heads, + use_qk_norm=cfg.use_qk_norm, + norm_eps=cfg.rms_norm_eps, + use_gated_attention=False, + layer_id=0, + ).eval() + converted = convert_hf_weights(hf_sd, cfg) + local = { + k.removeprefix("layers.0.attention."): v + for k, v in converted.items() + if k.startswith("layers.0.attention.") + } + attn.load_state_dict(local, strict=True) + head_dim = cfg.hidden_size // cfg.num_attention_heads + seq = x.shape[1] + rot = get_rotary_emb(head_dim, seq)[None, :seq].expand(x.shape[0], seq, -1, -1) + with torch.no_grad(): + return attn(x, rot, is_causal=True) + + +def test_hf_import_rope_permutation_matches_half_split_reference(): + torch.manual_seed(0) + n_heads, n_kv, head_dim = 4, 2, 8 + dim = n_heads * head_dim + Wq = torch.randn(n_heads * head_dim, dim) + Wk = torch.randn(n_kv * head_dim, dim) + Wv = torch.randn(n_kv * head_dim, dim) + Wo = torch.randn(dim, dim) + x = torch.randn(2, 16, dim) + + hf_sd = { + "model.layers.0.self_attn.q_proj.weight": Wq, + "model.layers.0.self_attn.k_proj.weight": Wk, + "model.layers.0.self_attn.v_proj.weight": Wv, + "model.layers.0.self_attn.o_proj.weight": Wo, + } + cfg = make_tiny_config( + hidden_size=dim, num_attention_heads=n_heads, num_key_value_heads=n_kv + ) + ref = _hf_reference_attn(x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim) + out = _run_converted_gqa(x, hf_sd, cfg) + torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4) + + +def test_hf_import_qk_norm_matches_norm_before_rope_reference(): + torch.manual_seed(1) + n_heads, n_kv, head_dim = 4, 2, 8 + dim = n_heads * head_dim + Wq = torch.randn(n_heads * head_dim, dim) + Wk = torch.randn(n_kv * head_dim, dim) + Wv = torch.randn(n_kv * head_dim, dim) + Wo = torch.randn(dim, dim) + gq = torch.randn(head_dim) + gk = torch.randn(head_dim) + x = torch.randn(2, 16, dim) + + hf_sd = { + "model.layers.0.self_attn.q_proj.weight": Wq, + "model.layers.0.self_attn.k_proj.weight": Wk, + "model.layers.0.self_attn.v_proj.weight": Wv, + "model.layers.0.self_attn.o_proj.weight": Wo, + "model.layers.0.self_attn.q_norm.weight": gq, + "model.layers.0.self_attn.k_norm.weight": gk, + } + cfg = make_tiny_config( + hidden_size=dim, + num_attention_heads=n_heads, + num_key_value_heads=n_kv, + use_qk_norm=True, + ) + ref = _hf_reference_attn( + x, Wq, Wk, Wv, Wo, n_heads, n_kv, head_dim, q_norm_w=gq, k_norm_w=gk + ) + out = _run_converted_gqa(x, hf_sd, cfg) + torch.testing.assert_close(out, ref, rtol=1e-4, atol=1e-4) + + def test_from_pretrained_hf_directory_with_moe(tmp_path): cfg = make_tiny_config( ffn_type="moe", @@ -354,7 +518,9 @@ def test_from_pretrained_hf_directory_with_moe(tmp_path): model = AutoRegressiveLM(cfg).eval() save_model( config=MOE_RAW, - state_dict=to_hf_keys(model.state_dict()), + state_dict=to_hf_keys( + model.state_dict(), cfg.hidden_size // cfg.num_attention_heads + ), save_directory=str(tmp_path), ) loaded = AutoModel.from_pretrained(tmp_path).eval() diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py index 85613ae..ad97881 100644 --- a/tests/trainer/test_callbacks.py +++ b/tests/trainer/test_callbacks.py @@ -4,7 +4,11 @@ import torch from astrai.model.components.decoder_block import DecoderBlock from astrai.serialization import Checkpoint -from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback +from astrai.trainer.train_callback import ( + GradientCheckpointingCallback, + TrainCallback, + _copy_tokenizer_files, +) from astrai.trainer.trainer import Trainer from tests.helpers import RandomTokenDataset @@ -174,3 +178,38 @@ def test_checkpoint_captures_completed_optimizer_step( assert ( Path(base_test_env["test_dir"]) / "epoch_0_step_1" / "metric.jsonl" ).is_file() + + +def test_checkpoint_snapshots_tokenizer_files( + base_test_env, train_config_factory, device, tmp_path +): + """Checkpoints copy tokenizer files from param_path so resume works.""" + param_dir = tmp_path / "model" + param_dir.mkdir() + (param_dir / "tokenizer.json").write_text("{}") + (param_dir / "tokenizer_config.json").write_text("{}") + + train_config = train_config_factory( + model_fn=lambda: base_test_env["model"], + dataset=RandomTokenDataset(length=2), + test_dir=base_test_env["test_dir"], + device=device, + batch_per_device=2, + ckpt_interval=1, + ) + + Trainer(train_config).train(param_path=str(param_dir)) + + ckpt_dir = Path(base_test_env["test_dir"]) / "epoch_0_step_1" + assert (ckpt_dir / "tokenizer.json").is_file() + assert (ckpt_dir / "tokenizer_config.json").is_file() + + # Resuming with param_path == checkpoint dir must not raise + # (samefile guard). + _copy_tokenizer_files(str(ckpt_dir), str(ckpt_dir)) + + +def test_copy_tokenizer_files_skips_missing_and_none(tmp_path): + _copy_tokenizer_files(None, str(tmp_path)) + _copy_tokenizer_files(str(tmp_path), str(tmp_path / "out")) + assert not (tmp_path / "out").exists() or not any((tmp_path / "out").iterdir()) diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index 473871b..b6b880b 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/tests/trainer/test_grpo_strategy.py @@ -162,3 +162,34 @@ def test_grpo_sync_old_model(grpo_strategy): if k in old_sd_after ) assert matches + + +def test_grpo_optimizer_step_syncs_old_model(grpo_strategy): + """optimizer_step must refresh old_model after each update.""" + strategy, device = grpo_strategy + + class _SteppedOptimizer: + def step(self): + with torch.no_grad(): + for p in strategy.model.parameters(): + p.add_(0.05) + + strategy.optimizer_step(_SteppedOptimizer()) + + policy_sd = strategy.model.state_dict() + old_sd = strategy.old_model.state_dict() + assert all( + torch.allclose(policy_sd[k], old_sd[k]) for k in policy_sd if k in old_sd + ) + + +def test_online_grpo_optimizer_step_skips_sync(grpo_strategy): + """old_model=None (online) must not attempt a sync.""" + strategy, device = grpo_strategy + strategy.old_model = None + + class _SteppedOptimizer: + def step(self): + return None + + strategy.optimizer_step(_SteppedOptimizer()) diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index dbb2d0f..f711153 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -45,6 +45,7 @@ class _RecordingRunner: self._fresh = True self.policy_version = result.policy_version self.weight_updates = [] + self.eval_calls = 0 def __call__(self, batch): self.calls += 1 @@ -52,6 +53,12 @@ class _RecordingRunner: self._fresh = False return self.result, fresh + def evaluate(self, batch): + # Mirrors RolloutRunner.evaluate: one-off scoring that never + # touches the replay cache or freshness state. + self.eval_calls += 1 + return self.result + def step(self): self.step_calls += 1 @@ -360,6 +367,26 @@ def test_loss_is_differentiable_dpo(device): assert has_grad +def test_validate_online_returns_none_without_runner(device): + strat = _make_grpo(device) + batch = {"input_ids": torch.randint(3, 200, (2, 4), device=device)} + assert strat.validate_online(batch) is None + + +def test_validate_online_uses_one_off_rollout_not_replay_cache(device): + strat = _make_grpo(device) + runner = _RecordingRunner(_make_rollout_result(device=device)) + strat.set_rollout_runner(runner) + + out = strat.validate_online( + {"input_ids": torch.randint(3, 200, (2, 4), device=device)} + ) + + assert torch.isfinite(out["loss"]).item() + assert runner.eval_calls == 1 + assert runner.calls == 0 # replay cache path untouched + + def test_ref_model_not_updated_by_backward_dpo(device): strat = _make_dpo(device) strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device))) diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index 385e368..411bb30 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -388,6 +388,21 @@ def test_rollout_runner_cache_returns_stale_flag(device): assert fresh2 is False +def test_rollout_runner_evaluate_leaves_cache_untouched(device): + runner, _ = _make_runner(device, rollout_interval=10) + batch = _make_instruction_batch() + cached, _ = runner(batch) + + eval_batch = _make_instruction_batch(n=1) + result = runner.evaluate(eval_batch) + + assert result.rewards.shape == result.responses.shape[:2] + replayed, fresh = runner(batch) + assert replayed is cached + assert fresh is False + assert runner._steps_since_rollout == 0 + + def test_rollout_runner_tags_generation_version_and_preserves_cached_behavior(device): runner, _ = _make_runner(device, rollout_interval=100) batch = _make_instruction_batch(n=1)