| """Main training loop for SAGE.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from typing import Optional |
|
|
| import torch |
| from torch.utils.data import DataLoader |
| import yaml |
|
|
| from data.dataset import DatasetConfig, PackedDataset |
| from eval.perplexity import evaluate_perplexity |
| from model.config import ModelConfig |
| from model.model import SageTransformer |
| from train.checkpoint import load_latest_checkpoint, save_checkpoint |
| from train.hardware import HardwareConfig |
| from train.loss import masked_cross_entropy |
| from train.optimizer import ScheduleConfig, create_optimizer, create_scheduler |
|
|
|
|
| @dataclass |
| class TrainerConfig: |
| """High-level trainer settings.""" |
|
|
| output_dir: str = "runs/default" |
| checkpoint_interval: int = 1000 |
| log_interval: int = 10 |
| eval_interval: int = 1000 |
| total_steps: int = 25_000 |
| seed: int = 42 |
| use_wandb: bool = True |
|
|
|
|
| def collate_batch(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: |
| """Stack packed dataset examples into a batch.""" |
| keys = batch[0].keys() |
| return {key: torch.stack([item[key] for item in batch], dim=0) for key in keys} |
|
|
|
|
| def create_dataloader(dataset: PackedDataset, batch_size: int) -> DataLoader: |
| """Create the training DataLoader.""" |
| return DataLoader(dataset, batch_size=batch_size, collate_fn=collate_batch) |
|
|
|
|
| def train( |
| model: SageTransformer, |
| train_dataset: PackedDataset, |
| validation_dataset: PackedDataset | None, |
| model_config: ModelConfig, |
| schedule_config: ScheduleConfig, |
| trainer_config: TrainerConfig, |
| ) -> dict[str, object]: |
| """Run the training loop and return the final summary.""" |
| torch.manual_seed(trainer_config.seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(trainer_config.seed) |
| hw = HardwareConfig(model_size_b=1.0, context_length=model_config.context_length) |
| device = torch.device(hw.device) |
| model = model.to(device) |
| optimizer = create_optimizer(model, schedule_config) |
| scheduler = create_scheduler(optimizer, schedule_config) |
| scaler = torch.GradScaler("cuda", enabled=(hw.device == "cuda" and hw.dtype == torch.float16)) |
| start_step = load_latest_checkpoint(model, optimizer, scheduler, scaler, trainer_config.output_dir, device) |
| train_dataset.skip(start_step * hw.grad_accum) |
| train_loader = create_dataloader(train_dataset, batch_size=hw.micro_batch) |
| train_iter = iter(train_loader) |
|
|
| Path(trainer_config.output_dir).mkdir(parents=True, exist_ok=True) |
| metrics_path = Path(trainer_config.output_dir) / "metrics.jsonl" |
| tokens_seen = start_step * hw.micro_batch * model_config.context_length |
| last_log_time = time.perf_counter() |
| wandb_run = _init_wandb(trainer_config, model_config, schedule_config, hw.summary()) |
|
|
| model.train() |
| for step in range(start_step, trainer_config.total_steps): |
| optimizer.zero_grad(set_to_none=True) |
| step_loss = 0.0 |
| for _ in range(hw.grad_accum): |
| try: |
| batch = next(train_iter) |
| except StopIteration: |
| train_iter = iter(train_loader) |
| batch = next(train_iter) |
| input_ids = batch["input_ids"].to(device) |
| labels = batch["labels"].to(device) |
| loss_mask = batch["loss_mask"].to(device) |
| if hw.use_amp: |
| with torch.autocast(device_type=hw.device, dtype=hw.dtype): |
| logits, _ = model(input_ids) |
| loss = masked_cross_entropy(logits, labels, loss_mask) / hw.grad_accum |
| else: |
| logits, _ = model(input_ids) |
| loss = masked_cross_entropy(logits, labels, loss_mask) / hw.grad_accum |
| scaler.scale(loss).backward() |
| step_loss += loss.item() |
| tokens_seen += int(input_ids.numel()) |
|
|
| scaler.unscale_(optimizer) |
| grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) |
| scaler.step(optimizer) |
| scaler.update() |
| scheduler.step() |
|
|
| if (step + 1) % trainer_config.log_interval == 0: |
| now = time.perf_counter() |
| elapsed = max(now - last_log_time, 1.0e-6) |
| tokens_per_second = (hw.micro_batch * hw.grad_accum * model_config.context_length) / elapsed |
| metrics = { |
| "step": step + 1, |
| "loss": step_loss, |
| "learning_rate": scheduler.get_last_lr()[0], |
| "tokens_seen": tokens_seen, |
| "tokens_per_second": tokens_per_second, |
| "grad_norm": float(grad_norm), |
| } |
| with metrics_path.open("a", encoding="utf-8") as handle: |
| handle.write(json.dumps(metrics) + "\n") |
| if wandb_run is not None: |
| wandb_run.log(metrics, step=step + 1) |
| last_log_time = now |
|
|
| if (step + 1) % trainer_config.eval_interval == 0 and validation_dataset is not None: |
| val_loader = create_dataloader(validation_dataset, batch_size=1) |
| evaluation = evaluate_perplexity(model, val_loader, device=device, dtype=hw.dtype if hw.use_amp else None) |
| with metrics_path.open("a", encoding="utf-8") as handle: |
| handle.write(json.dumps({"step": step + 1, **evaluation}) + "\n") |
| if wandb_run is not None: |
| wandb_run.log(evaluation, step=step + 1) |
|
|
| if (step + 1) % trainer_config.checkpoint_interval == 0: |
| save_checkpoint( |
| model=model, |
| optimizer=optimizer, |
| scheduler=scheduler, |
| scaler=scaler, |
| step=step + 1, |
| config={"model": model_config.to_dict(), "schedule": asdict(schedule_config), "trainer": asdict(trainer_config)}, |
| output_dir=trainer_config.output_dir, |
| ) |
|
|
| if wandb_run is not None: |
| wandb_run.finish() |
| return {"output_dir": trainer_config.output_dir, "tokens_seen": tokens_seen, "hardware": hw.summary()} |
|
|
|
|
| def _init_wandb( |
| trainer_config: TrainerConfig, |
| model_config: ModelConfig, |
| schedule_config: ScheduleConfig, |
| hardware_summary: dict[str, object], |
| ): |
| """Start a wandb run when available and enabled.""" |
| if not trainer_config.use_wandb: |
| return None |
| try: |
| import wandb |
| except ImportError: |
| return None |
| return wandb.init( |
| project="sage-llm", |
| name=Path(trainer_config.output_dir).name, |
| config={ |
| "model": model_config.to_dict(), |
| "schedule": asdict(schedule_config), |
| "trainer": asdict(trainer_config), |
| "hardware": hardware_summary, |
| }, |
| mode="offline", |
| ) |
|
|
|
|
| def build_argparser() -> argparse.ArgumentParser: |
| """Build the trainer CLI.""" |
| parser = argparse.ArgumentParser(description="Train the SAGE dense language model.") |
| parser.add_argument("--model-config", default="configs/model/1b.yaml") |
| parser.add_argument("--schedule-config", default="configs/train/schedule.yaml") |
| parser.add_argument("--train-shards", nargs="+", default=[]) |
| parser.add_argument("--validation-shards", nargs="*", default=[]) |
| parser.add_argument("--output-dir", default="runs/default") |
| parser.add_argument("--steps", type=int, default=None) |
| parser.add_argument("--disable-wandb", action="store_true") |
| return parser |
|
|
|
|
| def main(argv: Optional[list[str]] = None) -> None: |
| """CLI entrypoint for local training runs.""" |
| parser = build_argparser() |
| args = parser.parse_args(argv) |
| model_config = ModelConfig.from_yaml(args.model_config) |
| schedule_payload = yaml.safe_load(Path(args.schedule_config).read_text(encoding="utf-8")) |
| schedule = ScheduleConfig( |
| peak_learning_rate=schedule_payload["peak_learning_rate"], |
| min_learning_rate=schedule_payload["min_learning_rate"], |
| warmup_steps=schedule_payload["warmup_steps"], |
| weight_decay=schedule_payload["weight_decay"], |
| betas=tuple(schedule_payload["betas"]), |
| adam_eps=schedule_payload["adam_eps"], |
| total_steps=args.steps or schedule_payload["total_steps"] if "total_steps" in schedule_payload else (args.steps or 25_000), |
| ) |
| trainer_config = TrainerConfig( |
| output_dir=args.output_dir, |
| checkpoint_interval=schedule_payload.get("checkpoint_interval", 1000), |
| log_interval=schedule_payload.get("log_interval", 10), |
| eval_interval=schedule_payload.get("eval_interval", 1000), |
| total_steps=args.steps or schedule_payload.get("total_steps", 25_000), |
| seed=schedule_payload.get("seed", 42), |
| use_wandb=not args.disable_wandb, |
| ) |
| if not args.train_shards: |
| print("No training shards provided. The trainer entrypoint is configured correctly but requires shard paths to run.") |
| return |
| train_dataset = PackedDataset(DatasetConfig(tuple(args.train_shards), model_config.context_length, split="train")) |
| validation_dataset = None |
| if args.validation_shards: |
| validation_dataset = PackedDataset(DatasetConfig(tuple(args.validation_shards), model_config.context_length, split="validation")) |
| model = SageTransformer(model_config) |
| summary = train(model, train_dataset, validation_dataset, model_config, schedule, trainer_config) |
| print(json.dumps(summary, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
