File size: 992 Bytes
7c5e40e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """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",
]
|