diff --git a/astrai/config/__init__.py b/astrai/config/__init__.py index 1b8cada..8134526 100644 --- a/astrai/config/__init__.py +++ b/astrai/config/__init__.py @@ -1,3 +1,11 @@ +from astrai.config.cli import ( + GroupedCommand, + GroupedOption, + OptSpec, + apply_specs, + merge_yaml_into_kwargs, + opt, +) from astrai.config.model_config import ( AutoRegressiveLMConfig, BaseModelConfig, @@ -22,4 +30,10 @@ __all__ = [ "OutputConfig", "PipelineConfig", "ProcessingConfig", + "GroupedCommand", + "GroupedOption", + "OptSpec", + "apply_specs", + "merge_yaml_into_kwargs", + "opt", ] diff --git a/astrai/config/cli.py b/astrai/config/cli.py new file mode 100644 index 0000000..196e061 --- /dev/null +++ b/astrai/config/cli.py @@ -0,0 +1,315 @@ +"""Generate grouped click CLIs from pydantic config fields. + +A config class alone does not make a CLI: some options need defaults that +differ from the config defaults (e.g. ``num_workers``), choices come from +factory registries or frozenset validators, and values merge across three +layers (option defaults -> YAML -> explicit CLI flags). ``OptSpec`` records +those overrides in a declarative table; types and defaults are inferred +from the backing config field wherever the spec leaves them ``AUTO``. +""" + +import dataclasses +import re +import types +import typing as t +from collections import OrderedDict +from collections.abc import Sequence + +import click +import yaml + + +class GroupedOption(click.Option): + """A ``click.Option`` that carries a ``group`` label for help output.""" + + def __init__(self, *args, group: str = "Options", **kwargs): + super().__init__(*args, **kwargs) + self.group = group + + +class GroupedCommand(click.Command): + """A ``click.Command`` that renders options grouped by their ``group``.""" + + def format_options(self, ctx, formatter): + groups: OrderedDict[str, list] = OrderedDict() + for param in self.get_params(ctx): + record = param.get_help_record(ctx) + if record is None: + continue + group = getattr(param, "group", "Options") + groups.setdefault(group, []).append(record) + for group_name, records in groups.items(): + with formatter.section(group_name): + formatter.write_dl(records) + + +def opt(*param_decls, group: str, **kwargs): + """Shorthand for ``click.option`` that tags the option with a group.""" + kwargs.setdefault("cls", GroupedOption) + kwargs["group"] = group + return click.option(*param_decls, **kwargs) + + +class _Auto: + def __repr__(self) -> str: + return "AUTO" + + +AUTO = _Auto() + +_YAML_FLOAT_PATTERN = re.compile( + r"""^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)? + |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+) + |[-+]?\.(?:inf|Inf|INF) + |\.(?:nan|NaN|NAN))$""", + re.X, +) + +DEFAULT_YAML_SECTIONS = ("model", "data", "parallel", "training", "ckpt", "log") + + +def _enable_yaml12_floats() -> None: + """PyYAML implements YAML 1.1, where ``2e-5`` parses as a string; switch its + float resolver to the YAML 1.2 core schema so scientific notation works.""" + yaml.SafeLoader.add_implicit_resolver( + "tag:yaml.org,2002:float", _YAML_FLOAT_PATTERN, list("-+0123456789.") + ) + + +def merge_yaml_into_kwargs( + config_path: str, + passed_kwargs: dict, + explicit_keys: set[str] | None = None, + sections: Sequence[str] = DEFAULT_YAML_SECTIONS, + allowed_keys: Sequence[str] | None = None, +) -> dict: + """Merge option defaults, YAML values, then explicit CLI values. + + ``sections`` selects which top-level YAML mappings feed the flat kwargs + namespace. ``allowed_keys`` optionally restricts the accepted keys across + those sections: unknown keys are warned about once and dropped. + """ + _enable_yaml12_floats() + + with open(config_path) as f: + cfg = yaml.safe_load(f) or {} + if not isinstance(cfg, dict): + raise click.UsageError(f"config must be a mapping: {config_path}") + + merged = dict(passed_kwargs) + seen: set[str] = set() + for section in sections: + values = cfg.get(section) or {} + if not isinstance(values, dict): + raise click.UsageError(f"top-level {section} section must be a mapping") + if allowed_keys is None: + merged.update(values) + seen.update(values) + else: + merged.update({k: v for k, v in values.items() if k in allowed_keys}) + seen.update(values) + + if allowed_keys is not None: + unknown = sorted(seen - set(allowed_keys)) + if unknown: + click.echo( + f"Warning: ignoring unknown config keys: {', '.join(unknown)}", + err=True, + ) + + if explicit_keys is None: + explicit_keys = set(passed_kwargs) + for key in explicit_keys: + if key in passed_kwargs: + merged[key] = passed_kwargs[key] + + return merged + + +@dataclasses.dataclass(frozen=True) +class OptSpec: + """One CLI option, optionally backed by a field of a config class. + + ``type``/``default`` left as ``AUTO`` are inferred from the backing + config field (bool becomes a ``--x/--no-x`` pair, lists become + repeatable options). Standalone specs for CLI-only options must carry + ``type`` or ``default`` explicitly. ``choices`` overrides the inferred + type with ``click.Choice``. + """ + + name: str + group: str + type: t.Any = AUTO + default: t.Any = AUTO + help: str | None = None + choices: Sequence[str] | None = None + multiple: bool = False + is_flag: bool = False + required: bool = False + param_decls: tuple[str, ...] = () + + +def _resolve_hints(config_cls: type) -> dict[str, t.Any]: + """Resolve field annotations with ``typing.get_type_hints``. + + Raw ``Field.type`` stays a string under PEP 563 (``from __future__ + import annotations``) and never mentions ``types.UnionType``; resolved + hints cover old-style ``Optional[X]``/``Union[X, None]``, new-style + ``X | None``, and stringified forward references alike. + """ + try: + return t.get_type_hints(config_cls) + except Exception: + return {} + + +def _unwrap_optional(annotation: t.Any) -> t.Any | None: + """Return the single non-None member of an optional annotation.""" + if annotation is None or isinstance(annotation, str): + return None + origin = t.get_origin(annotation) + if origin is not t.Union and origin is not types.UnionType: + return None + args = [a for a in t.get_args(annotation) if a is not type(None)] + if len(args) == 1: + return args[0] + return None + + +def _annotation( + spec: OptSpec, + field: dataclasses.Field | None, + hints: dict[str, t.Any] | None = None, +) -> t.Any: + if field is not None: + if hints and spec.name in hints: + return hints[spec.name] + return field.type + if spec.type is not AUTO: + return spec.type + return None + + +def _click_type( + spec: OptSpec, + field: dataclasses.Field | None, + hints: dict[str, t.Any] | None = None, + annotation: t.Any = None, +) -> t.Any: + if spec.choices is not None: + return click.Choice(list(spec.choices)) + if annotation is None: + if spec.type is not AUTO and spec.type is not bool: + return spec.type + annotation = _annotation(spec, field, hints) + inner = _unwrap_optional(annotation) + if inner is not None: + return _click_type(spec, field, hints, annotation=inner) + origin = t.get_origin(annotation) + if annotation is bool or spec.type is bool: + return click.BOOL + if annotation is int: + return click.INT + if annotation is float: + return click.FLOAT + if annotation is str: + return click.STRING + if origin in (list, tuple) or annotation in (list, tuple): + return click.STRING + raise TypeError( + f"cannot infer a click type for option {spec.name!r}; set OptSpec.type" + ) + + +def _resolve_default(spec: OptSpec, field: dataclasses.Field | None) -> t.Any: + if spec.default is not AUTO: + return spec.default + if field is not None: + if field.default is not dataclasses.MISSING: + return field.default + if field.default_factory is not dataclasses.MISSING: # type: ignore[misc] + return field.default_factory() # type: ignore[misc] + return None + + +def _is_flag_pair( + spec: OptSpec, + field: dataclasses.Field | None, + hints: dict[str, t.Any] | None = None, +) -> bool: + if spec.is_flag: + return False + annotation = _annotation(spec, field, hints) + return annotation is bool + + +def option_from_spec( + spec: OptSpec, + fields_by_name: dict[str, dataclasses.Field] | None = None, + hints: dict[str, t.Any] | None = None, +) -> t.Callable: + """Build a grouped ``click.option`` decorator from one spec.""" + fields_by_name = fields_by_name or {} + hints = hints or {} + field = fields_by_name.get(spec.name) + + kwargs: dict[str, t.Any] = { + "cls": GroupedOption, + "group": spec.group, + } + if spec.help is not None: + kwargs["help"] = spec.help + if spec.required: + kwargs["required"] = True + + if spec.is_flag: + decls = spec.param_decls or (f"--{spec.name}",) + kwargs["is_flag"] = True + kwargs["default"] = _resolve_default(spec, field) + elif _is_flag_pair(spec, field, hints): + if spec.param_decls: + raise ValueError( + f"flag pair {spec.name!r} does not support custom param_decls" + ) + decls = (f"--{spec.name}/--no-{spec.name}",) + kwargs["type"] = click.BOOL + kwargs["default"] = _resolve_default(spec, field) + else: + decls = spec.param_decls or (f"--{spec.name}",) + kwargs["type"] = _click_type(spec, field, hints) + kwargs["default"] = _resolve_default(spec, field) + multiple = spec.multiple + if not multiple and field is not None: + annotation = _annotation(spec, field, hints) + origin = t.get_origin(annotation) + multiple = origin in (list, tuple) or annotation in (list, tuple) + if multiple: + kwargs["multiple"] = True + if isinstance(kwargs["default"], list): + kwargs["default"] = tuple(kwargs["default"]) + return opt(*decls, **kwargs) + + +def apply_specs( + specs: Sequence[OptSpec], + config_cls: type | None = None, +) -> t.Callable: + """Apply a table of specs as click options, in table order. + + Options render top-to-bottom in the command's help output following the + table order (click reverses decorator application, so specs are applied + reversed). ``config_cls`` supplies type/default inference for specs whose + name matches one of its fields. + """ + fields_by_name: dict[str, dataclasses.Field] = {} + hints: dict[str, t.Any] = {} + if config_cls is not None: + fields_by_name = {f.name: f for f in dataclasses.fields(config_cls)} + hints = _resolve_hints(config_cls) + + def decorator(func): + for spec in reversed(specs): + func = option_from_spec(spec, fields_by_name, hints)(func) + return func + + return decorator diff --git a/scripts/tools/server.py b/scripts/tools/server.py index 82b537a..b4a4fb9 100644 --- a/scripts/tools/server.py +++ b/scripts/tools/server.py @@ -2,9 +2,14 @@ from pathlib import Path import click import torch -import yaml from click.core import ParameterSource +from astrai.config.cli import ( + GroupedCommand, + OptSpec, + apply_specs, + merge_yaml_into_kwargs, +) from astrai.inference import run_server _DTYPES = ["bfloat16", "float16", "float32"] @@ -26,29 +31,65 @@ def _merge_yaml_into_kwargs( explicit_keys: set[str] | None = None, ) -> dict: """Merge Click defaults, YAML server values, then explicit CLI values.""" - with open(config_path, encoding="utf-8") as file: - config = yaml.safe_load(file) or {} - if not isinstance(config, dict): - raise click.UsageError(f"Serving config must be a mapping: {config_path}") - server = config.get("server") or {} - if not isinstance(server, dict): - raise click.UsageError("top-level server section must be a mapping") + return merge_yaml_into_kwargs( + config_path, + passed_kwargs, + explicit_keys, + sections=("server",), + allowed_keys=_SERVER_KEYS, + ) - unknown = sorted(set(server) - set(_SERVER_KEYS)) - if unknown: - click.echo( - f"Warning: ignoring unknown server config keys: {', '.join(unknown)}", - err=True, - ) - merged = dict(passed_kwargs) - merged.update({key: server[key] for key in _SERVER_KEYS if key in server}) - if explicit_keys is None: - explicit_keys = set(passed_kwargs) - for key in explicit_keys: - if key in passed_kwargs: - merged[key] = passed_kwargs[key] - return merged +_SPECS = [ + OptSpec( + "config_path", + "Server", + type=click.Path(exists=True, dir_okay=False), + param_decls=("--config", "-c", "config_path"), + help="Serving YAML config. CLI flags override YAML values.", + ), + OptSpec("host", "Server", type=str, default="0.0.0.0", help="Host address."), + OptSpec("port", "Server", type=int, default=8000, help="Port number."), + OptSpec( + "reload", + "Server", + is_flag=True, + default=False, + help="Enable auto-reload for development.", + ), + OptSpec( + "param_path", + "Model", + type=click.Path(exists=True), + default=None, + help="Path to model parameters.", + ), + OptSpec( + "device", "Model", type=str, default="cuda", help="Device to load model on." + ), + OptSpec( + "dtype", + "Model", + choices=_DTYPES, + default="bfloat16", + help="Data type for model weights.", + ), + OptSpec( + "max_batch_size", + "Performance", + type=int, + default=16, + help="Maximum batch size for continuous batching.", + ), + OptSpec( + "max_seq_len", + "Performance", + type=int, + default=None, + help="Maximum sequence length (KV cache size + prompt truncation). " + "Uses model config if not set.", + ), +] def _as_int(value, name: str) -> int | None: @@ -87,84 +128,22 @@ def _resolve_server_config( return resolved -@click.command(name="serve", help="Launch inference server (OpenAI-compatible API).") -@click.option( - "--config", - "-c", - "config_path", - type=click.Path(exists=True, dir_okay=False), - default=None, - help="Serving YAML config. CLI flags override YAML values.", -) -@click.option("--host", default="0.0.0.0", help="Host address.") -@click.option("--port", type=int, default=8000, help="Port number.") -@click.option( - "--reload", is_flag=True, default=False, help="Enable auto-reload for development." -) -@click.option( - "--param_path", - type=click.Path(exists=True), - default=None, - help="Path to model parameters.", -) -@click.option("--device", default="cuda", help="Device to load model on.") -@click.option( - "--dtype", - type=click.Choice(_DTYPES), - default="bfloat16", - help="Data type for model weights.", -) -@click.option( - "--max_batch_size", - type=int, - default=16, - help="Maximum batch size for continuous batching.", -) -@click.option( - "--max_seq_len", - type=int, - default=None, - help="Maximum sequence length (KV cache size + prompt truncation). Uses model config if not set.", +@click.command( + name="serve", + cls=GroupedCommand, + help="Launch inference server (OpenAI-compatible API).", ) +@apply_specs(_SPECS) @click.pass_context -def server_command( - ctx, - config_path, - host, - port, - reload, - param_path, - device, - dtype, - max_batch_size, - max_seq_len, -): +def server_command(ctx, config_path, **kwargs): """Launch inference server (OpenAI-compatible API).""" if config_path: - passed_kwargs = { - "host": host, - "port": port, - "reload": reload, - "param_path": param_path, - "device": device, - "dtype": dtype, - "max_batch_size": max_batch_size, - "max_seq_len": max_seq_len, - } explicit_keys = { key - for key in passed_kwargs + for key in kwargs if ctx.get_parameter_source(key) is ParameterSource.COMMANDLINE } - resolved = _resolve_server_config(config_path, passed_kwargs, explicit_keys) - host = resolved["host"] - port = resolved["port"] - reload = resolved["reload"] - param_path = resolved["param_path"] - device = resolved["device"] - dtype = resolved["dtype"] - max_batch_size = resolved["max_batch_size"] - max_seq_len = resolved["max_seq_len"] + kwargs = _resolve_server_config(config_path, kwargs, explicit_keys) click.echo(f"Config: {config_path}") dtype_map = { @@ -173,19 +152,22 @@ def server_command( "float32": torch.float32, } project_root = Path(__file__).parent.parent.parent - param_path = param_path or str(project_root / "params") + kwargs["param_path"] = kwargs.get("param_path") or str(project_root / "params") - click.echo(f"Starting server on http://{host}:{port}") - click.echo(f"Model: {param_path} | Device: {device} | Dtype: {dtype}") + click.echo(f"Starting server on http://{kwargs['host']}:{kwargs['port']}") + click.echo( + f"Model: {kwargs['param_path']} | " + f"Device: {kwargs['device']} | Dtype: {kwargs['dtype']}" + ) run_server( - host=host, - port=port, - reload=reload, - device=device, - dtype=dtype_map[dtype], - param_path=Path(param_path), - max_batch_size=max_batch_size, - max_seq_len=max_seq_len, + host=kwargs["host"], + port=kwargs["port"], + reload=kwargs["reload"], + device=kwargs["device"], + dtype=dtype_map[kwargs["dtype"]], + param_path=Path(kwargs["param_path"]), + max_batch_size=kwargs["max_batch_size"], + max_seq_len=kwargs["max_seq_len"], ) diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 2f15631..b0d6759 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -1,16 +1,19 @@ import os -import re -from collections import OrderedDict from collections.abc import Callable from functools import partial import click import torch -import yaml 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, @@ -24,79 +27,8 @@ from astrai.optim import OptimizerFactory from astrai.trainer import SchedulerFactory, Trainer from astrai.trainer.rollout import BaseRewardModel - -class GroupedOption(click.Option): - """A ``click.Option`` that carries a ``group`` label for help output.""" - - def __init__(self, *args, group: str = "Options", **kwargs): - super().__init__(*args, **kwargs) - self.group = group - - -class GroupedCommand(click.Command): - """A ``click.Command`` that renders options grouped by their ``group``.""" - - def format_options(self, ctx, formatter): - groups: OrderedDict[str, list] = OrderedDict() - for param in self.get_params(ctx): - record = param.get_help_record(ctx) - if record is None: - continue - group = getattr(param, "group", "Options") - groups.setdefault(group, []).append(record) - for group_name, records in groups.items(): - with formatter.section(group_name): - formatter.write_dl(records) - - -def opt(*param_decls, group: str, **kwargs): - """Shorthand for ``click.option`` that tags the option with a group.""" - kwargs.setdefault("cls", GroupedOption) - kwargs["group"] = group - return click.option(*param_decls, **kwargs) - - -_YAML_FLOAT_PATTERN = re.compile( - r"""^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)? - |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+) - |[-+]?\.(?:inf|Inf|INF) - |\.(?:nan|NaN|NAN))$""", - re.X, -) - - -def _enable_yaml12_floats() -> None: - """PyYAML implements YAML 1.1, where ``2e-5`` parses as a string; switch its - float resolver to the YAML 1.2 core schema so scientific notation works.""" - yaml.SafeLoader.add_implicit_resolver( - "tag:yaml.org,2002:float", _YAML_FLOAT_PATTERN, list("-+0123456789.") - ) - - -def _merge_yaml_into_kwargs( - config_path: str, - passed_kwargs: dict, - explicit_keys: set[str] | None = None, -) -> dict: - """Merge Click defaults, YAML values, then explicit CLI values.""" - _enable_yaml12_floats() - - with open(config_path) as f: - cfg = yaml.safe_load(f) or {} - - merged = dict(passed_kwargs) - for section in ("model", "data", "parallel", "training", "ckpt", "log"): - if section in cfg: - merged.update(cfg[section]) - - if explicit_keys is None: - explicit_keys = set(passed_kwargs) - for key in explicit_keys: - if key in passed_kwargs: - merged[key] = passed_kwargs[key] - - return merged - +# 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) @@ -105,6 +37,306 @@ _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", @@ -112,425 +344,7 @@ _START_METHODS = sorted(START_METHODS) help="Start model training (pretrain / SFT / DPO / GRPO).", context_settings={"show_default": True}, ) -@opt( - "--config", - "-c", - "config_path", - type=click.Path(exists=True), - group="Paths & Setup", - help="YAML config file. CLI flags override YAML values.", -) -@opt( - "--train_type", - type=click.Choice(_TRAIN_TYPE), - required=False, - group="Paths & Setup", - help="Training type.", -) -@opt( - "--data_root_path", - type=click.Path(exists=True), - group="Paths & Setup", - help="Root directory of the dataset.", -) -@opt( - "--param_path", - type=click.Path(exists=True), - group="Paths & Setup", - help="Path to model parameters or resume checkpoint.", -) -@opt( - "--resume", - is_flag=True, - default=False, - group="Paths & Setup", - help="Resume from checkpoint.", -) -@opt("--n_epoch", type=int, default=1, group="Training", help="Number of epochs.") -@opt( - "--batch_per_device", - type=int, - default=1, - group="Training", - help="Batch size per GPU.", -) -@opt( - "--grad_accum_steps", - type=int, - default=1, - group="Training", - help="Gradient accumulation steps.", -) -@opt( - "--warmup_ratio", - type=float, - default=0.05, - group="LR Schedule", - help="Fraction of total steps for LR warmup.", -) -@opt( - "--max_lr", - type=float, - default=3e-4, - group="Optimizer", - help="Max learning rate.", -) -@opt( - "--optimizer", - type=click.Choice(_OPTIMIZERS), - default="muon_adamw", - group="Optimizer", - help="Built-in optimizer.", -) -@opt( - "--max_grad_norm", - type=float, - default=1.0, - group="Training", - help="Max gradient norm for clipping.", -) -@opt( - "--weight_decay", - type=float, - default=0.1, - group="Optimizer", - help="Weight decay for eligible optimizer parameters.", -) -@opt( - "--nora_lr", type=float, default=5e-3, group="Optimizer", help="Nora learning rate." -) -@opt( - "--nora_beta", type=float, default=0.95, group="Optimizer", help="Nora EMA factor." -) -@opt( - "--nora_momentum", - type=float, - default=0.95, - group="Optimizer", - help="Nora update momentum.", -) -@opt( - "--nora_weight_decay", - type=float, - default=0.0, - group="Optimizer", - help="Nora weight decay.", -) -@opt( - "--muon_momentum", - type=float, - default=0.95, - group="Optimizer", - help="Muon momentum factor.", -) -@opt( - "--muon_nesterov/--no-muon_nesterov", - default=True, - group="Optimizer", - help="Muon Nesterov.", -) -@opt( - "--muon_ns_steps", - type=int, - default=5, - group="Optimizer", - help="Muon Newton-Schulz steps.", -) -@opt( - "--muon_adjust_lr", - type=click.Choice(["original", "match_rms_adamw"]), - default="match_rms_adamw", - group="Optimizer", - help="Muon LR adjustment strategy.", -) -@opt( - "--mano_momentum", - type=float, - default=0.95, - group="Optimizer", - help="Mano momentum factor.", -) -@opt( - "--mano_nesterov/--no-mano_nesterov", - default=True, - group="Optimizer", - help="Mano Nesterov momentum.", -) -@opt( - "--random_seed", - type=int, - default=3407, - group="Data Loading", - help="Random seed.", -) -@opt( - "--num_workers", - type=int, - default=4, - group="Data Loading", - help="DataLoader workers.", -) -@opt( - "--pin_memory/--no-pin_memory", - default=True, - group="Data Loading", - help="Pin memory.", -) -@opt( - "--persistent_workers/--no-persistent_workers", - default=True, - group="Data Loading", - help="Keep DataLoader workers alive between epochs.", -) -@opt( - "--window_size", - type=int, - default=None, - group="Data Loading", - help="Max input sequence length.", -) -@opt( - "--stride", - type=int, - default=None, - group="Data Loading", - help="Step size for sliding window.", -) -@opt("--dpo_beta", type=float, default=0.1, group="Algorithm", help="DPO beta.") -@opt("--group_size", type=int, default=4, group="Algorithm", help="GRPO group size.") -@opt( - "--grpo_clip_eps", - type=float, - default=0.2, - group="Algorithm", - help="GRPO clip epsilon.", -) -@opt( - "--grpo_kl_coef", - type=float, - default=0.01, - group="Algorithm", - help="GRPO KL penalty coefficient.", -) -@opt( - "--label_smoothing", - type=float, - default=0.0, - group="Data Loading", - help="Label smoothing.", -) -@opt( - "--moe_aux_loss_coef", - type=float, - default=0.01, - group="Algorithm", - help="MoE load balancing auxiliary loss coefficient (0=disable).", -) -@opt( - "--rollout_interval", - type=int, - default=512, - group="Algorithm", - help="Steps between rollouts.", -) -@opt( - "--rollout_max_policy_lag", - type=int, - default=None, - group="Algorithm", - help="Maximum accepted rollout/live policy-version gap.", -) -@opt( - "--rollout_temperature", - type=float, - default=0.7, - group="Algorithm", - help="Rollout temperature.", -) -@opt( - "--rollout_top_k", - type=int, - default=0, - group="Algorithm", - help="Rollout top-k (0=disable).", -) -@opt( - "--rollout_top_p", - type=float, - default=0.9, - group="Algorithm", - help="Rollout top-p.", -) -@opt( - "--rollout_max_tokens", - type=int, - default=1024, - group="Algorithm", - help="Max tokens per rollout response.", -) -@opt( - "--gradient_checkpointing/--no-gradient_checkpointing", - default=False, - group="Misc", - help="Enable activation checkpointing.", -) -@opt( - "--compile", - "compile_mode", - type=click.Choice(["default", "reduce-overhead", "max-autotune"]), - default=None, - group="Misc", - help="torch.compile mode. Omit to disable.", -) -@opt( - "--ckpt_interval", - type=int, - default=5000, - group="Checkpoint", - help="Steps between checkpoints.", -) -@opt( - "--ckpt_dir", - type=click.Path(), - default="checkpoint", - group="Checkpoint", - help="Checkpoint directory.", -) -@opt( - "--val_split", - type=float, - default=None, - group="Validation", - help="Validation split ratio.", -) -@opt( - "--val_step", - type=int, - default=1000, - group="Validation", - help="Steps between validation runs.", -) -@opt( - "--metrics", - multiple=True, - default=("loss", "lr", "grad_norm", "grad_snr"), - group="Validation", - help="Metrics to log (repeatable).", -) -@opt("--start_epoch", type=int, default=0, group="Checkpoint", help="Start epoch.") -@opt( - "--start_samples", - type=int, - default=0, - group="Checkpoint", - help="Start samples (per rank).", -) -@opt( - "--master_addr", - type=str, - default="localhost", - group="Distributed", - help="Master node address.", -) -@opt( - "--master_port", - type=str, - default="29500", - group="Distributed", - help="Master node port.", -) -@opt( - "--backend", - type=click.Choice(_BACKENDS), - default="nccl", - group="Distributed", - help="Distributed backend.", -) -@opt("--nprocs", type=int, default=1, group="Distributed", help="Number of GPUs.") -@opt( - "--parallel_mode", - type=click.Choice(_PARALLEL), - default="fsdp", - group="Distributed", - help="Parallel strategy.", -) -@opt( - "--device_type", - type=str, - default="cuda", - group="Distributed", - help="Device type.", -) -@opt( - "--start_method", - type=click.Choice(_START_METHODS), - default="spawn", - group="Distributed", - help="Multiprocessing start method.", -) -@opt( - "--neftune_alpha", - type=float, - default=0.0, - group="Algorithm", - help="NEFTune noise alpha.", -) -@opt( - "--schedule_type", - type=click.Choice(_SCHEDULES), - default="cosine", - group="LR Schedule", - help="LR scheduler.", -) -@opt( - "--min_rate", - type=float, - default=None, - group="LR Schedule", - help="Minimum LR as fraction of base LR.", -) -@opt( - "--cycle_length", - type=int, - default=None, - group="LR Schedule", - help="SGDR first cycle length.", -) -@opt( - "--t_mult", - type=int, - default=2, - group="LR Schedule", - help="SGDR cycle length multiplier.", -) -@opt( - "--stable_steps", - type=int, - default=None, - group="LR Schedule", - help="WSD stable plateau steps.", -) -@opt( - "--decay_steps", - type=int, - default=None, - group="LR Schedule", - help="WSD decay steps.", -) -@opt( - "--tp_size", - type=int, - default=None, - group="Distributed", - help="Tensor parallelism (future).", -) -@opt( - "--dry-run", - is_flag=True, - default=False, - group="Misc", - help="Validate config and print plan, do not train.", -) +@apply_specs(_SPECS, TrainConfig) @click.pass_context def train_command(ctx, config_path, dry_run, metrics, **kwargs): """Start model training (pretrain / SFT / DPO / GRPO).""" diff --git a/tests/config/test_cli.py b/tests/config/test_cli.py new file mode 100644 index 0000000..a438d39 --- /dev/null +++ b/tests/config/test_cli.py @@ -0,0 +1,230 @@ +import dataclasses +import typing as t + +import click +import pytest +from click.testing import CliRunner + +from astrai.config import TrainConfig, merge_yaml_into_kwargs +from astrai.config.cli import ( + GroupedCommand, + OptSpec, + apply_specs, + option_from_spec, +) + + +def test_merge_yaml_overrides_defaults_but_not_explicit_cli(tmp_path): + config_path = tmp_path / "train.yaml" + config_path.write_text( + "training:\n" + " optimizer: nora_nadamw\n" + " max_lr: 2e-4\n" + " nora_lr: 0.004\n" + " batch_per_device: 8\n", + encoding="utf-8", + ) + click_values = { + "optimizer": "nora_nadamw", + "max_lr": 3e-4, + "nora_lr": 5e-3, + "batch_per_device": 16, + } + + merged = merge_yaml_into_kwargs( + str(config_path), click_values, explicit_keys={"batch_per_device"} + ) + + assert merged["max_lr"] == 2e-4 + assert merged["nora_lr"] == 4e-3 + assert merged["batch_per_device"] == 16 + + +def test_merge_yaml_parses_scientific_notation_as_float(tmp_path): + config_path = tmp_path / "train.yaml" + config_path.write_text("training:\n max_lr: 2e-5\n", encoding="utf-8") + + merged = merge_yaml_into_kwargs( + str(config_path), {"max_lr": 3e-4}, explicit_keys=set() + ) + + assert merged["max_lr"] == 2e-5 + assert isinstance(merged["max_lr"], float) + + +def test_merge_yaml_ignores_unknown_sections(tmp_path): + config_path = tmp_path / "train.yaml" + config_path.write_text("unknown:\n max_lr: 1.0\n", encoding="utf-8") + + merged = merge_yaml_into_kwargs(str(config_path), {"max_lr": 3e-4}) + + assert merged["max_lr"] == 3e-4 + + +def test_type_inference_from_config_fields(): + @click.command() + @apply_specs( + [ + OptSpec("n_epoch", "G"), + OptSpec("rollout_max_policy_lag", "G"), + OptSpec("pin_memory", "G"), + OptSpec("metrics", "G"), + ], + TrainConfig, + ) + def cmd(**kwargs): + pass + + params = {p.name: p for p in cmd.params} + assert params["n_epoch"].default == 1 + assert params["n_epoch"].type.name == "integer" + assert params["rollout_max_policy_lag"].default is None + assert params["rollout_max_policy_lag"].type.name == "integer" + assert params["pin_memory"].secondary_opts == ["--no-pin_memory"] + assert params["pin_memory"].default is False # config default, no override + assert params["metrics"].multiple + assert params["metrics"].default == ("loss", "lr", "grad_norm") + + +def test_spec_overrides_beat_config_defaults(): + @click.command() + @apply_specs( + [ + OptSpec("num_workers", "G", default=4), + OptSpec("parallel_mode", "G", default="fsdp"), + ], + TrainConfig, + ) + def cmd(**kwargs): + pass + + params = {p.name: p for p in cmd.params} + assert params["num_workers"].default == 4 + assert params["parallel_mode"].default == "fsdp" + + +def test_cli_only_flag_pair_and_one_way_flag(): + @click.command() + @apply_specs( + [ + OptSpec("muon_nesterov", "G", type=bool, default=True), + OptSpec("dry_run", "G", is_flag=True, default=False), + ] + ) + def cmd(**kwargs): + pass + + params = {p.name: p for p in cmd.params} + assert params["muon_nesterov"].secondary_opts == ["--no-muon_nesterov"] + assert params["muon_nesterov"].default is True + assert params["dry_run"].is_flag + assert not params["dry_run"].secondary_opts + + +def test_uninferrable_type_raises_without_spec_type(): + with pytest.raises(TypeError, match="cannot infer"): + option_from_spec(OptSpec("mystery", "G"), {}) + + +def test_annotation_styles_old_and_new(): + """Optional[X], Union[X, None], X | None, List[str], and list[str] all + infer the same click types.""" + + @dataclasses.dataclass + class OldStyle: + opt_int: t.Optional[int] = None + opt_float: t.Optional[float] = None + union_int: t.Union[int, None] = None + names: t.List[str] = dataclasses.field(default_factory=lambda: ["loss", "lr"]) + flag: bool = True + + @dataclasses.dataclass + class NewStyle: + opt_int: int | None = None + names: list[str] = dataclasses.field(default_factory=lambda: ["loss"]) + flag: bool = False + + @click.command() + @apply_specs( + [ + OptSpec("opt_int", "G"), + OptSpec("opt_float", "G"), + OptSpec("union_int", "G"), + OptSpec("names", "G"), + OptSpec("flag", "G"), + ], + OldStyle, + ) + def old_cmd(**kwargs): + pass + + params = {p.name: p for p in old_cmd.params} + assert params["opt_int"].type.name == "integer" + assert params["opt_int"].default is None + assert params["opt_float"].type.name == "float" + assert params["union_int"].type.name == "integer" + assert params["names"].multiple + assert params["names"].type.name == "text" + assert params["flag"].secondary_opts == ["--no-flag"] + assert params["flag"].default is True + + @click.command() + @apply_specs( + [ + OptSpec("opt_int", "G"), + OptSpec("names", "G"), + OptSpec("flag", "G"), + ], + NewStyle, + ) + def new_cmd(**kwargs): + pass + + params = {p.name: p for p in new_cmd.params} + assert params["opt_int"].type.name == "integer" + assert params["names"].multiple + assert params["names"].default == ("loss",) + assert params["flag"].secondary_opts == ["--no-flag"] + assert params["flag"].default is False + + +def test_stringified_annotations_resolve_via_get_type_hints(): + """PEP 563 modules (``from __future__ import annotations``) leave + ``Field.type`` as a string; hints resolution still infers types.""" + + @dataclasses.dataclass + class FutureStyle: + opt_int: "t.Optional[int]" = None + count: "int" = 3 + + @click.command() + @apply_specs( + [ + OptSpec("opt_int", "G"), + OptSpec("count", "G"), + ], + FutureStyle, + ) + def cmd(**kwargs): + pass + + params = {p.name: p for p in cmd.params} + assert params["opt_int"].type.name == "integer" + assert params["count"].type.name == "integer" + assert params["count"].default == 3 + + +def test_help_order_follows_spec_table(): + @click.command(cls=GroupedCommand) + @apply_specs( + [ + OptSpec("first", "G", type=int, default=1), + OptSpec("second", "G", type=int, default=2), + ] + ) + def cmd(**kwargs): + pass + + result = CliRunner().invoke(cmd, ["--help"]) + assert result.exit_code == 0 + assert result.output.index("--first") < result.output.index("--second")