File size: 2,568 Bytes
69928d6 | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """``softasp-train`` — launch a training run (supports torchrun for DDP).
Usage::
# Single-GPU
uv run --locked softasp-train instantiation=sudoku model.pretrained_from=backbone.pt
# Multi-GPU via torchrun
uv run --locked torchrun --nproc_per_node=4 -m softasp.cli.train \\
instantiation=my_distributed_task paths=local
"""
from __future__ import annotations
import logging
from pathlib import Path
import hydra
from omegaconf import DictConfig, OmegaConf
from softasp_core.training.schedule import CORRECTED_PROTOCOL, LEGACY_PROTOCOL
from softasp_core.training.seed import seed_everything
from softasp_core.training.trainer import Trainer
from softasp_core.utils.logging import setup_logging
from softasp.configuration import config_directory
from softasp.instantiations.checkpoint import read_checkpoint, validate_checkpoint
from softasp.instantiations.runtime import provider
logger = logging.getLogger(__name__)
_CONFIGS = config_directory()
@hydra.main(version_base="1.3", config_path=_CONFIGS, config_name="config")
def main(cfg: DictConfig) -> None:
run_training(cfg)
def run_training(cfg: DictConfig) -> Path:
"""Train from an explicit configuration and return the exact run directory."""
setup_logging(level=cfg.logging.get("level", "INFO"))
logger.info("=== softasp-train ===")
# Resolve all interpolations
OmegaConf.resolve(cfg)
protocol = cfg.trainer.get("protocol", LEGACY_PROTOCOL)
if protocol not in (LEGACY_PROTOCOL, CORRECTED_PROTOCOL):
raise ValueError(f"Unknown training protocol: {protocol}")
if protocol == CORRECTED_PROTOCOL:
# Includes release-model construction and all subsequent stochastic setup.
seed_everything(cfg.trainer.seed, deterministic=cfg.trainer.deterministic)
OmegaConf.update(cfg, "initialization_seeded_before_model", True, force_add=True)
model_provider = provider(cfg, "model")
# Preparation precedes data construction to preserve initialization order.
prepared = model_provider.prepare(cfg)
train_ds, val_ds = provider(cfg, "dataset").build(cfg)
model = prepared if prepared is not None else model_provider.build(cfg)
evaluator = provider(cfg, "metrics").build(cfg)
# Train
trainer = Trainer(
cfg=cfg, model=model, evaluator=evaluator,
checkpoint_validator=lambda path: validate_checkpoint(cfg, read_checkpoint(path)),
)
run_dir = trainer.fit(train_ds, val_ds)
logger.info("Training complete.")
return run_dir
if __name__ == "__main__":
main()
|