src/training_schedule.py

from __future__ import annotations

import json
import math

import torch
from torch.utils.data import Dataset, Subset

def compute_training_schedule(
    dataset_size: int,
    micro_batch_size: int,
    grad_accum: int,
    num_epochs: int,
    use_ds: bool,
    drop_last: bool = True,
) -> dict[str, int | bool]:
    if dataset_size < 0:
        raise ValueError("dataset_size must be >= 0")
    if micro_batch_size <= 0 or grad_accum <= 0 or num_epochs <= 0:
        raise ValueError("micro_batch_size, grad_accum, and num_epochs must all be positive")

    if drop_last:
        batches_per_epoch = dataset_size // micro_batch_size
        used_samples_per_epoch = batches_per_epoch * micro_batch_size
        dropped_samples_per_epoch = dataset_size - used_samples_per_epoch
    else:
        batches_per_epoch = math.ceil(dataset_size / micro_batch_size) if dataset_size else 0
        used_samples_per_epoch = dataset_size
        dropped_samples_per_epoch = 0

    total_micro_batches = batches_per_epoch * num_epochs
    remainder_batches = batches_per_epoch % grad_accum if batches_per_epoch else 0
    has_remainder = remainder_batches != 0

    if use_ds:
        steps_per_epoch = batches_per_epoch // grad_accum
        total_steps = total_micro_batches // grad_accum
        final_remainder = total_micro_batches % grad_accum
    else:
        steps_per_epoch = batches_per_epoch // grad_accum + (1 if has_remainder and batches_per_epoch else 0)
        total_steps = steps_per_epoch * num_epochs
        final_remainder = 0

    return {
        "batches_per_epoch": batches_per_epoch,
        "used_samples_per_epoch": used_samples_per_epoch,
        "dropped_samples_per_epoch": dropped_samples_per_epoch,
        "remainder_batches": remainder_batches,
        "has_remainder": has_remainder,
        "total_micro_batches": total_micro_batches,
        "steps_per_epoch": steps_per_epoch,
        "total_steps": total_steps,
        "final_remainder": final_remainder,
        "dropped_samples_total": final_remainder * micro_batch_size if use_ds else 0,
    }

def choose_validation_size(
    dataset_size: int,
    validation_ratio: float,
    micro_batch_size: int,
    grad_accum: int,
    num_epochs: int,
    use_ds: bool,
) -> int:
    if not 0.0 <= validation_ratio < 1.0:
        raise ValueError(f"validation_ratio must be in [0, 1), got {validation_ratio}")
    if dataset_size < 2 or validation_ratio <= 0:
        return 0

    desired_val_size = max(1, int(round(dataset_size * validation_ratio)))
    aligned_candidates: list[tuple[int, int]] = []
    fallback_candidates: list[tuple[int, int]] = []
    for val_size in range(1, dataset_size):
        train_size = dataset_size - val_size
        schedule = compute_training_schedule(
            dataset_size=train_size,
            micro_batch_size=micro_batch_size,
            grad_accum=grad_accum,
            num_epochs=num_epochs,
            use_ds=use_ds,
            drop_last=True,
        )
        if int(schedule["batches_per_epoch"]) == 0:
            continue
        if int(schedule["dropped_samples_per_epoch"]) != 0:
            continue
        candidate = (abs(val_size - desired_val_size), val_size)
        if int(schedule["remainder_batches"]) == 0 and int(schedule["final_remainder"]) == 0:
            aligned_candidates.append(candidate)
        else:
            fallback_candidates.append(candidate)

    if aligned_candidates:
        return min(aligned_candidates)[1]
    if fallback_candidates:
        return min(fallback_candidates)[1]
    return min(desired_val_size, dataset_size - 1)

def build_train_validation_subsets(
    dataset: Dataset,
    validation_ratio: float,
    split_seed: int,
    micro_batch_size: int,
    grad_accum: int,
    num_epochs: int,
    use_ds: bool,
) -> tuple[Dataset, Dataset | None, dict[str, float | int | bool]]:
    dataset_size = len(dataset)
    validation_size = choose_validation_size(
        dataset_size=dataset_size,
        validation_ratio=validation_ratio,
        micro_batch_size=micro_batch_size,
        grad_accum=grad_accum,
        num_epochs=num_epochs,
        use_ds=use_ds,
    )
    requested_validation_size = max(1, int(round(dataset_size * validation_ratio))) if validation_ratio > 0 else 0
    metadata: dict[str, float | int | bool] = {
        "dataset_size": dataset_size,
        "requested_validation_size": requested_validation_size,
        "validation_size": validation_size,
        "train_size": dataset_size - validation_size,
        "requested_validation_ratio": validation_ratio,
        "actual_validation_ratio": (validation_size / dataset_size) if dataset_size else 0.0,
        "adjusted": validation_size != requested_validation_size,
    }
    train_schedule = compute_training_schedule(
        dataset_size=dataset_size - validation_size,
        micro_batch_size=micro_batch_size,
        grad_accum=grad_accum,
        num_epochs=num_epochs,
        use_ds=use_ds,
        drop_last=True,
    )
    metadata.update(
        {
            "effective_batch_size": micro_batch_size * grad_accum,
            "train_batches_per_epoch": int(train_schedule["batches_per_epoch"]),
            "train_remainder_batches": int(train_schedule["remainder_batches"]),
            "train_dropped_samples_per_epoch": int(train_schedule["dropped_samples_per_epoch"]),
            "accumulation_aligned": int(train_schedule["remainder_batches"]) == 0
            and int(train_schedule["final_remainder"]) == 0,
        }
    )

    if validation_size == 0:
        return dataset, None, metadata

    generator = torch.Generator().manual_seed(split_seed)
    permutation = torch.randperm(dataset_size, generator=generator).tolist()
    val_indices = sorted(permutation[:validation_size])
    train_indices = sorted(permutation[validation_size:])
    return Subset(dataset, train_indices), Subset(dataset, val_indices), metadata

def load_deepspeed_runtime_config(config_path: str, micro_batch_size: int, grad_accum: int) -> dict:
    with open(config_path, "r", encoding="utf-8") as f:
        ds_cfg = json.load(f)

    if not isinstance(ds_cfg, dict):
        raise ValueError(f"DeepSpeed config in {config_path} must be a JSON object.")

    runtime_cfg = dict(ds_cfg)
    runtime_cfg["train_micro_batch_size_per_gpu"] = micro_batch_size
    runtime_cfg["gradient_accumulation_steps"] = grad_accum
    return runtime_cfg