ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
76f1c10feb8dd68ac7c60b2a8150f458a6b69b50
AstrAI/scripts/tools/train.py
T
ViperEkura 76f1c10feb refactor: generate train and serve CLIs from config-backed option specs 
- Add astrai/config/cli.py: OptSpec tables plus apply_specs infer click types and defaults from config fields, covering Optional[X], Union[X, None], PEP 604 X | None, stringified PEP 563 annotations, bool flag pairs, and repeatable list options
- Move GroupedCommand/GroupedOption and the three-layer YAML merge (option defaults < YAML < explicit CLI) into the config package, adding unknown-key warning and mapping validation
- Replace ~420 lines of hand-written @opt decorators in scripts/tools/train.py with a 66-entry spec table; option names, defaults, flag styles, and YAML semantics verified unchanged
- Migrate scripts/tools/server.py to the same mechanism with its section binding, integer coercion, and dtype validation preserved locally
- Add tests/config/test_cli.py covering type inference across annotation styles, default overrides, flag pairs, merge precedence, scientific notation, and help ordering
2026-09-03 20:56:11 +08:00

680 lines
20 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
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(
"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_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"),
}
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
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"):
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,
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