ML Training¶

Training is configuration-driven. The entry point owns data loading, model and optimizer construction, validation, logging, distributed setup, and checkpoint persistence.

Run a single process with:

python scripts/train.py --config configs/ml_dry_run_cpu.json

Distributed training is enabled by the environment that torchrun creates and currently requires CUDA and NCCL:

OMP_NUM_THREADS=1 uv run torchrun --standalone --nproc-per-node=2 scripts/train.py \
  --config configs/ml_baseline.json

File-backed runs create the following structure. W&B runs use the W&B run ID for the nested checkpoint directory; other loggers use the local run name.

<run-dir>/
  config.json
  logs/
    metrics.jsonl
    payloads.jsonl
  checkpoints/
    <run-id-or-name>/
      latest.pt
      best_<metric>.pt
      final.pt

Files are conditional on the selected logger and checkpoint flags. Best checkpoints minimize their configured metric. run.init_from initializes compatible model weights only; it does not restore optimizer, scheduler, AMP scaler, or training cursor state.

latest and best_<metric> are updated only when validation runs. A missing or misspelled monitor emits a warning and does not produce a best checkpoint. Common monitor names include val/tf/loss_ce, val/tf/rmse, val/rollout/loss_ce, and val/rollout/rmse.

Warning: If run.name resolves to an existing directory under run.output_dir, the directory is deleted before training starts.

train_from_config ¶

train_from_config(path: str | Path) -> Path | None

Execute a complete training run from an experiment JSON file.

The entry point validates the configuration, builds one shared index and the train/validation loaders, constructs the model and optimizer, trains to the explicit or epoch-derived step limit, validates, logs, and saves configured checkpoints. Distributed execution is activated by torchrun environment variables and requires CUDA/NCCL.

Parameters:

path (str | Path) –

Experiment JSON path.

Returns:

Path | None –

The created run directory, or None for output-free stdout runs.

Raises:

OSError –

If configuration, data, output, or checkpoint files cannot be accessed.
TypeError –

If the JSON schema has incorrect value types.
ValueError –

If configuration, manifests, runtime, or distributed settings are invalid.

Examples:

Run the bundled CPU smoke experiment:

from batgrad.ml.train import train_from_config

run_dir = train_from_config("configs/ml_dry_run_cpu.json")

The equivalent repository command is:

python scripts/train.py --config configs/ml_dry_run_cpu.json

Warning

When run.name resolves to an existing directory under run.output_dir, that directory is removed before training. The run.init_from option initializes compatible model weights only; it does not resume optimizer, scheduler, AMP scaler, or training cursor state.

Note

AMP is CUDA-only. Cross-batch Mamba state carry is unavailable under DDP.