wehub-resource-sync
159 lines
6.3 KiB
Python
159 lines
6.3 KiB
Python
"""
|
|
Pretrain the mid-size (~400M) base model from scratch on the Pile HDF5 corpus.
|
|
|
|
This is the shared starting checkpoint for every post-training stage. It upgrades the
|
|
original ``train_transformer.py`` recipe with the things needed to actually train a
|
|
mid-size model on 2x H100: DistributedDataParallel, bf16 autocast, gradient accumulation,
|
|
a cosine LR schedule with warmup, weight-decay param groups, and periodic checkpointing.
|
|
The original ``train_transformer.py`` is left untouched.
|
|
|
|
Single GPU:
|
|
PYTHONPATH=. python scripts/pretrain_base.py
|
|
Both GPUs:
|
|
PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py
|
|
|
|
Override any config field from the CLI, e.g. ``--batch_size 16 --train_steps 50000``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from config.post_training_config import PretrainConfig
|
|
from data_loader.data_loader import get_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.utils import (
|
|
amp_autocast, build_model_from_config, save_stage_ckpt, set_seed, unwrap,
|
|
)
|
|
|
|
|
|
@torch.no_grad()
|
|
def estimate_loss(model, cfg, ctx, iters: int) -> dict[str, float]:
|
|
model.eval()
|
|
out = {}
|
|
for split, path in [("train", cfg.train_path), ("dev", cfg.dev_path)]:
|
|
if not os.path.exists(path):
|
|
continue
|
|
it = get_batch_iterator(path, cfg.batch_size, cfg.context_length, device=ctx.device)
|
|
losses = torch.zeros(iters)
|
|
for k in range(iters):
|
|
xb, yb = next(it)
|
|
with amp_autocast(cfg.amp_dtype, ctx.device):
|
|
_, loss = model(xb, yb)
|
|
losses[k] = loss.item()
|
|
out[split] = losses.mean().item()
|
|
model.train()
|
|
return out
|
|
|
|
|
|
def main():
|
|
cfg, extras = parse_config_with_json(
|
|
PretrainConfig, "configs/pretrain.json",
|
|
extra={"--resume": dict(type=str, default=None, help="checkpoint to resume from")})
|
|
resume = extras.resume
|
|
ctx = ddp_setup(cfg.device)
|
|
# Different data shuffle per rank (the loader shuffles via numpy global RNG).
|
|
set_seed(cfg.seed + ctx.rank)
|
|
|
|
model = build_model_from_config(cfg).to(ctx.device)
|
|
start_step = 0
|
|
if resume and os.path.exists(resume):
|
|
ck = torch.load(resume, map_location="cpu", weights_only=False)
|
|
unwrap(model).load_state_dict(ck["model_state_dict"])
|
|
start_step = ck.get("step", 0)
|
|
if ctx.is_main:
|
|
print(f"Resumed from {resume} at step {start_step}")
|
|
|
|
if cfg.compile:
|
|
model = torch.compile(model)
|
|
model = ddp_wrap(model, ctx)
|
|
|
|
optimizer = configure_optimizer(unwrap(model), cfg.lr, cfg.weight_decay)
|
|
if resume and os.path.exists(resume):
|
|
ck = torch.load(resume, map_location="cpu", weights_only=False)
|
|
if ck.get("optimizer_state_dict"):
|
|
optimizer.load_state_dict(ck["optimizer_state_dict"])
|
|
|
|
logger = None
|
|
if ctx.is_main:
|
|
n_params = sum(p.numel() for p in unwrap(model).parameters())
|
|
print(f"Model parameters: {n_params:,} (~{n_params/1e6:.0f}M) | world_size={ctx.world_size}")
|
|
print(f"Effective batch = {cfg.batch_size}*{cfg.grad_accum}*{ctx.world_size} "
|
|
f"= {cfg.batch_size*cfg.grad_accum*ctx.world_size} seqs/step")
|
|
logger = MetricsLogger("pretrain", cfg.log_dir, use_wandb=cfg.use_wandb,
|
|
wandb_project=cfg.wandb_project, config=vars(cfg).copy() if hasattr(cfg, "__dict__") else None)
|
|
|
|
batch_iter = get_batch_iterator(cfg.train_path, cfg.batch_size, cfg.context_length, device=ctx.device)
|
|
tokens_per_step = cfg.batch_size * cfg.context_length * cfg.grad_accum * ctx.world_size
|
|
|
|
model.train()
|
|
t0 = time.perf_counter()
|
|
for step in range(start_step, cfg.train_steps):
|
|
lr = cosine_lr(step, warmup_steps=cfg.warmup_steps, max_steps=cfg.train_steps,
|
|
lr=cfg.lr, min_lr=cfg.min_lr)
|
|
for g in optimizer.param_groups:
|
|
g["lr"] = lr
|
|
|
|
optimizer.zero_grad(set_to_none=True)
|
|
accum_loss = 0.0
|
|
for micro in range(cfg.grad_accum):
|
|
xb, yb = next(batch_iter)
|
|
# Only sync grads on the last micro-step (DDP optimization).
|
|
sync = (micro == cfg.grad_accum - 1) or not ctx.enabled
|
|
cm = model.no_sync() if (ctx.enabled and not sync) else _nullcm()
|
|
with cm, amp_autocast(cfg.amp_dtype, ctx.device):
|
|
_, loss = model(xb, yb)
|
|
loss = loss / cfg.grad_accum
|
|
loss.backward()
|
|
accum_loss += loss.item()
|
|
|
|
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
|
|
tok_s = tokens_per_step * 20 / dt if step > start_step else 0.0
|
|
t0 = time.perf_counter()
|
|
print(f"step {step} | loss {accum_loss:.4f} | lr {lr:.2e} | {tok_s:,.0f} tok/s")
|
|
if logger:
|
|
logger.log(step, {"train_loss": accum_loss, "lr": lr, "tok_per_s": tok_s})
|
|
|
|
if step > start_step and step % cfg.eval_steps == 0:
|
|
ev = estimate_loss(model, cfg, ctx, cfg.eval_iters)
|
|
ev = {k: reduce_scalar(v, ctx) for k, v in ev.items()}
|
|
if ctx.is_main:
|
|
print(f" [eval] step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in ev.items()))
|
|
if logger:
|
|
logger.log(step, {f"eval_{k}": v for k, v in ev.items()})
|
|
|
|
if ctx.is_main and step > start_step and step % cfg.save_every == 0:
|
|
save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="pretrain",
|
|
cfg=cfg, step=step, metrics={"train_loss": accum_loss})
|
|
print(f" saved checkpoint -> {cfg.out_ckpt} (step {step})")
|
|
|
|
if ctx.is_main:
|
|
save_stage_ckpt(cfg.out_ckpt, model, optimizer, stage="pretrain", cfg=cfg,
|
|
step=cfg.train_steps, metrics={"train_loss": accum_loss})
|
|
print(f"Done. Final checkpoint -> {cfg.out_ckpt}")
|
|
if logger:
|
|
logger.close()
|
|
cleanup(ctx)
|
|
|
|
|
|
import contextlib
|
|
|
|
|
|
def _nullcm():
|
|
return contextlib.nullcontext()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|