Source code for kempnerforge.config.training

"""Training configuration."""

from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum
from typing import TYPE_CHECKING, Literal

if TYPE_CHECKING:
    import torch

FP32_CROSS_ENTROPY_LOSSES = ("cross_entropy", "chunked_cross_entropy")

_FP16_UNSUPPORTED = (
    "train.mixed_precision='fp16' is not supported: it would train in pure fp16 "
    "with no loss scaling, so small gradients underflow to zero. Use 'bf16', or "
    "'fp32' on GPUs without bf16."
)




class ActivationCheckpointing(StrEnum):
    none = "none"
    full = "full"
    selective = "selective"





@dataclass
class TrainConfig:
    """Training hyperparameters.

    ``"fp16"`` stays in the ``mixed_precision`` ``Literal`` only so the config loader
    passes it on to ``__post_init__``, which rejects it with the reason; without it the
    loader would fail first with a generic Literal error.
    """

    batch_size: int = 8
    seq_len: int = 2048
    max_steps: int = 100000
    grad_accum_steps: int = 1
    grad_clip_norm: float = 1.0
    seed: int = 42
    data_seed: int | None = None
    compile_model: bool = True
    mixed_precision: Literal["bf16", "fp16", "fp32", "fp8"] = "bf16"
    fp32_master_weights: bool = False
    tf32: bool = False
    activation_checkpointing: ActivationCheckpointing = ActivationCheckpointing.none
    loss_fn: str = "cross_entropy"
    z_loss_weight: float = 0.0
    ce_chunk_size: int = 0
    fp32_cross_entropy: bool = False
    shutdown_timeout_sec: float = 600.0
    nccl_health_check_interval: int = 0

    def __post_init__(self) -> None:
        if self.batch_size <= 0:
            raise ValueError("batch_size must be positive")
        if self.seq_len <= 0:
            raise ValueError("seq_len must be positive")
        if self.max_steps <= 0:
            raise ValueError("max_steps must be positive")
        if self.grad_accum_steps <= 0:
            raise ValueError("grad_accum_steps must be positive")
        if self.grad_clip_norm <= 0:
            raise ValueError("grad_clip_norm must be positive")
        if self.mixed_precision == "fp16":
            raise ValueError(_FP16_UNSUPPORTED)
        if self.fp32_master_weights and self.mixed_precision == "fp32":
            raise ValueError(
                "train.fp32_master_weights has no effect with mixed_precision='fp32': the "
                "weights are already fp32. Use it with 'bf16' or 'fp8', or turn it off."
            )
        if self.fp32_cross_entropy and self.loss_fn not in FP32_CROSS_ENTROPY_LOSSES:
            raise ValueError(
                f"train.fp32_cross_entropy applies to {', '.join(FP32_CROSS_ENTROPY_LOSSES)}, "
                f"not loss_fn={self.loss_fn!r}"
            )

    @property
    def param_dtype(self) -> torch.dtype:
        """Resolve mixed_precision to the compute dtype, also the master weight dtype by default.

        FP8 resolves to bf16 -- FP8 is a compute mode, not a storage dtype. The storage
        dtype is :attr:`master_dtype`, which is fp32 with ``fp32_master_weights``.
        A config switched to ``"fp16"`` after construction skips ``__post_init__``,
        so ``"fp16"`` raises the same ``ValueError`` here.
        """
        import torch

        if self.mixed_precision == "fp16":
            raise ValueError(_FP16_UNSUPPORTED)
        _DTYPE_MAP = {
            "bf16": torch.bfloat16,
            "fp32": torch.float32,
            "fp8": torch.bfloat16,
        }
        return _DTYPE_MAP[self.mixed_precision]

    @property
    def master_dtype(self) -> torch.dtype:
        """Dtype the parameters, their gradients and the optimizer state are kept in."""
        import torch

        return torch.float32 if self.fp32_master_weights else self.param_dtype

    @property
    def is_fp8(self) -> bool:
        """Whether FP8 mixed precision is enabled."""
        return self.mixed_precision == "fp8"

    @property
    def effective_data_seed(self) -> int:
        """Seed for data shuffling / batch composition.

        Falls back to ``seed`` when ``data_seed`` is unset, so existing configs
        reproduce their current trajectory. Kept independent from ``seed``
        (parameter init) so stability studies can vary batch order while holding
        initialization fixed. Must stay identical across data-parallel ranks so
        the global shuffle is consistent before rank partitioning.
        """
        return self.data_seed if self.data_seed is not None else self.seed