from __future__ import annotations import argparse import csv import json import math import os import sys import time from functools import partial from pathlib import Path _REPO_ROOT = Path(__file__).resolve().parents[1] if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) import torch from torch.utils.data import DataLoader, Dataset, Subset from transformers import AutoModelForCausalLM, AutoTokenizer, get_cosine_schedule_with_warmup from configs import cfg, emit_log_spacing, setup_logger from src.checkpoints import ( find_latest_training_checkpoint, load_trainer_state, maybe_upload_checkpoint, packing_checkpoint_metadata, read_env_flag, save_checkpoint, validate_resume_packing_state, ) from src.kd_contracts import build_tokenizer_contract from src.losses import compute_loss_for_phase from src.optim import build_adamw_optimizer from src.provenance import resolve_model_vocab_size, validate_provenance from src.sequence_packing import SequencePackedDataset, collate_packed_fn from src.training_data import ( DistillationDataset, collate_fn, extract_shard_id_range, move_batch_to_device, resolve_dataloader_runtime, torch_load_cpu, ) from src.training_schedule import ( build_train_validation_subsets, compute_training_schedule, load_deepspeed_runtime_config, ) from src.transformers_compat import format_model_load_error, resolve_attention_backend from src.validation import evaluate_validation_loss def _log_gpu(logger) -> None: if torch.cuda.is_available(): device = torch.cuda.current_device() alloc = torch.cuda.max_memory_allocated(device) / (1024**3) reserved = torch.cuda.max_memory_reserved(device) / (1024**3) total = torch.cuda.get_device_properties(device).total_memory / (1024**3) pct = alloc / total * 100 logger.info(f"[GPU] {alloc:.1f}/{total:.0f} GiB ({pct:.0f}%) peak alloc, {reserved:.1f} GiB peak reserved") def main() -> None: parser = argparse.ArgumentParser(description="Quintus training (SFT / KD)") packing_cfg = getattr(cfg.training, "sequence_packing", None) sequence_packing_default = bool(getattr(packing_cfg, "enabled", False)) pack_length_default = int(getattr(packing_cfg, "pack_length", cfg.data.max_seq_len)) mask_first_after_separator = bool(getattr(packing_cfg, "mask_first_token_after_separator", True)) parser.add_argument("--num_samples", type=int, default=cfg.data.num_samples) parser.add_argument("--phase", type=str, choices=["sft", "kd", "online_kd"], default="online_kd", help="Training phase") parser.add_argument("--resume_from_checkpoint", action="store_true", help="Resume from latest epoch in current output directory") parser.add_argument("--init_from_checkpoint", type=str, default=None, help="Initialize weights from a specific path before training") parser.add_argument( "--compile_model", action="store_true", default=bool(getattr(cfg.training, "compile_model", False)), help="Enable torch.compile after checkpoint loading. Off by default for KD memory safety.", ) parser.add_argument("--local_rank", type=int, default=-1, help=argparse.SUPPRESS) parser.add_argument("--deepspeed", type=str, default=None, help="Enable DeepSpeed with the given config path.") parser.add_argument("--no_deepspeed", action="store_true", help="Run without DeepSpeed.") parser.add_argument( "--allow_partial_final_window", action="store_true", help="Allow DeepSpeed to drop a final incomplete accumulation window during smoke tests.", ) parser.add_argument("--teacher_model", type=str, default=cfg.model.teacher) parser.add_argument("--teacher_revision", type=str, default=cfg.model.teacher_revision) parser.add_argument("--student_model", type=str, default=cfg.model.student) parser.add_argument("--student_revision", type=str, default=cfg.model.student_revision) parser.add_argument("--tokenizer_model", type=str, default=getattr(cfg.model, "tokenizer", cfg.model.student)) parser.add_argument("--tokenizer_revision", type=str, default=getattr(cfg.model, "tokenizer_revision", cfg.model.student_revision)) parser.add_argument("--student_dir", type=str, default=cfg.paths.student_dir) parser.add_argument("--tokenizer_dir", type=str, default=getattr(cfg.paths, "tokenizer_dir", cfg.paths.student_dir)) parser.add_argument("--distilled_dir", type=str, default=cfg.paths.distilled_dir) parser.add_argument("--num_epochs", type=int, default=cfg.training.num_epochs) parser.add_argument("--max_steps", type=int, default=-1, help="Stop after this many optimizer steps. -1 = no limit.") parser.add_argument("--learning_rate", type=float, default=float(cfg.training.learning_rate)) parser.add_argument("--alpha", type=float, default=cfg.training.alpha) parser.add_argument("--temperature", type=float, default=cfg.training.temperature) parser.add_argument( "--online_kd_token_chunk_size", type=int, default=int(getattr(cfg.training, "online_kd_token_chunk_size", 2048)), help="Token chunk size for full-vocabulary online KD loss.", ) parser.add_argument("--micro_batch_size", type=int, default=cfg.training.micro_batch_size) parser.add_argument("--grad_accum_steps", type=int, default=cfg.training.grad_accum_steps) parser.add_argument("--sequence_packing", action="store_true", default=False, help="Enable sequence packing for online_kd.") parser.add_argument("--no_sequence_packing", action="store_true", default=False, help="Disable sequence packing.") parser.add_argument("--pack_length", type=int, default=None, help="Packed sequence length.") parser.add_argument("--disable_checkpointing", action="store_true", default=False, help="Disable intermediate epoch/step/best checkpoint saves.") parser.add_argument("--gradient_checkpointing", action="store_true", default=bool(cfg.training.gradient_checkpointing), help="Enable gradient checkpointing (activation checkpointing).") parser.add_argument("--upload_kd_checkpoints", action="store_true", default=False) parser.add_argument("--upload_step_checkpoints", action="store_true", default=False) parser.add_argument( "--upload_last_checkpoint", action="store_true", default=False, help="Upload the final 'last' checkpoint to the Hub. Off by default.", ) parser.add_argument( "--hub_upload_strict", action="store_true", default=read_env_flag("QUINTUS_HUB_UPLOAD_STRICT", False), help="Fail training if a requested Hub checkpoint upload fails.", ) parser.add_argument("--hub_repo_id", type=str, default=f"{cfg.hub.username}/{cfg.hub.repo_name}") parser.add_argument("--ckpt_path_in_repo", type=str, default="models/online_kd_8b_17b_ep1_B200_20260608_alpha0.3") parser.add_argument("--commit_message_prefix", type=str, default="Online KD 8B->1.7B B200 Run (alpha=0.3)") args = parser.parse_args() if args.sequence_packing and args.no_sequence_packing: parser.error("Use either --sequence_packing or --no_sequence_packing, not both.") sequence_packing_enabled = sequence_packing_default if args.sequence_packing: sequence_packing_enabled = True elif args.no_sequence_packing: sequence_packing_enabled = False pack_length = int(args.pack_length if args.pack_length is not None else pack_length_default) if pack_length <= 0: parser.error(f"--pack_length must be positive, got {pack_length}.") if pack_length > int(cfg.data.max_seq_len): parser.error(f"--pack_length must be <= data.max_seq_len ({int(cfg.data.max_seq_len)}), got {pack_length}.") if sequence_packing_enabled and args.phase != "online_kd": parser.error("--sequence_packing is supported only with --phase online_kd.") if args.online_kd_token_chunk_size <= 0: parser.error( f"--online_kd_token_chunk_size must be positive, got {args.online_kd_token_chunk_size}." ) cfg.model.teacher = args.teacher_model cfg.model.teacher_revision = args.teacher_revision cfg.model.student = args.student_model cfg.model.student_revision = args.student_revision cfg.model.tokenizer = args.tokenizer_model cfg.model.tokenizer_revision = args.tokenizer_revision cfg.paths.student_dir = args.student_dir cfg.paths.tokenizer_dir = args.tokenizer_dir cfg.paths.distilled_dir = args.distilled_dir cfg.training.num_epochs = args.num_epochs cfg.training.learning_rate = args.learning_rate cfg.training.alpha = args.alpha cfg.training.temperature = args.temperature cfg.training.online_kd_token_chunk_size = int(args.online_kd_token_chunk_size) cfg.training.micro_batch_size = args.micro_batch_size cfg.training.grad_accum_steps = args.grad_accum_steps cfg.training.gradient_checkpointing = args.gradient_checkpointing cfg.training.disable_checkpointing = args.disable_checkpointing cfg.training.sequence_packing.enabled = sequence_packing_enabled cfg.training.sequence_packing.pack_length = pack_length cfg.training.sequence_packing.mask_first_token_after_separator = mask_first_after_separator cfg.data.num_samples = args.num_samples from omegaconf import OmegaConf if not hasattr(cfg, "hub"): cfg.hub = OmegaConf.create() cfg.hub.upload_kd_checkpoints = args.upload_kd_checkpoints cfg.hub.upload_step_checkpoints = args.upload_step_checkpoints cfg.hub.upload_last_checkpoint = args.upload_last_checkpoint cfg.hub.hub_upload_strict = args.hub_upload_strict cfg.hub.repo_id = args.hub_repo_id cfg.hub.ckpt_path_in_repo = args.ckpt_path_in_repo cfg.hub.commit_message_prefix = args.commit_message_prefix rank = int(os.environ.get("LOCAL_RANK", args.local_rank)) log = setup_logger("TRAIN", rank=rank) log.info("=" * 70) log.info("Quintus Training") log.info("=" * 70) tokenizer_dir = getattr(cfg.paths, "tokenizer_dir", cfg.paths.student_dir) tokenizer_model = getattr(cfg.model, "tokenizer", cfg.model.student) log.info(f" Student: {cfg.paths.student_dir}") log.info(f" Student id: {cfg.model.student}") log.info(f" Tokenizer: {tokenizer_dir}") log.info(f" Tokenizer id:{tokenizer_model}") log.info(f" Num samples: {args.num_samples:,}") log.info(f" Epochs: {cfg.training.num_epochs}") log.info(f" LR: {cfg.training.learning_rate}") log.info(f" Phase: {args.phase}") if args.phase in ("kd", "online_kd"): log.info(f" CE weight: {cfg.training.alpha}") log.info(f" Temperature: {cfg.training.temperature}") if args.phase == "online_kd": log.info(f" KD chunk: {cfg.training.online_kd_token_chunk_size} tokens") log.info(f" Micro batch: {cfg.training.micro_batch_size}") log.info(f" Grad accum: {cfg.training.grad_accum_steps}") log.info(f" Eff. batch: {cfg.training.micro_batch_size * cfg.training.grad_accum_steps}") log.info(f" Val ratio: {cfg.training.validation_ratio:.2%}") log.info(f" Remote code: {cfg.model.allow_remote_code}") log.info(f" Output dir: {cfg.paths.distilled_dir}") log.info(f" Log file: {cfg.paths.log_file}") log.info(f" Fused AdamW: {bool(getattr(cfg.training, 'fused_adamw', False))}") log.info( f" HF upload: regular={cfg.hub.upload_kd_checkpoints} " f"steps={cfg.hub.upload_step_checkpoints} " f"last={cfg.hub.upload_last_checkpoint} " f"strict={cfg.hub.hub_upload_strict}" ) log.info( f" HF target: {cfg.hub.repo_id}/" f"{cfg.hub.ckpt_path_in_repo}" ) if torch.cuda.is_available(): log.info(f" GPU: {torch.cuda.get_device_name(0)}") try: t_dir = tokenizer_dir if not os.path.exists(t_dir): log.warning(f"Tokenizer directory '{t_dir}' not found. Falling back to downloading '{tokenizer_model}' from HF Hub.") t_dir = tokenizer_model tokenizer = AutoTokenizer.from_pretrained( t_dir, trust_remote_code=cfg.model.allow_remote_code, ) except Exception as exc: log.error(format_model_load_error("Student tokenizer load", exc)) sys.exit(1) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token if sequence_packing_enabled: if tokenizer.eos_token_id is None: log.error("Sequence packing requires tokenizer.eos_token_id.") sys.exit(1) if tokenizer.pad_token_id is None: log.error("Sequence packing requires tokenizer.pad_token_id.") sys.exit(1) student_tokenizer_contract = build_tokenizer_contract(tokenizer) student_tokenizer_vocab_size = student_tokenizer_contract["full_vocab_size"] if args.phase == "kd": _prov_path_for_teacher = os.path.join(cfg.paths.logits_dir, "_provenance.json") if os.path.exists(_prov_path_for_teacher): with open(_prov_path_for_teacher, "r", encoding="utf-8") as _pf: _prov_data = json.load(_pf) _teacher_prov = _prov_data.get("teacher", {}) teacher_tokenizer_contract = { "full_vocab_size": _teacher_prov.get("tokenizer_size"), "fingerprint": _teacher_prov.get("tokenizer_fingerprint"), } log.info( f" Teacher contract read from provenance: " f"vocab={teacher_tokenizer_contract['full_vocab_size']}, " f"fingerprint={teacher_tokenizer_contract['fingerprint'][:12]}..." ) else: try: teacher_tokenizer = AutoTokenizer.from_pretrained( cfg.paths.teacher_dir if os.path.exists(cfg.paths.teacher_dir) else cfg.model.teacher, trust_remote_code=cfg.model.allow_remote_code, ) except Exception as exc: log.error(format_model_load_error("Teacher tokenizer load", exc)) sys.exit(1) teacher_tokenizer_contract = build_tokenizer_contract(teacher_tokenizer) del teacher_tokenizer else: teacher_tokenizer_contract = None attn_impl = resolve_attention_backend(log) log.info(f" Attention: {attn_impl}") try: from liger_kernel.transformers import apply_liger_kernel_to_qwen3 apply_liger_kernel_to_qwen3( rope=True, swiglu=True, rms_norm=True, cross_entropy=False, fused_linear_cross_entropy=False, ) log.info(" Liger: enabled") except ImportError: if cfg.training.micro_batch_size >= 6: log.error(" Liger: missing; install liger-kernel or lower micro_batch_size.") raise RuntimeError("liger_kernel is required for micro_batch_size >= 6.") else: log.warning(" Liger: not installed") try: s_dir = cfg.paths.student_dir if not os.path.exists(s_dir): log.warning(f"Student model directory '{s_dir}' not found. Falling back to downloading '{cfg.model.student}' from HF Hub.") s_dir = cfg.model.student model = AutoModelForCausalLM.from_pretrained( s_dir, dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=cfg.model.allow_remote_code, attn_implementation=attn_impl, ) except Exception as exc: log.error(format_model_load_error("Student model load", exc)) sys.exit(1) model.config.use_cache = False if getattr(cfg.training, "gradient_checkpointing", False): model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) log.info(" Grad ckpt: enabled") else: log.info(" Grad ckpt: disabled") start_epoch = 0 resume_state: dict = {} if args.resume_from_checkpoint and args.init_from_checkpoint: log.error("Use either --init_from_checkpoint or --resume_from_checkpoint, not both.") sys.exit(1) checkpoint_to_load = args.init_from_checkpoint if args.resume_from_checkpoint: latest_ckpt = find_latest_training_checkpoint(cfg.paths.distilled_dir) if latest_ckpt is None: log.error( f"--resume_from_checkpoint was set, but no epoch_* or step_* checkpoints were found in " f"{cfg.paths.distilled_dir}. Use --init_from_checkpoint for the first KD run." ) sys.exit(1) checkpoint_to_load = latest_ckpt resume_state = load_trainer_state(latest_ckpt, log) checkpoint_type = resume_state.get("checkpoint_type", os.path.basename(latest_ckpt).split("_")[0]) start_epoch = int(resume_state.get("start_epoch", 0) or 0) if checkpoint_type == "epoch": log.info(f"Interrupted run detected. Resuming after completed epoch {start_epoch}") else: log.info( f"Interrupted run detected. Resuming from {os.path.basename(latest_ckpt)} " f"at epoch_index={start_epoch}, next_batch_in_epoch=" f"{int(resume_state.get('next_batch_in_epoch', 0) or 0)}" ) validate_resume_packing_state( resume_state, enabled=sequence_packing_enabled, pack_length=pack_length, max_seq_len=int(cfg.data.max_seq_len), log=log, ) if checkpoint_to_load: log.info(f"Loading weights from: {checkpoint_to_load}") try: from safetensors.torch import load_file ckpt_file = os.path.join(checkpoint_to_load, "model.safetensors") if not os.path.exists(ckpt_file): ckpt_file = os.path.join(checkpoint_to_load, "pytorch_model.bin") if ckpt_file.endswith(".safetensors"): state_dict = load_file(ckpt_file) else: state_dict = torch.load(ckpt_file, map_location="cpu") new_state_dict = {} for k, v in state_dict.items(): if k.startswith("_orig_mod."): new_state_dict[k[len("_orig_mod."):]] = v else: new_state_dict[k] = v model.load_state_dict(new_state_dict) log.info("Weights loaded.") except Exception as e: log.error(f"Failed to load weights: {e}") sys.exit(1) model.train() if args.compile_model: log.info(" Compile: enabled") model = torch.compile(model, dynamic=True) else: log.info(" Compile: disabled") torch.set_float32_matmul_precision("high") student_model_vocab_size = resolve_model_vocab_size(model, tokenizer, "Student", log) log.info( f" Student V: tokenizer={student_tokenizer_vocab_size:,} " f"model={student_model_vocab_size:,}" ) _log_gpu(log) if args.phase == "kd": shard0 = os.path.join(cfg.paths.logits_dir, "shard_000000.pt") if os.path.exists(shard0): test_shard = torch_load_cpu(shard0) try: min_id, max_id = extract_shard_id_range(test_shard, shard0) except (KeyError, ValueError) as exc: log.error(str(exc)) sys.exit(1) if min_id < 0: log.error(f" Negative IDs (min={min_id}); int16 overflow.") log.error(" Regenerate shards.") sys.exit(1) if max_id >= student_tokenizer_vocab_size: log.error( f"VOCAB MISMATCH: shard max_id={max_id} >= " f"student tokenizer vocab={student_tokenizer_vocab_size}" ) sys.exit(1) log.info( f" Vocab check: PASS (ids in [{min_id}, {max_id}], " f"reachable tokenizer V={student_tokenizer_vocab_size})" ) else: log.warning(f" Shard {shard0} not found; skipping vocab check") data_path = os.path.join(cfg.paths.tokenized_dir, "train.jsonl") dataset = DistillationDataset(data_path, cfg.paths.logits_dir, cfg.data.max_seq_len, args.num_samples, args.phase) log.info(f" Dataset: {len(dataset):,} samples") if args.phase == "kd": prov_path = os.path.join(cfg.paths.logits_dir, "_provenance.json") validate_provenance( prov_path=prov_path, data_path=data_path, dataset=dataset, teacher_tokenizer_contract=teacher_tokenizer_contract, student_tokenizer_contract=student_tokenizer_contract, log=log, ) pad_id = tokenizer.pad_token_id if args.no_deepspeed: args.deepspeed = None use_ds = args.deepspeed is not None world_size = int(os.environ.get("WORLD_SIZE", 1)) if world_size != 1: log.error("This training path is single-GPU only. Re-run with NUM_GPUS=1.") sys.exit(1) is_main = rank in (-1, 0) ds_runtime_config = None if use_ds: try: ds_runtime_config = load_deepspeed_runtime_config( args.deepspeed, micro_batch_size=cfg.training.micro_batch_size, grad_accum=cfg.training.grad_accum_steps, ) except (OSError, ValueError, json.JSONDecodeError) as exc: log.error(str(exc)) sys.exit(1) train_dataset, val_dataset, split_meta = build_train_validation_subsets( dataset=dataset, validation_ratio=float(cfg.training.validation_ratio), split_seed=int(cfg.training.split_seed), micro_batch_size=cfg.training.micro_batch_size, grad_accum=cfg.training.grad_accum_steps, num_epochs=cfg.training.num_epochs, use_ds=use_ds, ) log.info( f" Train split: {len(train_dataset):,} samples | " f"Val split: {int(split_meta['validation_size']):,} samples" ) if bool(split_meta["accumulation_aligned"]): log.info( f" Accum align: train split is divisible by effective batch " f"{int(split_meta['effective_batch_size']):,}" ) else: if use_ds: log.warning( f" Accum align: train split leaves " f"{int(split_meta['train_remainder_batches'])} partial accumulation batches per epoch; " "DeepSpeed will carry partial accumulation across epoch boundaries" ) else: log.warning( f" Accum align: train split leaves " f"{int(split_meta['train_remainder_batches'])} partial accumulation batches per epoch; " "the fallback flush path will rescale gradients correctly" ) if bool(split_meta["adjusted"]): log.info( f" Val align: requested {int(split_meta['requested_validation_size']):,} " f"({float(split_meta['requested_validation_ratio']) * 100:.2f}%), " f"using {int(split_meta['validation_size']):,} " f"({float(split_meta['actual_validation_ratio']) * 100:.2f}%) " "to preserve the training schedule" ) elif val_dataset is not None: log.info( f" Val split: using {float(split_meta['actual_validation_ratio']) * 100:.2f}% " f"held out with split_seed={cfg.training.split_seed}" ) else: log.warning(" Validation disabled; tracking training loss.") effective_train_dataset: Dataset = train_dataset train_collate = partial(collate_fn, pad_token_id=pad_id) val_collate = partial(collate_fn, pad_token_id=pad_id) if sequence_packing_enabled: if isinstance(train_dataset, Subset): source_dataset = train_dataset.dataset train_source_indices = [int(index) for index in train_dataset.indices] else: source_dataset = train_dataset train_source_indices = list(range(len(train_dataset))) if not isinstance(source_dataset, DistillationDataset): log.error("Sequence packing requires DistillationDataset as the split source.") sys.exit(1) val_source_indices: set[int] = set() if isinstance(val_dataset, Subset) and val_dataset.dataset is source_dataset: val_source_indices = {int(index) for index in val_dataset.indices} try: packed_train_dataset = SequencePackedDataset( source=source_dataset, source_indices=train_source_indices, pack_length=pack_length, eos_token_id=int(tokenizer.eos_token_id), pad_token_id=int(tokenizer.pad_token_id), mask_first_after_separator=mask_first_after_separator, ) except (IndexError, ValueError) as exc: log.error(str(exc)) sys.exit(1) overlap = packed_train_dataset.source_index_set.intersection(val_source_indices) if overlap: first_overlap = min(overlap) log.error(f"Sequence packing split error: validation sample #{first_overlap} appears in training bins.") sys.exit(1) effective_train_dataset = packed_train_dataset train_collate = partial(collate_packed_fn, pad_token_id=pad_id) log.info(" Packing: enabled") log.info(f" Pack length: {packed_train_dataset.pack_length:,}") log.info(f" Train bins: {packed_train_dataset.bin_count:,}") log.info(f" Train rows: {packed_train_dataset.source_sample_count:,}") log.info(f" Avg samples: {packed_train_dataset.average_samples_per_bin:.2f} per bin") log.info(f" Original tokens: {packed_train_dataset.original_token_count:,}") log.info(f" Separator tokens: {packed_train_dataset.separator_token_count:,}") log.info(f" Pad tokens: {packed_train_dataset.pad_token_count:,}") log.info(f" Utilization: {packed_train_dataset.utilization * 100:.1f}%") else: log.info(" Packing: disabled") dataloader_runtime = resolve_dataloader_runtime() log.info( " DataLoader: " f"workers={int(dataloader_runtime['num_workers'])} " f"pin_memory={bool(dataloader_runtime['pin_memory'])} " f"persistent={bool(dataloader_runtime.get('persistent_workers', False))}" ) dataloader = DataLoader( effective_train_dataset, batch_size=cfg.training.micro_batch_size, shuffle=(args.phase != "kd"), collate_fn=train_collate, drop_last=True, **dataloader_runtime, ) if args.phase == "kd": log.info(" KD sampler: sequential shard-local order (split membership remains randomized)") val_dataloader = None if val_dataset is not None: val_dataloader = DataLoader( val_dataset, batch_size=cfg.training.micro_batch_size, shuffle=False, collate_fn=val_collate, drop_last=False, **dataloader_runtime, ) grad_accum = cfg.training.grad_accum_steps schedule = compute_training_schedule( dataset_size=len(effective_train_dataset), micro_batch_size=cfg.training.micro_batch_size, grad_accum=grad_accum, num_epochs=cfg.training.num_epochs, use_ds=use_ds, drop_last=True, ) batches_per_epoch = int(schedule["batches_per_epoch"]) remainder_batches = int(schedule["remainder_batches"]) has_remainder = bool(schedule["has_remainder"]) total_micro_batches = int(schedule["total_micro_batches"]) steps_per_epoch = int(schedule["steps_per_epoch"]) total_steps = int(schedule["total_steps"]) final_remainder = int(schedule["final_remainder"]) if batches_per_epoch == 0: schedule_unit = "packed bins" if sequence_packing_enabled else "samples" log.error( f"Dataset too small for micro_batch_size={cfg.training.micro_batch_size}. " f"Train split has {len(effective_train_dataset)} {schedule_unit} and drop_last=True would produce 0 batches." ) sys.exit(1) dropped_samples_per_epoch = int(schedule["dropped_samples_per_epoch"]) if dropped_samples_per_epoch: schedule_unit = "packed bins" if sequence_packing_enabled else "samples" log.warning( f" drop_last=True will discard {dropped_samples_per_epoch} {schedule_unit} per epoch " "before gradient accumulation begins" ) if use_ds and final_remainder: dropped_total = int(schedule["dropped_samples_total"]) schedule_unit = "packed bins" if sequence_packing_enabled else "samples" message = ( f"DeepSpeed would drop the final {final_remainder} micro-batches " f"({dropped_total} {schedule_unit} total) because {batches_per_epoch} batches per epoch " f"across {cfg.training.num_epochs} epochs yields {total_micro_batches} micro-batches, " f"which is not divisible by grad_accum={grad_accum}." ) if not args.allow_partial_final_window: log.error(message) log.error( "Adjust num_samples, micro_batch_size, grad_accum_steps, or num_epochs " "so total micro-batches is divisible by grad_accum, or rerun with " "--allow_partial_final_window for a smoke test." ) sys.exit(1) log.warning(message) log.warning("Proceeding because --allow_partial_final_window was set.") warmup_steps = int(total_steps * cfg.training.warmup_ratio) if has_remainder: if use_ds: log.info( f" NOTE: {batches_per_epoch} batches are not divisible by grad_accum={grad_accum}; " f"DeepSpeed carries {remainder_batches} leftover micro-batches across epoch boundaries" ) if final_remainder and args.allow_partial_final_window: log.info( f" NOTE: only the final {final_remainder} micro-batches of " "the last epoch are dropped because they never reach a full accumulation window" ) else: log.info( f" NOTE: {batches_per_epoch} batches are not divisible by grad_accum={grad_accum}; " f"the training loop will flush {remainder_batches} leftover micro-batches each epoch" ) if not use_ds: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) optimizer = build_adamw_optimizer(list(model.parameters()), log, allow_fused=not use_ds) scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps) resume_global_step = int(resume_state.get("global_step", 0) or 0) if args.resume_from_checkpoint else 0 saved_run_epochs = int(resume_state.get("num_epochs", cfg.training.num_epochs) or cfg.training.num_epochs) extending_completed_run = ( args.resume_from_checkpoint and saved_run_epochs < cfg.training.num_epochs and start_epoch >= saved_run_epochs ) scheduler_state_path = os.path.join(checkpoint_to_load, "scheduler.pt") if checkpoint_to_load else None if ( extending_completed_run and read_env_flag("QUINTUS_FRESH_SCHEDULER_ON_EXTEND", True) ): remaining_steps = max(1, total_steps - resume_global_step) extension_warmup_steps = int(remaining_steps * cfg.training.warmup_ratio) scheduler = get_cosine_schedule_with_warmup(optimizer, extension_warmup_steps, remaining_steps) log.info( f" Scheduler: fresh extension schedule " f"({remaining_steps:,} remaining steps, {extension_warmup_steps:,} warmup); " f"checkpoint was saved for {saved_run_epochs} epochs" ) elif args.resume_from_checkpoint and scheduler_state_path and os.path.exists(scheduler_state_path): try: scheduler.load_state_dict(torch.load(scheduler_state_path, map_location="cpu")) for param_group, lr in zip(optimizer.param_groups, scheduler.get_last_lr()): param_group["lr"] = lr log.info(f" Scheduler: restored from {scheduler_state_path}") except Exception as exc: log.warning(f" Scheduler restore failed ({exc}); continuing with a fresh schedule") log.info(f" Batches/ep: {batches_per_epoch:,}") step_label = "Steps/ep" step_note = "" if has_remainder: if use_ds: step_label = "Steps/ep*" step_note = " (floor; cross-epoch carry shifts exact epoch boundaries)" else: step_note = " (includes remainder flush)" log.info(f" {step_label}: {steps_per_epoch:,}{step_note}") log.info(f" Steps total: {total_steps:,} ({warmup_steps:,} warmup)") log.info( " Best ckpt: held-out validation loss" if val_dataloader is not None else " Best ckpt: training loss (validation disabled)" ) if use_ds: import deepspeed model, optimizer, _, scheduler = deepspeed.initialize( model=model, optimizer=optimizer, lr_scheduler=scheduler, config=ds_runtime_config, ) device = model.device log.info("[DS] DeepSpeed ZeRO-2 initialized") log.info(f"[DS] DeepSpeed will accumulate over {grad_accum} micro-batches internally") else: log.info(f" Device: {device}") _log_gpu(log) teacher_model = None if args.phase == "online_kd": teacher_source = cfg.paths.teacher_dir if os.path.exists(cfg.paths.teacher_dir) else cfg.model.teacher if teacher_source != cfg.model.teacher: log.info(f"Loading frozen teacher model from local directory '{teacher_source}' on device {device}...") else: log.info(f"Loading frozen teacher model '{teacher_source}' on device {device}...") try: teacher_model = AutoModelForCausalLM.from_pretrained( teacher_source, dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=cfg.model.allow_remote_code, attn_implementation=attn_impl, ).to(device) for p in teacher_model.parameters(): p.requires_grad = False teacher_model.eval() log.info(f"Teacher model '{teacher_source}' loaded and frozen.") except Exception as exc: log.error(f"Failed to load teacher model: {exc}") sys.exit(1) checkpoint_packing_metadata = packing_checkpoint_metadata( enabled=sequence_packing_enabled, pack_length=pack_length, max_seq_len=int(cfg.data.max_seq_len), ) os.makedirs(cfg.paths.distilled_dir, exist_ok=True) loss_log: list[dict] = [] global_step = resume_global_step micro_step_global = int(resume_state.get("micro_step_global", 0) or 0) if args.resume_from_checkpoint else 0 best_metric_name = "validation loss" if val_dataloader is not None else "training loss" best_selection_loss = float("inf") if args.resume_from_checkpoint and "best_selection_loss" in resume_state: try: best_selection_loss = float(resume_state["best_selection_loss"]) log.info(f" Best resume: restored prior best {best_metric_name}={best_selection_loss:.4f}") except (TypeError, ValueError): log.warning(" Best resume: prior best_selection_loss was unreadable; recomputing from this run") best_checkpoint_tag = resume_state.get("best_checkpoint_tag") best_ckpt_path = os.path.join(cfg.paths.distilled_dir, "best") if not os.path.isdir(best_ckpt_path): best_ckpt_path = None if best_checkpoint_tag: candidate_best_path = os.path.join(cfg.paths.distilled_dir, str(best_checkpoint_tag)) if os.path.isdir(candidate_best_path): best_ckpt_path = candidate_best_path log.info(f" Best resume: using {best_checkpoint_tag} as the current best checkpoint") t_start = time.time() alpha = cfg.training.alpha temperature = cfg.training.temperature log_every = max(1, min(50, total_steps // 20)) checkpoint_every_steps = max(0, int(os.environ.get("TRAIN_CHECKPOINT_EVERY_STEPS", "2000"))) if getattr(cfg.training, "disable_checkpointing", False): checkpoint_every_steps = 0 running_loss = 0.0 running_ce = 0.0 running_kd = 0.0 running_count = 0 emit_log_spacing(log) log.info("-" * 70) log.info("Training Start") if checkpoint_every_steps: log.info(f" Mid-epoch checkpoint interval: every {checkpoint_every_steps:,} optimizer steps") else: log.info(" Mid-epoch checkpoints disabled") log.info("-" * 70) window_tokens = 0 window_t_start = time.time() _gpu_loss_accum = torch.zeros(1, device=device) _gpu_ce_accum = torch.zeros(1, device=device) _gpu_kd_accum = torch.zeros(1, device=device) _gpu_tokens_accum = torch.zeros(1, dtype=torch.long, device=device) training_complete = False for epoch in range(start_epoch, cfg.training.num_epochs): if training_complete: break t_epoch = time.time() epoch_loss = 0.0 epoch_ce = 0.0 epoch_kd = 0.0 epoch_steps = 0 epoch_tokens = 0 micro_in_epoch = 0 resume_batch_offset = 0 if args.resume_from_checkpoint and epoch == start_epoch: resume_batch_offset = int(resume_state.get("next_batch_in_epoch", 0) or 0) if resume_batch_offset: log.info(f" Resume: skipping {resume_batch_offset:,} already-processed batches in epoch {epoch + 1}") for batch_idx, batch in enumerate(dataloader): if resume_batch_offset and batch_idx < resume_batch_offset: continue batch = move_batch_to_device(batch, device) input_ids = batch["input_ids"] attention_mask = batch["attention_mask"] labels = batch["labels"] loss_mask = batch["loss_mask"] logits = model(input_ids=input_ids, attention_mask=attention_mask).logits if args.phase == "online_kd" and teacher_model is not None: with torch.no_grad(): teacher_logits = teacher_model(input_ids=input_ids, attention_mask=attention_mask).logits else: teacher_logits = None loss, ce, kd = compute_loss_for_phase( args.phase, logits, labels, loss_mask, batch, alpha, temperature, teacher_logits=teacher_logits, online_kd_token_chunk_size=int(cfg.training.online_kd_token_chunk_size), ) if not torch.isfinite(loss): log.error( f"Non-finite loss in phase={args.phase}: " f"loss={loss.item()} ce={ce.item()} kd={kd.item()}" ) if args.phase == "kd": log.error("Action: regenerate teacher logits.") else: log.error("Action: check dataset / reduce LR.") sys.exit(1) micro_in_epoch += 1 micro_step_global += 1 _gpu_loss_accum += loss.detach() _gpu_ce_accum += ce.detach() _gpu_kd_accum += kd.detach() _gpu_tokens_accum += attention_mask.sum() if use_ds: model.backward(loss) model.step() else: scaled = loss / grad_accum scaled.backward() if micro_in_epoch % grad_accum == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) is_optim_step = ( (micro_step_global % grad_accum == 0) if use_ds else (micro_in_epoch % grad_accum == 0) ) if is_optim_step: global_step += 1 epoch_steps += 1 running_count += 1 step_loss = _gpu_loss_accum.item() / grad_accum step_ce = _gpu_ce_accum.item() / grad_accum step_kd = _gpu_kd_accum.item() / grad_accum step_tokens = _gpu_tokens_accum.item() _gpu_loss_accum.zero_() _gpu_ce_accum.zero_() _gpu_kd_accum.zero_() _gpu_tokens_accum.zero_() epoch_tokens += step_tokens window_tokens += step_tokens epoch_loss += step_loss epoch_ce += step_ce epoch_kd += step_kd running_loss += step_loss running_ce += step_ce running_kd += step_kd if global_step % log_every == 0 or global_step == total_steps: avg_loss = running_loss / max(running_count, 1) avg_ce = running_ce / max(running_count, 1) avg_kd = running_kd / max(running_count, 1) try: lr = scheduler.get_last_lr()[0] except Exception: lr = cfg.training.learning_rate window_elapsed = max(time.time() - window_t_start, 0.1) rolling_tok_s = window_tokens / window_elapsed rolling_eta_s = (window_elapsed / max(running_count, 1)) * (total_steps - global_step) / log_every * running_count cum_tok_s = epoch_tokens / max(time.time() - t_epoch, 1) log.info( f" E{epoch + 1}/{cfg.training.num_epochs} " f"S{global_step:>4}/{total_steps} | " f"loss={avg_loss:.4f} ce={avg_ce:.4f} kd={avg_kd:.4f} | " f"lr={lr:.2e} | {rolling_tok_s:,.0f} tok/s (avg {cum_tok_s:,.0f}) | ETA {rolling_eta_s / 60:.1f}m" ) loss_log.append( { "step": global_step, "epoch": epoch + 1, "loss_total": round(avg_loss, 5), "loss_ce": round(avg_ce, 5), "loss_kd": round(avg_kd, 5), "lr": lr, "tok_per_sec": round(rolling_tok_s, 0), "tok_per_sec_cumulative": round(cum_tok_s, 0), } ) window_tokens = 0 window_t_start = time.time() running_loss = 0.0 running_ce = 0.0 running_kd = 0.0 running_count = 0 if checkpoint_every_steps and global_step % checkpoint_every_steps == 0 and is_main: log.info(f" Saving mid-epoch checkpoint at step {global_step}...") step_tag = f"step_{global_step}" step_ckpt_path = save_checkpoint( model, tokenizer, cfg.paths.distilled_dir, step_tag, log, scheduler=scheduler, trainer_state={ **checkpoint_packing_metadata, "checkpoint_type": "step", "phase": args.phase, "epoch_index": epoch, "start_epoch": epoch, "global_step": global_step, "micro_step_global": micro_step_global, "next_batch_in_epoch": micro_in_epoch, "num_epochs": cfg.training.num_epochs, "micro_batch_size": cfg.training.micro_batch_size, "grad_accum_steps": grad_accum, }, ) maybe_upload_checkpoint(step_ckpt_path, step_tag, log) if args.max_steps > 0 and global_step >= args.max_steps: log.info(f"Reached max_steps={args.max_steps}. Stopping training.") training_complete = True break if training_complete: break if not use_ds: remainder = micro_in_epoch % grad_accum if remainder != 0: flush_scale = grad_accum / remainder for parameter in model.parameters(): if parameter.grad is not None: parameter.grad.mul_(flush_scale) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad(set_to_none=True) global_step += 1 epoch_steps += 1 step_loss = _gpu_loss_accum.item() / remainder step_ce = _gpu_ce_accum.item() / remainder step_kd = _gpu_kd_accum.item() / remainder step_tokens = _gpu_tokens_accum.item() _gpu_loss_accum.zero_() _gpu_ce_accum.zero_() _gpu_kd_accum.zero_() _gpu_tokens_accum.zero_() epoch_tokens += step_tokens window_tokens += step_tokens running_loss += step_loss running_ce += step_ce running_kd += step_kd running_count += 1 avg_loss = running_loss / max(running_count, 1) avg_ce = running_ce / max(running_count, 1) avg_kd = running_kd / max(running_count, 1) epoch_loss += step_loss epoch_ce += step_ce epoch_kd += step_kd running_loss = 0.0 running_ce = 0.0 running_kd = 0.0 running_count = 0 elapsed = time.time() - t_start try: lr = scheduler.get_last_lr()[0] except Exception: lr = cfg.training.learning_rate tok_s = epoch_tokens / max(time.time() - t_epoch, 1) eta_s = (elapsed / max(global_step, 1)) * (total_steps - global_step) log.info( f" E{epoch + 1}/{cfg.training.num_epochs} " f"S{global_step:>4}/{total_steps} | " f"loss={avg_loss:.4f} ce={avg_ce:.4f} kd={avg_kd:.4f} | " f"lr={lr:.2e} | {tok_s:,.0f} tok/s | ETA {eta_s / 60:.1f}m [flush]" ) loss_log.append( { "step": global_step, "epoch": epoch + 1, "loss_total": round(avg_loss, 5), "loss_ce": round(avg_ce, 5), "loss_kd": round(avg_kd, 5), "lr": lr, "tok_per_sec": round(tok_s, 0), } ) window_tokens = 0 window_t_start = time.time() log.info(f" Epoch {epoch + 1}: flushed {remainder} leftover micro-batches") else: optimizer.zero_grad(set_to_none=True) elif (micro_step_global % grad_accum) != 0 and epoch < cfg.training.num_epochs - 1: carry = micro_step_global % grad_accum log.info(f" Epoch {epoch + 1}: carrying {carry} micro-batches into the next epoch") avg_epoch_loss = epoch_loss / max(epoch_steps, 1) avg_epoch_ce = epoch_ce / max(epoch_steps, 1) avg_epoch_kd = epoch_kd / max(epoch_steps, 1) epoch_elapsed = time.time() - t_epoch log.info( f" Epoch {epoch + 1} done | " f"avg_loss={avg_epoch_loss:.4f} ce={avg_epoch_ce:.4f} kd={avg_epoch_kd:.4f} | " f"{epoch_tokens:,} tok | {epoch_elapsed / 60:.1f}m" ) _log_gpu(log) val_metrics = None if val_dataloader is not None: val_start = time.time() val_limit = min(20, len(val_dataloader)) if args.max_steps > 0 else -1 if val_limit > 0: log.info(f" Validation start | capping at {val_limit} batches for dry run (total {len(val_dataloader)} batches)") else: log.info(f" Validation start | {len(val_dataloader):,} batches") val_metrics = evaluate_validation_loss( phase=args.phase, model=model, dataloader=val_dataloader, device=device, alpha=alpha, temperature=temperature, online_kd_token_chunk_size=int(cfg.training.online_kd_token_chunk_size), teacher_model=teacher_model, max_batches=val_limit, ) log.info( f" Validation | loss={val_metrics['loss']:.4f} ce={val_metrics['ce']:.4f} " f"kd={val_metrics['kd']:.4f} | {int(val_metrics['batches'])} batches | " f"{(time.time() - val_start) / 60:.1f}m" ) if is_main: selection_loss = val_metrics["loss"] if val_metrics is not None else avg_epoch_loss is_new_best = selection_loss < best_selection_loss epoch_tag = f"epoch_{epoch + 1}" if is_new_best: best_selection_loss = selection_loss best_checkpoint_tag = epoch_tag log.info(f" Best update: {best_metric_name}={best_selection_loss:.4f} from {epoch_tag}") else: log.info( f" Best unchanged: current {best_metric_name}={selection_loss:.4f}; " f"best={best_selection_loss:.4f} from {best_checkpoint_tag}" ) epoch_state = { **checkpoint_packing_metadata, "checkpoint_type": "epoch", "phase": args.phase, "epoch_index": epoch, "start_epoch": epoch + 1, "global_step": global_step, "micro_step_global": micro_step_global, "next_batch_in_epoch": 0, "num_epochs": cfg.training.num_epochs, "micro_batch_size": cfg.training.micro_batch_size, "grad_accum_steps": grad_accum, "selection_loss": float(selection_loss), "best_selection_loss": float(best_selection_loss), "best_metric_name": best_metric_name, "best_checkpoint_tag": best_checkpoint_tag, } if read_env_flag("QUINTUS_SAVE_EPOCH_CHECKPOINTS", True) and not getattr(cfg.training, "disable_checkpointing", False): epoch_ckpt_path = save_checkpoint( model, tokenizer, cfg.paths.distilled_dir, epoch_tag, log, scheduler=scheduler, trainer_state=epoch_state, ) maybe_upload_checkpoint(epoch_ckpt_path, epoch_tag, log) else: log.info(f" Skipping intermediate {epoch_tag} save") if is_new_best and not getattr(cfg.training, "disable_checkpointing", False): best_ckpt_path = save_checkpoint( model, tokenizer, cfg.paths.distilled_dir, "best", log, scheduler=scheduler, trainer_state=dict(epoch_state, checkpoint_type="best"), ) if use_ds and final_remainder: model.zero_grad() running_loss = 0.0 running_ce = 0.0 running_kd = 0.0 running_count = 0 log.warning(f" Training end: dropped final {final_remainder} leftover micro-batches") if is_main: if best_ckpt_path and os.path.isdir(best_ckpt_path) and not getattr(cfg.training, "disable_checkpointing", False): maybe_upload_checkpoint(best_ckpt_path, "best", log) last_ckpt_path = save_checkpoint( model, tokenizer, cfg.paths.distilled_dir, "last", log, scheduler=scheduler, trainer_state={ **checkpoint_packing_metadata, "checkpoint_type": "last", "phase": args.phase, "start_epoch": cfg.training.num_epochs, "global_step": global_step, "micro_step_global": micro_step_global, "next_batch_in_epoch": 0, "num_epochs": cfg.training.num_epochs, "micro_batch_size": cfg.training.micro_batch_size, "grad_accum_steps": grad_accum, "best_selection_loss": float(best_selection_loss) if math.isfinite(best_selection_loss) else None, "best_metric_name": best_metric_name, "best_checkpoint_tag": best_checkpoint_tag, }, ) maybe_upload_checkpoint(last_ckpt_path, "last", log) csv_path = os.path.join(cfg.paths.distilled_dir, cfg.paths.loss_csv) if loss_log and is_main: with open(csv_path, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=loss_log[0].keys()) writer.writeheader() writer.writerows(loss_log) log.info(f"Loss CSV -> {csv_path}") total_elapsed = time.time() - t_start emit_log_spacing(log) log.info("=" * 70) log.info("Training complete") log.info(f" Wall time: {total_elapsed / 3600:.2f}h ({total_elapsed / 60:.1f}m)") log.info(f" Optim steps: {global_step}") log.info(f" Micro steps: {micro_step_global}") log.info(f" Best {best_metric_name}: {best_selection_loss:.4f}") log.info(f" Best ckpt: {best_ckpt_path}") log.info(f" Output dir: {cfg.paths.distilled_dir}/") log.info("=" * 70) if __name__ == "__main__": try: main() except Exception: try: setup_logger("TRAIN").exception("Uncaught training failure") except Exception: pass raise