1063 lines
39 KiB
Python
1063 lines
39 KiB
Python
"""
|
||
train/sft.py — Supervised Fine-Tuning (SFT) entry point.
|
||
|
||
Loads a pretrained checkpoint and fine-tunes it on instruction/conversation
|
||
data using SFTDataset, which masks prompt tokens with ignore_index=-1 so only
|
||
the assistant response tokens contribute to the loss.
|
||
|
||
Launch single-GPU:
|
||
python train/sft.py \\
|
||
--base_checkpoint checkpoints/korean_1b_fp8_run1/checkpoint-0034000 \\
|
||
--sft_data data/sft/train.jsonl \\
|
||
--device cuda:0
|
||
|
||
Launch multi-GPU (DDP via torchrun):
|
||
torchrun --nproc_per_node=8 train/sft.py \\
|
||
--base_checkpoint checkpoints/korean_1b_fp8_run1/checkpoint-0034000 \\
|
||
--sft_data data/sft/train.jsonl
|
||
|
||
KEY DIFFERENCES from pretrain.py:
|
||
- Loads weights from a pretrained checkpoint via LLM.from_pretrained()
|
||
- Uses SFTDataset (JSONL instruction data) instead of PackedDataset
|
||
- Lower default learning rate (2e-5 vs 2e-4)
|
||
- Fewer default steps (3000 vs 100000)
|
||
- Copies tokenizer.json to checkpoint_dir for easy deployment
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import random
|
||
import signal
|
||
import shutil
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch.utils.data import DataLoader, DistributedSampler, RandomSampler
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Data Mixing: Interleave SFT + Pretrain batches for forgetting prevention
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class MixingDataLoader:
|
||
"""
|
||
Wraps two DataLoaders and yields batches from one or the other
|
||
based on a probability ratio.
|
||
|
||
With ``pretrain_ratio=0.3``, 30% of batches come from the pretrain
|
||
loader and 70% from the SFT loader. Both loaders cycle infinitely.
|
||
|
||
This is duck-type compatible with DataLoader for the Trainer's needs:
|
||
- ``__iter__`` yields batches
|
||
- ``__len__`` returns the SFT loader length (used for epoch estimation)
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
sft_loader: DataLoader,
|
||
pretrain_loader: DataLoader,
|
||
pretrain_ratio: float = 0.3,
|
||
sft_sampler: DistributedSampler | RandomSampler | None = None,
|
||
pretrain_sampler: DistributedSampler | RandomSampler | None = None,
|
||
) -> None:
|
||
self.sft_loader = sft_loader
|
||
self.pretrain_loader = pretrain_loader
|
||
self.pretrain_ratio = pretrain_ratio
|
||
self.sft_sampler = sft_sampler
|
||
self.pretrain_sampler = pretrain_sampler
|
||
self._epoch = 0
|
||
|
||
def __len__(self) -> int:
|
||
return len(self.sft_loader)
|
||
|
||
def __iter__(self):
|
||
sft_iter = iter(self.sft_loader)
|
||
pt_iter = iter(self.pretrain_loader)
|
||
|
||
while True:
|
||
use_pretrain = random.random() < self.pretrain_ratio
|
||
try:
|
||
if use_pretrain:
|
||
batch = next(pt_iter)
|
||
else:
|
||
batch = next(sft_iter)
|
||
except StopIteration:
|
||
# Whichever exhausted, restart it
|
||
if use_pretrain:
|
||
self._epoch += 1
|
||
if self.pretrain_sampler is not None and hasattr(self.pretrain_sampler, 'set_epoch'):
|
||
self.pretrain_sampler.set_epoch(self._epoch)
|
||
pt_iter = iter(self.pretrain_loader)
|
||
try:
|
||
batch = next(pt_iter)
|
||
except StopIteration:
|
||
raise RuntimeError(
|
||
"Pretrain DataLoader is empty after restart. "
|
||
"Check pretrain_data path and drop_last settings."
|
||
)
|
||
else:
|
||
self._epoch += 1
|
||
if self.sft_sampler is not None and hasattr(self.sft_sampler, 'set_epoch'):
|
||
self.sft_sampler.set_epoch(self._epoch)
|
||
sft_iter = iter(self.sft_loader)
|
||
try:
|
||
batch = next(sft_iter)
|
||
except StopIteration:
|
||
raise RuntimeError(
|
||
"SFT DataLoader is empty after restart. "
|
||
"Check sft_data path and drop_last settings."
|
||
)
|
||
yield batch
|
||
|
||
# B200 Tensor Core 최대 활용: TF32 matmul + cuDNN
|
||
torch.backends.cuda.matmul.allow_tf32 = True
|
||
torch.backends.cudnn.allow_tf32 = True
|
||
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 model import LLM
|
||
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="Supervised Fine-Tuning (SFT) of a pretrained decoder-only LLM.",
|
||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||
)
|
||
|
||
# --- Required paths -----------------------------------------------------
|
||
parser.add_argument(
|
||
"--base_checkpoint",
|
||
type=Path,
|
||
required=True,
|
||
help=(
|
||
"Path to the pretrained checkpoint directory. "
|
||
"Must contain model.pt and config.yaml (produced by save_checkpoint)."
|
||
),
|
||
)
|
||
parser.add_argument(
|
||
"--sft_data",
|
||
type=Path,
|
||
required=True,
|
||
help="Path to the JSONL SFT training data file.",
|
||
)
|
||
|
||
# --- Optional paths -----------------------------------------------------
|
||
parser.add_argument(
|
||
"--val_data",
|
||
type=Path,
|
||
default=None,
|
||
help="Optional path to JSONL SFT validation data file.",
|
||
)
|
||
parser.add_argument(
|
||
"--checkpoint_dir",
|
||
type=Path,
|
||
default=Path("checkpoints/korean_1b_sft"),
|
||
help="Root directory for saving SFT checkpoints.",
|
||
)
|
||
parser.add_argument(
|
||
"--resume",
|
||
type=Path,
|
||
default=None,
|
||
help="Path to an SFT checkpoint directory to resume fine-tuning from.",
|
||
)
|
||
parser.add_argument(
|
||
"--tokenizer",
|
||
type=Path,
|
||
default=None,
|
||
help=(
|
||
"Override path to tokenizer.json. "
|
||
"Defaults to <base_checkpoint>/tokenizer.json, "
|
||
"then falls back to tokenizer/korean_sp/tokenizer.json."
|
||
),
|
||
)
|
||
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."
|
||
),
|
||
)
|
||
|
||
# --- Training hyper-parameters ------------------------------------------
|
||
parser.add_argument(
|
||
"--max_steps",
|
||
type=int,
|
||
default=3000,
|
||
help="Total number of optimiser steps.",
|
||
)
|
||
parser.add_argument(
|
||
"--batch_size",
|
||
type=int,
|
||
default=4,
|
||
help="Per-GPU micro-batch size.",
|
||
)
|
||
parser.add_argument(
|
||
"--lr",
|
||
type=float,
|
||
default=2e-5,
|
||
help=(
|
||
"Peak learning rate. "
|
||
"SFT uses a much lower lr than pretraining (2e-5 vs 2e-4) "
|
||
"to preserve pretrained representations."
|
||
),
|
||
)
|
||
parser.add_argument(
|
||
"--weight_decay",
|
||
type=float,
|
||
default=0.01,
|
||
help="AdamW weight decay. Lower than pretrain (0.01 vs 0.1).",
|
||
)
|
||
parser.add_argument(
|
||
"--warmup_steps",
|
||
type=int,
|
||
default=100,
|
||
help="Number of linear LR warmup steps.",
|
||
)
|
||
parser.add_argument(
|
||
"--grad_accum",
|
||
type=int,
|
||
default=2,
|
||
help="Gradient accumulation steps.",
|
||
)
|
||
parser.add_argument(
|
||
"--seed",
|
||
type=int,
|
||
default=42,
|
||
help="Base random seed (rank offset is added automatically in DDP).",
|
||
)
|
||
parser.add_argument(
|
||
"--use_fp8",
|
||
action="store_true",
|
||
default=False,
|
||
help=(
|
||
"Enable TransformerEngine FP8 training "
|
||
"(requires B200/H100, uses MXFP8BlockScaling)."
|
||
),
|
||
)
|
||
|
||
# --- Single-GPU device override (ignored when using torchrun) -----------
|
||
parser.add_argument(
|
||
"--device",
|
||
type=str,
|
||
default=None,
|
||
help=(
|
||
"Explicit device string (e.g. 'cuda:0'). "
|
||
"Ignored when running under torchrun (DDP auto-assigns devices)."
|
||
),
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--config", type=Path, default=None,
|
||
help="YAML config file. Values under 'train:' section are used as CLI defaults.",
|
||
)
|
||
parser.add_argument("--save_interval", type=int, default=500, help="Checkpoint save interval (steps).")
|
||
parser.add_argument("--eval_interval", type=int, default=250, help="Validation eval interval (steps).")
|
||
parser.add_argument("--neftune_alpha", type=float, default=5.0, help="NEFTune noise magnitude (0 to disable).")
|
||
parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Maximum gradient L2 norm for clipping.")
|
||
|
||
# --- Data mixing (forgetting prevention) --------------------------------
|
||
parser.add_argument(
|
||
"--pretrain_data",
|
||
type=Path,
|
||
default=None,
|
||
help="Path to pretrain .bin file for data mixing. Enables SFT+pretrain interleaving.",
|
||
)
|
||
parser.add_argument(
|
||
"--pretrain_mix_ratio",
|
||
type=float,
|
||
default=0.3,
|
||
help="Fraction of batches from pretrain data (0.3 = 30%% pretrain, 70%% SFT).",
|
||
)
|
||
|
||
# First pass: just get --config
|
||
args, remaining = parser.parse_known_args()
|
||
|
||
# Load YAML config and apply values as defaults
|
||
if args.config is not None:
|
||
if not args.config.exists():
|
||
raise FileNotFoundError(f"Config file not found: {args.config}")
|
||
import yaml
|
||
with open(args.config, "r") as f:
|
||
yaml_cfg = yaml.safe_load(f)
|
||
train_section = yaml_cfg.get("train", {})
|
||
yaml_to_arg = {
|
||
"max_steps": "max_steps",
|
||
"batch_size": "batch_size",
|
||
"lr": "lr",
|
||
"weight_decay": "weight_decay",
|
||
"warmup_steps": "warmup_steps",
|
||
"grad_accum_steps": "grad_accum",
|
||
"save_interval": "save_interval",
|
||
"eval_interval": "eval_interval",
|
||
"neftune_alpha": "neftune_alpha",
|
||
"pretrain_mix_ratio": "pretrain_mix_ratio",
|
||
"max_grad_norm": "max_grad_norm",
|
||
}
|
||
# pretrain_data is a Path, handle separately
|
||
if "pretrain_data" in train_section:
|
||
parser.set_defaults(pretrain_data=Path(train_section["pretrain_data"]))
|
||
new_defaults = {}
|
||
for yaml_key, arg_name in yaml_to_arg.items():
|
||
if yaml_key in train_section:
|
||
new_defaults[arg_name] = train_section[yaml_key]
|
||
if new_defaults:
|
||
parser.set_defaults(**new_defaults)
|
||
|
||
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
|
||
# (Copied from pretrain.py to avoid circular import; identical logic)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
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 (Linear, etc.)
|
||
- 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] = []
|
||
|
||
# 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'.
|
||
no_decay_name_suffixes = ("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},
|
||
]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tokenizer resolution helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _resolve_tokenizer_path(args: argparse.Namespace) -> Path:
|
||
"""
|
||
Determine the tokenizer path in priority order:
|
||
1. Explicit --tokenizer argument
|
||
2. tokenizer.json inside the base_checkpoint directory
|
||
3. Project default: tokenizer/korean_sp/tokenizer.json
|
||
"""
|
||
if args.tokenizer is not None:
|
||
p = Path(args.tokenizer)
|
||
if not p.exists():
|
||
raise FileNotFoundError(f"Tokenizer not found at --tokenizer path: {p}")
|
||
return p
|
||
|
||
ckpt_tok = args.base_checkpoint / "tokenizer.json"
|
||
if ckpt_tok.exists():
|
||
return ckpt_tok
|
||
|
||
default_tok = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
|
||
if default_tok.exists():
|
||
return default_tok
|
||
|
||
raise FileNotFoundError(
|
||
"Could not locate tokenizer.json. Tried:\n"
|
||
f" 1. {ckpt_tok}\n"
|
||
f" 2. {default_tok}\n"
|
||
"Use --tokenizer to specify an explicit path."
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Dynamic padding collate function
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def dynamic_collate_fn(batch: list) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
"""
|
||
Collate function that pads each batch to its own maximum sequence length
|
||
instead of a fixed global max_seq_len. This reduces wasted FLOPs on
|
||
short sequences and speeds up SFT which tends to have highly variable
|
||
response lengths.
|
||
|
||
Pads to the batch-local max, aligned to 64 tokens (for Flash Attention
|
||
efficiency), with a floor of 512 tokens so micro-batches are not too short.
|
||
|
||
Args:
|
||
batch: List of ``(input_ids, labels)`` tuples from SFTDataset.
|
||
|
||
Returns:
|
||
Tuple of ``(input_ids, labels, attention_mask)`` tensors shaped
|
||
``[B, max_len]``.
|
||
``input_ids`` is right-padded with 0 (pad token).
|
||
``labels`` is right-padded with -1 (cross-entropy ignore_index).
|
||
``attention_mask`` is 1 for real tokens, 0 for padding.
|
||
"""
|
||
# 64-token alignment + minimum 512 floor
|
||
raw_max = max(item[0].size(0) for item in batch)
|
||
max_len = max(512, ((raw_max + 63) // 64) * 64)
|
||
|
||
input_ids_list, labels_list, mask_list = [], [], []
|
||
for ids, labs in batch:
|
||
pad_len = max_len - ids.size(0)
|
||
input_ids_list.append(F.pad(ids, (0, pad_len), value=0))
|
||
labels_list.append(F.pad(labs, (0, pad_len), value=-1))
|
||
mask_list.append(
|
||
F.pad(torch.ones(ids.size(0), dtype=torch.long), (0, pad_len), value=0)
|
||
)
|
||
|
||
return (
|
||
torch.stack(input_ids_list),
|
||
torch.stack(labels_list),
|
||
torch.stack(mask_list),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# NEFTune helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def add_neftune_hook(model: torch.nn.Module, noise_alpha: float = 10.0):
|
||
"""
|
||
Register a forward hook on the model's input embedding layer that adds
|
||
uniform noise scaled by noise_alpha during training (NEFTune).
|
||
|
||
Reference: "NEFTune: Noisy Embeddings Improve Instruction Finetuning"
|
||
(Jain et al., 2023). https://arxiv.org/abs/2310.05914
|
||
|
||
Args:
|
||
model: Raw (non-DDP) model instance.
|
||
noise_alpha: Noise magnitude parameter (paper default: 10).
|
||
|
||
Returns:
|
||
The hook handle (call ``handle.remove()`` to deactivate), or None if
|
||
the embedding layer could not be located.
|
||
"""
|
||
# Unwrap DDP if needed
|
||
raw = model.module if hasattr(model, "module") else model
|
||
|
||
# 1) Try the standard HuggingFace accessor first.
|
||
embedding: torch.nn.Embedding | None = None
|
||
if hasattr(raw, "get_input_embeddings"):
|
||
try:
|
||
emb = raw.get_input_embeddings()
|
||
if isinstance(emb, torch.nn.Embedding):
|
||
embedding = emb
|
||
except Exception:
|
||
pass
|
||
|
||
# 2) Fallback: walk common attribute paths found in open-source LLMs.
|
||
if embedding is None:
|
||
for attr_path in [
|
||
"embedding",
|
||
"embed_tokens",
|
||
"token_embedding",
|
||
"wte",
|
||
"word_embeddings",
|
||
"tok_embeddings",
|
||
"transformer.wte",
|
||
"model.embed_tokens",
|
||
"model.embedding",
|
||
]:
|
||
obj = raw
|
||
for part in attr_path.split("."):
|
||
obj = getattr(obj, part, None)
|
||
if obj is None:
|
||
break
|
||
if obj is not None and isinstance(obj, torch.nn.Embedding):
|
||
embedding = obj
|
||
break
|
||
|
||
if embedding is None:
|
||
print("[WARN] NEFTune: embedding layer을 찾지 못함, NEFTune 비활성화")
|
||
return None
|
||
|
||
print(
|
||
f"[INFO] NEFTune: {type(embedding).__name__} hook 등록 "
|
||
f"(shape={tuple(embedding.weight.shape)}, alpha={noise_alpha})"
|
||
)
|
||
|
||
def _hook(
|
||
module: torch.nn.Module,
|
||
inp: tuple,
|
||
out: torch.Tensor,
|
||
) -> torch.Tensor:
|
||
if module.training:
|
||
# out shape: [B, seq_len, d_model]
|
||
mag = noise_alpha / ((out.size(1) * out.size(2)) ** 0.5)
|
||
out = out + torch.empty_like(out).uniform_(-mag, mag)
|
||
return out
|
||
|
||
return embedding.register_forward_hook(_hook)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
|
||
# ---- Distributed setup -------------------------------------------------
|
||
is_ddp = "RANK" in os.environ
|
||
rank = 0
|
||
local_rank = 0
|
||
world_size = 1
|
||
|
||
if is_ddp:
|
||
rank, local_rank, world_size, device = setup_ddp()
|
||
else:
|
||
# Single-GPU: honour --device flag, else pick cuda:0 or cpu.
|
||
if args.device is not None:
|
||
device = torch.device(args.device)
|
||
elif torch.cuda.is_available():
|
||
device = torch.device("cuda:0")
|
||
else:
|
||
device = torch.device("cpu")
|
||
|
||
# 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)
|
||
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}")
|
||
|
||
# ---- Validate base checkpoint ------------------------------------------
|
||
if not args.base_checkpoint.exists():
|
||
raise FileNotFoundError(
|
||
f"Base checkpoint directory not found: {args.base_checkpoint}"
|
||
)
|
||
for required_file in ("model.pt", "config.yaml"):
|
||
if not (args.base_checkpoint / required_file).exists():
|
||
raise FileNotFoundError(
|
||
f"Expected {required_file} inside base checkpoint: {args.base_checkpoint}"
|
||
)
|
||
|
||
# ---- Load pretrained model ---------------------------------------------
|
||
# LLM.from_pretrained() reads config.yaml + model.pt and returns the model on CPU.
|
||
# We move it to the target device immediately after loading.
|
||
#
|
||
# NOTE: fp8_model_init() is intentionally NOT used here (same as pretrain.py).
|
||
# MXFP8Tensor weights are incompatible with DDP's _broadcast_coalesced.
|
||
# Weights stay in float32; TransformerEngine quantizes on-the-fly inside fp8_autocast.
|
||
model = LLM.from_pretrained(args.base_checkpoint)
|
||
|
||
# When FP8 flag is passed at SFT time, enable it on the loaded config.
|
||
# This is useful if the pretrained model was trained without FP8 but you
|
||
# want to fine-tune with FP8 precision (the TE layers must exist in the model).
|
||
if args.use_fp8:
|
||
model.config.use_fp8 = True
|
||
|
||
# Move model to target device in bfloat16 (more memory-efficient than fp32
|
||
# for fine-tuning, and required when BF16 autocast + TE are active).
|
||
model = model.to(device=device, dtype=torch.bfloat16)
|
||
|
||
# ---- Gradient checkpointing ----------------------------------------
|
||
# Trades activation memory for recomputation during backward pass.
|
||
# Especially useful for large models / long sequences in SFT.
|
||
if hasattr(model, 'gradient_checkpointing_enable'):
|
||
model.gradient_checkpointing_enable()
|
||
if rank == 0:
|
||
print("[INFO] Gradient checkpointing enabled")
|
||
|
||
# FP8 alignment check: (batch_size × seq_len) must be divisible by 8.
|
||
if model.config.use_fp8:
|
||
seq_len = model.config.max_seq_len
|
||
if (args.batch_size * seq_len) % 8 != 0:
|
||
raise ValueError(
|
||
f"FP8: batch_size × max_seq_len = {args.batch_size} × {seq_len} "
|
||
f"= {args.batch_size * seq_len} must be divisible by 8."
|
||
)
|
||
|
||
if is_main_process():
|
||
total_params = sum(p.numel() for p in model.parameters())
|
||
print(f"Pretrained model loaded: {total_params:,} parameters")
|
||
print(f"LMConfig: {model.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,
|
||
bucket_cap_mb=800,
|
||
find_unused_parameters=False,
|
||
)
|
||
|
||
# ---- Tokenizer ---------------------------------------------------------
|
||
tokenizer_path = _resolve_tokenizer_path(args)
|
||
if is_main_process():
|
||
print(f"Loading tokenizer from: {tokenizer_path}")
|
||
|
||
# Use the fast tokenizers library (same as the rest of the project).
|
||
from tokenizers import Tokenizer # type: ignore[import]
|
||
tokenizer = Tokenizer.from_file(str(tokenizer_path))
|
||
|
||
# ---- Dataset & DataLoader ----------------------------------------------
|
||
# Import SFTDataset (created separately alongside this file).
|
||
# SFTDataset returns (input_ids, targets) where prompt token positions in
|
||
# targets are filled with -1. The Trainer._compute_loss already uses
|
||
# ignore_index=-1, so only response tokens contribute to the gradient.
|
||
from data.sft_dataset import SFTDataset # type: ignore[import]
|
||
|
||
max_seq_len_cfg = (
|
||
model.config.max_seq_len
|
||
if not isinstance(model, torch.nn.parallel.DistributedDataParallel)
|
||
else model.module.config.max_seq_len
|
||
)
|
||
|
||
# DDP optimization: rank 0 does tokenization + cache, other ranks load cache.
|
||
# This avoids 8× redundant work and 8× memory usage.
|
||
# Rank 0 gets all 64 CPU cores for parallel tokenization (reserve 8 for system)
|
||
tok_workers = 64 if is_main_process() else 0
|
||
if is_ddp:
|
||
if is_main_process():
|
||
# Rank 0: full tokenization with parallel workers + save cache
|
||
train_dataset = SFTDataset(
|
||
data_path=args.sft_data,
|
||
tokenizer=tokenizer,
|
||
max_seq_len=max_seq_len_cfg,
|
||
tokenizer_path=tokenizer_path,
|
||
num_workers=tok_workers,
|
||
)
|
||
# Barrier: wait for rank 0 to finish tokenization and save cache
|
||
torch.distributed.barrier()
|
||
if not is_main_process():
|
||
# Other ranks: load from cache (rank 0 already saved it)
|
||
train_dataset = SFTDataset(
|
||
data_path=args.sft_data,
|
||
tokenizer=tokenizer,
|
||
max_seq_len=max_seq_len_cfg,
|
||
)
|
||
else:
|
||
train_dataset = SFTDataset(
|
||
data_path=args.sft_data,
|
||
tokenizer=tokenizer,
|
||
max_seq_len=max_seq_len_cfg,
|
||
tokenizer_path=tokenizer_path,
|
||
num_workers=tok_workers,
|
||
)
|
||
|
||
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,
|
||
# SFT datasets are typically small enough that 2–4 workers suffice.
|
||
# We use 4 to balance I/O with CPU parsing overhead from JSONL.
|
||
num_workers=4,
|
||
pin_memory=True,
|
||
drop_last=True,
|
||
prefetch_factor=2,
|
||
persistent_workers=True,
|
||
collate_fn=dynamic_collate_fn,
|
||
)
|
||
|
||
# ---- Pretrain Data Mixing (forgetting prevention) -----------------------
|
||
# When --pretrain_data is specified, create a second DataLoader for pretrain
|
||
# data and wrap both in MixingDataLoader. This interleaves SFT and pretrain
|
||
# batches (default 70/30 ratio) so the model retains pretrained knowledge.
|
||
pretrain_sampler = None
|
||
if args.pretrain_data is not None:
|
||
if not args.pretrain_data.exists():
|
||
raise FileNotFoundError(f"Pretrain data not found: {args.pretrain_data}")
|
||
|
||
from data import PackedDataset
|
||
|
||
max_seq_len = (
|
||
model.config.max_seq_len
|
||
if not isinstance(model, torch.nn.parallel.DistributedDataParallel)
|
||
else model.module.config.max_seq_len
|
||
)
|
||
pretrain_dataset = PackedDataset(args.pretrain_data, seq_len=max_seq_len)
|
||
|
||
if is_ddp:
|
||
pretrain_sampler = DistributedSampler(
|
||
pretrain_dataset,
|
||
num_replicas=world_size,
|
||
rank=rank,
|
||
shuffle=True,
|
||
seed=args.seed + 1000, # different seed from SFT
|
||
)
|
||
else:
|
||
pretrain_sampler = RandomSampler(pretrain_dataset)
|
||
|
||
pretrain_loader = DataLoader(
|
||
pretrain_dataset,
|
||
batch_size=args.batch_size,
|
||
sampler=pretrain_sampler,
|
||
num_workers=4,
|
||
pin_memory=True,
|
||
drop_last=True,
|
||
prefetch_factor=2,
|
||
persistent_workers=True,
|
||
)
|
||
|
||
# Wrap both loaders in MixingDataLoader
|
||
effective_loader = MixingDataLoader(
|
||
sft_loader=train_loader,
|
||
pretrain_loader=pretrain_loader,
|
||
pretrain_ratio=args.pretrain_mix_ratio,
|
||
sft_sampler=train_sampler if is_ddp else None,
|
||
pretrain_sampler=pretrain_sampler if is_ddp else None,
|
||
)
|
||
|
||
if is_main_process():
|
||
print(
|
||
f"[INFO] Data mixing enabled: "
|
||
f"{(1 - args.pretrain_mix_ratio) * 100:.0f}% SFT + "
|
||
f"{args.pretrain_mix_ratio * 100:.0f}% pretrain"
|
||
)
|
||
print(f"[INFO] Pretrain data: {args.pretrain_data} ({len(pretrain_dataset):,} samples)")
|
||
else:
|
||
effective_loader = train_loader
|
||
|
||
# Optional validation loader.
|
||
# NOTE: The current Trainer implementation does not yet accept a val_loader
|
||
# argument; the eval_interval config field is reserved for future use.
|
||
# We construct the loader here so that once Trainer gains eval support,
|
||
# wiring it in requires only passing val_loader=val_loader below.
|
||
val_loader: DataLoader | None = None
|
||
if args.val_data is not None:
|
||
if not args.val_data.exists():
|
||
raise FileNotFoundError(f"Validation data not found: {args.val_data}")
|
||
if is_ddp:
|
||
if is_main_process():
|
||
val_dataset = SFTDataset(
|
||
data_path=args.val_data,
|
||
tokenizer=tokenizer,
|
||
max_seq_len=train_dataset.max_seq_len,
|
||
tokenizer_path=tokenizer_path,
|
||
num_workers=tok_workers,
|
||
)
|
||
torch.distributed.barrier()
|
||
if not is_main_process():
|
||
val_dataset = SFTDataset(
|
||
data_path=args.val_data,
|
||
tokenizer=tokenizer,
|
||
max_seq_len=train_dataset.max_seq_len,
|
||
)
|
||
else:
|
||
val_dataset = SFTDataset(
|
||
data_path=args.val_data,
|
||
tokenizer=tokenizer,
|
||
max_seq_len=train_dataset.max_seq_len,
|
||
tokenizer_path=tokenizer_path,
|
||
num_workers=tok_workers,
|
||
)
|
||
val_loader = DataLoader(
|
||
val_dataset,
|
||
batch_size=args.batch_size,
|
||
shuffle=False,
|
||
num_workers=2,
|
||
pin_memory=True,
|
||
drop_last=False,
|
||
collate_fn=dynamic_collate_fn,
|
||
)
|
||
if is_main_process():
|
||
print(f"Validation dataset: {len(val_dataset):,} samples")
|
||
|
||
# ---- Optimizer ---------------------------------------------------------
|
||
# Use the same two-group split (weight_decay / no weight_decay) as pretrain.
|
||
# Unwrap DDP to get the raw model's parameters.
|
||
raw_model = getattr(model, "module", model)
|
||
param_groups = build_optimizer_param_groups(raw_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.
|
||
)
|
||
|
||
# ---- TrainConfig -------------------------------------------------------
|
||
# Set use_fp8 from the (possibly overridden) model config so Trainer builds
|
||
# the correct FP8 recipe and wraps forward passes in fp8_autocast.
|
||
use_fp8 = raw_model.config.use_fp8
|
||
|
||
train_config = TrainConfig(
|
||
max_steps=args.max_steps,
|
||
checkpoint_dir=str(args.checkpoint_dir),
|
||
grad_accum_steps=args.grad_accum,
|
||
use_fp8=use_fp8,
|
||
log_file=str(args.log_file) if args.log_file is not None else None,
|
||
save_interval=args.save_interval,
|
||
log_interval=10,
|
||
eval_interval=args.eval_interval,
|
||
max_grad_norm=args.max_grad_norm,
|
||
)
|
||
|
||
# ---- LR Scheduler ------------------------------------------------------
|
||
scheduler = get_cosine_schedule_with_warmup(
|
||
optimizer=optimizer,
|
||
warmup_steps=args.warmup_steps,
|
||
total_steps=train_config.max_steps,
|
||
)
|
||
|
||
# ---- Resume from SFT checkpoint ----------------------------------------
|
||
# When --resume is given we restore the SFT optimizer/scheduler state as
|
||
# well so learning rate, momentum buffers, etc. are correctly restored.
|
||
# NOTE: This resumes SFT training, NOT the pretrain checkpoint.
|
||
# The pretrain weights were already loaded above via from_pretrained().
|
||
start_step = 0
|
||
if args.resume is not None:
|
||
if not args.resume.exists():
|
||
raise FileNotFoundError(f"Resume checkpoint 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 SFT from {args.resume} at step {start_step} (loss={resume_loss:.4f})")
|
||
|
||
if args.resume is not None and isinstance(train_sampler, DistributedSampler):
|
||
steps_per_epoch = len(train_loader)
|
||
approx_epoch = start_step // steps_per_epoch if steps_per_epoch > 0 else 0
|
||
train_sampler.set_epoch(approx_epoch)
|
||
if is_main_process():
|
||
print(f"[INFO] Resume: sampler epoch set to {approx_epoch}")
|
||
|
||
# ---- Checkpoint directory ----------------------------------------------
|
||
args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# ---- Copy tokenizer to checkpoint dir for easy deployment later --------
|
||
# This mirrors the tokenizer into the SFT checkpoint root so that the
|
||
# final checkpoint directory is self-contained for convert_to_hf.py, etc.
|
||
if is_main_process():
|
||
dest_tok = args.checkpoint_dir / "tokenizer.json"
|
||
if not dest_tok.exists():
|
||
shutil.copy2(str(tokenizer_path), str(dest_tok))
|
||
print(f"Tokenizer copied to {dest_tok}")
|
||
|
||
# ---- Trainer -----------------------------------------------------------
|
||
# When data mixing is active, pass effective_loader (MixingDataLoader).
|
||
# MixingDataLoader handles its own epoch cycling, so no external sampler needed.
|
||
trainer = Trainer(
|
||
model=model,
|
||
train_loader=effective_loader,
|
||
optimizer=optimizer,
|
||
scheduler=scheduler,
|
||
config=train_config,
|
||
device=device,
|
||
rank=rank,
|
||
sampler=train_sampler if is_ddp and args.pretrain_data is None else None,
|
||
val_loader=val_loader,
|
||
)
|
||
|
||
# ---- Signal handlers for graceful shutdown ----------------------------
|
||
import signal as _signal_mod
|
||
|
||
_trainer_ref = trainer
|
||
|
||
def _graceful_shutdown_handler(signum, frame):
|
||
sig_name = _signal_mod.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_mod.SIGHUP, _signal_mod.SIGTERM):
|
||
_signal_mod.signal(_sig, _graceful_shutdown_handler)
|
||
|
||
# ---- SFT banner --------------------------------------------------------
|
||
if is_main_process():
|
||
import datetime
|
||
|
||
inner_config = raw_model.config
|
||
eff_batch_seqs = args.batch_size * args.grad_accum * world_size
|
||
eff_tokens_per_step = eff_batch_seqs * inner_config.max_seq_len
|
||
train_samples = len(train_dataset)
|
||
precision_label = "FP8 (MXFP8BlockScaling)" if use_fp8 else "BF16"
|
||
nccl_debug = os.environ.get("NCCL_DEBUG", "not set")
|
||
omp_threads = os.environ.get("OMP_NUM_THREADS", "not set")
|
||
|
||
mix_label = "none"
|
||
if args.pretrain_data is not None:
|
||
mix_label = (
|
||
f"{(1 - args.pretrain_mix_ratio) * 100:.0f}% SFT + "
|
||
f"{args.pretrain_mix_ratio * 100:.0f}% pretrain"
|
||
)
|
||
|
||
print(
|
||
f"\n{'='*70}\n"
|
||
f" LLM Supervised Fine-Tuning — "
|
||
f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
|
||
f"{'='*70}\n"
|
||
f" base ckpt : {args.base_checkpoint}\n"
|
||
f" sft data : {args.sft_data} ({train_samples:,} samples)\n"
|
||
f" data mix : {mix_label}\n"
|
||
f" model : {inner_config.num_params:,} params | "
|
||
f"d_model={inner_config.d_model} n_layers={inner_config.n_layers}\n"
|
||
f" precision : {precision_label}\n"
|
||
f" GPUs : {world_size} | batch/GPU={args.batch_size} "
|
||
f"grad_accum={args.grad_accum}\n"
|
||
f" eff_batch : {eff_batch_seqs} seqs "
|
||
f"= {eff_tokens_per_step:,} tok/step\n"
|
||
f" max_steps : {train_config.max_steps:,}\n"
|
||
f" lr : {args.lr:.2e} "
|
||
f"warmup={args.warmup_steps} weight_decay={args.weight_decay}\n"
|
||
f" ckpt_dir : {args.checkpoint_dir}\n"
|
||
f" env : OMP_NUM_THREADS={omp_threads} NCCL_DEBUG={nccl_debug}\n"
|
||
f"{'='*70}\n"
|
||
)
|
||
|
||
# ---- NEFTune -----------------------------------------------------------
|
||
# Add uniform noise to embeddings during training to improve instruction
|
||
# following (Jain et al., 2023). Hook is registered on the raw (non-DDP)
|
||
# model so it survives DDP's internal module wrapping.
|
||
neftune_alpha = getattr(args, 'neftune_alpha', 5.0)
|
||
neftune_handle = add_neftune_hook(raw_model, noise_alpha=neftune_alpha)
|
||
if rank == 0:
|
||
if neftune_handle is not None:
|
||
print(f"[INFO] NEFTune enabled (noise_alpha={neftune_alpha})")
|
||
else:
|
||
print("[WARN] NEFTune disabled - embedding layer not found")
|
||
|
||
# ---- Train -------------------------------------------------------------
|
||
try:
|
||
trainer.train(start_step=start_step)
|
||
except KeyboardInterrupt:
|
||
if is_main_process():
|
||
print("\n[INFO] SFT interrupted by user (KeyboardInterrupt).")
|
||
except Exception as e:
|
||
import traceback
|
||
if is_main_process():
|
||
tb = traceback.format_exc()
|
||
print(f"\n[ERROR] SFT failed at rank {rank}:\n{tb}")
|
||
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')}] "
|
||
f"[FATAL] {tb}\n"
|
||
)
|
||
raise
|
||
finally:
|
||
# Remove NEFTune hook so the model is clean for inference/saving.
|
||
if neftune_handle is not None:
|
||
neftune_handle.remove()
|
||
if is_ddp:
|
||
cleanup_ddp()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|