pathcosmos's picture
Upload folder using huggingface_hub (#14)
17758b3
raw
history blame contribute delete
880 Bytes
"""
train — LLM pretraining package.
Public API:
TrainConfig : Dataclass of training hyper-parameters.
Trainer : Core training loop with gradient accumulation, AMP, and logging.
Utility functions (re-exported from train.utils):
get_cosine_schedule_with_warmup
save_checkpoint
load_checkpoint
get_grad_norm
setup_ddp
cleanup_ddp
is_main_process
"""
from train.trainer import TrainConfig, Trainer
from train.utils import (
cleanup_ddp,
get_cosine_schedule_with_warmup,
get_grad_norm,
is_main_process,
load_checkpoint,
save_checkpoint,
setup_ddp,
)
__all__ = [
# Core classes
"TrainConfig",
"Trainer",
# Utility functions
"get_cosine_schedule_with_warmup",
"save_checkpoint",
"load_checkpoint",
"get_grad_norm",
"setup_ddp",
"cleanup_ddp",
"is_main_process",
]