初始化项目，由ModelHub XC社区提供模型

Model: pathcosmos/frankenstallm
Source: Original Platform

source/train/__init__.py

@@ -0,0 +1,41 @@
+"""
+train — LLM pretraining package.
+
+Public API:
+    TrainConfig   : Dataclass of training hyper-parameters.
+    Trainer       : Core training loop with gradient accumulation, AMP, and logging.
+
+Utility functions (re-exported from train.utils):
+    get_cosine_schedule_with_warmup
+    save_checkpoint
+    load_checkpoint
+    get_grad_norm
+    setup_ddp
+    cleanup_ddp
+    is_main_process
+"""
+
+from train.trainer import TrainConfig, Trainer
+from train.utils import (
+    cleanup_ddp,
+    get_cosine_schedule_with_warmup,
+    get_grad_norm,
+    is_main_process,
+    load_checkpoint,
+    save_checkpoint,
+    setup_ddp,
+)
+
+__all__ = [
+    # Core classes
+    "TrainConfig",
+    "Trainer",
+    # Utility functions
+    "get_cosine_schedule_with_warmup",
+    "save_checkpoint",
+    "load_checkpoint",
+    "get_grad_norm",
+    "setup_ddp",
+    "cleanup_ddp",
+    "is_main_process",
+]

source/train/orpo.py

