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
This commit is contained in:
+135
-226
@@ -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()
|
||||
|
||||
+45
-94
@@ -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()
|
||||
|
||||
+36
-34
@@ -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()
|
||||
|
||||
+39
-53
@@ -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()
|
||||
|
||||
+240
-301
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user