pathcosmos/EVAFRILL-Mo-3B
130
1"""2train/sft.py — Supervised Fine-Tuning (SFT) entry point.3 4Loads a pretrained checkpoint and fine-tunes it on instruction/conversation5data using SFTDataset, which masks prompt tokens with ignore_index=-1 so only6the assistant response tokens contribute to the loss.7 8Launch single-GPU:9 python train/sft.py \\10 --base_checkpoint checkpoints/korean_1b_fp8_run1/checkpoint-0034000 \\11 --sft_data data/sft/train.jsonl \\12 --device cuda:013 14Launch multi-GPU (DDP via torchrun, 7 GPU):15 torchrun --nproc_per_node=7 train/sft.py \\16 --base_checkpoint checkpoints/3b_final/checkpoint-0319772 \\17 --sft_data data/sft_combined/train_filtered.jsonl18 19KEY DIFFERENCES from pretrain.py:20 - Loads weights from a pretrained checkpoint via LLM.from_pretrained()21 - Uses SFTDataset (JSONL instruction data) instead of PackedDataset22 - Lower default learning rate (2e-5 vs 2e-4)23 - Fewer default steps (3000 vs 100000)24 - Copies tokenizer.json to checkpoint_dir for easy deployment25"""26 27from __future__ import annotations28 29import argparse30import os31import random32import signal33import shutil34import sys35from pathlib import Path36 37import numpy as np38import torch39import torch.nn.functional as F40from torch.utils.data import DataLoader, DistributedSampler, RandomSampler41 42# B200 Tensor Core 최대 활용: TF32 matmul + cuDNN43torch.backends.cuda.matmul.allow_tf32 = True44torch.backends.cudnn.allow_tf32 = True45torch.set_float32_matmul_precision("high") # TF32 precision for fp32 matmul46 47# Allow imports from the project root regardless of working directory.48_PROJECT_ROOT = Path(__file__).resolve().parent.parent49if str(_PROJECT_ROOT) not in sys.path:50 sys.path.insert(0, str(_PROJECT_ROOT))51 52from model import LLM53from train.trainer import TrainConfig, Trainer54from train.utils import (55 cleanup_ddp,56 get_cosine_schedule_with_warmup,57 is_main_process,58 load_checkpoint,59 setup_ddp,60)61 62# ---------------------------------------------------------------------------63# Optional TransformerEngine import (FP8 support)64# ---------------------------------------------------------------------------65try:66 import transformer_engine.pytorch as te # type: ignore[import]67 HAS_TE = True68except ImportError:69 te = None # type: ignore[assignment]70 HAS_TE = False71 72 73# ---------------------------------------------------------------------------74# Argument parsing75# ---------------------------------------------------------------------------76 77 78def parse_args() -> argparse.Namespace:79 parser = argparse.ArgumentParser(80 description="Supervised Fine-Tuning (SFT) of a pretrained decoder-only LLM.",81 formatter_class=argparse.ArgumentDefaultsHelpFormatter,82 )83 84 # --- Required paths -----------------------------------------------------85 parser.add_argument(86 "--base_checkpoint",87 type=Path,88 required=True,89 help=(90 "Path to the pretrained checkpoint directory. "91 "Must contain model.pt and config.yaml (produced by save_checkpoint)."92 ),93 )94 parser.add_argument(95 "--sft_data",96 type=Path,97 required=True,98 help="Path to the JSONL SFT training data file.",99 )100 101 # --- Optional paths -----------------------------------------------------102 parser.add_argument(103 "--val_data",104 type=Path,105 default=None,106 help="Optional path to JSONL SFT validation data file.",107 )108 parser.add_argument(109 "--checkpoint_dir",110 type=Path,111 default=Path("checkpoints/korean_1b_sft"),112 help="Root directory for saving SFT checkpoints.",113 )114 parser.add_argument(115 "--resume",116 type=Path,117 default=None,118 help="Path to an SFT checkpoint directory to resume fine-tuning from.",119 )120 parser.add_argument(121 "--tokenizer",122 type=Path,123 default=None,124 help=(125 "Override path to tokenizer.json. "126 "Defaults to <base_checkpoint>/tokenizer.json, "127 "then falls back to tokenizer/korean_sp/tokenizer.json."128 ),129 )130 parser.add_argument(131 "--log_file",132 type=Path,133 default=None,134 help=(135 "Path to a text file for structured training logs (rank-0 only). "136 "If omitted, logs go only to stdout."137 ),138 )139 140 # --- Training hyper-parameters ------------------------------------------141 parser.add_argument(142 "--max_steps",143 type=int,144 default=3000,145 help="Total number of optimiser steps.",146 )147 parser.add_argument(148 "--batch_size",149 type=int,150 default=4,151 help="Per-GPU micro-batch size.",152 )153 parser.add_argument(154 "--lr",155 type=float,156 default=2e-5,157 help=(158 "Peak learning rate. "159 "SFT uses a much lower lr than pretraining (2e-5 vs 2e-4) "160 "to preserve pretrained representations."161 ),162 )163 parser.add_argument(164 "--weight_decay",165 type=float,166 default=0.01,167 help="AdamW weight decay. Lower than pretrain (0.01 vs 0.1).",168 )169 parser.add_argument(170 "--warmup_steps",171 type=int,172 default=100,173 help="Number of linear LR warmup steps.",174 )175 parser.add_argument(176 "--grad_accum",177 type=int,178 default=2,179 help="Gradient accumulation steps.",180 )181 parser.add_argument(182 "--seed",183 type=int,184 default=42,185 help="Base random seed (rank offset is added automatically in DDP).",186 )187 parser.add_argument(188 "--use_fp8",189 action="store_true",190 default=False,191 help=(192 "Enable TransformerEngine FP8 training "193 "(requires B200/H100, uses MXFP8BlockScaling)."194 ),195 )196 197 # --- Single-GPU device override (ignored when using torchrun) -----------198 parser.add_argument(199 "--device",200 type=str,201 default=None,202 help=(203 "Explicit device string (e.g. 'cuda:0'). "204 "Ignored when running under torchrun (DDP auto-assigns devices)."205 ),206 )207 208 parser.add_argument(209 "--config", type=Path, default=None,210 help="YAML config file. Values under 'train:' section are used as CLI defaults.",211 )212 parser.add_argument("--save_interval", type=int, default=500, help="Checkpoint save interval (steps).")213 parser.add_argument("--eval_interval", type=int, default=250, help="Validation eval interval (steps).")214 parser.add_argument("--neftune_alpha", type=float, default=5.0, help="NEFTune noise magnitude (0 to disable).")215 parser.add_argument("--no_fp8", action="store_true", default=False, help="Force disable FP8 even if pretrained config has use_fp8=True.")216 parser.add_argument("--num_workers", type=int, default=4, help="Number of DataLoader worker processes.")217 parser.add_argument("--max_val_batches", type=int, default=0, help="Max validation batches (0=unlimited).")218 219 # First pass: just get --config220 args, remaining = parser.parse_known_args()221 222 # Load YAML config and apply values as defaults223 if args.config is not None:224 if not args.config.exists():225 raise FileNotFoundError(f"Config file not found: {args.config}")226 import yaml227 with open(args.config, "r") as f:228 yaml_cfg = yaml.safe_load(f)229 train_section = yaml_cfg.get("train", {})230 yaml_to_arg = {231 "max_steps": "max_steps",232 "batch_size": "batch_size",233 "lr": "lr",234 "weight_decay": "weight_decay",235 "warmup_steps": "warmup_steps",236 "grad_accum_steps": "grad_accum",237 "save_interval": "save_interval",238 "eval_interval": "eval_interval",239 "neftune_alpha": "neftune_alpha",240 "max_val_batches": "max_val_batches",241 }242 new_defaults = {}243 for yaml_key, arg_name in yaml_to_arg.items():244 if yaml_key in train_section:245 new_defaults[arg_name] = train_section[yaml_key]246 if new_defaults:247 parser.set_defaults(**new_defaults)248 249 return parser.parse_args()250 251 252# ---------------------------------------------------------------------------253# Seed helper254# ---------------------------------------------------------------------------255 256 257def set_seed(seed: int) -> None:258 """Set deterministic seeds for Python, NumPy, and PyTorch."""259 random.seed(seed)260 np.random.seed(seed)261 torch.manual_seed(seed)262 torch.cuda.manual_seed_all(seed)263 264 265# ---------------------------------------------------------------------------266# Optimizer parameter groups267# (Copied from pretrain.py to avoid circular import; identical logic)268# ---------------------------------------------------------------------------269 270 271def build_optimizer_param_groups(272 model: torch.nn.Module,273 weight_decay: float,274) -> list[dict]:275 """276 Split parameters into two groups:277 - decay group : weight tensors with ndim >= 2 (Linear, etc.)278 - no-decay group: bias, LayerNorm/RMSNorm weights, and embedding weights279 280 This follows standard practice (e.g. GPT-style training).281 """282 decay_params: list[torch.nn.Parameter] = []283 no_decay_params: list[torch.nn.Parameter] = []284 285 # Module types whose parameters should never be decayed.286 no_decay_module_types = (287 torch.nn.Embedding,288 torch.nn.LayerNorm,289 )290 # Also skip any parameter whose name ends with '.bias'.291 no_decay_name_suffixes = ("bias",)292 293 # Collect module-level exclusions.294 no_decay_module_params: set[int] = set()295 for module in model.modules():296 if isinstance(module, no_decay_module_types):297 for param in module.parameters(recurse=False):298 no_decay_module_params.add(id(param))299 300 seen: set[int] = set()301 for name, param in model.named_parameters():302 if not param.requires_grad:303 continue304 if id(param) in seen:305 continue306 seen.add(id(param))307 308 if (309 id(param) in no_decay_module_params310 or any(name.endswith(sfx) for sfx in no_decay_name_suffixes)311 or param.ndim < 2312 ):313 no_decay_params.append(param)314 else:315 decay_params.append(param)316 317 return [318 {"params": decay_params, "weight_decay": weight_decay},319 {"params": no_decay_params, "weight_decay": 0.0},320 ]321 322 323# ---------------------------------------------------------------------------324# Tokenizer resolution helper325# ---------------------------------------------------------------------------326 327 328def _resolve_tokenizer_path(args: argparse.Namespace) -> Path:329 """330 Determine the tokenizer path in priority order:331 1. Explicit --tokenizer argument332 2. tokenizer.json inside the base_checkpoint directory333 3. Project default: tokenizer/korean_sp/tokenizer.json334 """335 if args.tokenizer is not None:336 p = Path(args.tokenizer)337 if not p.exists():338 raise FileNotFoundError(f"Tokenizer not found at --tokenizer path: {p}")339 return p340 341 ckpt_tok = args.base_checkpoint / "tokenizer.json"342 if ckpt_tok.exists():343 return ckpt_tok344 345 default_tok = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"346 if default_tok.exists():347 return default_tok348 349 raise FileNotFoundError(350 "Could not locate tokenizer.json. Tried:\n"351 f" 1. {ckpt_tok}\n"352 f" 2. {default_tok}\n"353 "Use --tokenizer to specify an explicit path."354 )355 356 357# ---------------------------------------------------------------------------358# Dynamic padding collate function359# ---------------------------------------------------------------------------360 361 362def dynamic_collate_fn(batch: list) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:363 """364 Collate function that pads each batch to its own maximum sequence length365 instead of a fixed global max_seq_len. This reduces wasted FLOPs on366 short sequences and speeds up SFT which tends to have highly variable367 response lengths.368 369 Pads to the batch-local max, aligned to 64 tokens (for Flash Attention370 efficiency), with a floor of 512 tokens so micro-batches are not too short.371 372 Args:373 batch: List of ``(input_ids, labels)`` tuples from SFTDataset.374 375 Returns:376 Tuple of ``(input_ids, labels, attention_mask)`` tensors shaped377 ``[B, max_len]``.378 ``input_ids`` is right-padded with 0 (pad token).379 ``labels`` is right-padded with -1 (cross-entropy ignore_index).380 ``attention_mask`` is 1 for real tokens, 0 for padding.381 """382 # 64-token alignment + minimum 512 floor383 raw_max = max(item[0].size(0) for item in batch)384 max_len = max(512, ((raw_max + 63) // 64) * 64)385 386 input_ids_list, labels_list, mask_list = [], [], []387 for ids, labs in batch:388 pad_len = max_len - ids.size(0)389 input_ids_list.append(F.pad(ids, (0, pad_len), value=0))390 labels_list.append(F.pad(labs, (0, pad_len), value=-1))391 mask_list.append(392 F.pad(torch.ones(ids.size(0), dtype=torch.long), (0, pad_len), value=0)393 )394 395 return (396 torch.stack(input_ids_list),397 torch.stack(labels_list),398 torch.stack(mask_list),399 )400 401 402# ---------------------------------------------------------------------------403# NEFTune helper404# ---------------------------------------------------------------------------405 406 407def add_neftune_hook(model: torch.nn.Module, noise_alpha: float = 10.0):408 """409 Register a forward hook on the model's input embedding layer that adds410 uniform noise scaled by noise_alpha during training (NEFTune).411 412 Reference: "NEFTune: Noisy Embeddings Improve Instruction Finetuning"413 (Jain et al., 2023). https://arxiv.org/abs/2310.05914414 415 Args:416 model: Raw (non-DDP) model instance.417 noise_alpha: Noise magnitude parameter (paper default: 10).418 419 Returns:420 The hook handle (call ``handle.remove()`` to deactivate), or None if421 the embedding layer could not be located.422 """423 # Unwrap DDP if needed424 raw = model.module if hasattr(model, "module") else model425 426 # 1) Try the standard HuggingFace accessor first.427 embedding: torch.nn.Embedding | None = None428 if hasattr(raw, "get_input_embeddings"):429 try:430 emb = raw.get_input_embeddings()431 if isinstance(emb, torch.nn.Embedding):432 embedding = emb433 except Exception:434 pass435 436 # 2) Fallback: walk common attribute paths found in open-source LLMs.437 if embedding is None:438 for attr_path in [439 "embedding",440 "embed_tokens",441 "token_embedding",442 "wte",443 "word_embeddings",444 "tok_embeddings",445 "transformer.wte",446 "model.embed_tokens",447 "model.embedding",448 ]:449 obj = raw450 for part in attr_path.split("."):451 obj = getattr(obj, part, None)452 if obj is None:453 break454 if obj is not None and isinstance(obj, torch.nn.Embedding):455 embedding = obj456 break457 458 if embedding is None:459 print("[WARN] NEFTune: embedding layer을 찾지 못함, NEFTune 비활성화")460 return None461 462 print(463 f"[INFO] NEFTune: {type(embedding).__name__} hook 등록 "464 f"(shape={tuple(embedding.weight.shape)}, alpha={noise_alpha})"465 )466 467 def _hook(468 module: torch.nn.Module,469 inp: tuple,470 out: torch.Tensor,471 ) -> torch.Tensor:472 if module.training:473 # out shape: [B, seq_len, d_model]474 mag = noise_alpha / ((out.size(1) * out.size(2)) ** 0.5)475 out = out + torch.empty_like(out).uniform_(-mag, mag)476 return out477 478 return embedding.register_forward_hook(_hook)479 480 481# ---------------------------------------------------------------------------482# Main483# ---------------------------------------------------------------------------484 485 486def main() -> None:487 args = parse_args()488 489 # ---- Distributed setup -------------------------------------------------490 is_ddp = "RANK" in os.environ491 rank = 0492 local_rank = 0493 world_size = 1494 495 if is_ddp:496 rank, local_rank, world_size, device = setup_ddp()497 else:498 # Single-GPU: honour --device flag, else pick cuda:0 or cpu.499 if args.device is not None:500 device = torch.device(args.device)501 elif torch.cuda.is_available():502 device = torch.device("cuda:0")503 else:504 device = torch.device("cpu")505 506 # Per-rank seed so data shuffling differs across replicas.507 set_seed(args.seed + rank)508 509 # ---- NUMA affinity for optimal GPU↔CPU memory locality ---------------510 # B200 topology: GPU 0-3 → NUMA node 0 (cores 0-35)511 # GPU 4-6 → NUMA node 1 (cores 36-71) [7 GPU 환경]512 try:513 if local_rank < 4:514 os.sched_setaffinity(0, set(range(0, 36))) # NUMA node 0515 else:516 os.sched_setaffinity(0, set(range(36, 72))) # NUMA node 1517 if is_main_process():518 print(f"NUMA affinity: rank {rank} (GPU {local_rank}) → "519 f"{'NUMA0 cores 0-35' if local_rank < 4 else 'NUMA1 cores 36-71'}")520 except (AttributeError, OSError) as e:521 if is_main_process():522 print(f"[WARN] NUMA affinity failed: {e}")523 524 # ---- Validate base checkpoint ------------------------------------------525 if not args.base_checkpoint.exists():526 raise FileNotFoundError(527 f"Base checkpoint directory not found: {args.base_checkpoint}"528 )529 for required_file in ("model.pt", "config.yaml"):530 if not (args.base_checkpoint / required_file).exists():531 raise FileNotFoundError(532 f"Expected {required_file} inside base checkpoint: {args.base_checkpoint}"533 )534 535 # ---- Load pretrained model ---------------------------------------------536 # LLM.from_pretrained() reads config.yaml + model.pt and returns the model on CPU.537 # We move it to the target device immediately after loading.538 #539 # NOTE: fp8_model_init() is intentionally NOT used here (same as pretrain.py).540 # MXFP8Tensor weights are incompatible with DDP's _broadcast_coalesced.541 # Weights stay in float32; TransformerEngine quantizes on-the-fly inside fp8_autocast.542 model = LLM.from_pretrained(args.base_checkpoint)543 544 # FP8 override: --no_fp8 forces BF16 even if pretrained config had use_fp8=True.545 # --use_fp8 enables FP8 if pretrained config had it disabled.546 if args.no_fp8:547 model.config.use_fp8 = False548 elif args.use_fp8:549 model.config.use_fp8 = True550 551 # Move model to target device in bfloat16 (more memory-efficient than fp32552 # for fine-tuning, and required when BF16 autocast + TE are active).553 model = model.to(device=device, dtype=torch.bfloat16)554 555 # ---- Gradient checkpointing ----------------------------------------556 # Trades activation memory for recomputation during backward pass.557 # Especially useful for large models / long sequences in SFT.558 if hasattr(model, 'gradient_checkpointing_enable'):559 model.gradient_checkpointing_enable()560 if rank == 0:561 print("[INFO] Gradient checkpointing enabled")562 563 # FP8 alignment check: (batch_size × seq_len) must be divisible by 8.564 if model.config.use_fp8:565 seq_len = model.config.max_seq_len566 if (args.batch_size * seq_len) % 8 != 0:567 raise ValueError(568 f"FP8: batch_size × max_seq_len = {args.batch_size} × {seq_len} "569 f"= {args.batch_size * seq_len} must be divisible by 8."570 )571 572 if is_main_process():573 total_params = sum(p.numel() for p in model.parameters())574 print(f"Pretrained model loaded: {total_params:,} parameters")575 print(f"LMConfig: {model.config}")576 577 # ---- Wrap in DDP -------------------------------------------------------578 if is_ddp:579 from torch.nn.parallel import DistributedDataParallel as DDP580 581 model = DDP(582 model,583 device_ids=[local_rank],584 output_device=local_rank,585 gradient_as_bucket_view=True,586 bucket_cap_mb=800,587 find_unused_parameters=False,588 )589 590 # ---- Tokenizer ---------------------------------------------------------591 tokenizer_path = _resolve_tokenizer_path(args)592 if is_main_process():593 print(f"Loading tokenizer from: {tokenizer_path}")594 595 # Use the fast tokenizers library (same as the rest of the project).596 from tokenizers import Tokenizer # type: ignore[import]597 tokenizer = Tokenizer.from_file(str(tokenizer_path))598 599 # ---- Dataset & DataLoader ----------------------------------------------600 # Import SFTDataset (created separately alongside this file).601 # SFTDataset returns (input_ids, targets) where prompt token positions in602 # targets are filled with -1. The Trainer._compute_loss already uses603 # ignore_index=-1, so only response tokens contribute to the gradient.604 from data.sft_dataset import SFTDataset # type: ignore[import]605 606 train_dataset = SFTDataset(607 data_path=args.sft_data,608 tokenizer=tokenizer,609 max_seq_len=model.config.max_seq_len610 if not isinstance(model, torch.nn.parallel.DistributedDataParallel)611 else model.module.config.max_seq_len,612 )613 614 if is_ddp:615 train_sampler: DistributedSampler | RandomSampler = DistributedSampler(616 train_dataset,617 num_replicas=world_size,618 rank=rank,619 shuffle=True,620 seed=args.seed,621 )622 shuffle = False623 else:624 train_sampler = RandomSampler(train_dataset)625 shuffle = False # Sampler is provided; DataLoader must not also shuffle.626 627 train_loader = DataLoader(628 train_dataset,629 batch_size=args.batch_size,630 sampler=train_sampler,631 # SFT datasets are typically small enough that 2–4 workers suffice.632 # We use 4 to balance I/O with CPU parsing overhead from JSONL.633 num_workers=args.num_workers,634 pin_memory=True,635 drop_last=True,636 prefetch_factor=2,637 persistent_workers=True,638 collate_fn=dynamic_collate_fn,639 )640 641 # Optional validation loader.642 # NOTE: The current Trainer implementation does not yet accept a val_loader643 # argument; the eval_interval config field is reserved for future use.644 # We construct the loader here so that once Trainer gains eval support,645 # wiring it in requires only passing val_loader=val_loader below.646 val_loader: DataLoader | None = None647 if args.val_data is not None:648 if not args.val_data.exists():649 raise FileNotFoundError(f"Validation data not found: {args.val_data}")650 val_dataset = SFTDataset(651 data_path=args.val_data,652 tokenizer=tokenizer,653 max_seq_len=train_dataset.max_seq_len,654 )655 val_loader = DataLoader(656 val_dataset,657 batch_size=args.batch_size,658 shuffle=False,659 num_workers=2,660 pin_memory=True,661 drop_last=False,662 collate_fn=dynamic_collate_fn,663 )664 if is_main_process():665 print(f"Validation dataset: {len(val_dataset):,} samples")666 667 # ---- Optimizer ---------------------------------------------------------668 # Use the same two-group split (weight_decay / no weight_decay) as pretrain.669 # Unwrap DDP to get the raw model's parameters.670 raw_model = getattr(model, "module", model)671 param_groups = build_optimizer_param_groups(raw_model, args.weight_decay)672 optimizer = torch.optim.AdamW(673 param_groups,674 lr=args.lr,675 betas=(0.9, 0.95),676 eps=1e-8,677 fused=torch.cuda.is_available(), # Use fused kernel when on CUDA.678 )679 680 # ---- TrainConfig -------------------------------------------------------681 # Set use_fp8 from the (possibly overridden) model config so Trainer builds682 # the correct FP8 recipe and wraps forward passes in fp8_autocast.683 use_fp8 = raw_model.config.use_fp8684 685 train_config = TrainConfig(686 max_steps=args.max_steps,687 checkpoint_dir=str(args.checkpoint_dir),688 grad_accum_steps=args.grad_accum,689 use_fp8=use_fp8,690 log_file=str(args.log_file) if args.log_file is not None else None,691 save_interval=args.save_interval,692 log_interval=10,693 eval_interval=args.eval_interval,694 max_val_batches=args.max_val_batches,695 )696 697 # ---- LR Scheduler ------------------------------------------------------698 scheduler = get_cosine_schedule_with_warmup(699 optimizer=optimizer,700 warmup_steps=args.warmup_steps,701 total_steps=train_config.max_steps,702 )703 704 # ---- Resume from SFT checkpoint ----------------------------------------705 # When --resume is given we restore the SFT optimizer/scheduler state as706 # well so learning rate, momentum buffers, etc. are correctly restored.707 # NOTE: This resumes SFT training, NOT the pretrain checkpoint.708 # The pretrain weights were already loaded above via from_pretrained().709 start_step = 0710 if args.resume is not None:711 if not args.resume.exists():712 raise FileNotFoundError(f"Resume checkpoint not found: {args.resume}")713 start_step, resume_loss = load_checkpoint(714 path=args.resume,715 model=model,716 optimizer=optimizer,717 scheduler=scheduler,718 )719 if is_main_process():720 print(f"Resumed SFT from {args.resume} at step {start_step} (loss={resume_loss:.4f})")721 722 if args.resume is not None and isinstance(train_sampler, DistributedSampler):723 steps_per_epoch = len(train_loader)724 approx_epoch = start_step // steps_per_epoch if steps_per_epoch > 0 else 0725 train_sampler.set_epoch(approx_epoch)726 if is_main_process():727 print(f"[INFO] Resume: sampler epoch set to {approx_epoch}")728 729 # ---- Checkpoint directory ----------------------------------------------730 args.checkpoint_dir.mkdir(parents=True, exist_ok=True)731 732 # ---- Copy tokenizer to checkpoint dir for easy deployment later --------733 # This mirrors the tokenizer into the SFT checkpoint root so that the734 # final checkpoint directory is self-contained for convert_to_hf.py, etc.735 if is_main_process():736 dest_tok = args.checkpoint_dir / "tokenizer.json"737 if not dest_tok.exists():738 shutil.copy2(str(tokenizer_path), str(dest_tok))739 print(f"Tokenizer copied to {dest_tok}")740 741 # ---- Trainer -----------------------------------------------------------742 trainer = Trainer(743 model=model,744 train_loader=train_loader,745 optimizer=optimizer,746 scheduler=scheduler,747 config=train_config,748 device=device,749 rank=rank,750 sampler=train_sampler if is_ddp else None,751 val_loader=val_loader,752 )753 754 # ---- Signal handlers for graceful shutdown ----------------------------755 import signal as _signal_mod756 757 _trainer_ref = trainer758 759 def _graceful_shutdown_handler(signum, frame):760 sig_name = _signal_mod.Signals(signum).name761 if is_main_process():762 import datetime as _dt763 ts = _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")764 msg = (765 f"[{ts}] [SIGNAL] Received {sig_name} (signum={signum}). "766 f"Initiating graceful shutdown..."767 )768 print(f"\n{msg}")769 if args.log_file is not None:770 try:771 with open(args.log_file, "a", encoding="utf-8") as f:772 f.write(msg + "\n")773 except Exception:774 pass775 _trainer_ref.request_shutdown(sig_name)776 777 for _sig in (_signal_mod.SIGHUP, _signal_mod.SIGTERM):778 _signal_mod.signal(_sig, _graceful_shutdown_handler)779 780 # ---- SFT banner --------------------------------------------------------781 if is_main_process():782 import datetime783 784 inner_config = raw_model.config785 eff_batch_seqs = args.batch_size * args.grad_accum * world_size786 eff_tokens_per_step = eff_batch_seqs * inner_config.max_seq_len787 train_samples = len(train_dataset)788 precision_label = "FP8 (MXFP8BlockScaling)" if use_fp8 else "BF16"789 nccl_debug = os.environ.get("NCCL_DEBUG", "not set")790 omp_threads = os.environ.get("OMP_NUM_THREADS", "not set")791 792 print(793 f"\n{'='*70}\n"794 f" LLM Supervised Fine-Tuning — "795 f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"796 f"{'='*70}\n"797 f" base ckpt : {args.base_checkpoint}\n"798 f" sft data : {args.sft_data} ({train_samples:,} samples)\n"799 f" model : {inner_config.num_params:,} params | "800 f"d_model={inner_config.d_model} n_layers={inner_config.n_layers}\n"801 f" precision : {precision_label}\n"802 f" GPUs : {world_size} | batch/GPU={args.batch_size} "803 f"grad_accum={args.grad_accum}\n"804 f" eff_batch : {eff_batch_seqs} seqs "805 f"= {eff_tokens_per_step:,} tok/step\n"806 f" max_steps : {train_config.max_steps:,}\n"807 f" lr : {args.lr:.2e} "808 f"warmup={args.warmup_steps} weight_decay={args.weight_decay}\n"809 f" ckpt_dir : {args.checkpoint_dir}\n"810 f" env : OMP_NUM_THREADS={omp_threads} NCCL_DEBUG={nccl_debug}\n"811 f"{'='*70}\n"812 )813 814 # ---- NEFTune -----------------------------------------------------------815 # Add uniform noise to embeddings during training to improve instruction816 # following (Jain et al., 2023). Hook is registered on the raw (non-DDP)817 # model so it survives DDP's internal module wrapping.818 neftune_alpha = getattr(args, 'neftune_alpha', 5.0)819 neftune_handle = add_neftune_hook(raw_model, noise_alpha=neftune_alpha)820 if rank == 0:821 if neftune_handle is not None:822 print(f"[INFO] NEFTune enabled (noise_alpha={neftune_alpha})")823 else:824 print("[WARN] NEFTune disabled - embedding layer not found")825 826 # ---- Train -------------------------------------------------------------827 try:828 trainer.train(start_step=start_step)829 except KeyboardInterrupt:830 if is_main_process():831 print("\n[INFO] SFT interrupted by user (KeyboardInterrupt).")832 except Exception as e:833 import traceback834 if is_main_process():835 tb = traceback.format_exc()836 print(f"\n[ERROR] SFT failed at rank {rank}:\n{tb}")837 if args.log_file is not None:838 with open(args.log_file, "a", encoding="utf-8") as f:839 import datetime840 f.write(841 f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "842 f"[FATAL] {tb}\n"843 )844 raise845 finally:846 # Remove NEFTune hook so the model is clean for inference/saving.847 if neftune_handle is not None:848 neftune_handle.remove()849 if is_ddp:850 cleanup_ddp()851 852 853if __name__ == "__main__":854 main()855