@@ -0,0 +1,586 @@
+"""
+ORPO (Odds Ratio Preference Optimization) training script.
+Uses TRL 0.29.0 ORPOTrainer/ORPOConfig (trl.experimental.orpo).
+Optimized for 8x NVIDIA B200 GPUs (183GB VRAM each, ~1.47TB total).
+
+Usage:
+    # Full training (8 GPU DDP)
+    torchrun --nproc_per_node=8 train/orpo.py \
+        --config configs/korean_3b_orpo.yaml
+
+    # Quick test (200 steps)
+    python train/orpo.py --config configs/korean_3b_orpo.yaml --max_steps 200
+
+    # Single GPU test
+    python train/orpo.py --config configs/korean_3b_orpo.yaml --device cuda:0
+
+Prerequisites:
+    pip install trl==0.29.0 transformers accelerate peft datasets
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import json
+import logging
+import os
+import signal as _signal_mod
+import sys
+import time
+import traceback
+from pathlib import Path
+
+import torch
+from datasets import Dataset, load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    EarlyStoppingCallback,
+    TrainerCallback,
+)
+
+# TRL imports -- ORPOTrainer/ORPOConfig (TRL 0.29.0, experimental path)
+try:
+    from trl.experimental.orpo import ORPOConfig, ORPOTrainer
+except ImportError:
+    print("ERROR: trl not installed or outdated. Run: pip install trl==0.29.0")
+    sys.exit(1)
+
+# Telegram notifications
+try:
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+    from scripts.telegram_notify import send_telegram_safe
+    HAS_TELEGRAM = True
+except ImportError:
+    HAS_TELEGRAM = False
+    def send_telegram_safe(msg, **kw): return False
+
+# ---------------------------------------------------------------------------
+# Logging setup
+# ---------------------------------------------------------------------------
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+log = logging.getLogger("orpo")
+
+
+# ---------------------------------------------------------------------------
+# Custom callback for detailed monitoring
+# ---------------------------------------------------------------------------
+class ORPOMonitorCallback(TrainerCallback):
+    """Monitors ORPO-specific metrics and sends alerts on anomalies."""
+
+    def __init__(self, alert_fn=send_telegram_safe):
+        self.alert_fn = alert_fn
+        self.start_time = None
+        self.last_eval_loss = None
+        self.eval_loss_increases = 0
+        self.negative_margin_streak = 0
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self.start_time = time.time()
+        log.info("ORPO training begin -- monitoring active")
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs is None:
+            return
+        step = state.global_step
+
+        # Monitor rewards/margins
+        margin = logs.get("rewards/margins")
+        if margin is not None:
+            if margin < 0:
+                self.negative_margin_streak += 1
+                if self.negative_margin_streak >= 10:
+                    msg = (f"[ORPO ALERT] rewards/margins negative for "
+                           f"{self.negative_margin_streak} consecutive logs at step {step} "
+                           f"(margin={margin:.4f})")
+                    log.warning(msg)
+                    self.alert_fn(msg)
+            else:
+                self.negative_margin_streak = 0
+
+        # Log key metrics every logging step
+        loss = logs.get("loss")
+        chosen = logs.get("rewards/chosen")
+        rejected = logs.get("rewards/rejected")
+        if loss is not None:
+            elapsed = time.time() - self.start_time if self.start_time else 0
+            log.info(
+                f"step={step} loss={loss:.4f} "
+                f"margin={margin if margin is not None else 'N/A'} "
+                f"chosen={chosen if chosen is not None else 'N/A'} "
+                f"rejected={rejected if rejected is not None else 'N/A'} "
+                f"elapsed={elapsed/3600:.1f}h"
+            )
+
+        # Check for NaN/Inf
+        if loss is not None and (not isinstance(loss, (int, float)) or loss != loss):
+            msg = f"[ORPO CRITICAL] NaN/Inf loss detected at step {step}!"
+            log.error(msg)
+            self.alert_fn(msg)
+
+    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+        if metrics is None:
+            return
+        eval_loss = metrics.get("eval_loss")
+        step = state.global_step
+
+        if eval_loss is not None:
+            log.info(f"[EVAL] step={step} eval_loss={eval_loss:.4f}")
+            if self.last_eval_loss is not None and eval_loss > self.last_eval_loss:
+                self.eval_loss_increases += 1
+                log.warning(
+                    f"[EVAL] eval_loss increased: {self.last_eval_loss:.4f} -> {eval_loss:.4f} "
+                    f"({self.eval_loss_increases}/3 before early stop)"
+                )
+            else:
+                self.eval_loss_increases = 0
+            self.last_eval_loss = eval_loss
+
+    def on_train_end(self, args, state, control, **kwargs):
+        elapsed = time.time() - self.start_time if self.start_time else 0
+        log.info(f"ORPO training ended -- total time: {elapsed/3600:.2f}h, "
+                 f"total steps: {state.global_step}")
+
+    def on_save(self, args, state, control, **kwargs):
+        log.info(f"Checkpoint saved at step {state.global_step}")
+
+
+class VRAMMonitorCallback(TrainerCallback):
+    """Measures peak VRAM usage across all GPUs during training."""
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        if torch.cuda.is_available():
+            for i in range(torch.cuda.device_count()):
+                torch.cuda.reset_peak_memory_stats(i)
+            log.info("[VRAM] Peak memory stats reset for all GPUs")
+
+    def on_train_end(self, args, state, control, **kwargs):
+        if torch.cuda.is_available():
+            for i in range(torch.cuda.device_count()):
+                peak_mb = torch.cuda.max_memory_allocated(i) / (1024**2)
+                log.info(f"[VRAM] GPU {i} peak: {peak_mb:.0f} MiB")
+
+
+def load_hf_preference_dataset(dataset_name: str, token: str | None = None) -> Dataset:
+    """Load and normalize a HuggingFace preference dataset to {prompt, chosen, rejected}."""
+    ds = load_dataset(dataset_name, split="train", token=token)
+
+    # kuotient/orca-math-korean-dpo-pairs format: {system, question, chosen, rejected}
+    if "question" in ds.column_names and "chosen" in ds.column_names:
+        def normalize(example):
+            prompt = example.get("system", "") + "\n" + example["question"]
+            return {"prompt": prompt.strip(), "chosen": example["chosen"], "rejected": example["rejected"]}
+        return ds.map(normalize, remove_columns=ds.column_names)
+
+    # nayohan/preference-collection-ko-full format: {response_A, response_B, orig_preference}
+    if "orig_preference" in ds.column_names:
+        def normalize_pref(example):
+            prompt = example.get("orig_instruction", example.get("instruction", ""))
+            if example["orig_preference"] == "B":
+                return {"prompt": prompt, "chosen": example["orig_response_B"], "rejected": example["orig_response_A"]}
+            else:
+                return {"prompt": prompt, "chosen": example["orig_response_A"], "rejected": example["orig_response_B"]}
+        return ds.map(normalize_pref, remove_columns=ds.column_names)
+
+    # Already in {prompt, chosen, rejected} format
+    if all(c in ds.column_names for c in ["prompt", "chosen", "rejected"]):
+        return ds
+
+    raise ValueError(f"Unknown dataset format. Columns: {ds.column_names}")
+
+
+def load_custom_jsonl(path: str) -> Dataset:
+    """Load custom JSONL with {prompt, chosen, rejected} fields."""
+    data = []
+    with open(path) as f:
+        for line in f:
+            data.append(json.loads(line))
+    return Dataset.from_list(data)
+
+
+def load_yaml_config(path: str) -> dict:
+    """Load YAML config and return as dict."""
+    import yaml
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ORPO Training (TRL 0.29.0 -- 8xB200 optimized)")
+    parser.add_argument("--config", type=str, default=None, help="YAML config file path")
+    parser.add_argument("--model_path", type=str, default=None, help="HF format model path")
+    parser.add_argument("--dataset", type=str, default="kuotient/orca-math-korean-dpo-pairs")
+    parser.add_argument("--custom_data_path", type=str, default=None, help="Custom JSONL preference data")
+    parser.add_argument("--output_dir", type=str, default="checkpoints/korean_3b_orpo")
+    parser.add_argument("--hf_token", type=str, default=None)
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--lr", type=float, default=5e-6)
+    parser.add_argument("--beta", type=float, default=0.1, help="ORPO beta (odds ratio weight)")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=4)
+    parser.add_argument("--max_length", type=int, default=1536)
+    parser.add_argument("--bf16", action="store_true", default=True)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--eval_split_ratio", type=float, default=0.05, help="Fraction of data for eval")
+    parser.add_argument("--eval_steps", type=int, default=500)
+    parser.add_argument("--early_stopping_patience", type=int, default=3)
+    parser.add_argument("--max_steps", type=int, default=-1, help="Override max steps (for quick test)")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--save_total_limit", type=int, default=5)
+    parser.add_argument("--warmup_ratio", type=float, default=0.05)
+    parser.add_argument("--lr_scheduler_type", type=str, default="cosine")
+    parser.add_argument("--logging_steps", type=int, default=10)
+    parser.add_argument("--save_steps", type=int, default=500)
+    parser.add_argument("--gradient_checkpointing", action="store_true", default=True)
+    parser.add_argument("--report_to", type=str, default="none")
+    parser.add_argument("--dataset_num_proc", type=int, default=8,
+                        help="Number of processes for parallel tokenization in ORPOTrainer")
+    parser.add_argument("--dataloader_num_workers", type=int, default=4,
+                        help="Number of dataloader worker processes")
+    parser.add_argument("--no_load_best", action="store_true", default=False,
+                        help="Disable load_best_model_at_end (for sweep/quick tests)")
+    parser.add_argument("--max_samples", type=int, default=0,
+                        help="Limit dataset size (0=use all, >0=subset for benchmarking)")
+    parser.add_argument("--skip_filter", action="store_true", default=False,
+                        help="Skip NaN-prevention filter (for benchmarking only)")
+    args = parser.parse_args()
+
+    # Override CLI defaults with YAML config values
+    if args.config:
+        cfg = load_yaml_config(args.config)
+        for key, value in cfg.items():
+            if hasattr(args, key):
+                setattr(args, key, value)
+
+    if not args.model_path:
+        parser.error("--model_path is required (or set model_path in YAML config)")
+
+    # Log all resolved config
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    is_main = local_rank == 0
+    if is_main:
+        log.info("=" * 70)
+        log.info("ORPO Training Configuration (8xB200 optimized)")
+        log.info("=" * 70)
+        for k, v in sorted(vars(args).items()):
+            log.info(f"  {k}: {v}")
+        log.info("=" * 70)
+
+        # GPU info
+        if torch.cuda.is_available():
+            for i in range(torch.cuda.device_count()):
+                mem = torch.cuda.get_device_properties(i).total_memory / 1e9
+                log.info(f"  GPU {i}: {torch.cuda.get_device_name(i)} ({mem:.1f} GB)")
+
+    # Validate paths
+    if not Path(args.model_path).exists():
+        raise FileNotFoundError(f"Model path not found: {args.model_path}")
+    if args.custom_data_path and not Path(args.custom_data_path).exists():
+        raise FileNotFoundError(f"Data path not found: {args.custom_data_path}")
+
+    # NCCL/DDP environment diagnostics
+    if is_main:
+        log.info("--- DDP/NCCL Environment ---")
+        for env_key in ["RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT",
+                        "NCCL_IB_DISABLE", "NCCL_BUFFSIZE", "NCCL_P2P_LEVEL",
+                        "OMP_NUM_THREADS", "PYTORCH_CUDA_ALLOC_CONF"]:
+            log.info(f"  {env_key}={os.environ.get(env_key, '(not set)')}")
+        log.info(f"  torch.distributed.is_available={torch.distributed.is_available()}")
+        if torch.distributed.is_initialized():
+            log.info(f"  world_size={torch.distributed.get_world_size()}, "
+                     f"rank={torch.distributed.get_rank()}")
+
+    # Load model (bfloat16 + flash_attention_2 for B200)
+    log.info(f"Loading model from {args.model_path}...")
+    t0 = time.time()
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            attn_implementation="flash_attention_2",
+        )
+    except Exception as e:
+        log.error(f"Model loading failed: {e}")
+        send_telegram_safe(f"[ORPO FATAL] Model load failed: {e}")
+        raise
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    if is_main:
+        n_params = sum(p.numel() for p in model.parameters())
+        log.info(f"Model loaded: {n_params:,} params in {time.time()-t0:.1f}s")
+        log.info(f"Tokenizer: vocab_size={tokenizer.vocab_size}, "
+                 f"pad_token='{tokenizer.pad_token}', eos_token='{tokenizer.eos_token}'")
+
+    # Load dataset
+    t0 = time.time()
+    try:
+        if args.custom_data_path:
+            log.info(f"Loading custom data from {args.custom_data_path}...")
+            fsize_mb = Path(args.custom_data_path).stat().st_size / 1e6
+            log.info(f"  File size: {fsize_mb:.1f} MB")
+            dataset = load_custom_jsonl(args.custom_data_path)
+        else:
+            log.info(f"Loading dataset {args.dataset}...")
+            dataset = load_hf_preference_dataset(args.dataset, token=args.hf_token)
+    except Exception as e:
+        log.error(f"Dataset loading failed: {e}")
+        send_telegram_safe(f"[ORPO FATAL] Data load failed: {e}")
+        raise
+
+    # Subset for benchmarking (skip tokenization bottleneck)
+    if args.max_samples > 0 and len(dataset) > args.max_samples:
+        dataset = dataset.select(range(args.max_samples))
+        if is_main:
+            log.info(f"[BENCH] Dataset subset: {args.max_samples:,} samples")
+
+    if is_main:
+        log.info(f"Dataset loaded: {len(dataset)} pairs in {time.time()-t0:.1f}s")
+        # Data quality check
+        sample = dataset[0]
+        log.info(f"Sample keys: {list(sample.keys())}")
+        for key in ["prompt", "chosen", "rejected"]:
+            if key not in sample:
+                raise ValueError(f"Dataset missing required column: {key}")
+            val = sample[key]
+            log.info(f"  {key}: {str(val)[:100]}...")
+
+        # Length distribution check (sample first 1000)
+        sample_size = min(1000, len(dataset))
+        lengths = [len(str(dataset[i]["prompt"])) + max(len(str(dataset[i]["chosen"])),
+                   len(str(dataset[i]["rejected"]))) for i in range(sample_size)]
+        avg_len = sum(lengths) / len(lengths)
+        max_len = max(lengths)
+        log.info(f"  Char lengths (sample {sample_size}): avg={avg_len:.0f}, max={max_len}")
+
+    # Filter out samples where prompt is too long for the response to fit in max_length.
+    # Without this, samples with 0 response tokens cause NaN in ORPO log-probability computation
+    # (division by zero in average_log_prob when loss_mask is all-zero).
+    # Also catches TRL truncation bug: tokenize_row uses longer_response_length = max(chosen_len, rejected_len)
+    # and truncates BOTH responses to [:max_length - longer_response_length]. When longer >= max_length,
+    # the shorter response becomes EMPTY → NaN.
+    if args.skip_filter:
+        if is_main:
+            log.info("[BENCH] Skipping NaN-prevention filter (--skip_filter)")
+    else:
+        pre_filter = len(dataset)
+        def _has_response_room(example):
+            prompt_tok_len = len(tokenizer.encode(example["prompt"], add_special_tokens=False))
+            chosen_tok_len = len(tokenizer.encode(example["chosen"], add_special_tokens=False))
+            rejected_tok_len = len(tokenizer.encode(example["rejected"], add_special_tokens=False))
+
+            # 1. Prompt must leave room for at least 16 response tokens
+            if prompt_tok_len + 16 > args.max_length:
+                return False
+
+            # 2. Each response independently must fit with prompt
+            # (TRL adds BOS/EOS, so use +2 margin)
+            if prompt_tok_len + chosen_tok_len + 2 > args.max_length * 2:
+                return False  # extremely long, will cause issues
+            if prompt_tok_len + rejected_tok_len + 2 > args.max_length * 2:
+                return False
+
+            # 3. The longer response must not exceed max_length alone
+            # (TRL bug: both responses truncated by max(chosen_len, rejected_len))
+            longer = max(chosen_tok_len, rejected_tok_len)
+            if longer >= args.max_length:
+                return False
+
+            return True
+        dataset = dataset.filter(_has_response_room, num_proc=min(args.dataset_num_proc, 32) if is_main else 1)
+        if is_main:
+            log.info(f"Filtered: {pre_filter:,} -> {len(dataset):,} "
+                     f"(removed {pre_filter - len(dataset):,} samples with prompt > max_length-16 or TRL truncation risk)")
+
+    # Train/eval split
+    split = dataset.train_test_split(test_size=args.eval_split_ratio, seed=args.seed)
+    train_dataset = split["train"]
+    eval_dataset = split["test"]
+    log.info(f"Train: {len(train_dataset):,}, Eval: {len(eval_dataset):,}")
+
+    # Compute training stats for warmup_steps calculation
+    n_gpus = max(torch.cuda.device_count(), 1) if torch.cuda.is_available() else 1
+    eff_batch = args.batch_size * args.gradient_accumulation_steps * n_gpus
+    steps_per_epoch = len(train_dataset) // eff_batch
+    total_steps = args.max_steps if args.max_steps > 0 else steps_per_epoch * args.epochs
+    computed_warmup_steps = int(total_steps * args.warmup_ratio)
+    if is_main:
+        log.info(f"Training plan: eff_batch={eff_batch}, steps/epoch={steps_per_epoch:,}, "
+                 f"total={total_steps:,}, warmup={computed_warmup_steps}")
+
+    # DDP tokenization strategy:
+    # TRL ORPOTrainer uses main_process_first() — rank 0 tokenizes first, then ranks 1-7.
+    # With multiprocessing (num_proc>1), ranks 1-7 all spawn workers simultaneously,
+    # causing CPU/memory oversubscription (e.g. 7 ranks × 8 workers = 56 processes).
+    # Fix: rank 0 uses full num_proc for speed, other ranks use 1 (should hit cache).
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    if world_size > 1:
+        effective_num_proc = args.dataset_num_proc if is_main else 1
+        if is_main:
+            log.info(f"DDP tokenization: rank 0 uses num_proc={args.dataset_num_proc}, "
+                     f"other {world_size-1} ranks use num_proc=1 (cache)")
+    else:
+        effective_num_proc = args.dataset_num_proc
+
+    # ORPOConfig (TRL 0.29.0) -- optimized for 8x B200
+    orpo_config = ORPOConfig(
+        output_dir=args.output_dir,
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.lr,
+        beta=args.beta,
+        lr_scheduler_type=args.lr_scheduler_type,
+        warmup_steps=computed_warmup_steps,
+        weight_decay=args.weight_decay,
+        bf16=args.bf16,
+        logging_steps=args.logging_steps,
+        save_steps=args.save_steps,
+        save_total_limit=args.save_total_limit,
+        max_length=args.max_length,
+        gradient_checkpointing=args.gradient_checkpointing,
+        report_to=args.report_to,
+        remove_unused_columns=False,
+        eval_strategy="steps",
+        eval_steps=args.eval_steps,
+        metric_for_best_model="eval_loss" if not args.no_load_best else None,
+        load_best_model_at_end=not args.no_load_best,
+        greater_is_better=False if not args.no_load_best else None,
+        max_steps=args.max_steps,
+        seed=args.seed,
+        # B200 hardware optimizations
+        dataloader_num_workers=args.dataloader_num_workers,
+        dataloader_pin_memory=True,
+        ddp_find_unused_parameters=False,
+        ddp_timeout=7200,  # 2h — tokenization takes ~30min on 683K samples
+        dataset_num_proc=effective_num_proc,
+    )
+
+    # ORPOTrainer (no reference model needed)
+    log.info("Initializing ORPOTrainer (tokenization will happen here — may take a while)...")
+    t0 = time.time()
+    monitor = ORPOMonitorCallback()
+    try:
+        trainer = ORPOTrainer(
+            model=model,
+            args=orpo_config,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=tokenizer,
+            callbacks=[cb for cb in [
+                EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience)
+                    if not args.no_load_best else None,
+                monitor,
+                VRAMMonitorCallback(),
+            ] if cb is not None],
+        )
+    except Exception as e:
+        log.error(f"ORPOTrainer init failed: {e}\n{traceback.format_exc()}")
+        send_telegram_safe(f"[ORPO FATAL] Trainer init failed: {e}")
+        raise
+    if is_main:
+        log.info(f"ORPOTrainer initialized in {time.time()-t0:.1f}s "
+                 f"(dataset_num_proc={args.dataset_num_proc})")
+
+    # SIGHUP/SIGTERM defense -- graceful shutdown with emergency checkpoint
+    def _graceful_shutdown_handler(signum, frame):
+        sig_name = _signal_mod.Signals(signum).name
+        log.warning(f"Received {sig_name}. Saving emergency checkpoint...")
+        try:
+            emergency_path = os.path.join(args.output_dir, "emergency_checkpoint")
+            trainer.save_model(emergency_path)
+            log.info(f"Emergency checkpoint saved to {emergency_path}")
+            send_telegram_safe(
+                f"[ORPO] Signal {sig_name} received at step {trainer.state.global_step}. "
+                f"Emergency checkpoint saved."
+            )
+        except Exception as e:
+            log.error(f"Emergency save failed: {e}")
+            send_telegram_safe(f"[ORPO] Emergency save FAILED after {sig_name}: {e}")
+        sys.exit(1)
+
+    for _sig in (_signal_mod.SIGHUP, _signal_mod.SIGTERM):
+        _signal_mod.signal(_sig, _graceful_shutdown_handler)
+
+    # Pre-training VRAM report
+    if is_main and torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+        alloc = torch.cuda.memory_allocated() / 1e9
+        reserved = torch.cuda.memory_reserved() / 1e9
+        log.info(f"Pre-train VRAM: allocated={alloc:.1f}GB, reserved={reserved:.1f}GB")
+
+    start_msg = (
+        f"[ORPO] Training started\n"
+        f"  model: {args.model_path}\n"
+        f"  beta: {args.beta}, lr: {args.lr}\n"
+        f"  train: {len(train_dataset):,}, eval: {len(eval_dataset):,}\n"
+        f"  eff_batch: {eff_batch}, steps/epoch: {steps_per_epoch:,}, total: {total_steps:,}\n"
+        f"  warmup: {computed_warmup_steps} steps ({args.warmup_ratio*100:.0f}%)\n"
+        f"  max_length: {args.max_length}, max_steps: {args.max_steps}\n"
+        f"  dataset_num_proc: {args.dataset_num_proc}, dl_workers: {args.dataloader_num_workers}"
+    )
+    log.info(start_msg.replace("[ORPO] ", ""))
+    send_telegram_safe(start_msg)
+
+    try:
+        trainer.train()
+
+        # Post-training VRAM report
+        if is_main and torch.cuda.is_available():
+            peak = torch.cuda.max_memory_allocated() / 1e9
+            log.info(f"Peak VRAM usage: {peak:.1f}GB")
+
+        trainer.save_model(args.output_dir)
+        log.info(f"Model saved to {args.output_dir}")
+
+        # Extract final metrics
+        final_metrics = {}
+        for entry in reversed(trainer.state.log_history):
+            if "loss" in entry and "loss" not in final_metrics:
+                final_metrics["loss"] = entry["loss"]
+            if "eval_loss" in entry and "eval_loss" not in final_metrics:
+                final_metrics["eval_loss"] = entry["eval_loss"]
+            if "rewards/margins" in entry and "rewards/margins" not in final_metrics:
+                final_metrics["rewards/margins"] = entry["rewards/margins"]
+            if len(final_metrics) >= 3:
+                break
+
+        done_msg = (
+            f"[ORPO] Training complete!\n"
+            f"  output: {args.output_dir}\n"
+            f"  steps: {trainer.state.global_step}\n"
+            f"  final loss: {final_metrics.get('loss', 'N/A')}\n"
+            f"  final eval_loss: {final_metrics.get('eval_loss', 'N/A')}\n"
+            f"  final margins: {final_metrics.get('rewards/margins', 'N/A')}"
+        )
+        log.info(done_msg.replace("[ORPO] ", ""))
+        send_telegram_safe(done_msg)
+
+    except KeyboardInterrupt:
+        log.warning("Training interrupted by user (KeyboardInterrupt)")
+        send_telegram_safe(f"[ORPO] Training interrupted at step {trainer.state.global_step}")
+        trainer.save_model(os.path.join(args.output_dir, "interrupted_checkpoint"))
+        log.info("Interrupted checkpoint saved.")
+
+    except Exception as e:
+        tb = traceback.format_exc()
+        error_msg = f"[ORPO] Training FAILED at step {trainer.state.global_step}: {e}"
+        log.error(f"{error_msg}\n{tb}")
+        send_telegram_safe(f"{error_msg}\n{tb[:500]}")
+        # Try emergency save
+        try:
+            trainer.save_model(os.path.join(args.output_dir, "error_checkpoint"))
+            log.info("Error checkpoint saved.")
+        except Exception:
+            log.error("Error checkpoint save also failed.")
+        raise
+
+
+if __name__ == "__main__":
+    main()

source/train/pretrain.py

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

source/train/trainer.py

@@ -0,0 +1,594 @@
+"""
+train/trainer.py — Core training loop.
+
+Provides:
+    TrainConfig : Dataclass of all training hyper-parameters.
+    Trainer     : Orchestrates gradient accumulation, AMP, gradient clipping,
+                  tensorboard logging, and checkpoint saving.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import math
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+try:
+    from torch.utils.tensorboard import SummaryWriter
+    HAS_TENSORBOARD = True
+except (ImportError, AttributeError):
+    SummaryWriter = None  # type: ignore[misc,assignment]
+    HAS_TENSORBOARD = False
+
+from train.utils import get_grad_norm, is_main_process, save_checkpoint
+
+
+# ---------------------------------------------------------------------------
+# Optional TransformerEngine import (FP8 support)
+# ---------------------------------------------------------------------------
+try:
+    import transformer_engine.pytorch as te  # type: ignore[import]
+    from transformer_engine.common.recipe import DelayedScaling, Format  # type: ignore[import]
+    HAS_TE = True
+except ImportError:
+    te = None  # type: ignore[assignment]
+    HAS_TE = False
+
+
+# ---------------------------------------------------------------------------
+# Configuration dataclass
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TrainConfig:
+    """Hyper-parameters that control the training loop."""
+
+    # Total number of optimiser update steps.
+    max_steps: int = 100_000
+
+    # Number of forward passes accumulated before each optimiser step.
+    grad_accum_steps: int = 1
+
+    # Maximum global gradient L2 norm; clips if exceeded (0 = disabled).
+    max_grad_norm: float = 1.0
+
+    # Log training metrics every this many *optimiser* steps.
+    log_interval: int = 10
+
+    # Save a checkpoint every this many optimiser steps.
+    save_interval: int = 1000
+
+    # Run validation (if val_loader provided) every this many optimiser steps.
+    eval_interval: int = 500
+
+    # Root directory where checkpoint sub-folders are written.
+    checkpoint_dir: str = "checkpoints"
+
+    # Use bf16 autocast during the forward pass (no GradScaler needed for bf16).
+    use_amp: bool = True
+
+    # Pass model through torch.compile() before training.
+    compile_model: bool = False
+
+    # FP8 (TransformerEngine) settings — only relevant when use_fp8=True.
+    use_fp8: bool = False
+    fp8_amax_history_len: int = 16
+    fp8_amax_compute_algo: str = "max"   # "max" | "most_recent"
+    fp8_format: str = "MXFP8"            # "MXFP8" (B200 block scaling) | "HYBRID" (E4M3+E5M2)
+
+    # Path to a text log file (rank-0 only). None = stdout only.
+    log_file: Optional[str] = None
+
+    # grad_norm을 파일에 기록하는 간격 (0=비활성)
+    log_grad_norm_interval: int = 100
+    # GPU 메모리를 파일에 기록하는 간격 (0=비활성)
+    log_memory_interval: int = 100
+
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+
+
+class Trainer:
+    """
+    Manages the full pretraining loop for a decoder-only LLM.
+
+    Supports:
+        - Gradient accumulation over ``config.grad_accum_steps`` micro-batches.
+        - bf16 mixed-precision via ``torch.autocast`` (no GradScaler required).
+        - Global gradient norm clipping.
+        - Tensorboard logging on the main process.
+        - Periodic checkpoint saving via :func:`train.utils.save_checkpoint`.
+        - Optional ``torch.compile`` acceleration.
+
+    Args:
+        model:        The LLM (plain ``nn.Module`` or DDP-wrapped).
+        train_loader: DataLoader yielding ``(input_ids, targets)`` batches.
+        optimizer:    AdamW (or any ``Optimizer``) configured externally.
+        scheduler:    LR scheduler produced by the caller.
+        config:       ``TrainConfig`` instance controlling all loop behaviour.
+        device:       Target device for data and model.
+        rank:         Process rank (used to suppress logging on non-main ranks).
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        train_loader: DataLoader,
+        optimizer: Optimizer,
+        scheduler: LambdaLR,
+        config: TrainConfig,
+        device: torch.device,
+        rank: int = 0,
+        sampler: Optional[DistributedSampler] = None,
+        val_loader: Optional[DataLoader] = None,
+    ) -> None:
+        self.model = model
+        self.train_loader = train_loader
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.config = config
+        self.device = device
+        self.rank = rank
+        self._is_main = is_main_process()
+        self._sampler = sampler   # for set_epoch() on each data pass
+        self._epoch = 0
+        self._val_loader = val_loader
+        self._best_val_loss: float = float("inf")
+        self._val_patience_counter: int = 0
+        self._val_patience_limit: int = 10  # early stopping patience (v2: 5→10, warmup 후 충분한 학습 보장)
+
+        # Graceful shutdown support — signal handler에서 flag 설정,
+        # 학습 루프가 각 step 완료 후 확인하여 비상 체크포인트 저장 후 종료
+        self._shutdown_requested = False
+        self._shutdown_signal = ""
+
+        # Build FP8 recipe once (reused every step) ----------------------
+        self._fp8_recipe = None
+        if config.use_fp8 and HAS_TE:
+            if config.fp8_format == "MXFP8":
+                from transformer_engine.common.recipe import MXFP8BlockScaling  # type: ignore[import]
+                self._fp8_recipe = MXFP8BlockScaling()
+            else:
+                self._fp8_recipe = DelayedScaling(
+                    fp8_format=getattr(Format, config.fp8_format),
+                    amax_history_len=config.fp8_amax_history_len,
+                    amax_compute_algo=config.fp8_amax_compute_algo,
+                )
+
+        # Optionally compile the model (unwrap DDP first to compile the inner module).
+        if config.compile_model:
+            inner: nn.Module = getattr(self.model, "module", self.model)
+            compiled = torch.compile(inner)
+            if hasattr(self.model, "module"):
+                self.model.module = compiled  # type: ignore[assignment]
+            else:
+                self.model = compiled  # type: ignore[assignment]
+
+        # Tensorboard writer — only on rank 0.
+        self._writer: Optional[SummaryWriter] = None
+        self._log_fh = None  # optional file handle for structured text log
+        if self._is_main:
+            if HAS_TENSORBOARD:
+                log_dir = Path(config.checkpoint_dir) / "tensorboard"
+                self._writer = SummaryWriter(log_dir=str(log_dir))
+            if config.log_file is not None:
+                Path(config.log_file).parent.mkdir(parents=True, exist_ok=True)
+                self._log_fh = open(config.log_file, "a", encoding="utf-8", buffering=1)
+
+        # 학습 시작 시각 기록 (통계 요약 로그에 사용)
+        import datetime
+        self._train_start_time = datetime.datetime.now()
+
+        # Infinite iterator over the DataLoader.
+        self._loader_iter = iter(self.train_loader)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def request_shutdown(self, signal_name: str = "UNKNOWN") -> None:
+        """Request graceful shutdown after the current training step.
+
+        Called from signal handlers (SIGHUP, SIGTERM). Sets a flag
+        that the training loop checks after each optimizer step.
+        The loop will save an emergency checkpoint and exit cleanly.
+        """
+        self._shutdown_requested = True
+        self._shutdown_signal = signal_name
+
+    def train(self, start_step: int = 0) -> None:
+        """
+        Run the main training loop from ``start_step`` to ``config.max_steps``.
+
+        Args:
+            start_step: First optimiser step index (non-zero when resuming).
+        """
+        cfg = self.config
+        model = self.model
+
+        model.train()
+
+        # Timing state for tokens/sec estimation.
+        t0 = time.perf_counter()
+        running_loss = 0.0
+        log_step_count = 0
+        accum_loss = torch.tensor(0.0, device=self.device)  # initialise so end-of-training save is safe on empty loops
+
+        for step in range(start_step, cfg.max_steps):
+            # ---- Gradient accumulation loop --------------------------------
+            self.optimizer.zero_grad(set_to_none=True)
+            # Accumulate loss on GPU to avoid one GPU-CPU sync per micro-step.
+            accum_loss = torch.zeros(1, device=self.device)
+
+            for micro_step in range(cfg.grad_accum_steps):
+                batch = self._next_batch()
+                # Suppress DDP all-reduce on all but the last micro-step (Bug 3).
+                is_last_micro = micro_step == cfg.grad_accum_steps - 1
+                sync_ctx = (
+                    contextlib.nullcontext()
+                    if not isinstance(model, DDP) or is_last_micro
+                    else model.no_sync()
+                )
+                try:
+                    with sync_ctx:
+                        micro_loss = self._step(batch)  # returns detached GPU tensor
+                except torch.cuda.OutOfMemoryError as e:
+                    torch.cuda.empty_cache()
+                    mem_total = torch.cuda.get_device_properties(self.device).total_memory / 1e9
+                    mem_alloc = torch.cuda.memory_allocated() / 1e9
+                    raise RuntimeError(
+                        f"CUDA OOM at step {step}, micro_step {micro_step}. "
+                        f"GPU mem: {mem_alloc:.1f}/{mem_total:.1f} GB. "
+                        f"Try reducing batch_size or grad_accum_steps."
+                    ) from e
+                except RuntimeError as e:
+                    self._log(f"RuntimeError at step {step}, micro_step {micro_step}: {e}", level="ERROR")
+                    raise
+                accum_loss += micro_loss  # GPU-side accumulation, no CPU sync
+
+            # Single GPU-CPU sync per optimizer step (was one sync per micro-step).
+            avg_loss = accum_loss.item() / cfg.grad_accum_steps
+
+            # Detect NaN/Inf loss — indicates numerical instability.
+            if not math.isfinite(avg_loss):
+                mem_gb = torch.cuda.memory_allocated() / 1e9
+                mem_total = torch.cuda.get_device_properties(self.device).total_memory / 1e9
+                raise RuntimeError(
+                    f"Non-finite loss detected: {avg_loss}. "
+                    f"GPU mem: {mem_gb:.1f}/{mem_total:.1f} GB. "
+                    f"Check lr, grad clipping, FP8 amax history. "
+                    f"Try: lower lr, increase fp8_amax_history_len, or switch to BF16."
+                )
+
+            # ---- Gradient clipping -----------------------------------------
+            # clip_grad_norm_ already computes the global norm internally.
+            # Reuse its return value to avoid a second pass of ~50 GPU-CPU syncs.
+            if cfg.max_grad_norm > 0.0:
+                grad_norm = torch.nn.utils.clip_grad_norm_(
+                    model.parameters(), cfg.max_grad_norm
+                ).item()
+            else:
+                grad_norm = get_grad_norm(model)
+
+            # ---- Optimiser + scheduler step ---------------------------------
+            self.optimizer.step()
+            self.scheduler.step()
+
+            # ---- Graceful shutdown check -----------------------------------
+            # Signal handler가 request_shutdown()을 호출하면 이 flag가 True.
+            # 현재 step의 optimizer 업데이트가 완료된 시점에서 체크하므로
+            # 모델 가중치는 항상 일관된 상태로 저장됩니다.
+            if self._shutdown_requested:
+                self._log(
+                    f"Graceful shutdown initiated (signal: {self._shutdown_signal}) "
+                    f"at step {step + 1}, loss={avg_loss:.4f}",
+                    level="WARN",
+                )
+                if self._is_main:
+                    ckpt_path = save_checkpoint(
+                        model=self.model,
+                        optimizer=self.optimizer,
+                        scheduler=self.scheduler,
+                        step=step + 1,
+                        loss=avg_loss,
+                        path=cfg.checkpoint_dir,
+                    )
+                    self._log(f"Emergency checkpoint saved → {ckpt_path}", level="WARN")
+                # DDP 동기화: 모든 rank가 함께 종료하도록 barrier
+                try:
+                    if torch.distributed.is_initialized():
+                        torch.distributed.barrier()
+                except Exception:
+                    pass  # DDP 미사용 또는 이미 해체된 경우 무시
+                self._log("Shutdown complete. Exiting training loop.", level="WARN")
+                if self._writer is not None:
+                    self._writer.close()
+                if self._log_fh is not None:
+                    self._log_fh.flush()
+                return
+
+            running_loss += avg_loss
+            log_step_count += 1
+
+            # ---- Logging ---------------------------------------------------
+            if (step + 1) % cfg.log_interval == 0 and self._is_main:
+                t1 = time.perf_counter()
+                elapsed = t1 - t0
+
+                avg_loss = running_loss / log_step_count
+
+                # Estimate throughput: tokens processed during this log window.
+                batch_size, seq_len = self._last_batch_shape
+                tokens_per_sec = (
+                    batch_size * seq_len * cfg.grad_accum_steps * cfg.log_interval
+                ) / max(elapsed, 1e-9)
+
+                current_lr = self.scheduler.get_last_lr()[0]
+                global_step = step + 1
+
+                mem_gb = torch.cuda.memory_allocated() / 1e9
+                self._log(
+                    f"step {global_step:>7d} | "
+                    f"loss {avg_loss:.4f} | "
+                    f"lr {current_lr:.2e} | "
+                    f"gnorm {grad_norm:.3f} | "
+                    f"tok/s {tokens_per_sec:,.0f} | "
+                    f"mem {mem_gb:.1f}GB | "
+                    f"epoch {self._epoch}"
+                )
+
+                if self._writer is not None:
+                    self._writer.add_scalar("train/loss", avg_loss, global_step)
+                    self._writer.add_scalar("train/lr", current_lr, global_step)
+                    self._writer.add_scalar("train/grad_norm", grad_norm, global_step)
+                    self._writer.add_scalar("train/tokens_per_sec", tokens_per_sec, global_step)
+
+                # Reset accumulators.
+                running_loss = 0.0
+                log_step_count = 0
+                t0 = t1
+
+            # ---- Validation ------------------------------------------------
+            if (step + 1) % cfg.eval_interval == 0 and self._val_loader is not None:
+                val_loss = self._run_validation()
+                # Determine early stopping on rank 0, broadcast to all ranks
+                # so every DDP rank exits together (prevents hang).
+                should_stop = False
+                if self._is_main:
+                    self._log(f"step {step + 1:>7d} | val_loss {val_loss:.4f}")
+                    if self._writer is not None:
+                        self._writer.add_scalar("val/loss", val_loss, step + 1)
+                    # Save best checkpoint when val loss improves.
+                    if val_loss < self._best_val_loss:
+                        self._best_val_loss = val_loss
+                        self._val_patience_counter = 0
+                        best_path = save_checkpoint(
+                            model=self.model,
+                            optimizer=self.optimizer,
+                            scheduler=self.scheduler,
+                            step=step + 1,
+                            loss=val_loss,
+                            path=cfg.checkpoint_dir,
+                            suffix="best",
+                        )
+                        self._log(
+                            f"New best val_loss={val_loss:.4f} → {best_path}"
+                        )
+                    else:
+                        self._val_patience_counter += 1
+                        self._log(
+                            f"val_loss {val_loss:.4f} did not improve "
+                            f"(best={self._best_val_loss:.4f}, "
+                            f"patience={self._val_patience_counter}/{self._val_patience_limit})"
+                        )
+                        if self._val_patience_counter >= self._val_patience_limit:
+                            self._log(
+                                f"Early stopping triggered at step {step + 1} "
+                                f"(patience {self._val_patience_limit} exhausted)"
+                            )
+                            should_stop = True
+                # Broadcast early stopping decision to all DDP ranks.
+                if torch.distributed.is_initialized():
+                    stop_tensor = torch.tensor(
+                        [1 if should_stop else 0], dtype=torch.int32,
+                        device=self.device,
+                    )
+                    torch.distributed.broadcast(stop_tensor, src=0)
+                    should_stop = stop_tensor.item() == 1
+                if should_stop:
+                    return
+
+            # ---- Checkpoint save -------------------------------------------
+            if (step + 1) % cfg.save_interval == 0 and self._is_main:
+                ckpt_path = save_checkpoint(
+                    model=self.model,
+                    optimizer=self.optimizer,
+                    scheduler=self.scheduler,
+                    step=step + 1,
+                    loss=avg_loss,
+                    path=cfg.checkpoint_dir,
+                )
+                self._log(f"Checkpoint saved → {ckpt_path}")
+
+        # ---- End of training cleanup ---------------------------------------
+        if self._is_main:
+            # Save final checkpoint.
+            final_path = save_checkpoint(
+                model=self.model,
+                optimizer=self.optimizer,
+                scheduler=self.scheduler,
+                step=cfg.max_steps,
+                loss=avg_loss,
+                path=cfg.checkpoint_dir,
+            )
+            self._log(f"Training complete. Final checkpoint → {final_path}")
+
+            import datetime
+            elapsed = (datetime.datetime.now() - self._train_start_time).total_seconds()
+            total_steps_done = cfg.max_steps - start_step
+            self._log(
+                f"Training summary: {total_steps_done} steps, "
+                f"{elapsed/3600:.2f}h elapsed, "
+                f"avg {total_steps_done/elapsed:.1f} steps/s"
+            )
+
+            if self._writer is not None:
+                self._writer.close()
+            if self._log_fh is not None:
+                self._log_fh.close()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def _run_validation(self) -> float:
+        """
+        Evaluate the model on the entire validation set and return the mean loss.
+
+        Temporarily switches the model to eval mode and back to train mode
+        afterwards so that dropout / NEFTune hooks are inactive during eval.
+        """
+        model = self.model
+        model.eval()
+        total_loss = 0.0
+        total_batches = 0
+
+        for batch in self._val_loader:  # type: ignore[union-attr]
+            input_ids = batch[0].to(self.device, dtype=torch.long, non_blocking=True)
+            targets   = batch[1].to(self.device, dtype=torch.long, non_blocking=True)
+            # Consume attention_mask if provided (model does not use it yet).
+            _attn_mask = batch[2].to(self.device, non_blocking=True) if len(batch) > 2 else None  # noqa: F841
+
+            device_type = self.device.type
+            with contextlib.ExitStack() as stack:
+                if self.config.use_fp8 and self._fp8_recipe is not None:
+                    stack.enter_context(
+                        torch.autocast(device_type=device_type, dtype=torch.bfloat16)
+                    )
+                    stack.enter_context(
+                        te.fp8_autocast(enabled=True, fp8_recipe=self._fp8_recipe)
+                    )
+                elif self.config.use_amp:
+                    stack.enter_context(
+                        torch.autocast(device_type=device_type, dtype=torch.bfloat16)
+                    )
+                logits, _ = model(input_ids)
+                loss = self._compute_loss(logits, targets)
+
+            total_loss += loss.item()
+            total_batches += 1
+
+        model.train()
+        if total_batches == 0:
+            self._log("Validation set is empty — returning inf", level="WARN")
+            return float("inf")
+        return total_loss / total_batches
+
+    def _log(self, msg: str, level: str = "INFO") -> None:
+        """Print to stdout and optionally write to the log file."""
+        import datetime
+        ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        line = f"[{ts}] [{level}] {msg}"
+        print(line)
+        if self._log_fh is not None:
+            self._log_fh.write(line + "\n")
+
+    def _step(self, batch: tuple) -> torch.Tensor:
+        """
+        Execute one forward + backward pass for a single micro-batch.
+
+        The loss is divided by ``grad_accum_steps`` so that gradients
+        accumulated over multiple micro-batches sum to the correct scale.
+
+        Args:
+            batch: ``(input_ids, targets)`` or ``(input_ids, targets, attention_mask)``
+                   tensors on CPU; moved to device here.
+
+        Returns:
+            Raw (un-scaled) loss as a detached GPU tensor (no CPU sync).
+            The caller is responsible for calling .item() once per optimizer step.
+        """
+        input_ids = batch[0].to(self.device, dtype=torch.long, non_blocking=True)
+        targets   = batch[1].to(self.device, dtype=torch.long, non_blocking=True)
+        # Consume attention_mask if the dataset provides it (future-proof).
+        # Current model forward(input_ids, targets=None) does not accept
+        # attention_mask, so we read it but do not forward it yet.
+        _attn_mask = batch[2].to(self.device, non_blocking=True) if len(batch) > 2 else None  # noqa: F841
+
+        # Store for tokens/sec calculation.
+        self._last_batch_shape = (input_ids.shape[0], input_ids.shape[1])
+
+        device_type = self.device.type
+        # te.fp8_autocast must be combined with torch.autocast(bfloat16) so that
+        # all tensors entering TE modules are in BF16 (not FP32 master weights).
+        # te.fp8_autocast only affects TE modules (te.Linear, te.LayerNormMLP).
+        # Hybrid Mamba-2 layers use nn.Linear → stay in bf16 under torch.autocast.
+        with contextlib.ExitStack() as stack:
+            if self.config.use_fp8 and self._fp8_recipe is not None:
+                stack.enter_context(
+                    torch.autocast(device_type=device_type, dtype=torch.bfloat16)
+                )
+                stack.enter_context(
+                    te.fp8_autocast(enabled=True, fp8_recipe=self._fp8_recipe)
+                )
+            elif self.config.use_amp:
+                stack.enter_context(
+                    torch.autocast(device_type=device_type, dtype=torch.bfloat16)
+                )
+            logits, _ = self.model(input_ids)
+            loss = self._compute_loss(logits, targets)
+
+        # Scale loss for gradient accumulation before backward.
+        scaled_loss = loss / self.config.grad_accum_steps
+        scaled_loss.backward()
+
+        # Return detached GPU tensor — no CPU sync here.
+        # Caller accumulates on GPU and calls .item() once per optimizer step.
+        return loss.detach()
+
+    @staticmethod
+    def _compute_loss(
+        logits: torch.Tensor, targets: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Compute cross-entropy loss, ignoring target positions equal to -1.
+
+        Args:
+            logits:  ``[B, T, vocab_size]`` float tensor.
+            targets: ``[B, T]`` long tensor (may contain -1 as ignore index).
+
+        Returns:
+            Scalar loss tensor.
+        """
+        B, T, V = logits.shape
+        return nn.functional.cross_entropy(
+            logits.view(B * T, V),
+            targets.view(B * T),
+            ignore_index=-1,
+        )
+
+    def _next_batch(self) -> tuple:
+        """Return the next batch, restarting the DataLoader iterator if exhausted."""
+        try:
+            return next(self._loader_iter)
+        except StopIteration:
+            self._epoch += 1
+            # Advance DistributedSampler epoch so each pass has a fresh shuffle.
+            if self._sampler is not None:
+                self._sampler.set_epoch(self._epoch)
+            self._loader_iter = iter(self.train_loader)
+            return next(self._loader_iter)

