| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from rich.console import Console |
|
|
| from .config import SFTConfig |
| from .io import write_json |
| from .schemas import RewardRecord |
| from .training_common import ( |
| evaluate_loader, |
| make_loader, |
| make_model, |
| move_batch, |
| optimizer_parameters, |
| seed_everything, |
| ) |
|
|
| console = Console() |
|
|
|
|
| def soft_label_loss(logits: torch.Tensor, rewards: torch.Tensor, temperature: float) -> torch.Tensor: |
| if temperature <= 0: |
| raise ValueError("soft_target_temperature must be positive") |
| targets = F.softmax(rewards / temperature, dim=-1) |
| return -(targets * F.log_softmax(logits, dim=-1)).sum(dim=-1).mean() |
|
|
|
|
| def train_sft( |
| records: list[RewardRecord], |
| config: SFTConfig, |
| output_dir: str | Path, |
| checkpoint: str | None = None, |
| ) -> dict: |
| seed_everything(config.training.seed) |
| worker_ids = records[0].worker_ids |
| model = make_model(config.model, worker_ids, checkpoint) |
| train_loader = make_loader( |
| records, model, "train", config.training.batch_size, shuffle=True |
| ) |
| try: |
| validation_loader = make_loader( |
| records, model, "validation", config.training.batch_size, shuffle=False |
| ) |
| except ValueError: |
| validation_loader = None |
|
|
| trainable, total = model.trainable_parameter_counts() |
| console.print( |
| f"Device={model.device_ref}; trainable={trainable:,}/{total:,} " |
| f"({100 * trainable / total:.4f}%)" |
| ) |
| optimizer = torch.optim.AdamW( |
| optimizer_parameters(model), |
| lr=config.training.learning_rate, |
| weight_decay=config.training.weight_decay, |
| ) |
| history: list[dict] = [] |
| global_step = 0 |
| optimizer.zero_grad(set_to_none=True) |
| for epoch in range(config.training.epochs): |
| model.train() |
| epoch_losses = [] |
| for step, batch in enumerate(train_loader): |
| inputs, rewards = move_batch(batch, model.device_ref) |
| logits = model(**inputs) |
| loss = soft_label_loss(logits, rewards, config.training.soft_target_temperature) |
| (loss / config.training.gradient_accumulation_steps).backward() |
| epoch_losses.append(float(loss.detach())) |
| should_step = ( |
| (step + 1) % config.training.gradient_accumulation_steps == 0 |
| or step + 1 == len(train_loader) |
| ) |
| if should_step: |
| torch.nn.utils.clip_grad_norm_( |
| list(optimizer_parameters(model)), config.training.max_grad_norm |
| ) |
| optimizer.step() |
| optimizer.zero_grad(set_to_none=True) |
| global_step += 1 |
| if global_step % config.training.log_every == 0: |
| console.print(f"SFT step={global_step} loss={epoch_losses[-1]:.4f}") |
| epoch_report = { |
| "epoch": epoch + 1, |
| "train_loss": sum(epoch_losses) / max(1, len(epoch_losses)), |
| } |
| if validation_loader is not None: |
| epoch_report["validation"] = evaluate_loader(model, validation_loader) |
| history.append(epoch_report) |
| console.print(epoch_report) |
|
|
| output = Path(output_dir) |
| model.save_checkpoint( |
| output, |
| metadata={"stage": "sft", "global_step": global_step, "history": history}, |
| ) |
| report = { |
| "stage": "sft", |
| "checkpoint": str(output), |
| "worker_ids": worker_ids, |
| "global_step": global_step, |
| "history": history, |
| } |
| write_json(output / "training_report.json", report) |
| return report |
|
|
|
|
