src/train.py

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