ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
350e4a18493a4222f84b93abfb88d52bdd25a1b1
AstrAI/scripts/tools/train.py
T
ViperEkura 350e4a1849 feat: add online ppo with value-model critic and gae advantages 
- 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
2026-09-05 01:59:50 +08:00

714 lines
21 KiB
Python
Raw Blame History

import os
from collections.abc import Callable
from functools import partial
import click
import torch
from click.core import ParameterSource
from torch import optim
from astrai.config import AutoRegressiveLMConfig, TrainConfig
from astrai.config.cli import (
GroupedCommand,
OptSpec,
apply_specs,
merge_yaml_into_kwargs,
)
from astrai.config.train_config import (
BACKENDS,
PARALLEL_MODES,
START_METHODS,
TRAIN_TYPES,
)
from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
from astrai.model import AutoRegressiveLM, ValueModel
from astrai.model.components.decoder_block import DecoderBlock
from astrai.optim import OptimizerFactory
from astrai.trainer import SchedulerFactory, Trainer
from astrai.trainer.rollout import BaseRewardModel
# Re-exported under its historical name for tests importing it from here.
_merge_yaml_into_kwargs = merge_yaml_into_kwargs
_TRAIN_TYPE = sorted(TRAIN_TYPES)
_PARALLEL = sorted(PARALLEL_MODES)
_SCHEDULES = ["cosine", "sgdr", "wsd"]
_OPTIMIZERS = OptimizerFactory.list_registered()
_BACKENDS = sorted(BACKENDS)
_START_METHODS = sorted(START_METHODS)
# Option table: types/defaults marked AUTO are inferred from TrainConfig
# fields; everything else (CLI-only options and default overrides) is
# declared inline. Table order is the --help order.
_SPECS = [
OptSpec(
"config_path",
"Paths & Setup",
type=click.Path(exists=True),
param_decls=("--config", "-c", "config_path"),
help="YAML config file. CLI flags override YAML values.",
),
OptSpec("train_type", "Paths & Setup", choices=_TRAIN_TYPE, help="Training type."),
OptSpec(
"data_root_path",
"Paths & Setup",
type=click.Path(exists=True),
help="Root directory of the dataset.",
),
OptSpec(
"param_path",
"Paths & Setup",
type=click.Path(exists=True),
help="Path to model parameters or resume checkpoint.",
),
OptSpec(
"resume",
"Paths & Setup",
is_flag=True,
default=False,
help="Resume from checkpoint.",
),
OptSpec("n_epoch", "Training", help="Number of epochs."),
OptSpec("batch_per_device", "Training", default=1, help="Batch size per GPU."),
OptSpec(
"grad_accum_steps",
"Training",
help="Gradient accumulation steps.",
),
OptSpec("max_grad_norm", "Training", help="Max gradient norm for clipping."),
OptSpec(
"warmup_ratio",
"LR Schedule",
type=float,
default=0.05,
help="Fraction of total steps for LR warmup.",
),
OptSpec(
"max_lr",
"Optimizer",
type=float,
default=3e-4,
help="Max learning rate.",
),
OptSpec(
"optimizer",
"Optimizer",
choices=_OPTIMIZERS,
default="muon_adamw",
help="Built-in optimizer.",
),
OptSpec(
"weight_decay",
"Optimizer",
type=float,
default=0.1,
help="Weight decay for eligible optimizer parameters.",
),
OptSpec(
"nora_lr", "Optimizer", type=float, default=5e-3, help="Nora learning rate."
),
OptSpec(
"nora_beta", "Optimizer", type=float, default=0.95, help="Nora EMA factor."
),
OptSpec(
"nora_momentum",
"Optimizer",
type=float,
default=0.95,
help="Nora update momentum.",
),
OptSpec(
"nora_weight_decay",
"Optimizer",
type=float,
default=0.0,
help="Nora weight decay.",
),
OptSpec(
"muon_momentum",
"Optimizer",
type=float,
default=0.95,
help="Muon momentum factor.",
),
OptSpec(
"muon_nesterov",
"Optimizer",
type=bool,
default=True,
help="Muon Nesterov.",
),
OptSpec(
"muon_ns_steps",
"Optimizer",
type=int,
default=5,
help="Muon Newton-Schulz steps.",
),
OptSpec(
"muon_adjust_lr",
"Optimizer",
choices=["original", "match_rms_adamw"],
default="match_rms_adamw",
help="Muon LR adjustment strategy.",
),
OptSpec(
"mano_momentum",
"Optimizer",
type=float,
default=0.95,
help="Mano momentum factor.",
),
OptSpec(
"mano_nesterov",
"Optimizer",
type=bool,
default=True,
help="Mano Nesterov momentum.",
),
OptSpec("random_seed", "Data Loading", help="Random seed."),
OptSpec("num_workers", "Data Loading", default=4, help="DataLoader workers."),
OptSpec("pin_memory", "Data Loading", default=True, help="Pin memory."),
OptSpec(
"persistent_workers",
"Data Loading",
default=True,
help="Keep DataLoader workers alive between epochs.",
),
OptSpec(
"window_size",
"Data Loading",
type=int,
default=None,
help="Max input sequence length.",
),
OptSpec(
"stride",
"Data Loading",
type=int,
default=None,
help="Step size for sliding window.",
),
OptSpec(
"label_smoothing",
"Data Loading",
type=float,
default=0.0,
help="Label smoothing.",
),
OptSpec("dpo_beta", "Algorithm", type=float, default=0.1, help="DPO beta."),
OptSpec("group_size", "Algorithm", type=int, default=4, help="GRPO group size."),
OptSpec(
"grpo_clip_eps",
"Algorithm",
type=float,
default=0.2,
help="GRPO clip epsilon.",
),
OptSpec(
"grpo_kl_coef",
"Algorithm",
type=float,
default=0.01,
help="GRPO KL penalty coefficient.",
),
OptSpec(
"ppo_gamma",
"Algorithm",
type=float,
default=1.0,
help="PPO reward discount factor.",
),
OptSpec(
"ppo_gae_lambda",
"Algorithm",
type=float,
default=0.95,
help="PPO GAE bias/variance trade-off.",
),
OptSpec(
"ppo_vf_coef",
"Algorithm",
type=float,
default=0.5,
help="PPO value-loss coefficient.",
),
OptSpec(
"moe_aux_loss_coef",
"Algorithm",
help="MoE load balancing auxiliary loss coefficient (0=disable).",
),
OptSpec("rollout_interval", "Algorithm", help="Steps between rollouts."),
OptSpec(
"rollout_max_policy_lag",
"Algorithm",
help="Maximum accepted rollout/live policy-version gap.",
),
OptSpec("rollout_temperature", "Algorithm", help="Rollout temperature."),
OptSpec("rollout_top_k", "Algorithm", help="Rollout top-k (0=disable)."),
OptSpec("rollout_top_p", "Algorithm", help="Rollout top-p."),
OptSpec("rollout_max_tokens", "Algorithm", help="Max tokens per rollout response."),
OptSpec("neftune_alpha", "Algorithm", help="NEFTune noise alpha."),
OptSpec("val_split", "Validation", help="Validation split ratio."),
OptSpec("val_step", "Validation", help="Steps between validation runs."),
OptSpec(
"metrics",
"Validation",
default=("loss", "lr", "grad_norm", "grad_snr"),
help="Metrics to log (repeatable).",
),
OptSpec("ckpt_interval", "Checkpoint", help="Steps between checkpoints."),
OptSpec(
"ckpt_dir",
"Checkpoint",
type=click.Path(),
default="checkpoint",
help="Checkpoint directory.",
),
OptSpec("start_epoch", "Checkpoint", help="Start epoch."),
OptSpec("start_samples", "Checkpoint", help="Start samples (per rank)."),
OptSpec("master_addr", "Distributed", help="Master node address."),
OptSpec("master_port", "Distributed", help="Master node port."),
OptSpec("backend", "Distributed", choices=_BACKENDS, help="Distributed backend."),
OptSpec("nprocs", "Distributed", help="Number of GPUs."),
OptSpec(
"parallel_mode",
"Distributed",
choices=_PARALLEL,
default="fsdp",
help="Parallel strategy.",
),
OptSpec("device_type", "Distributed", help="Device type."),
OptSpec(
"start_method",
"Distributed",
choices=_START_METHODS,
help="Multiprocessing start method.",
),
OptSpec(
"tp_size",
"Distributed",
type=int,
default=None,
help="Tensor parallelism (future).",
),
OptSpec(
"gradient_checkpointing",
"Misc",
type=bool,
default=False,
help="Enable activation checkpointing.",
),
OptSpec(
"compile_mode",
"Misc",
choices=["default", "reduce-overhead", "max-autotune"],
param_decls=("--compile", "compile_mode"),
help="torch.compile mode. Omit to disable.",
),
OptSpec(
"dry_run",
"Misc",
is_flag=True,
default=False,
param_decls=("--dry-run",),
help="Validate config and print plan, do not train.",
),
OptSpec(
"schedule_type",
"LR Schedule",
choices=_SCHEDULES,
default="cosine",
help="LR scheduler.",
),
OptSpec(
"min_rate",
"LR Schedule",
type=float,
default=None,
help="Minimum LR as fraction of base LR.",
),
OptSpec(
"cycle_length",
"LR Schedule",
type=int,
default=None,
help="SGDR first cycle length.",
),
OptSpec(
"t_mult",
"LR Schedule",
type=int,
default=2,
help="SGDR cycle length multiplier.",
),
OptSpec(
"stable_steps",
"LR Schedule",
type=int,
default=None,
help="WSD stable plateau steps.",
),
OptSpec(
"decay_steps",
"LR Schedule",
type=int,
default=None,
help="WSD decay steps.",
),
]
@click.command(
name="train",
cls=GroupedCommand,
help="Start model training (pretrain / SFT / DPO / GRPO).",
context_settings={"show_default": True},
)
@apply_specs(_SPECS, TrainConfig)
@click.pass_context
def train_command(ctx, config_path, dry_run, metrics, **kwargs):
"""Start model training (pretrain / SFT / DPO / GRPO)."""
kwargs["metrics"] = metrics
if config_path:
explicit_keys = {
key
for key in kwargs
if ctx.get_parameter_source(key) is ParameterSource.COMMANDLINE
}
kwargs = _merge_yaml_into_kwargs(config_path, kwargs, explicit_keys)
required = ["train_type", "data_root_path", "param_path"]
missing = [k for k in required if kwargs.get(k) is None]
if missing:
raise click.UsageError(
f"Missing required options: {', '.join(missing)}. "
f"Use --config YAML or provide them directly."
)
# Convert tuple back to list
kwargs["metrics"] = list(kwargs["metrics"])
# Remove tp_size (not yet wired)
kwargs.pop("tp_size", None)
if dry_run:
_print_dry_run(kwargs)
return
train(**kwargs)
def _print_dry_run(kwargs: dict) -> None:
"""Print training plan summary."""
rows = [
("Train type", kwargs.get("train_type")),
("Model path", kwargs.get("param_path")),
("Data path", kwargs.get("data_root_path")),
("Parallel mode", kwargs.get("parallel_mode", "none")),
("GPUs", str(kwargs.get("nprocs", 1))),
("Epochs", str(kwargs.get("n_epoch", 1))),
("Batch/device", str(kwargs.get("batch_per_device", 1))),
("Grad accum", str(kwargs.get("grad_accum_steps", 1))),
("Optimizer", str(kwargs.get("optimizer", "muon_adamw"))),
("Max LR", str(kwargs.get("max_lr", "?"))),
("Schedule", str(kwargs.get("schedule_type", "cosine"))),
("Warmup ratio", str(kwargs.get("warmup_ratio", 0.05))),
("Window size", str(kwargs.get("window_size", "config default"))),
("Checkpoint dir", str(kwargs.get("ckpt_dir", "checkpoint"))),
("Checkpoint interval", str(kwargs.get("ckpt_interval", 5000))),
("Resume", str(kwargs.get("resume", False))),
]
max_len = max(len(k) for k, _ in rows)
click.secho("\n=== Training Plan (dry-run) ===", fg="cyan", bold=True)
for key, val in rows:
click.echo(f" {key:<{max_len}s} : {val}")
click.secho("=" * 40, fg="cyan")
def create_model(config):
return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
def create_value_model(config):
return ValueModel(config).to(dtype=torch.bfloat16)
def create_optimizer(
model, optimizer_name: str = "muon_adamw", **kwargs
) -> optim.Optimizer:
return OptimizerFactory.create(optimizer_name, model, **kwargs)
def create_scheduler(
optimizer: optim.Optimizer, **kwargs
) -> optim.lr_scheduler.LRScheduler:
schedule_type = kwargs.pop("schedule_type")
return SchedulerFactory.create(schedule_type, optimizer, **kwargs)
def compute_total_steps(
dataset_len: int,
n_epoch: int,
batch_per_device: int,
nprocs: int,
grad_accum_steps: int,
) -> int:
def ceil_div(a: int, b: int) -> int:
return (a + b - 1) // b
samples_per_replica = ceil_div(dataset_len, nprocs)
batches_per_replica = ceil_div(samples_per_replica, batch_per_device)
total_steps = (batches_per_replica // grad_accum_steps) * n_epoch
return total_steps
def train(
train_type: str,
param_path: str,
data_root_path: str,
resume: bool,
n_epoch: int,
batch_per_device: int,
start_epoch: int,
start_samples: int,
grad_accum_steps: int,
warmup_ratio: float,
ckpt_interval: int,
ckpt_dir: str,
val_split: float,
val_step: int,
metrics: list[str],
max_grad_norm: float,
random_seed: int,
num_workers: int,
pin_memory: bool,
persistent_workers: bool,
gradient_checkpointing: bool,
window_size: int,
stride: int,
nprocs: int,
parallel_mode: str,
device_type: str,
backend: str,
master_addr: str,
master_port: str,
start_method: str,
neftune_alpha: float,
schedule_type: str,
min_rate: float,
cycle_length: int,
t_mult: int,
stable_steps: int,
decay_steps: int,
**kwargs,
):
if train_type not in _TRAIN_TYPE:
raise ValueError(
f"Invalid train_type '{train_type}'. "
f"Must be one of: {', '.join(_TRAIN_TYPE)}"
)
if not os.path.exists(param_path):
raise FileNotFoundError(f"Model directory not found: {param_path}")
if nprocs > 1 and parallel_mode == "none":
raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
# Load config
config_path = os.path.join(param_path, "config.json")
config = AutoRegressiveLMConfig.from_file(config_path)
config.neftune_alpha = neftune_alpha
if window_size is None:
window_size = config.max_position_embeddings
strategy_kwargs = {
"beta": kwargs.pop("dpo_beta"),
"label_smoothing": kwargs.pop("label_smoothing"),
"clip_eps": kwargs.pop("grpo_clip_eps"),
"kl_coef": kwargs.pop("grpo_kl_coef"),
"group_size": kwargs.pop("group_size"),
"gamma": kwargs.pop("ppo_gamma"),
"gae_lambda": kwargs.pop("ppo_gae_lambda"),
"vf_coef": kwargs.pop("ppo_vf_coef"),
}
rollout_interval = kwargs.pop("rollout_interval", 512)
rollout_max_policy_lag = kwargs.pop("rollout_max_policy_lag", None)
rollout_temperature = kwargs.pop("rollout_temperature", 0.7)
rollout_top_k = kwargs.pop("rollout_top_k", 0)
rollout_top_p = kwargs.pop("rollout_top_p", 0.9)
rollout_max_tokens = kwargs.pop("rollout_max_tokens", 1024)
reward_model_fn: Callable[[], BaseRewardModel] | None = None
critic_model_fn = None
if train_type == "online_ppo":
# The optimizer defaults to the policy's; critic_optimizer_fn can
# override it in the TrainConfig.
critic_model_fn = partial(create_value_model, config)
executor_kwargs = {}
if parallel_mode == "ddp":
executor_kwargs.update(
gradient_as_bucket_view=True,
broadcast_buffers=False,
)
model_fn = partial(create_model, config)
dataset = DatasetFactory.load(
train_type=train_type,
load_path=data_root_path,
window_size=window_size,
stride=stride,
tokenizer_path=param_path,
)
optimizer_name = kwargs.pop("optimizer", "muon_adamw")
optimizer_kwargs = {
"lr": kwargs.pop("max_lr"),
"weight_decay": kwargs.pop("weight_decay"),
"nora_lr": kwargs.pop("nora_lr", 5e-3),
"nora_beta": kwargs.pop("nora_beta", 0.95),
"nora_momentum": kwargs.pop("nora_momentum", 0.95),
"nora_weight_decay": kwargs.pop("nora_weight_decay", 0.0),
"momentum": kwargs.pop("muon_momentum", 0.95),
"nesterov": kwargs.pop("muon_nesterov", True),
"ns_steps": kwargs.pop("muon_ns_steps", 5),
"adjust_lr_fn": kwargs.pop("muon_adjust_lr", "match_rms_adamw"),
"mano_momentum": kwargs.pop("mano_momentum", 0.95),
"mano_nesterov": kwargs.pop("mano_nesterov", True),
}
optimizer_fn = partial(
create_optimizer,
optimizer_name=optimizer_name,
**optimizer_kwargs,
)
if optimizer_name == "nora_nadamw":
optimizer_hyperparameters = {
key: optimizer_kwargs[key]
for key in (
"lr",
"weight_decay",
"nora_lr",
"nora_beta",
"nora_momentum",
"nora_weight_decay",
)
}
optimizer_hyperparameters.update(
{"nadamw_betas": [0.9, 0.999], "nadamw_eps": 1e-8, "nora_eps": 1e-10}
)
elif optimizer_name == "mano_adamw":
optimizer_hyperparameters = {
key: optimizer_kwargs[key]
for key in ("lr", "weight_decay", "mano_momentum", "mano_nesterov")
}
optimizer_hyperparameters.update(
{"adamw_betas": [0.9, 0.95], "adamw_eps": 1e-8}
)
else:
optimizer_hyperparameters = {
key: optimizer_kwargs[key]
for key in (
"lr",
"weight_decay",
"momentum",
"nesterov",
"ns_steps",
"adjust_lr_fn",
)
}
total_steps = compute_total_steps(
len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
)
warmup_steps = int(warmup_ratio * total_steps)
warmup_steps = min(warmup_steps, total_steps)
scheduler_kwargs = {"warmup_steps": warmup_steps}
if schedule_type == "cosine":
scheduler_kwargs["lr_decay_steps"] = total_steps - warmup_steps
elif schedule_type == "sgdr":
scheduler_kwargs["cycle_length"] = cycle_length or (total_steps - warmup_steps)
scheduler_kwargs["t_mult"] = t_mult
elif schedule_type == "wsd":
remaining = total_steps - warmup_steps
stable_steps_ = stable_steps or max(1, int(remaining * 0.8))
scheduler_kwargs["stable_steps"] = stable_steps_
scheduler_kwargs["decay_steps"] = max(
1, decay_steps or (remaining - stable_steps_)
)
if min_rate is not None:
scheduler_kwargs["min_rate"] = min_rate
scheduler_fn = partial(
create_scheduler,
schedule_type=schedule_type,
**scheduler_kwargs,
)
grad_ckpt_modules = [DecoderBlock] if gradient_checkpointing else []
compile_mode = kwargs.pop("compile_mode", None)
collate_fn = None
if train_type == "dpo":
collate_fn = dpo_collate_fn
elif train_type == "grpo":
collate_fn = grpo_collate_fn
elif train_type in ("online_grpo", "online_dpo", "online_ppo"):
collate_fn = None
train_config = TrainConfig(
model_fn=model_fn,
strategy=train_type,
dataset=dataset,
optimizer_fn=optimizer_fn,
scheduler_fn=scheduler_fn,
optimizer_name=optimizer_name,
optimizer_hyperparameters=optimizer_hyperparameters,
ckpt_dir=ckpt_dir,
n_epoch=n_epoch,
batch_per_device=batch_per_device,
start_epoch=start_epoch,
start_samples=start_samples,
ckpt_interval=ckpt_interval,
grad_accum_steps=grad_accum_steps,
max_grad_norm=max_grad_norm,
random_seed=random_seed,
num_workers=num_workers,
pin_memory=pin_memory,
persistent_workers=persistent_workers,
nprocs=nprocs,
backend=backend,
master_addr=master_addr,
master_port=master_port,
parallel_mode=parallel_mode,
device_type=device_type,
start_method=start_method,
val_split=val_split,
val_step=val_step,
metrics=metrics,
gradient_checkpointing_modules=grad_ckpt_modules,
compile_mode=compile_mode,
executor_kwargs=executor_kwargs,
strategy_kwargs=strategy_kwargs,
neftune_alpha=neftune_alpha,
collate_fn=collate_fn,
rollout_interval=rollout_interval,
rollout_max_policy_lag=rollout_max_policy_lag,
rollout_temperature=rollout_temperature,
rollout_top_k=rollout_top_k,
rollout_top_p=rollout_top_p,
rollout_max_tokens=rollout_max_tokens,
reward_model_fn=reward_model_fn,
critic_model_fn=critic_model_fn,
moe_aux_loss_coef=kwargs.pop("moe_aux_loss_coef", 0.01),
)
trainer = Trainer(train_config)
trainer.train(param_path=param_path, resume=resume)
if __name__ == "__main__":
train_command()
Reference in New Issue View Git Blame Copy Permalink