587 lines
26 KiB
Python
587 lines
26 KiB
Python
"""
|
||
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()
|