File size: 13,490 Bytes
e2bfccc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 | """Main CLI entry point."""
import sys
from pathlib import Path
from typing import Optional
import click
import torch
from taoTrain.config import load_config, load_tokenizer_config, TrainingModeEnum, PretrainConfig, SFTConfig, RLConfig
from taoTrain.utils import set_seed, get_device
from taoTrain.core import BaseModel, create_model, create_datasets
from taoTrain.data import get_dataloader
from taoTrain.training import PretrainTrainer, SFTTrainer, RLTrainer
from taoTrain.benchmarks import BenchmarkRunner
from taoTrain.tokenizers import TokenizerTrainer
@click.group()
def main():
"""TaoTrain: A clean, modular PyTorch LLM training framework."""
pass
@main.command()
@click.option(
"--config",
type=click.Path(exists=True),
required=True,
help="Path to training config file (YAML or JSON)",
)
def pretrain(config: str):
"""Pretrain a language model."""
_train_command(config, TrainingModeEnum.PRETRAIN)
@main.command()
@click.option(
"--config",
type=click.Path(exists=True),
required=True,
help="Path to training config file (YAML or JSON)",
)
def sft(config: str):
"""Supervised fine-tune a language model."""
_train_command(config, TrainingModeEnum.SFT)
@main.command()
@click.option(
"--config",
type=click.Path(exists=True),
required=True,
help="Path to training config file (YAML or JSON)",
)
def rl(config: str):
"""Train with reinforcement learning."""
_train_command(config, TrainingModeEnum.RL)
@main.command()
@click.option(
"--config",
type=click.Path(exists=True),
required=True,
help="Path to tokenizer config file (YAML or JSON)",
)
def train_tokenizer(config: str):
"""Train a SentencePiece tokenizer from a YAML/JSON config file."""
try:
click.echo("🚀 TaoTrain Tokenizer Trainer")
click.echo(f"{'=' * 50}")
# Load tokenizer config
click.echo(f"Loading config from {config}...")
tokenizer_config = load_tokenizer_config(config)
# Train tokenizer from config
result = TokenizerTrainer.train_from_config(tokenizer_config)
# Display results
click.echo(f"\n{'=' * 50}")
click.echo("✅ Tokenizer Training Complete!")
click.echo(f"\n📊 Configuration:")
click.echo(f" - Input file: {tokenizer_config.jsonl_path}")
click.echo(f" - Samples: {tokenizer_config.max_samples or 'all'}")
click.echo(f" - Output dir: {result['output_dir']}")
click.echo(f" - Vocab size: {result['vocab_size']}")
click.echo(f" - Model type: {result['model_type']}")
if tokenizer_config.special_tokens:
click.echo(f" - Special tokens: {tokenizer_config.special_tokens}")
click.echo(f"\n📁 Generated Files:")
click.echo(f" - Model: {result['model_file']}")
click.echo(f" - Vocab: {result['vocab_file']}")
click.echo(f"\n📝 Next Steps:")
click.echo(f" 1. Use this tokenizer in your pretraining config:")
click.echo(f" dataset:")
click.echo(f" local: true")
click.echo(f" jsonl_path: {tokenizer_config.jsonl_path}")
click.echo(f" tokenizer_path: {result['model_file']}")
click.echo(f"")
click.echo(f" 2. Run pretraining with:")
click.echo(f" train pretrain --config your_config.yaml")
except ImportError as e:
click.echo(f"❌ Error: {e}", err=True)
sys.exit(1)
except FileNotFoundError as e:
click.echo(f"❌ File Error: {e}", err=True)
sys.exit(1)
except ValueError as e:
click.echo(f"❌ Validation Error: {e}", err=True)
sys.exit(1)
except Exception as e:
click.echo(f"❌ Unexpected Error: {e}", err=True)
click.echo(f" Please report this issue.", err=True)
sys.exit(1)
@click.command()
@click.option(
"--jsonl-path",
type=click.Path(exists=True),
required=True,
help="Path to JSONL file containing training data",
)
@click.option(
"--output-dir",
type=click.Path(),
default="tokenizers",
help="Directory to save tokenizer files",
)
@click.option(
"--vocab-size",
type=int,
default=50000,
help="Vocabulary size for the tokenizer",
)
@click.option(
"--model-type",
type=click.Choice(["unigram", "bpe", "char", "word"]),
default="unigram",
help="SentencePiece model type",
)
@click.option(
"--character-coverage",
type=float,
default=0.9995,
help="Character coverage for SentencePiece",
)
@click.option(
"--tokenizer-prefix",
type=str,
default=None,
help="Prefix for tokenizer output files (default: model_type)",
)
def train_tokenizer_command(
jsonl_path: str,
output_dir: str,
vocab_size: int,
model_type: str,
character_coverage: float,
tokenizer_prefix: Optional[str],
):
"""Train a SentencePiece tokenizer from JSONL data."""
try:
click.echo("🚀 TaoTrain Tokenizer Trainer")
click.echo(f"{'=' * 50}")
# Train tokenizer
result = TokenizerTrainer.train_sentencepiece(
jsonl_path=jsonl_path,
output_dir=output_dir,
vocab_size=vocab_size,
model_type=model_type,
character_coverage=character_coverage,
tokenizer_prefix=tokenizer_prefix,
)
# Display results
click.echo(f"\n{'=' * 50}")
click.echo("✅ Tokenizer Training Complete!")
click.echo(f"\n📊 Configuration:")
click.echo(f" - Input file: {jsonl_path}")
click.echo(f" - Output dir: {result['output_dir']}")
click.echo(f" - Vocab size: {result['vocab_size']}")
click.echo(f" - Model type: {result['model_type']}")
click.echo(f"\n📁 Generated Files:")
click.echo(f" - Model: {result['model_file']}")
click.echo(f" - Vocab: {result['vocab_file']}")
click.echo(f"\n📝 Next Steps:")
click.echo(f" 1. Use this tokenizer in your pretraining config:")
click.echo(f" dataset:")
click.echo(f" local: true")
click.echo(f" jsonl_path: {jsonl_path}")
click.echo(f" tokenizer_path: {result['model_file']}")
click.echo(f"")
click.echo(f" 2. Run pretraining with:")
click.echo(f" train pretrain --config your_config.yaml")
except ImportError as e:
click.echo(f"❌ Error: {e}", err=True)
sys.exit(1)
except FileNotFoundError as e:
click.echo(f"❌ File Error: {e}", err=True)
sys.exit(1)
except ValueError as e:
click.echo(f"❌ Validation Error: {e}", err=True)
sys.exit(1)
except Exception as e:
click.echo(f"❌ Unexpected Error: {e}", err=True)
click.echo(f" Please report this issue.", err=True)
sys.exit(1)
# Keep legacy CLI command as train-tokenizer-legacy for backward compatibility
main.add_command(train_tokenizer_command, name="train-tokenizer-legacy")
def _train_command(config_path: str, mode: TrainingModeEnum):
"""Internal training command."""
try:
# Load config
click.echo(f"Loading config from {config_path}...")
train_config = load_config(config_path, mode)
# Set seed
set_seed(train_config.seed)
# Get device
device = get_device(train_config.device)
click.echo(f"Using device: {device}")
# Create model
click.echo("Creating model...")
model = create_model(train_config, device)
total_params, trainable_params = _count_params(model)
click.echo(f" - Total parameters: {total_params:,}")
click.echo(f" - Trainable parameters: {trainable_params:,}")
# Load pretrained checkpoint if provided (for SFT/RL)
if train_config.checkpoint_path:
click.echo(f"Loading pretrained checkpoint from {train_config.checkpoint_path}...")
from taoTrain.checkpointing.checkpoint import CheckpointManager
checkpoint_manager = CheckpointManager(train_config.checkpoint_dir)
checkpoint = checkpoint_manager.load(train_config.checkpoint_path, device=device)
# CheckpointManager.load() normalizes format and ensures 'model_state' key exists
if "model_state" in checkpoint:
model.load_state_dict(checkpoint["model_state"], strict=False)
click.echo(" ✓ Checkpoint loaded successfully")
else:
raise KeyError(f"Invalid checkpoint format: 'model_state' key not found. "
f"Available keys: {list(checkpoint.keys())}")
# Create datasets
click.echo("Loading datasets...")
train_dataset, val_dataset = create_datasets(train_config)
click.echo(f" - Train samples: {len(train_dataset)}")
if val_dataset:
click.echo(f" - Val samples: {len(val_dataset)}")
# Select trainer
if mode == TrainingModeEnum.PRETRAIN:
trainer_class = PretrainTrainer
elif mode == TrainingModeEnum.SFT:
trainer_class = SFTTrainer
elif mode == TrainingModeEnum.RL:
trainer_class = RLTrainer
else:
raise ValueError(f"Unknown training mode: {mode}")
# Create trainer
click.echo("Setting up trainer...")
trainer = trainer_class(
model=model,
train_dataset=train_dataset,
val_dataset=val_dataset,
config=train_config,
device=device,
)
# Training loop
click.echo("\nStarting training...\n")
for epoch in range(train_config.num_epochs):
if train_config.max_steps and trainer.global_step >= train_config.max_steps:
break
epoch_metrics = trainer.train_epoch()
click.echo(f"\nEpoch {epoch + 1} complete")
click.echo(f" - Loss: {epoch_metrics.get('loss', 'N/A')}")
click.echo(f" - Learning rate: {epoch_metrics.get('lr', 'N/A')}")
# Final checkpoint
final_path = Path(train_config.checkpoint_dir) / "final_model.pt"
trainer.save_checkpoint(final_path)
click.echo(f"\nTraining complete! Final model saved to {final_path}")
# Log finish
trainer.logger.finish()
except Exception as e:
click.echo(f"Error during training: {e}", err=True)
sys.exit(1)
@main.command()
@click.option(
"--model",
type=click.Path(exists=True),
required=True,
help="Path to model checkpoint",
)
@click.option(
"--benchmark-type",
type=click.Choice(["all", "perplexity", "throughput", "memory"]),
default="all",
help="Type of benchmark to run",
)
@click.option(
"--batch-size",
type=int,
default=32,
help="Batch size for benchmarking",
)
@click.option(
"--seq-length",
type=int,
default=1024,
help="Sequence length for benchmarking",
)
def benchmark(model: str, benchmark_type: str, batch_size: int, seq_length: int):
"""Benchmark a trained model."""
try:
click.echo(f"Loading model from {model}...")
device = get_device("cuda")
runner = BenchmarkRunner.load_from_checkpoint(model, device=device)
click.echo("Running benchmarks...\n")
if benchmark_type == "throughput" or benchmark_type == "all":
click.echo("Throughput benchmark:")
results = runner.benchmark_throughput(batch_size, seq_length)
for key, val in results.items():
click.echo(f" {key}: {val:.2f}")
if benchmark_type == "memory" or benchmark_type == "all":
click.echo("\nMemory benchmark:")
results = runner.benchmark_memory()
for key, val in results.items():
click.echo(f" {key}: {val:.2f}")
click.echo("\nBenchmarking complete!")
except Exception as e:
click.echo(f"Error during benchmarking: {e}", err=True)
sys.exit(1)
@main.command()
@click.option("--repo", type=str, default=".aim", help="AimStack repository path")
def view_logs(repo: str):
"""View training logs with AimStack."""
try:
import subprocess
click.echo(f"Opening AimStack dashboard for repo: {repo}")
subprocess.run(["aim", "up", "--repo", repo])
except FileNotFoundError:
click.echo("Error: 'aim' command not found. Install with: pip install aim", err=True)
sys.exit(1)
def _count_params(model: BaseModel) -> tuple[int, int]:
"""Count model parameters."""
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
return total, trainable
if __name__ == "__main__":
main()
|