486 lines
17 KiB
Python
486 lines
17 KiB
Python
"""
|
||
train/pretrain.py — Main pretraining entry point.
|
||
|
||
Launch single-GPU:
|
||
python train/pretrain.py --config configs/small.yaml --train_data data/train.bin
|
||
|
||
Launch multi-GPU with torchrun:
|
||
torchrun --nproc_per_node=8 train/pretrain.py --config configs/small.yaml \
|
||
--train_data data/train.bin
|
||
|
||
The script auto-detects whether it is running inside a torchrun launch by
|
||
checking for the RANK environment variable.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import random
|
||
import signal
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
from torch.utils.data import DataLoader, DistributedSampler, RandomSampler
|
||
|
||
# B200 Tensor Core 최대 활용: TF32 matmul + cuDNN
|
||
torch.backends.cuda.matmul.allow_tf32 = True
|
||
torch.backends.cudnn.allow_tf32 = True
|
||
torch.backends.cudnn.benchmark = True # fixed seq_len=4096 → safe to auto-tune
|
||
torch.set_float32_matmul_precision("high") # TF32 precision for fp32 matmul
|
||
|
||
# Allow imports from the project root regardless of working directory.
|
||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||
if str(_PROJECT_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||
|
||
from data import PackedDataset
|
||
from model import LLM, LMConfig
|
||
from train.trainer import TrainConfig, Trainer
|
||
from train.utils import (
|
||
cleanup_ddp,
|
||
get_cosine_schedule_with_warmup,
|
||
is_main_process,
|
||
load_checkpoint,
|
||
setup_ddp,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Optional TransformerEngine import (FP8 support)
|
||
# ---------------------------------------------------------------------------
|
||
try:
|
||
import transformer_engine.pytorch as te # type: ignore[import]
|
||
HAS_TE = True
|
||
except ImportError:
|
||
te = None # type: ignore[assignment]
|
||
HAS_TE = False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Argument parsing
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description="Pretrain a decoder-only LLM.",
|
||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||
)
|
||
|
||
# Paths
|
||
parser.add_argument(
|
||
"--config",
|
||
type=Path,
|
||
default=Path("configs/small.yaml"),
|
||
help="Path to the LMConfig YAML file.",
|
||
)
|
||
parser.add_argument(
|
||
"--train_data",
|
||
type=Path,
|
||
required=True,
|
||
help="Path to the training data .bin file (numpy uint16 memmap).",
|
||
)
|
||
parser.add_argument(
|
||
"--val_data",
|
||
type=Path,
|
||
default=None,
|
||
help="Optional path to validation data .bin file.",
|
||
)
|
||
parser.add_argument(
|
||
"--checkpoint_dir",
|
||
type=Path,
|
||
default=Path("checkpoints"),
|
||
help="Root directory for saving checkpoints.",
|
||
)
|
||
parser.add_argument(
|
||
"--resume",
|
||
type=Path,
|
||
default=None,
|
||
help="Path to a checkpoint directory to resume training from.",
|
||
)
|
||
|
||
# Training hyper-parameters
|
||
parser.add_argument(
|
||
"--max_steps",
|
||
type=int,
|
||
default=None,
|
||
help="Override the number of optimiser steps (default: TrainConfig.max_steps).",
|
||
)
|
||
parser.add_argument(
|
||
"--batch_size",
|
||
type=int,
|
||
default=8,
|
||
help="Per-GPU micro-batch size.",
|
||
)
|
||
parser.add_argument(
|
||
"--lr",
|
||
type=float,
|
||
default=3e-4,
|
||
help="Peak learning rate.",
|
||
)
|
||
parser.add_argument(
|
||
"--weight_decay",
|
||
type=float,
|
||
default=0.1,
|
||
help="AdamW weight decay coefficient.",
|
||
)
|
||
parser.add_argument(
|
||
"--warmup_steps",
|
||
type=int,
|
||
default=2000,
|
||
help="Number of linear warmup steps.",
|
||
)
|
||
parser.add_argument(
|
||
"--grad_accum",
|
||
type=int,
|
||
default=1,
|
||
help="Gradient accumulation steps.",
|
||
)
|
||
parser.add_argument(
|
||
"--seed",
|
||
type=int,
|
||
default=42,
|
||
help="Base random seed (rank offset is added automatically).",
|
||
)
|
||
parser.add_argument(
|
||
"--log_file",
|
||
type=Path,
|
||
default=None,
|
||
help="Path to a text file for structured training logs (rank-0 only). "
|
||
"If omitted, logs go only to stdout.",
|
||
)
|
||
parser.add_argument(
|
||
"--use_fp8",
|
||
action="store_true",
|
||
default=False,
|
||
help="Enable TransformerEngine FP8 training (overrides config; requires B200/H100).",
|
||
)
|
||
|
||
return parser.parse_args()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Seed helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def set_seed(seed: int) -> None:
|
||
"""Set deterministic seeds for Python, NumPy, and PyTorch."""
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
torch.cuda.manual_seed_all(seed)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Optimizer parameter groups
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def build_optimizer_param_groups(
|
||
model: torch.nn.Module,
|
||
weight_decay: float,
|
||
) -> list[dict]:
|
||
"""
|
||
Split parameters into two groups:
|
||
- decay group : weight tensors with ndim >= 2
|
||
- no-decay group: bias, LayerNorm/RMSNorm weights, and embedding weights
|
||
|
||
This follows standard practice (e.g. GPT-style training).
|
||
"""
|
||
decay_params: list[torch.nn.Parameter] = []
|
||
no_decay_params: list[torch.nn.Parameter] = []
|
||
|
||
# Names of module types whose parameters should never be decayed.
|
||
no_decay_module_types = (
|
||
torch.nn.Embedding,
|
||
torch.nn.LayerNorm,
|
||
)
|
||
# Also skip any parameter whose name ends with '.bias' or 'norm'.
|
||
# Mamba-2 SSM parameters that should never be decayed
|
||
no_decay_name_suffixes = ("bias", "A_log", "D", "dt_bias")
|
||
|
||
# Collect module-level exclusions.
|
||
no_decay_module_params: set[int] = set()
|
||
for module in model.modules():
|
||
if isinstance(module, no_decay_module_types):
|
||
for param in module.parameters(recurse=False):
|
||
no_decay_module_params.add(id(param))
|
||
|
||
seen: set[int] = set()
|
||
for name, param in model.named_parameters():
|
||
if not param.requires_grad:
|
||
continue
|
||
if id(param) in seen:
|
||
continue
|
||
seen.add(id(param))
|
||
|
||
if (
|
||
id(param) in no_decay_module_params
|
||
or any(name.endswith(sfx) for sfx in no_decay_name_suffixes)
|
||
or param.ndim < 2
|
||
):
|
||
no_decay_params.append(param)
|
||
else:
|
||
decay_params.append(param)
|
||
|
||
return [
|
||
{"params": decay_params, "weight_decay": weight_decay},
|
||
{"params": no_decay_params, "weight_decay": 0.0},
|
||
]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
|
||
# ---- Distributed setup -------------------------------------------------
|
||
is_ddp = "RANK" in os.environ
|
||
rank = 0
|
||
local_rank = 0
|
||
world_size = 1
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
if is_ddp:
|
||
rank, local_rank, world_size, device = setup_ddp()
|
||
|
||
# Per-rank seed so data shuffling differs across replicas.
|
||
set_seed(args.seed + rank)
|
||
|
||
# ---- NUMA affinity for optimal GPU↔CPU memory locality ---------------
|
||
# B200 topology: GPU 0-3 → NUMA node 0 (cores 0-35)
|
||
# GPU 4-7 → NUMA node 1 (cores 36-71)
|
||
# Without pinning, 5/8 ranks end up on wrong NUMA → 3.2x memory latency.
|
||
try:
|
||
if local_rank < 4:
|
||
os.sched_setaffinity(0, set(range(0, 36))) # NUMA node 0
|
||
else:
|
||
os.sched_setaffinity(0, set(range(36, 72))) # NUMA node 1
|
||
if is_main_process():
|
||
print(f"NUMA affinity: rank {rank} (GPU {local_rank}) → "
|
||
f"{'NUMA0 cores 0-35' if local_rank < 4 else 'NUMA1 cores 36-71'}")
|
||
except (AttributeError, OSError) as e:
|
||
if is_main_process():
|
||
print(f"[WARN] NUMA affinity failed: {e}")
|
||
|
||
# ---- Model -------------------------------------------------------------
|
||
if not args.config.exists():
|
||
raise FileNotFoundError(f"Config file not found: {args.config}")
|
||
|
||
lm_config = LMConfig.from_yaml(args.config)
|
||
|
||
# CLI --use_fp8 flag overrides whatever the config file says.
|
||
if args.use_fp8:
|
||
lm_config.use_fp8 = True
|
||
|
||
# FP8 alignment check: (batch_size × seq_len) must be divisible by 8.
|
||
if lm_config.use_fp8 and (args.batch_size * lm_config.max_seq_len) % 8 != 0:
|
||
raise ValueError(
|
||
f"FP8: batch_size × max_seq_len = {args.batch_size} × {lm_config.max_seq_len} "
|
||
f"= {args.batch_size * lm_config.max_seq_len} must be divisible by 8."
|
||
)
|
||
|
||
# Note: fp8_model_init() is intentionally omitted — MXFP8Tensor weights are
|
||
# incompatible with DDP's _broadcast_coalesced during multi-GPU init.
|
||
# Weights remain in float32; TE quantizes on-the-fly inside fp8_autocast.
|
||
model = LLM(lm_config).to(device)
|
||
|
||
if is_main_process():
|
||
total_params = sum(p.numel() for p in model.parameters())
|
||
print(f"Model parameters: {total_params:,}")
|
||
print(f"LMConfig: {lm_config}")
|
||
|
||
# ---- Wrap in DDP -------------------------------------------------------
|
||
if is_ddp:
|
||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||
|
||
model = DDP(
|
||
model,
|
||
device_ids=[local_rank],
|
||
output_device=local_rank,
|
||
gradient_as_bucket_view=True, # zero-copy gradient → NCCL buffer
|
||
bucket_cap_mb=800, # larger buckets for NVLS (was 400)
|
||
find_unused_parameters=False, # fixed graph, no traversal overhead
|
||
# NOTE: static_graph=True 제거 — TE FP8 레이어의 동적 autograd hooks와 충돌
|
||
)
|
||
|
||
# ---- Dataset & DataLoader ----------------------------------------------
|
||
# PackedDataset: non-overlapping stride=seq_len windows.
|
||
# Avoids 600M random-index mmap accesses from stride-1 TextDataset.
|
||
train_dataset = PackedDataset(args.train_data, seq_len=lm_config.max_seq_len)
|
||
|
||
if is_ddp:
|
||
train_sampler: DistributedSampler | RandomSampler = DistributedSampler(
|
||
train_dataset,
|
||
num_replicas=world_size,
|
||
rank=rank,
|
||
shuffle=True,
|
||
seed=args.seed,
|
||
)
|
||
shuffle = False
|
||
else:
|
||
train_sampler = RandomSampler(train_dataset)
|
||
shuffle = False # Sampler is provided; DataLoader must not also shuffle.
|
||
|
||
train_loader = DataLoader(
|
||
train_dataset,
|
||
batch_size=args.batch_size,
|
||
sampler=train_sampler,
|
||
num_workers=6, # 6×8=48 workers, fits 72-core budget with OMP=4
|
||
pin_memory=True,
|
||
drop_last=True,
|
||
prefetch_factor=4, # deeper pipeline for larger worker pool
|
||
persistent_workers=True, # keep workers alive across epochs — eliminates respawn stall
|
||
)
|
||
|
||
# ---- Optimizer ---------------------------------------------------------
|
||
param_groups = build_optimizer_param_groups(
|
||
getattr(model, "module", model), args.weight_decay
|
||
)
|
||
optimizer = torch.optim.AdamW(
|
||
param_groups,
|
||
lr=args.lr,
|
||
betas=(0.9, 0.95),
|
||
eps=1e-8,
|
||
fused=torch.cuda.is_available(), # Use fused kernel when on CUDA.
|
||
)
|
||
|
||
# ---- LR Scheduler ------------------------------------------------------
|
||
train_config = TrainConfig(
|
||
checkpoint_dir=str(args.checkpoint_dir),
|
||
grad_accum_steps=args.grad_accum,
|
||
use_fp8=lm_config.use_fp8,
|
||
log_file=str(args.log_file) if args.log_file is not None else None,
|
||
)
|
||
if args.max_steps is not None:
|
||
train_config.max_steps = args.max_steps
|
||
|
||
scheduler = get_cosine_schedule_with_warmup(
|
||
optimizer=optimizer,
|
||
warmup_steps=args.warmup_steps,
|
||
total_steps=train_config.max_steps,
|
||
)
|
||
|
||
# ---- Resume from checkpoint --------------------------------------------
|
||
start_step = 0
|
||
if args.resume is not None:
|
||
if not args.resume.exists():
|
||
raise FileNotFoundError(f"Checkpoint path not found: {args.resume}")
|
||
start_step, resume_loss = load_checkpoint(
|
||
path=args.resume,
|
||
model=model,
|
||
optimizer=optimizer,
|
||
scheduler=scheduler,
|
||
)
|
||
if is_main_process():
|
||
print(f"Resumed from {args.resume} at step {start_step} (loss={resume_loss:.4f})")
|
||
|
||
# ---- Checkpoint directory ----------------------------------------------
|
||
args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# ---- Trainer -----------------------------------------------------------
|
||
trainer = Trainer(
|
||
model=model,
|
||
train_loader=train_loader,
|
||
optimizer=optimizer,
|
||
scheduler=scheduler,
|
||
config=train_config,
|
||
device=device,
|
||
rank=rank,
|
||
sampler=train_sampler if is_ddp else None,
|
||
)
|
||
|
||
# ---- Signal handlers for graceful shutdown ----------------------------
|
||
# SIGHUP: SSH 세션 끊김 시 발생 → 이전에 학습을 죽인 주범
|
||
# SIGTERM: kill 명령 또는 시스템 종료 시 발생
|
||
# 핸들러가 trainer.request_shutdown()을 호출하면, 학습 루프가
|
||
# 현재 step 완료 후 비상 체크포인트를 저장하고 깨끗하게 종료합니다.
|
||
_trainer_ref = trainer
|
||
|
||
def _graceful_shutdown_handler(signum, frame):
|
||
sig_name = signal.Signals(signum).name
|
||
if is_main_process():
|
||
import datetime as _dt
|
||
ts = _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
msg = (
|
||
f"[{ts}] [SIGNAL] Received {sig_name} (signum={signum}). "
|
||
f"Initiating graceful shutdown..."
|
||
)
|
||
print(f"\n{msg}")
|
||
# 로그 파일에도 즉시 기록 (시그널 핸들러 내에서 안전하게)
|
||
if args.log_file is not None:
|
||
try:
|
||
with open(args.log_file, "a", encoding="utf-8") as f:
|
||
f.write(msg + "\n")
|
||
except Exception:
|
||
pass # 시그널 핸들러 내에서는 예외 무시
|
||
_trainer_ref.request_shutdown(sig_name)
|
||
|
||
for _sig in (signal.SIGHUP, signal.SIGTERM):
|
||
signal.signal(_sig, _graceful_shutdown_handler)
|
||
|
||
if is_main_process():
|
||
import datetime
|
||
eff_tokens_per_step = args.batch_size * lm_config.max_seq_len * args.grad_accum * world_size
|
||
nccl_debug = os.environ.get("NCCL_DEBUG", "not set")
|
||
omp_threads = os.environ.get("OMP_NUM_THREADS", "not set")
|
||
print(
|
||
f"\n{'='*70}\n"
|
||
f" LLM Pretraining — {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
|
||
f"{'='*70}\n"
|
||
f" model : {lm_config.num_params:,} params | "
|
||
f"d_model={lm_config.d_model} n_layers={lm_config.n_layers}\n"
|
||
f" precision : {'FP8 (MXFP8BlockScaling)' if lm_config.use_fp8 else 'BF16'}\n"
|
||
f" GPUs : {world_size} | batch/GPU={args.batch_size} "
|
||
f"grad_accum={args.grad_accum}\n"
|
||
f" eff_batch : {args.batch_size * args.grad_accum * world_size} seqs "
|
||
f"= {eff_tokens_per_step:,} tok/step\n"
|
||
f" max_steps : {train_config.max_steps:,} "
|
||
f"({train_config.max_steps * eff_tokens_per_step / 1e9:.1f}B tokens total)\n"
|
||
f" data : {args.train_data}\n"
|
||
f" ckpt_dir : {args.checkpoint_dir}\n"
|
||
f" env : OMP_NUM_THREADS={omp_threads} NCCL_DEBUG={nccl_debug}\n"
|
||
f"{'='*70}\n"
|
||
)
|
||
|
||
try:
|
||
trainer.train(start_step=start_step)
|
||
# 학습 완료 또는 graceful shutdown 후 상태 출력
|
||
if is_main_process():
|
||
if trainer._shutdown_requested:
|
||
print(
|
||
f"\n[INFO] Training gracefully shut down via {trainer._shutdown_signal}. "
|
||
f"Emergency checkpoint saved. Resume with same command."
|
||
)
|
||
else:
|
||
print("\n[INFO] Training completed successfully.")
|
||
except KeyboardInterrupt:
|
||
if is_main_process():
|
||
print("\n[INFO] Training interrupted by user (KeyboardInterrupt).")
|
||
except Exception as e:
|
||
import traceback
|
||
if is_main_process():
|
||
tb = traceback.format_exc()
|
||
print(f"\n[ERROR] Training failed at rank {rank}:\n{tb}")
|
||
# log_file에도 기록
|
||
if args.log_file is not None:
|
||
with open(args.log_file, "a", encoding="utf-8") as f:
|
||
import datetime
|
||
f.write(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] [FATAL] {tb}\n")
|
||
raise
|
||
finally:
|
||
if is_ddp:
|
||
cleanup_ddp()
|
||
|
||
# Note: DDP cleanup is handled in the try/finally block above.
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|