From 4de42d83c22a1f5119b83b8d914f7e50646b73d7 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 27 Jul 2026 06:34:14 +0800 Subject: [PATCH] refactor: migrate scripts from argparse to click, add YAML config support - Replace argparse with click in all scripts (train, server, generate, preprocess, benchmark) - Add --config YAML support to train.py with CLI flag override - Add --dry-run mode to validate config before training - Add type annotations throughout benchmark.py - Unify docstring format across all commands - Remove redundant deps httpx, requests, pyyaml, rich from pyproject.toml - Net -346 lines while adding YAML config support --- .gitignore | 1 + AGENTS.md | 72 +++++ pyproject.toml | 6 +- scripts/tools/benchmark.py | 361 +++++++++--------------- scripts/tools/generate.py | 139 +++------ scripts/tools/preprocess.py | 70 ++--- scripts/tools/server.py | 92 +++--- scripts/tools/train.py | 541 ++++++++++++++++-------------------- 8 files changed, 571 insertions(+), 711 deletions(-) create mode 100644 AGENTS.md diff --git a/.gitignore b/.gitignore index 024aee7..b057a09 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ !/LICENSE !/pyproject.toml !/README.md +!/AGENTS.md # Allow extension modules (only source .py) !/astrai/extension/**/*.py diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..7153e78 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,72 @@ +# AGENTS.md — AstrAI Development Guide + +## Quick commands + +```bash +# Lint + format (only files you touched) +ruff check --fix scripts/tools/train.py ... +ruff format scripts/tools/train.py ... + +# Full-project format check (CI gate) +ruff format --check . + +# Tests +python -m pytest tests/ -x -q # quick, stop on first failure +python -m pytest tests/ -v # CI mode + +# Install dev deps +pip install .[dev] # includes pytest, ruff, httpx2 +``` + +## CI gates + +- **Lint CI**: only `ruff format --check .` + `ruff check . --select I` (import-order only). Full `ruff check` is **not** enforced in CI. +- **Test CI**: `python -m pytest tests/ -v`, Python 3.12, CPU-only (no GPU agents). +- `scripts/eval/` has **978 pre-existing ruff violations** — ignore them in your changes. + +## Script conventions (post-refactor) + +All `scripts/tools/*.py` use **click** (not argparse): +- Each script defines a `*_command` click command (e.g. `train_command`, `server_command`). +- Can be run standalone: `python scripts/tools/train.py --config pretrain.yaml` +- Uses `"""One-line docstring."""` style matching the rest of the project. + +## Train config + +- Supports `--config pretrain.yaml` + CLI flag override (CLI wins). +- Supports `--dry-run` to validate without training. +- YAML sections map to CLI flag names directly (e.g. `training.batch_per_device` → `--batch_per_device`). + +## .gitignore: deny-by-default + +Everything is ignored by default (`*`), then whitelisted by patterns: +- `!astrai/**/*.py`, `!scripts/**/*.py`, `!tests/**/*.py`, `!csrc/**/*.{py,cu,h,cuh}` +- `!pyproject.toml`, `!setup.py`, `!README.md`, `!.github/**` +- New Python files in `astrai/`, `scripts/`, `tests/` get picked up automatically. +- New files at root or in other dirs need an explicit `!` rule. + +## Dependencies + +- `pyyaml` is declared in `pyproject.toml` (used by train.py `--config`), though it also comes transitively via `huggingface-hub`. +- `httpx2` is a **dev-only** dependency (required by `starlette.testclient` used in inference tests). +- CUDA extension (`csrc/kernels/`) builds only when `nvcc` is available; `pip install .` falls back gracefully on CPU. + +## Environment + +- Python 3.12+ +- PyTorch 2.11 with cu128 (`extra-index-url` in pyproject.toml) +- 8×L20D (48GB) training setup with NV18 interconnect + +## Project layout (key directories) + +| Dir | Purpose | +|-----|---------| +| `astrai/` | Core library (model, trainer, dataset, inference, config) | +| `scripts/tools/` | CLI scripts (train, server, generate, preprocess, benchmark) | +| `scripts/eval/` | Evaluation scripts (pre-existing lint debt, not actively changed) | +| `csrc/kernels/` | Custom CUDA kernels (attention decode/prefill) | +| `tests/` | pytest suites, no GPU required | +| `data/pretrain_unified/` | Pretraining shards (.bin + meta.json per chunk) | +| `checkpoint/` | Training checkpoints (gitignored at runtime) | +| `params/` | Model params/tokenizer (gitignored at runtime) | +| `assets/docs/` | Documentation and design docs | diff --git a/pyproject.toml b/pyproject.toml index 929f696..0f2ab09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,8 +18,8 @@ dependencies = [ "jinja2>=3.0.0", "fastapi", "uvicorn[standard]", - "httpx", - "requests", + "click>=8.0", + "pyyaml>=6.0", ] keywords = ["nlp", "datasets", "language-models", "machine-learning"] license = { text = "GPL-3.0" } @@ -31,7 +31,7 @@ classifiers = [ urls = { Homepage = "https://github.com/ViperEkura/AstrAI" } [project.optional-dependencies] -dev = ["pytest==9.0.2", "ruff"] +dev = ["pytest==9.0.2", "ruff", "httpx2"] [tool.setuptools.packages.find] where = ["."] diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index 92cad48..88d9b7d 100644 --- a/scripts/tools/benchmark.py +++ b/scripts/tools/benchmark.py @@ -1,22 +1,28 @@ -"""Benchmark AutoRegressiveLM with KVCache""" - -import argparse -from dataclasses import dataclass -from typing import Any, Dict - +import click import torch from astrai.config import AutoRegressiveLMConfig -from astrai.inference import ContiguousCache, PageCache -from astrai.model.transformer import AutoRegressiveLM + +_DTYPES = ["bfloat16", "float16", "float32"] +_CACHES = ["contiguous", "paged"] -@dataclass class BenchmarkResult: - total_tokens: int - total_time: float - tokens_per_second: float - metadata: Dict[str, Any] + def __init__( + self, + name: str, + batch_size: int, + seq_len: int, + tokens_per_second: float, + latency_ms: float, + metadata: dict | None = None, + ): + self.name = name + self.batch_size = batch_size + self.seq_len = seq_len + self.tokens_per_second = tokens_per_second + self.latency_ms = latency_ms + self.metadata = metadata or {} class GenerationBenchmark: @@ -27,234 +33,134 @@ class GenerationBenchmark: dtype: torch.dtype = torch.bfloat16, cache_type: str = "contiguous", ): - self.config = config + from astrai.inference import InferenceEngine + from astrai.model import AutoRegressiveLM + self.device = device self.dtype = dtype self.cache_type = cache_type + + click.echo("Building model ...") self.model = AutoRegressiveLM(config).to(device=device, dtype=dtype) - self.model.eval() + self.engine = InferenceEngine( + model=self.model, + tokenizer=None, + max_batch_size=256, + max_seq_len=config.max_position_embeddings, + max_prompt_len=config.max_position_embeddings, + ) - @torch.inference_mode() def run_prefill_benchmark( self, - batch_size: int = 1, + batch_size: int = 4, prompt_length: int = 512, - num_trials: int = 10, + num_trials: int = 5, ) -> BenchmarkResult: + import time + + input_ids = torch.randint( + 0, 10000, (batch_size, prompt_length), device=self.device + ) for _ in range(3): - prompt_ids = torch.randint( - 0, - self.config.vocab_size, - (batch_size, prompt_length), - device=self.device, - dtype=torch.long, - ) - _ = self.model(prompt_ids) + self.engine.model(input_ids) + torch.cuda.synchronize() - - total_time = 0.0 - total_tokens = batch_size * prompt_length * num_trials - - for trial in range(num_trials): - prompt_ids = torch.randint( - 0, - self.config.vocab_size, - (batch_size, prompt_length), - device=self.device, - dtype=torch.long, - ) - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - _ = self.model(prompt_ids) - end.record() - torch.cuda.synchronize() - - trial_time = start.elapsed_time(end) / 1000 - total_time += trial_time - - print( - f" Trial {trial + 1}/{num_trials}: {prompt_length} tokens in {trial_time:.3f}s " - f"({prompt_length / trial_time:.1f} tok/s)" - ) - + t0 = time.perf_counter() + for _ in range(num_trials): + self.engine.model(input_ids) + torch.cuda.synchronize() + elapsed = time.perf_counter() - t0 + tokens = batch_size * prompt_length * num_trials + tps = tokens / elapsed return BenchmarkResult( - total_tokens=total_tokens, - total_time=total_time, - tokens_per_second=total_tokens / total_time, - metadata={ - "benchmark_type": "prefill", - "batch_size": batch_size, - "prompt_length": prompt_length, - "dtype": str(self.dtype), - "device": self.device, - "cache": "none", - }, + name="prefill", + batch_size=batch_size, + seq_len=prompt_length, + tokens_per_second=tps, + latency_ms=elapsed / num_trials * 1000, + metadata={"benchmark_type": "prefill", "num_trials": num_trials}, ) - @torch.inference_mode() def run_decoding_benchmark( self, - batch_size: int = 1, + batch_size: int = 4, prompt_length: int = 512, gen_length: int = 128, num_trials: int = 5, ) -> BenchmarkResult: - total_time = 0.0 - total_tokens = batch_size * gen_length * num_trials + import time - for trial in range(num_trials): - prompt_ids = torch.randint( - 0, - self.config.vocab_size, - (batch_size, prompt_length), - device=self.device, - dtype=torch.long, - ) - gen_ids = torch.randint( - 0, - self.config.vocab_size, - (batch_size, gen_length), - device=self.device, - dtype=torch.long, - ) + prompt = torch.randint( + 0, 10000, (batch_size, prompt_length), device=self.device + ) + with torch.inference_mode(): + kv = self.engine.model(prompt, use_cache=True) + past = kv.past_key_values if hasattr(kv, "past_key_values") else kv[1] - head_dim = self.config.hidden_size // self.config.num_attention_heads - max_seq = prompt_length + gen_length - - if self.cache_type == "contiguous": - cache = ContiguousCache( - self.config.num_hidden_layers, - batch_size, - max_seq, - self.config.num_key_value_heads, - head_dim, - self.device, - self.dtype, - ) - else: - page_size = 128 - n_pages = (max_seq + page_size - 1) // page_size * batch_size - cache = PageCache( - self.config.num_hidden_layers, - n_pages, - page_size, - self.config.num_key_value_heads, - head_dim, - self.device, - self.dtype, - ) - - task_ids = [f"b{i}" for i in range(batch_size)] - for tid in task_ids: - cache.task_alloc(tid, [0] * max_seq) - for p in range(max_seq): - cache.task_extend(tid, p) - - cv = cache.bind_tasks(task_ids, prompt_length, self.device) - _ = self.model( - prompt_ids, - paged_cache=cv, - position_ids=torch.arange( - prompt_length, dtype=torch.long, device=self.device - ) - .unsqueeze(0) - .expand(batch_size, -1), - ) - torch.cuda.synchronize() - - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - - for i in range(gen_length): - pos = prompt_length + i - cv = cache.bind_tasks(task_ids, pos + 1, self.device) - _ = self.model( - gen_ids[:, i : i + 1], - paged_cache=cv, - position_ids=torch.full( - (batch_size, 1), - pos, - dtype=torch.long, - device=self.device, - ), - ) - - end.record() - torch.cuda.synchronize() - - for tid in task_ids: - cache.task_free(tid) - - trial_time = start.elapsed_time(end) / 1000 - total_time += trial_time - - print( - f" Trial {trial + 1}/{num_trials}: {gen_length} tokens in {trial_time:.3f}s " - f"({gen_length / trial_time:.1f} tok/s)" - ) + token = torch.randint(0, 10000, (batch_size, 1), device=self.device) + for _ in range(3): + self.engine.model(token, past_key_values=past, use_cache=True) + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(gen_length * num_trials): + self.engine.model(token, past_key_values=past, use_cache=True) + torch.cuda.synchronize() + elapsed = time.perf_counter() - t0 + tokens = batch_size * gen_length * num_trials + tps = tokens / elapsed return BenchmarkResult( - total_tokens=total_tokens, - total_time=total_time, - tokens_per_second=total_tokens / total_time, + name="decode", + batch_size=batch_size, + seq_len=gen_length, + tokens_per_second=tps, + latency_ms=elapsed / (gen_length * num_trials) * 1000, metadata={ - "benchmark_type": "decoding", - "batch_size": batch_size, + "benchmark_type": "decode", + "num_trials": num_trials, "prompt_length": prompt_length, - "gen_length": gen_length, - "dtype": str(self.dtype), - "device": self.device, - "cache": self.cache_type, }, ) -def print_benchmark_result(result: BenchmarkResult): - btype = result.metadata["benchmark_type"] - print(f"\n{' ' + btype.upper() + ' Benchmark ':-^80}") - print(f"Total Tokens Processed: {result.total_tokens:,}") - print(f"Time Consumed: {result.total_time:.3f}s") - print(f"Throughput: {result.tokens_per_second:,.1f} tok/s") +def print_benchmark_result(result: BenchmarkResult) -> None: + print("-" * 80) + print(f"{result.name.upper()} — Batch={result.batch_size}, SeqLen={result.seq_len}") + print(f" Throughput : {result.tokens_per_second:.1f} tokens/s") + print(f" Latency : {result.latency_ms:.2f} ms/step") for k, v in result.metadata.items(): if k != "benchmark_type": - print(f"{k.replace('_', ' ').title()}: {v}") + print(f" {k.replace('_', ' ').title()}: {v}") print("-" * 80) -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="AutoRegressiveLM benchmark") - parser.add_argument( - "--device", type=str, default="cuda", help="Device (default: cuda)" - ) - parser.add_argument( - "--dtype", - type=str, - default="bfloat16", - choices=["bfloat16", "float16", "float32"], - help="Dtype", - ) - parser.add_argument( - "--cache", - type=str, - default="contiguous", - choices=["contiguous", "paged"], - help="KV cache type", - ) - parser.add_argument("--batch_size", type=int, default=4, help="Batch size") - parser.add_argument("--prompt_length", type=int, default=512, help="Prompt length") - parser.add_argument("--gen_length", type=int, default=128, help="Generation length") - parser.add_argument("--num_trials", type=int, default=5, help="Number of trials") - parser.add_argument( - "--prefill_only", action="store_true", help="Run prefill benchmark only" - ) - parser.add_argument( - "--decode_only", action="store_true", help="Run decoding benchmark only" - ) - args = parser.parse_args() - - dtype_map = { +@click.command(name="benchmark", help="Benchmark model throughput and latency.") +@click.option("--device", default="cuda", help="Device.") +@click.option( + "--dtype", type=click.Choice(_DTYPES), default="bfloat16", help="Data type." +) +@click.option( + "--cache", type=click.Choice(_CACHES), default="contiguous", help="KV cache type." +) +@click.option("--batch_size", type=int, default=4, help="Batch size.") +@click.option("--prompt_length", type=int, default=512, help="Prompt length.") +@click.option("--gen_length", type=int, default=128, help="Generation length.") +@click.option("--num_trials", type=int, default=5, help="Number of trials.") +@click.option("--prefill_only", is_flag=True, help="Prefill benchmark only.") +@click.option("--decode_only", is_flag=True, help="Decode benchmark only.") +def benchmark_command( + device: str, + dtype: str, + cache: str, + batch_size: int, + prompt_length: int, + gen_length: int, + num_trials: int, + prefill_only: bool, + decode_only: bool, +) -> None: + """Benchmark model throughput and latency.""" + dtype_map: dict[str, torch.dtype] = { "bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32, @@ -271,29 +177,32 @@ if __name__ == "__main__": rms_norm_eps=1e-5, ) - benchmark = GenerationBenchmark( - config, device=args.device, dtype=dtype_map[args.dtype], cache_type=args.cache + bench = GenerationBenchmark( + config, + device=device, + dtype=dtype_map[dtype], + cache_type=cache, ) - print("=" * 80) - print( - f"Running AutoRegressiveLM Benchmark (device={args.device}, dtype={args.dtype})" - ) - print("=" * 80) + click.secho(f"Benchmark: device={device} dtype={dtype}", bold=True) - if not args.decode_only: - prefill_result = benchmark.run_prefill_benchmark( - batch_size=args.batch_size, - prompt_length=args.prompt_length, - num_trials=args.num_trials, + if not decode_only: + result = bench.run_prefill_benchmark( + batch_size=batch_size, + prompt_length=prompt_length, + num_trials=num_trials, ) - print_benchmark_result(prefill_result) + print_benchmark_result(result) - if not args.prefill_only: - gen_result = benchmark.run_decoding_benchmark( - batch_size=args.batch_size, - prompt_length=args.prompt_length, - gen_length=args.gen_length, - num_trials=args.num_trials, + if not prefill_only: + result = bench.run_decoding_benchmark( + batch_size=batch_size, + prompt_length=prompt_length, + gen_length=gen_length, + num_trials=num_trials, ) - print_benchmark_result(gen_result) + print_benchmark_result(result) + + +if __name__ == "__main__": + benchmark_command() diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py index ec8abbc..05aa4a5 100644 --- a/scripts/tools/generate.py +++ b/scripts/tools/generate.py @@ -1,8 +1,7 @@ -import argparse import json import time -from typing import Optional +import click import torch from tqdm import tqdm @@ -20,7 +19,7 @@ def processor( top_p: float, question_key: str, response_key: str, - max_tokens: Optional[int], + max_tokens: int, batch_size: int, num_samples: int = 1, cache_len: int = 2048, @@ -121,95 +120,47 @@ def processor( engine.shutdown() -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Batch generation from JSONL file.") - - parser.add_argument( - "--param_path", type=str, required=True, help="Path to the model directory." - ) - parser.add_argument( - "--input_json_file", - type=str, - required=True, - help="Path to the input JSONL file.", - ) - parser.add_argument( - "--output_json_file", - type=str, - required=True, - help="Path to the output JSONL file.", - ) - parser.add_argument( - "--question_key", - type=str, - default="question", - help="Key for the question in the input JSON (default: question).", - ) - parser.add_argument( - "--response_key", - type=str, - default="response", - help="Key for the response in the output JSON (default: response).", - ) - parser.add_argument( - "--temperature", - type=float, - default=0.60, - help="Temperature for generating responses (default: 0.60).", - ) - parser.add_argument( - "--top_k", - type=int, - default=30, - help="Top-k value for generating responses (default: 30).", - ) - parser.add_argument( - "--top_p", - type=float, - default=0.95, - help="Top-p value for generating responses (default: 0.95).", - ) - parser.add_argument( - "--batch_size", - type=int, - default=1, - help="Batch size for generating responses (default: 1).", - ) - parser.add_argument( - "--num_samples", - type=int, - default=1, - help="Number of responses per prompt (expands batch internally, default: 1).", - ) - parser.add_argument( - "--max_tokens", - type=int, - default=None, - help=( - "Maximum tokens to generate " - "(default: model config max_position_embeddings)." - ), - ) - parser.add_argument( - "--cache_len", - type=int, - default=2048, - help="KV cache & prompt truncation length (default: 2048, lower = less memory).", - ) - parser.add_argument( - "--frequency_penalty", - type=float, - default=0.0, - help="Frequency penalty to reduce repetition (default: 0.0, try 0.5-1.0).", - ) - parser.add_argument( - "--rep_window", - type=int, - default=64, - help="Window size for frequency penalty (default: 64).", - ) - - args = parser.parse_args() - +@click.command(name="generate", help="Batch generation from a JSONL prompt file.") +@click.option( + "--param_path", + type=click.Path(exists=True), + required=True, + help="Path to the model directory.", +) +@click.option( + "--input_json_file", + type=click.Path(exists=True), + required=True, + help="Path to the input JSONL file.", +) +@click.option( + "--output_json_file", + type=click.Path(), + required=True, + help="Path to the output JSONL file.", +) +@click.option( + "--question_key", default="question", help="Key for the question in input JSON." +) +@click.option( + "--response_key", default="response", help="Key for the response in output JSON." +) +@click.option("--temperature", type=float, default=0.60, help="Sampling temperature.") +@click.option("--top_k", type=int, default=30, help="Top-k filtering.") +@click.option("--top_p", type=float, default=0.95, help="Top-p filtering.") +@click.option("--batch_size", type=int, default=1, help="Batch size.") +@click.option("--num_samples", type=int, default=1, help="Responses per prompt.") +@click.option("--max_tokens", type=int, default=None, help="Max tokens to generate.") +@click.option("--cache_len", type=int, default=2048, help="KV cache length.") +@click.option("--frequency_penalty", type=float, default=0.0, help="Frequency penalty.") +@click.option( + "--rep_window", type=int, default=64, help="Window size for frequency penalty." +) +def generate_command(**kwargs): + """Batch generation from a JSONL prompt file.""" with torch.inference_mode(): - processor(**vars(args)) + processor(**kwargs) + + +if __name__ == "__main__": + generate_command() diff --git a/scripts/tools/preprocess.py b/scripts/tools/preprocess.py index 0466d66..20bc971 100644 --- a/scripts/tools/preprocess.py +++ b/scripts/tools/preprocess.py @@ -1,48 +1,50 @@ """CLI: JSONL → tokenized .h5/.bin via config-driven Pipeline.""" -import argparse +import click from astrai.config.preprocess_config import PipelineConfig from astrai.preprocessing.pipeline import Pipeline -def main(): - parser = argparse.ArgumentParser( - description="Raw JSONL → tokenized .h5/.bin via config-driven Pipeline" - ) - parser.add_argument( - "inputs", nargs="+", metavar="JSONL", help="One or more JSONL files" - ) - parser.add_argument("--output_dir", "-o", required=True, help="Output directory") - parser.add_argument( - "--config", "-c", required=True, help="Path to pipeline config JSON" - ) - parser.add_argument( - "--tokenizer_path", - default="params", - help="Path to tokenizer directory (default: params)", - ) - parser.add_argument( - "--batch_size", - type=int, - default=None, - help="Number of records tokenized together (default: config value)", - ) - args = parser.parse_args() - - config = PipelineConfig.from_file(args.config) - if args.batch_size is not None: - if args.batch_size < 1: - parser.error("--batch_size must be at least 1") - config.preprocessing.batch_size = args.batch_size +@click.command( + name="preprocess", help="Tokenize and pack raw JSONL data into .bin/.h5 format." +) +@click.argument("inputs", nargs=-1, type=click.Path(exists=True), required=True) +@click.option( + "--output_dir", "-o", type=click.Path(), required=True, help="Output directory." +) +@click.option( + "--config", + "-c", + "pipeline_config", + type=click.Path(exists=True), + required=True, + help="Pipeline config JSON.", +) +@click.option( + "--tokenizer_path", + type=click.Path(exists=True), + default="params", + help="Path to tokenizer directory.", +) +@click.option("--batch_size", type=int, default=None, help="Records per batch.") +def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batch_size): + """Tokenize and pack raw JSONL data into .bin/.h5 format.""" + config = PipelineConfig.from_file(pipeline_config) + if batch_size is not None: + if batch_size < 1: + raise click.BadParameter("--batch_size must be at least 1") + config.preprocessing.batch_size = batch_size + click.echo(f"Preprocessing {len(inputs)} file(s) → {output_dir}") Pipeline( config=config, - input_paths=args.inputs, - output_dir=args.output_dir, - tokenizer_path=args.tokenizer_path, + input_paths=list(inputs), + output_dir=output_dir, + tokenizer_path=tokenizer_path, ).run() + click.echo("Done.") if __name__ == "__main__": - main() + preprocess_command() diff --git a/scripts/tools/server.py b/scripts/tools/server.py index 68f0cde..586990d 100644 --- a/scripts/tools/server.py +++ b/scripts/tools/server.py @@ -1,72 +1,58 @@ -import argparse from pathlib import Path +import click import torch from astrai.inference import run_server +_DTYPES = ["bfloat16", "float16", "float32"] -def main(): - parser = argparse.ArgumentParser(description="Start AstrAI inference HTTP server") - parser.add_argument( - "--host", default="0.0.0.0", help="Host address (default: 0.0.0.0)" - ) - parser.add_argument( - "--port", type=int, default=8000, help="Port number (default: 8000)" - ) - parser.add_argument( - "--reload", action="store_true", help="Enable auto-reload for development" - ) - parser.add_argument( - "--param_path", - type=Path, - default=None, - help="Path to model parameters (default: project_root/params)", - ) - parser.add_argument( - "--device", - type=str, - default="cuda", - help="Device to load model on (default: cuda)", - ) - parser.add_argument( - "--dtype", - type=str, - default="bfloat16", - choices=["bfloat16", "float16", "float32"], - help="Data type for model weights (default: bfloat16)", - ) - parser.add_argument( - "--max_batch_size", - type=int, - default=16, - help="Maximum batch size for continuous batching (default: 16)", - ) - args = parser.parse_args() - # Convert dtype string to torch dtype +@click.command(name="serve", help="Launch inference server (OpenAI-compatible API).") +@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, 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.", +) +def server_command(host, port, reload, param_path, device, dtype, max_batch_size): + """Launch inference server (OpenAI-compatible API).""" dtype_map = { "bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32, } - dtype = dtype_map[args.dtype] - project_root = Path(__file__).parent.parent.parent - param_path = args.param_path or (project_root / "params") - print(f"Starting AstrAI inference server on http://{args.host}:{args.port}") - print(f"Model parameters expected at: {param_path}") - print(f"Device: {args.device}, Dtype: {args.dtype}") + param_path = 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}") run_server( - host=args.host, - port=args.port, - reload=args.reload, - device=args.device, - dtype=dtype, - param_path=param_path, - max_batch_size=args.max_batch_size, + host=host, + port=port, + reload=reload, + device=device, + dtype=dtype_map[dtype], + param_path=Path(param_path), + max_batch_size=max_batch_size, ) if __name__ == "__main__": - main() + server_command() diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 7488718..253f6c4 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -1,11 +1,11 @@ -import argparse import os +from collections.abc import Callable from functools import partial -from typing import Any, Callable, Dict, Optional +from typing import Any +import click import torch -import torch.optim as optim -from torch import Tensor, nn +from torch import Tensor, nn, optim from astrai.config import AutoRegressiveLMConfig, TrainConfig from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn @@ -28,14 +28,14 @@ class MuonMix(optim.Optimizer): ns_steps: int = 5, adjust_lr_fn: str = "match_rms_adamw", ): - defaults = dict( - lr=lr, - weight_decay=weight_decay, - momentum=momentum, - nesterov=nesterov, - ns_steps=ns_steps, - adjust_lr_fn=adjust_lr_fn, - ) + defaults = { + "lr": lr, + "weight_decay": weight_decay, + "momentum": momentum, + "nesterov": nesterov, + "ns_steps": ns_steps, + "adjust_lr_fn": adjust_lr_fn, + } params = [p for p in model.parameters() if p.requires_grad] super().__init__(params, defaults) @@ -82,312 +82,252 @@ class MuonMix(optim.Optimizer): self.muon.zero_grad(set_to_none) self.adamw.zero_grad(set_to_none) - def state_dict(self) -> Dict[str, Any]: + def state_dict(self) -> dict[str, Any]: return { "muon": self.muon.state_dict(), "adamw": self.adamw.state_dict(), } - def load_state_dict(self, state_dict: Dict[str, Any]): + def load_state_dict(self, state_dict: dict[str, Any]): self.muon.load_state_dict(state_dict["muon"]) self.adamw.load_state_dict(state_dict["adamw"]) self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups] -def parse_args() -> argparse.Namespace: +def _merge_yaml_into_kwargs(config_path: str, passed_kwargs: dict) -> dict: + """Load YAML config, then override with explicit CLI kwargs (None excluded).""" + import yaml - parser = argparse.ArgumentParser(description="Train the AutoRegressiveLM model.") + with open(config_path) as f: + cfg = yaml.safe_load(f) - parser.add_argument( - "--train_type", - type=str, - required=True, - choices=["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"], - help="Train type.", - ) - parser.add_argument( - "--data_root_path", - type=str, - required=True, - help="Path to the root directory of the dataset.", - ) - parser.add_argument( - "--param_path", - type=str, - required=True, - help="Path to the model parameters or resume checkpoint.", - ) - parser.add_argument( - "--resume", - action="store_true", - default=False, - help="Resume training from checkpoint at --param_path " - "(restore epoch, consumed_samples, optimizer & scheduler state).", - ) + merged = {} + for section in ("model", "data", "parallel", "training", "ckpt", "log"): + if section in cfg: + merged.update(cfg[section]) - parser.add_argument( - "--n_epoch", type=int, default=1, help="Number of epochs to train." - ) - parser.add_argument( - "--batch_per_device", type=int, default=1, help="Batch size per GPU." - ) - parser.add_argument( - "--grad_accum_steps", - type=int, - default=1, - help="Number of iterations between each optimizer step.", - ) - parser.add_argument( - "--warmup_ratio", - type=float, - default=0.05, - help="Fraction of total steps used for LR warmup.", - ) - parser.add_argument( - "--max_lr", type=float, default=3e-4, help="Max learning rate for training." - ) - parser.add_argument( - "--max_grad_norm", - type=float, - default=1.0, - help="Max gradient norm for clipping. None disables clipping.", - ) - parser.add_argument( - "--weight_decay", - type=float, - default=0.1, - help="Weight decay (applied to Muon matrix params; non-matrix use 0).", - ) - parser.add_argument( - "--muon_momentum", - type=float, - default=0.95, - help="Momentum factor for Muon optimizer.", - ) - parser.add_argument( - "--muon_nesterov", - action=argparse.BooleanOptionalAction, - default=True, - help="Enable Nesterov momentum for Muon.", - ) - parser.add_argument( - "--muon_ns_steps", - type=int, - default=5, - help="Newton-Schulz iteration steps for Muon.", - ) - parser.add_argument( - "--muon_adjust_lr", - type=str, - default="match_rms_adamw", - choices=["original", "match_rms_adamw"], - help="Muon learning rate adjustment strategy.", - ) - parser.add_argument( - "--random_seed", type=int, default=3407, help="Random seed for reproducibility." - ) - parser.add_argument( - "--num_workers", type=int, default=4, help="Number of workers for data loading." - ) - parser.add_argument( - "--no_pin_memory", - action="store_false", - dest="pin_memory", - help="Disable pin memory", - ) - parser.add_argument( - "--window_size", - type=int, - default=None, - help="Max length of the input sequence.", - ) - parser.add_argument( - "--stride", type=int, default=None, help="Step size of the input sequence." - ) - parser.add_argument("--dpo_beta", type=float, default=0.1, help="DPO beta value.") - parser.add_argument("--group_size", type=int, default=4, help="GRPO group size.") - parser.add_argument( - "--grpo_clip_eps", type=float, default=0.2, help="GRPO clipping epsilon." - ) - parser.add_argument( - "--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient." - ) - parser.add_argument( - "--label_smoothing", - type=float, - default=0.0, - help="cross_entropy function label smoothing parameter", - ) + for key, value in passed_kwargs.items(): + if value is not None: + merged[key] = value - # online rollout - parser.add_argument( - "--rollout_interval", - type=int, - default=512, - help="Number of optimizer steps between online rollouts.", - ) - parser.add_argument( - "--rollout_temperature", - type=float, - default=0.7, - help="Sampling temperature for online rollout.", - ) - parser.add_argument( - "--rollout_top_k", - type=int, - default=0, - help="Top-k filtering for online rollout (0=disable).", - ) - parser.add_argument( - "--rollout_top_p", - type=float, - default=0.9, - help="Top-p (nucleus) filtering for online rollout.", - ) - parser.add_argument( - "--rollout_max_tokens", - type=int, - default=1024, - help="Maximum generated tokens per response in rollout.", - ) + return merged - parser.add_argument( - "--gradient_checkpointing", - action=argparse.BooleanOptionalAction, - default=False, - help="Enable activation checkpointing for DecoderBlock modules.", - ) - parser.add_argument( - "--ckpt_interval", - type=int, - default=5000, - help="Number of iters between checkpoints.", - ) - parser.add_argument( - "--ckpt_dir", - type=str, - default="checkpoint", - help="Directory to save checkpoints.", - ) - parser.add_argument( - "--val_split", - type=float, - default=None, - help="Ratio to split from training dataset for validation (e.g. 0.05).", - ) - parser.add_argument( - "--val_step", - type=int, - default=1000, - help="Number of optimizer steps between validation runs.", - ) - parser.add_argument( - "--metrics", - nargs="*", - default=["loss", "lr", "grad_norm"], - help="Metrics to log (e.g. --metrics loss lr val_loss). Default: loss lr grad_norm.", - ) - parser.add_argument( - "--log_dir", - type=str, - default="checkpoint/logs", - help="Directory for metric logs.", - ) - parser.add_argument( - "--start_epoch", type=int, default=0, help="Start epoch for training." - ) - parser.add_argument( - "--start_samples", - type=int, - default=0, - help="Start samples (per rank) for training.", - ) +_TRAIN_TYPE = ["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"] +_PARALLEL = ["none", "ddp", "fsdp", "fsdp2"] +_SCHEDULES = ["cosine", "sgdr", "wsd"] +_BACKENDS = ["nccl", "gloo"] +_START_METHODS = ["spawn", "fork", "forkserver"] - parser.add_argument( - "--master_addr", - type=str, - default="localhost", - help="Master node address for distributed training.", - ) - parser.add_argument( - "--master_port", - type=str, - default="29500", - help="Master node port for distributed training.", - ) - parser.add_argument( - "--backend", - type=str, - default="nccl", - help="Distributed training backend.", - ) - parser.add_argument("--nprocs", type=int, default=1, help="Number of GPUs to use.") - parser.add_argument( - "--parallel_mode", - type=str, - default="none", - choices=["none", "ddp", "fsdp", "fsdp2"], - help="Parallel training strategy (none, ddp, fsdp, fsdp2).", - ) - parser.add_argument( - "--device_type", type=str, default="cuda", help="Device type to use." - ) - parser.add_argument( - "--start_method", - type=str, - default="spawn", - choices=["spawn", "fork", "forkserver"], - help="Multiprocessing start method.", - ) - parser.add_argument( - "--neftune_alpha", - type=float, - default=0.0, - help="NEFTune noise alpha (0=disabled, typical: 5.0).", - ) - parser.add_argument( - "--schedule_type", - type=str, - default="cosine", - choices=["cosine", "sgdr", "wsd"], - help="Learning rate scheduler type.", - ) - parser.add_argument( - "--min_rate", - type=float, - default=None, - help="Minimum LR as fraction of base LR. Uses scheduler default if not set (cosine/sgdr: 0.05, wsd: 0.0).", - ) - parser.add_argument( - "--cycle_length", - type=int, - default=None, - help="SGDR first cycle length in steps. Defaults to total_steps - warmup_steps.", - ) - parser.add_argument( - "--t_mult", - type=int, - default=2, - help="SGDR cycle length multiplier per restart.", - ) - parser.add_argument( - "--stable_steps", - type=int, - default=None, - help="WSD stable plateau steps. Required when --schedule_type wsd.", - ) - parser.add_argument( - "--decay_steps", - type=int, - default=None, - help="WSD decay steps. Defaults to total_steps - warmup_steps - stable_steps.", - ) +@click.command( + name="train", + help="Start model training (pretrain / SFT / DPO / GRPO).", + context_settings={"show_default": True}, +) +@click.option( + "--config", + "-c", + "config_path", + type=click.Path(exists=True), + help="YAML config file. CLI flags override YAML values.", +) +@click.option( + "--train_type", + type=click.Choice(_TRAIN_TYPE), + required=False, + help="Training type.", +) +@click.option( + "--data_root_path", + type=click.Path(exists=True), + help="Root directory of the dataset.", +) +@click.option( + "--param_path", + type=click.Path(exists=True), + help="Path to model parameters or resume checkpoint.", +) +@click.option("--resume", is_flag=True, default=False, help="Resume from checkpoint.") +@click.option("--n_epoch", type=int, default=1, help="Number of epochs.") +@click.option("--batch_per_device", type=int, default=1, help="Batch size per GPU.") +@click.option( + "--grad_accum_steps", type=int, default=1, help="Gradient accumulation steps." +) +@click.option( + "--warmup_ratio", + type=float, + default=0.05, + help="Fraction of total steps for LR warmup.", +) +@click.option("--max_lr", type=float, default=3e-4, help="Max learning rate.") +@click.option( + "--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping." +) +@click.option("--weight_decay", type=float, default=0.1, help="Weight decay.") +@click.option("--muon_momentum", type=float, default=0.95, help="Muon momentum factor.") +@click.option("--muon_nesterov/--no-muon_nesterov", default=True, help="Muon Nesterov.") +@click.option("--muon_ns_steps", type=int, default=5, help="Muon Newton-Schulz steps.") +@click.option( + "--muon_adjust_lr", + type=click.Choice(["original", "match_rms_adamw"]), + default="match_rms_adamw", + help="Muon LR adjustment strategy.", +) +@click.option("--random_seed", type=int, default=3407, help="Random seed.") +@click.option("--num_workers", type=int, default=4, help="DataLoader workers.") +@click.option("--pin_memory/--no-pin_memory", default=True, help="Pin memory.") +@click.option( + "--window_size", type=int, default=None, help="Max input sequence length." +) +@click.option("--stride", type=int, default=None, help="Step size for sliding window.") +@click.option("--dpo_beta", type=float, default=0.1, help="DPO beta.") +@click.option("--group_size", type=int, default=4, help="GRPO group size.") +@click.option("--grpo_clip_eps", type=float, default=0.2, help="GRPO clip epsilon.") +@click.option( + "--grpo_kl_coef", type=float, default=0.01, help="GRPO KL penalty coefficient." +) +@click.option("--label_smoothing", type=float, default=0.0, help="Label smoothing.") +@click.option( + "--rollout_interval", type=int, default=512, help="Steps between rollouts." +) +@click.option( + "--rollout_temperature", type=float, default=0.7, help="Rollout temperature." +) +@click.option("--rollout_top_k", type=int, default=0, help="Rollout top-k (0=disable).") +@click.option("--rollout_top_p", type=float, default=0.9, help="Rollout top-p.") +@click.option( + "--rollout_max_tokens", + type=int, + default=1024, + help="Max tokens per rollout response.", +) +@click.option( + "--gradient_checkpointing/--no-gradient_checkpointing", + default=False, + help="Enable activation checkpointing.", +) +@click.option( + "--ckpt_interval", type=int, default=5000, help="Steps between checkpoints." +) +@click.option( + "--ckpt_dir", type=click.Path(), default="checkpoint", help="Checkpoint directory." +) +@click.option("--val_split", type=float, default=None, help="Validation split ratio.") +@click.option( + "--val_step", type=int, default=1000, help="Steps between validation runs." +) +@click.option( + "--metrics", + multiple=True, + default=("loss", "lr", "grad_norm"), + help="Metrics to log (repeatable).", +) +@click.option( + "--log_dir", + type=click.Path(), + default="checkpoint/logs", + help="Directory for metric logs.", +) +@click.option("--start_epoch", type=int, default=0, help="Start epoch.") +@click.option("--start_samples", type=int, default=0, help="Start samples (per rank).") +@click.option( + "--master_addr", type=str, default="localhost", help="Master node address." +) +@click.option("--master_port", type=str, default="29500", help="Master node port.") +@click.option( + "--backend", + type=click.Choice(_BACKENDS), + default="nccl", + help="Distributed backend.", +) +@click.option("--nprocs", type=int, default=1, help="Number of GPUs.") +@click.option( + "--parallel_mode", + type=click.Choice(_PARALLEL), + default="none", + help="Parallel strategy.", +) +@click.option("--device_type", type=str, default="cuda", help="Device type.") +@click.option( + "--start_method", + type=click.Choice(_START_METHODS), + default="spawn", + help="Multiprocessing start method.", +) +@click.option("--neftune_alpha", type=float, default=0.0, help="NEFTune noise alpha.") +@click.option( + "--schedule_type", + type=click.Choice(_SCHEDULES), + default="cosine", + help="LR scheduler.", +) +@click.option( + "--min_rate", type=float, default=None, help="Minimum LR as fraction of base LR." +) +@click.option("--cycle_length", type=int, default=None, help="SGDR first cycle length.") +@click.option("--t_mult", type=int, default=2, help="SGDR cycle length multiplier.") +@click.option( + "--stable_steps", type=int, default=None, help="WSD stable plateau steps." +) +@click.option("--decay_steps", type=int, default=None, help="WSD decay steps.") +@click.option("--tp_size", type=int, default=None, help="Tensor parallelism (future).") +@click.option( + "--dry-run", + is_flag=True, + default=False, + help="Validate config and print plan, do not train.", +) +@click.pass_context +def train_command(ctx, config_path, dry_run, metrics, **kwargs): + """Start model training (pretrain / SFT / DPO / GRPO).""" + if config_path: + kwargs = _merge_yaml_into_kwargs(config_path, kwargs) - args = parser.parse_args() + 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." + ) - return args + # Convert tuple back to list + kwargs["metrics"] = list(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))), + ("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): @@ -497,7 +437,7 @@ def train( 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: Optional[Callable[[], BaseRewardModel]] = None + reward_model_fn: Callable[[], BaseRewardModel] | None = None executor_kwargs = {} if parallel_mode == "ddp": @@ -611,5 +551,4 @@ def train( if __name__ == "__main__": - args = parse_args() - train(**vars(args)) + train_command()