wehub-resource-sync
128 lines
5.1 KiB
Python
128 lines
5.1 KiB
Python
"""
|
|
Supervised Fine-Tuning of the pretrained base on packed instruction data.
|
|
|
|
Loads the base checkpoint, trains with the prompt-masked SFT loss, periodically reports
|
|
masked dev loss, and saves an SFT checkpoint. DDP + bf16, single code path for 1 or N GPUs.
|
|
|
|
Single GPU:
|
|
PYTHONPATH=. python scripts/train_sft.py
|
|
Both GPUs:
|
|
PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import math
|
|
import time
|
|
|
|
import torch
|
|
|
|
from config.post_training_config import SFTConfig
|
|
from data_loader.sft_dataset import get_sft_batch_iterator
|
|
from src.post_training.cli import parse_config_with_json
|
|
from src.post_training.distributed import ddp_setup, ddp_wrap, cleanup, reduce_scalar
|
|
from src.post_training.logging_utils import MetricsLogger
|
|
from src.post_training.optim import configure_optimizer, cosine_lr
|
|
from src.post_training.sft import sft_loss
|
|
from src.post_training.utils import amp_autocast, load_backbone_from_ckpt, save_stage_ckpt, set_seed, unwrap
|
|
|
|
DEV_PATH = "/ephemeral/data/sft_dev_packed.h5"
|
|
|
|
|
|
@torch.no_grad()
|
|
def eval_dev(model, cfg, ctx, dev_path: str, max_batches: int = 50) -> float:
|
|
model.eval()
|
|
it = get_sft_batch_iterator(dev_path, cfg.batch_size, device=ctx.device,
|
|
rank=ctx.rank, world_size=ctx.world_size, shuffle=False, infinite=False)
|
|
total, n = 0.0, 0
|
|
for tokens, mask, _ in it:
|
|
with amp_autocast(cfg.amp_dtype, ctx.device):
|
|
logits, _ = model(tokens)
|
|
loss = sft_loss(logits, tokens, mask)
|
|
total += loss.item(); n += 1
|
|
if n >= max_batches:
|
|
break
|
|
model.train()
|
|
return total / max(1, n)
|
|
|
|
|
|
def main():
|
|
cfg, _ = parse_config_with_json(SFTConfig, "configs/sft.json")
|
|
ctx = ddp_setup(cfg.device)
|
|
set_seed(cfg.seed + ctx.rank)
|
|
|
|
model = load_backbone_from_ckpt(cfg, cfg.pretrained_ckpt, ctx.device)
|
|
if cfg.compile:
|
|
model = torch.compile(model)
|
|
model = ddp_wrap(model, ctx)
|
|
optimizer = configure_optimizer(unwrap(model), cfg.lr, cfg.weight_decay)
|
|
|
|
logger = None
|
|
if ctx.is_main:
|
|
print(f"SFT from {cfg.pretrained_ckpt} | world_size={ctx.world_size}")
|
|
logger = MetricsLogger("sft", cfg.log_dir, use_wandb=cfg.use_wandb, wandb_project=cfg.wandb_project)
|
|
|
|
train_it = get_sft_batch_iterator(cfg.data_path, cfg.batch_size, device=ctx.device,
|
|
rank=ctx.rank, world_size=ctx.world_size, shuffle=True, infinite=True)
|
|
|
|
# Estimate total steps for the cosine schedule from dataset size.
|
|
import h5py
|
|
with h5py.File(cfg.data_path, "r") as f:
|
|
n_rows = f["tokens"].shape[0]
|
|
steps_per_epoch = max(1, n_rows // (cfg.batch_size * ctx.world_size))
|
|
total_steps = cfg.max_steps if cfg.max_steps > 0 else steps_per_epoch * cfg.epochs
|
|
if ctx.is_main:
|
|
print(f"{n_rows} packed rows | ~{steps_per_epoch} steps/epoch | total_steps={total_steps}")
|
|
|
|
model.train()
|
|
t0 = time.perf_counter()
|
|
for step in range(total_steps):
|
|
lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=total_steps, lr=cfg.lr, min_lr=cfg.min_lr)
|
|
for g in optimizer.param_groups:
|
|
g["lr"] = lr
|
|
|
|
tokens, mask, epoch = next(train_it)
|
|
if epoch >= cfg.epochs and cfg.max_steps <= 0:
|
|
break
|
|
optimizer.zero_grad(set_to_none=True)
|
|
with amp_autocast(cfg.amp_dtype, ctx.device):
|
|
logits, _ = model(tokens)
|
|
loss = sft_loss(logits, tokens, mask)
|
|
loss.backward()
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
|
|
optimizer.step()
|
|
|
|
if ctx.is_main and step % 20 == 0:
|
|
dt = time.perf_counter() - t0; t0 = time.perf_counter()
|
|
print(f"step {step}/{total_steps} | loss {loss.item():.4f} | ppl {math.exp(min(20, loss.item())):.2f} | lr {lr:.2e} | {dt:.1f}s/20")
|
|
if logger:
|
|
logger.log(step, {"train_loss": loss.item(), "lr": lr})
|
|
|
|
if step > 0 and step % cfg.eval_steps == 0:
|
|
dev = reduce_scalar(eval_dev(model, cfg, ctx, DEV_PATH), ctx)
|
|
if ctx.is_main:
|
|
print(f" [eval] step {step} | dev_loss {dev:.4f} | dev_ppl {math.exp(min(20, dev)):.2f}")
|
|
if logger:
|
|
logger.log(step, {"dev_loss": dev})
|
|
|
|
if ctx.is_main and step > 0 and step % cfg.save_every == 0:
|
|
save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="sft", cfg=cfg, step=step,
|
|
metrics={"train_loss": loss.item()})
|
|
|
|
if ctx.is_main:
|
|
# Use the unwrapped model for the final eval: the other ranks have already reached
|
|
# cleanup(), so calling the DDP-wrapped model here would launch a collective with no
|
|
# peer and hang (NCCL timeout). The periodic eval above runs on all ranks, so it is fine.
|
|
dev = eval_dev(unwrap(model), cfg, ctx, DEV_PATH)
|
|
save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="sft", cfg=cfg, step=total_steps,
|
|
metrics={"dev_loss": dev})
|
|
print(f"Done SFT. dev_loss {dev:.4f} -> {cfg.out_ckpt}")
|
|
if logger:
|
|
logger.close()
|
|
cleanup(ctx)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|