chore: import upstream snapshot with attribution

scripts/train_sft.py

@@ -0,0 +1,127 @@
+"""
+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()