source/train/utils.py

@@ -0,0 +1,331 @@
+"""
+train/utils.py — Training utility functions.
+
+Provides:
+    get_cosine_schedule_with_warmup  : LambdaLR scheduler with linear warmup + cosine decay
+    save_checkpoint                  : Persist model/optimizer/scheduler state to disk
+    load_checkpoint                  : Restore state from a saved checkpoint directory
+    get_grad_norm                    : Compute total L2 gradient norm across all parameters
+    setup_ddp                        : Initialise NCCL distributed process group
+    cleanup_ddp                      : Tear down distributed process group
+    is_main_process                  : True when this process is rank 0 (or non-distributed)
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import shutil
+from pathlib import Path
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import yaml
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+
+
+# ---------------------------------------------------------------------------
+# Learning-rate schedule
+# ---------------------------------------------------------------------------
+
+
+def get_cosine_schedule_with_warmup(
+    optimizer: Optimizer,
+    warmup_steps: int,
+    total_steps: int,
+    min_lr_ratio: float = 0.1,
+) -> LambdaLR:
+    """
+    Create a LambdaLR scheduler with:
+      - Linear warmup: lr scales from 0 → 1 over [0, warmup_steps)
+      - Cosine decay:  lr scales from 1 → min_lr_ratio over [warmup_steps, total_steps]
+
+    Args:
+        optimizer:     The wrapped optimizer.
+        warmup_steps:  Number of linear-warmup steps.
+        total_steps:   Total number of training steps.
+        min_lr_ratio:  Minimum lr as a fraction of the peak lr (default 0.1).
+
+    Returns:
+        A LambdaLR scheduler instance.
+    """
+    if warmup_steps < 0:
+        raise ValueError(f"warmup_steps must be >= 0, got {warmup_steps}")
+    if total_steps <= 0:
+        raise ValueError(f"total_steps must be > 0, got {total_steps}")
+    if not (0.0 <= min_lr_ratio <= 1.0):
+        raise ValueError(f"min_lr_ratio must be in [0, 1], got {min_lr_ratio}")
+
+    def lr_lambda(current_step: int) -> float:
+        # Linear warmup phase.
+        if current_step < warmup_steps:
+            return float(current_step) / float(max(1, warmup_steps))
+
+        # After total_steps, hold at min_lr_ratio.
+        if current_step >= total_steps:
+            return min_lr_ratio
+
+        # Cosine decay phase.
+        decay_steps = total_steps - warmup_steps
+        progress = float(current_step - warmup_steps) / float(max(1, decay_steps))
+        cosine_factor = 0.5 * (1.0 + math.cos(math.pi * progress))
+        # Scale cosine output from [0, 1] into [min_lr_ratio, 1].
+        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine_factor
+
+    return LambdaLR(optimizer, lr_lambda)
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint save / load
+# ---------------------------------------------------------------------------
+
+
+def save_checkpoint(
+    model: torch.nn.Module,
+    optimizer: Optimizer,
+    scheduler: LambdaLR,
+    step: int,
+    loss: float,
+    path: str | Path,
+    suffix: str | None = None,
+) -> Path:
+    """
+    Save a training checkpoint to ``path/checkpoint-{step:07d}/``.
+
+    Saves:
+        - model.pt        : model state_dict
+        - optimizer.pt    : optimizer state_dict
+        - scheduler.pt    : scheduler state_dict
+        - train_state.pt  : step and loss scalars
+        - config.yaml     : model LMConfig (if the model exposes a ``.config`` attribute)
+
+    Handles both plain ``nn.Module`` and DDP-wrapped models by unwrapping
+    via ``.module`` when present.
+
+    Args:
+        model:     The model (plain or DDP-wrapped).
+        optimizer: The optimizer.
+        scheduler: The LR scheduler.
+        step:      Current training step (used in directory name).
+        loss:      Current loss value (stored for reference).
+        path:      Root checkpoint directory.
+
+    Returns:
+        Path to the created checkpoint sub-directory.
+    """
+    dir_name = f"checkpoint-{suffix}" if suffix else f"checkpoint-{step:07d}"
+    ckpt_dir = Path(path) / dir_name
+    tmp_dir = Path(path) / f".tmp_{dir_name}"
+
+    # Write to temp directory first for crash safety
+    if tmp_dir.exists():
+        shutil.rmtree(tmp_dir)
+    tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    raw_model: torch.nn.Module = getattr(model, "module", model)
+
+    torch.save(raw_model.state_dict(), tmp_dir / "model.pt")
+    torch.save(optimizer.state_dict(), tmp_dir / "optimizer.pt")
+    torch.save(scheduler.state_dict(), tmp_dir / "scheduler.pt")
+
+    import random as _random
+    train_state = {
+        "step": step,
+        "loss": loss,
+        "rng_state": {
+            "python": _random.getstate(),
+            "numpy": np.random.get_state(),
+            "torch_cpu": torch.random.get_rng_state(),
+            "torch_cuda": torch.cuda.get_rng_state_all(),
+        },
+    }
+    torch.save(train_state, tmp_dir / "train_state.pt")
+
+    # Persist the model config when available.
+    if hasattr(raw_model, "config"):
+        cfg = raw_model.config
+        if hasattr(cfg, "to_dict"):
+            config_dict = cfg.to_dict()
+        else:
+            # Fallback: try __dict__ for plain dataclasses.
+            config_dict = {
+                k: v for k, v in vars(cfg).items() if not k.startswith("_")
+            }
+        with open(tmp_dir / "config.yaml", "w", encoding="utf-8") as f:
+            yaml.safe_dump(config_dict, f, default_flow_style=False, sort_keys=False)
+
+    # Atomic swap: rename old → trash, tmp → final, delete trash
+    trash_dir = Path(path) / f".trash_{dir_name}"
+    if trash_dir.exists():
+        shutil.rmtree(trash_dir)
+    if ckpt_dir.exists():
+        ckpt_dir.rename(trash_dir)
+    tmp_dir.rename(ckpt_dir)
+    if trash_dir.exists():
+        shutil.rmtree(trash_dir)
+
+    # Clean up old checkpoints (keep recent N + best)
+    cleanup_old_checkpoints(Path(path))
+
+    return ckpt_dir
+
+
+def cleanup_old_checkpoints(path: Path, keep: int = 5) -> None:
+    """Remove old checkpoints, keeping the most recent `keep` plus checkpoint-best."""
+    ckpts = sorted(
+        [d for d in path.glob("checkpoint-[0-9]*") if d.is_dir()],
+        key=lambda d: d.stat().st_mtime,
+    )
+    for old in ckpts[:-keep]:
+        shutil.rmtree(old)
+
+
+def load_checkpoint(
+    path: str | Path,
+    model: torch.nn.Module,
+    optimizer: Optional[Optimizer] = None,
+    scheduler: Optional[LambdaLR] = None,
+) -> Tuple[int, float]:
+    """
+    Load a checkpoint from a directory created by :func:`save_checkpoint`.
+
+    The model weights are always restored.  Optimizer and scheduler states are
+    only restored when the corresponding objects are provided.
+
+    Args:
+        path:      Path to the checkpoint directory (e.g. ``checkpoints/checkpoint-0001000``).
+        model:     Model to load weights into (plain or DDP-wrapped).
+        optimizer: Optional optimizer to restore state into.
+        scheduler: Optional LR scheduler to restore state into.
+
+    Returns:
+        ``(step, loss)`` — the training step and loss recorded at save time.
+    """
+    ckpt_dir = Path(path)
+    if not ckpt_dir.is_dir():
+        raise FileNotFoundError(f"Checkpoint directory not found: {ckpt_dir}")
+
+    # Unwrap DDP model if necessary.
+    raw_model: torch.nn.Module = getattr(model, "module", model)
+
+    # Determine the device the model lives on.
+    try:
+        device = next(raw_model.parameters()).device
+    except StopIteration:
+        device = torch.device("cpu")
+
+    raw_model.load_state_dict(
+        torch.load(ckpt_dir / "model.pt", map_location=device, weights_only=True)
+    )
+
+    if optimizer is not None:
+        optimizer.load_state_dict(
+            torch.load(ckpt_dir / "optimizer.pt", map_location=device, weights_only=True)
+        )
+
+    if scheduler is not None:
+        scheduler.load_state_dict(
+            torch.load(ckpt_dir / "scheduler.pt", map_location=device, weights_only=True)
+        )
+
+    train_state = torch.load(
+        ckpt_dir / "train_state.pt", map_location="cpu", weights_only=True
+    )
+    step: int = int(train_state["step"])
+    loss: float = float(train_state["loss"])
+
+    # Restore RNG states if available (for exact resume reproducibility)
+    rng_state = train_state.get("rng_state")
+    if rng_state is not None:
+        import random as _random
+        try:
+            _random.setstate(rng_state["python"])
+            np.random.set_state(rng_state["numpy"])
+            torch.random.set_rng_state(rng_state["torch_cpu"])
+            torch.cuda.set_rng_state_all(rng_state["torch_cuda"])
+        except Exception as e:
+            print(f"[WARN] RNG state restore failed (non-fatal): {e}")
+
+    return step, loss
+
+
+# ---------------------------------------------------------------------------
+# Gradient utilities
+# ---------------------------------------------------------------------------
+
+
+def get_grad_norm(model: torch.nn.Module) -> float:
+    """
+    Compute the total L2 norm of all parameter gradients.
+
+    Uses a single GPU kernel + one GPU-CPU sync instead of one sync per
+    parameter (the naive loop approach).  Only parameters with non-None
+    ``.grad`` attribute contribute.
+
+    Args:
+        model: The model (plain or DDP-wrapped).
+
+    Returns:
+        Scalar float — the global gradient L2 norm.
+    """
+    raw_model: torch.nn.Module = getattr(model, "module", model)
+    grads = [p.grad.detach().float() for p in raw_model.parameters() if p.grad is not None]
+    if not grads:
+        return 0.0
+    # Stack individual norms and compute the L2 norm of norms — single sync.
+    return torch.stack([g.norm(2) for g in grads]).norm(2).item()
+
+
+# ---------------------------------------------------------------------------
+# Distributed training helpers
+# ---------------------------------------------------------------------------
+
+
+def setup_ddp() -> Tuple[int, int, int, torch.device]:
+    """
+    Initialise the NCCL distributed process group for DDP training.
+
+    Reads ``RANK``, ``LOCAL_RANK``, and ``WORLD_SIZE`` from the environment
+    (set automatically by ``torchrun``).
+
+    Returns:
+        ``(rank, local_rank, world_size, device)``
+    """
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+
+    # Limit CPU thread count per process to avoid contention across 8 ranks.
+    # 72 cores / 8 ranks = 9; use 4 to leave headroom for DataLoader workers.
+    os.environ.setdefault("OMP_NUM_THREADS", "4")
+    os.environ.setdefault("MKL_NUM_THREADS", "4")
+
+    import datetime as _dt
+    dist.init_process_group(
+        backend="nccl",
+        timeout=_dt.timedelta(seconds=7200),  # 2h for large checkpoint loads
+    )
+
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+
+    return rank, local_rank, world_size, device
+
+
+def cleanup_ddp() -> None:
+    """Tear down the distributed process group (call at end of training)."""
+    if dist.is_available() and dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def is_main_process() -> bool:
+    """
+    Return ``True`` when this process is rank 0 or when running without DDP.
+
+    Reads the ``RANK`` environment variable; if it is absent the process is
+    assumed to be the sole process (rank 0).
+    """
+    return int(os.environ.get("RANK", "0")) == 0