ViperEkura
350e4a1849
- 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
38 lines
899 B
Python
38 lines
899 B
Python
from astrai.model.automodel import AutoModel
|
|
from astrai.model.components.attention import GQA
|
|
from astrai.model.components.decoder_block import DecoderBlock
|
|
from astrai.model.components.linear import Linear
|
|
from astrai.model.components.lora import (
|
|
LoRAConfig,
|
|
inject_lora,
|
|
load_lora,
|
|
merge_lora,
|
|
save_lora,
|
|
)
|
|
from astrai.model.components.mlp import MLP, DeepSeekMoE
|
|
from astrai.model.components.norm import RMSNorm
|
|
from astrai.model.encoder import EmbeddingEncoder
|
|
from astrai.model.transformer import AutoRegressiveLM
|
|
from astrai.model.value import ValueModel
|
|
|
|
__all__ = [
|
|
# Modules
|
|
"Linear",
|
|
"RMSNorm",
|
|
"MLP",
|
|
"DeepSeekMoE",
|
|
"GQA",
|
|
"DecoderBlock",
|
|
# Models
|
|
"AutoRegressiveLM",
|
|
"EmbeddingEncoder",
|
|
"AutoModel",
|
|
"ValueModel",
|
|
# LoRA
|
|
"LoRAConfig",
|
|
"inject_lora",
|
|
"merge_lora",
|
|
"save_lora",
|
|
"load_lora",
|
|
]
|