source/src/strata/training/__init__.py

"""STRATA training: packed-LM data and a pure-DDP training loop.

This subpackage deliberately contains only plain-PyTorch + DistributedDataParallel
logic. It has no DeepSpeed, ZeRO, or CPU/parameter-offload behaviour, per the
project constraints. Distributed execution is opt-in and driven entirely by the
standard ``torchrun`` environment variables.
"""

from __future__ import annotations

from strata.training.distributed import (
    DistInfo,
    cleanup_distributed,
    is_main_process,
    seed_everything,
    setup_distributed,
)
from strata.training.lm_data import (
    PackedLMDataset,
    PackedCorpusMeta,
    build_packed_corpus,
    choose_token_dtype,
)
from strata.training.loop import TrainConfig, train_lm

__all__ = [
    "DistInfo",
    "setup_distributed",
    "cleanup_distributed",
    "is_main_process",
    "seed_everything",
    "PackedLMDataset",
    "PackedCorpusMeta",
    "build_packed_corpus",
    "choose_token_dtype",
    "TrainConfig",
    "train_lm",
]