42 lines
880 B
Python
42 lines
880 B
Python
|
"""
|
||
|
|
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",
|
||
|
|
]
|