- register online_ppo train type backed by PPOStrategy: token-level clipped surrogate over GAE advantages plus masked value regression against rollout-pinned returns, with explained-variance metrics
- fold the reference-KL penalty (k3 estimator) into per-token rewards before GAE and pin advantages/returns on RolloutResult so replayed gradient steps optimize fixed targets
- add self-contained ValueModel critic with a zero-initialized value head and backbone warm-started from policy weights; AutoRegressiveLM stays untouched and trunk parity is pinned by tests
- step the critic's own optimizer outside the policy-version lock with the same max_grad_norm clipping as the policy
- persist critic state as value_model.pt/value_optimizer.pt checkpoint extras; resume restores it, fails loudly when missing, and the train.sh completeness check requires the extras for online_ppo configs
- extract shared rollout sequence/logprob helpers from GRPO (behavior unchanged) and add ppo_gamma/ppo_gae_lambda/ppo_vf_coef CLI options
- 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
- re-register the linear family with the operator dispatcher (ASTR_OPS / op_backend / resolve)
- fix bf16 gemv misaligned-address faults and element mispairing for offset weights
- reject misaligned bf16_swiglu inputs with a clear error and fall back in the backend gate
- make the rollout reuse decision, validation, and return atomic under one policy snapshot
- add the documented post-scoring rollout version check
- derive live+1 under the scheduler lock in optimizer_step via apply_weight_update(None, ...)
- reject rollout_max_policy_lag below rollout_interval - 1 at config time
- sync gemv stream-test inputs before switching streams; drop dead loader imports
Feed sampler-aligned behavior log-probabilities directly into online GRPO instead of allocating, synchronizing, and forwarding a duplicate old-policy model. Keep the old-model path as an offline compatibility fallback and validate supplied rollout tensors before loss computation.
Track a monotonic policy version across optimizer steps, scheduler updates, and rollout results. Serialize synchronous generation with weight acknowledgements and invalidate reusable prefix KV entries so cached samples remain attributable to the behavior policy that generated them.
- avoid constructing model_fn more than once when reading config
- keep inference package exports focused on public entry points
- rename extra strategy arguments to strategy_kwargs
- Converts detached strategy metrics before returning loss output
- Removes redundant item conversion from the trainer loop
- Updates the documented contract and regression tests
- Propagates MoE load-balancing loss through model outputs
- Logs task, auxiliary, and weighted losses across strategies
- Computes only explicitly requested callback metrics
- Preserves tensor compute_loss API and adds regression tests
- Add broadcast_state_dict to sync state_dict from rank-0 to all ranks
- Fix create_ref_model returning None on non-rank-0 under FSDP
- Fix sync_old_model only updating old_model on rank-0 under FSDP
- Split skip_no_cuda/skip_no_kernel markers and hoist to top-level conftest
- Add distributed tests for broadcast_state_dict and create_ref_model
- Extract _resolve_base_type and _validate_component as module-level helpers
- Replace ForwardRef._evaluate private API with eval in module namespace
- Remove broad except Exception in __init_subclass__, _component_base always set
- Replace direct _entries mutation in strategy.py with register() call form
- Remove dead TOKENIZER_CLASSES registry from AutoTokenizer
- Extract ModelFactory(BaseFactory[nn.Module]) as pure factory
- AutoModel now inherits only nn.Module, no factory state
- Move @AutoModel.register to @ModelFactory.register in transformer.py and encoder.py
- Move create_ref_model from astrai/parallel/utils.py into executor.py
- Remove unused ColumnParallelLinear/RowParallelLinear (module.py)
- Update imports in strategy.py and train_context.py
- Drop unused astrai.parallel.utils and astrai.parallel.module
- unshard() and full_tensor() are collective ops, all ranks must participate
- Old code returned None on non-rank-0 before calling unshard, causing deadlock
- Fix: all ranks unshard/full_tensor, only rank-0 keeps the result
- Move create_ref_model to parallel/utils.py, accept executor+model directly
- Guard create_ref_model and sync_old_model against None on non-rank-0
- Fix pre-existing ruff lint warnings (F401, F541, F841, E741)
- Exclude .md/.json/.yml from ruff format check
- Unblock SIGTERM/SIGINT via pthread_sigmask in early signal handler
- Do not restore SIG_DFL on unregister to prevent pending signal kills
- add loss_mask parameter to get_logprobs to decouple attention from loss masking
- DPO/GRPO strategies compose key-padding + causal mask before model forward
- prevents prompt tokens from being masked out of attention and missing causal masking
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
- Drop BaseStrategy.model_fn (stored but never read)
- Drop model_fn= passed to StrategyFactory.create in train_context
- Simplify FSDPExecutor.clip_grad_norm None branch to delegate to super()
- Remove DDPExecutor._gather_state_dict override (identical to base)
- DPO requires sequence-level sum of token logprobs, not per-token mean
- mean reduction made beta*ratio_diff ~0.03 (near-zero gradient)
- loss stalled at 0.6931 because logsigmoid(0.03) has vanishing grad
- sum gives beta*ratio_diff ~10 with meaningful gradients
- Mask prompt tokens to 0 so their logprobs excluded from ratio/KL
- Switch to token-level ratio + PPO clipping via reduction='none'
- Slice response token logprobs from full sequence output
- Replace k3 KL estimator with non-negative k1 estimator
- Fix epsilon from finifo.eps (~1e-38) to 1e-8
- Remove unused 'reduction' param from GRPOStrategy.__init__
- Clarify offline batch semantics in docstring
- Add 11 unit tests for masking, advantage, KL, sync, clipping
- Sync training.md and architecture.md docs
ViperEkura
0z5a