Training Utilities

Helper functions for training orchestration and result collection.

Overview

The training utilities module provides reusable components for the training workflow, including hardware detection, configuration parsing, checkpoint management, and summary artifact generation.

Usage

from pathlib import Path

from naics_embedder.utils.training import (
    detect_hardware,
    parse_config_overrides,
    resolve_checkpoint,
)

hardware = detect_hardware(log_info=True, cuda_precision='bf16-mixed')
overrides, invalid = parse_config_overrides([
    'training.learning_rate=1e-4',
    'data_loader.queries_per_step=128',
])
checkpoint_info = resolve_checkpoint('last', Path('checkpoints'), 'reference')

TrainingResult.best_score is outcome validation MRR, not an in-sample loss. The four guards refuse used fresh directories, cross-directory resumes, different settings/seed, and a run that early stopping ended. A spent epoch budget is a no-op. Exact resume uses last, with its monitor and epoch-summary files. refuse_other_constructor_settings separately checks saved LoRA rank/alpha/dropout and active MoE expert/routing/balancing controls against the current config, refusing missing required values. Inactive MoE controls are ignored. run_settings records 21 settings including effective accelerator/precision, accumulation, clipping and epoch budget.

Data Classes

HardwareInfo

Container for detected hardware configuration.

CheckpointInfo

Resolved checkpoint path and metadata.

TrainingResult

Structured result from a completed training run.

API Reference

Training orchestration utilities for NAICS Embedder.

This module provides helper functions to simplify training setup by extracting common operations into reusable components. These utilities handle hardware detection, configuration parsing, checkpoint management, and trainer creation.

Functions:

Name Description
detect_hardware

Detect the accelerator and the precision the trainer runs at.
get_gpu_memory_info

Query current GPU memory usage.
parse_config_overrides

Parse and validate command-line config overrides.
effective_precision

The precision a run trains at on an accelerator (P31).
run_settings

A run's free settings, epoch budget, fusion, accumulation and clipping (P21).
resolve_checkpoint

Resolve checkpoint path from user input.
read_checkpoint

A Lightning checkpoint's contents, on the CPU.
refuse_a_fresh_start_into_a_used_directory

P19's guard on a fresh run's directory.
refuse_a_resume_from_another_directory

P19's guard on a resume's directory.
refuse_a_resume_under_other_settings

P19's guard on a resume's settings and seed.
refuse_a_resume_of_a_stopped_run

P19's guard on a resume of a run early stopping ended.
outcome_checkpoint

The ModelCheckpoint the monitor's MRR drives (P17).
outcome_early_stopping

The EarlyStopping the monitor's MRR drives (P17).
create_trainer

Create a configured PyTorch Lightning Trainer.
TrainingResult

Structured result from a training run.

HardwareInfo dataclass

Hardware configuration detected for training.

Attributes:

Name Type Description
accelerator str

The accelerator type (cuda, mps, cpu).
precision str

The precision the trainer runs at: the configured CUDA precision on CUDA, 32-true elsewhere.
num_devices int

Number of available devices.
gpu_memory Optional[Dict[str, float]]

Optional GPU memory information dictionary.

CheckpointInfo dataclass

Resolved checkpoint information.

Attributes:

Name Type Description
path Optional[str]

Resolved filesystem path to the checkpoint, or None.
is_same_stage bool

Whether checkpoint is from the same experiment stage.
exists bool

Whether the checkpoint file exists.

TrainingResult dataclass

Structured result from a training run.

Provides a clean interface for accessing training outputs, metrics, and paths for downstream processing or testing.

Attributes:

Name Type Description
best_checkpoint_path Optional[str]

Path to the kept checkpoint: the earliest epoch with the highest val/outcome_mrr.
last_checkpoint_path Optional[str]

Path to the last model checkpoint.
config_path Optional[str]

Path to the saved configuration file.
best_score Optional[float]

The kept epoch's val/outcome_mrr, the monitor's validation MRR.
stopped_epoch int

Epoch at which training stopped (early stopping or max).
early_stopped bool

Whether early stopping was triggered.
metrics Dict[str, Any]

Dictionary of final metrics.

detect_hardware(log_info=False, *, cuda_precision='bf16-mixed')

Detect available hardware and the precision the trainer runs at.

Queries the system for CUDA, MPS, or CPU availability and returns the accelerator and precision settings for PyTorch Lightning.

Parameters:

Name Type Description Default
log_info bool

If True, log detailed hardware information.

 False
cuda_precision str

The precision on CUDA; train passes training.trainer.precision (spec 4.2). Off CUDA the trainer runs at 32-true.

 'bf16-mixed'

Returns:

Type Description
HardwareInfo

HardwareInfo with detected accelerator, precision, device count,
HardwareInfo

and optional GPU memory information.
Example

hw = detect_hardware(log_info=True, cuda_precision='bf16-mixed') print(f'Training on {hw.accelerator} with {hw.precision} precision')

get_gpu_memory_info()

Query current GPU memory usage.

Returns memory statistics including total, reserved, allocated, and free memory in gigabytes, along with utilization percentage.

Returns:

Type Description
Optional[Dict[str, float]]

Dictionary with memory statistics, or None if CUDA unavailable.
Example

info = get_gpu_memory_info() if info: ... print(f'GPU Memory: {info["free_gb"]:.1f} GB free')

parse_config_overrides(overrides)

Parse and validate command-line configuration overrides.

Converts a list of key=value strings into a dictionary suitable for use with Config.override(). Invalid overrides are collected and returned for warning messages.

Parameters:

Name Type Description Default
overrides Optional[List[str]]

List of override strings like training.learning_rate=1e-4.

 required

Returns:

Type Description
Tuple[Dict[str, Any], List[str]]

Tuple of (valid_overrides_dict, list_of_invalid_override_strings).
Example

overrides = ['training.learning_rate=1e-4', 'invalid', 'batch_size=32'] valid, invalid = parse_config_overrides(overrides) print(valid) # {'training.learning_rate': 0.0001, 'batch_size': 32} print(invalid) # ['invalid']

effective_precision(cfg, accelerator)

The precision a run trains at on accelerator (P4, P31).

The configured training.trainer.precision on CUDA, and 32-true everywhere else: the rule detect_hardware resolves for the trainer, so the precision a run records is the one it trains at, and tools sweep names a run's precision from its accelerator alone.

Parameters:

Name Type Description Default
cfg Config

The run's configuration.

 required
accelerator str

cuda, mps or cpu.

 required

Returns:

Type Description
str

A Lightning precision.

run_settings(cfg, *, accelerator, precision)

A run's free settings (R7), its epoch budget, its fusion and the trainer's accumulation and clipping, with where it trains (P21).

train records them in the model's hyperparameters, so an exact resume under other settings is refused (P19), and tools sweep writes them as the arm's settings. The two trainer settings change a run mid-way too: the warmup counts optimizer steps, so another accumulate_grad_batches would stretch it, and another gradient_clip_val would change every later step. Every value is a JSON type, the logit-scale range a list, so an arm record's settings equal a run's after the record's round trip.

Parameters:

Name Type Description Default
cfg Config

The run's configuration.

 required
accelerator str

The accelerator the run trains on.

 required
precision str

The precision it trains at (effective_precision).

 required

Returns:

Type Description
Dict[str, Any]

The settings, in a fixed order.

resolve_checkpoint(ckpt_path, checkpoint_dir, experiment_name)

Resolve a checkpoint path from user input.

Handles three cases: 1. None: No checkpoint specified, start fresh. 2. "last" or "last.ckpt": Auto-detect last checkpoint in experiment dir. 3. Explicit path: Validate and resolve the provided path.

Parameters:

Name Type Description Default
ckpt_path Optional[str]

User-provided checkpoint path or keyword.

 required
checkpoint_dir Path

Base directory for checkpoints.

 required
experiment_name str

Name of the current experiment.

 required

Returns:

Type Description
CheckpointInfo

CheckpointInfo with resolved path and metadata.
Example

info = resolve_checkpoint('last', Path('checkpoints'), '01_text') if info.exists: ... print(f'Resuming from {info.path}')

read_checkpoint(path)

A Lightning checkpoint's contents, on the CPU.

Lightning checkpoints carry pickled hyperparameters and loop state, so they load with weights_only=False: they are trusted artifacts of this project's own training runs.

refuse_a_fresh_start_into_a_used_directory(checkpoint_dir)

Refuse a fresh run into a checkpoint directory that exists and is not empty (P19).

The fresh run's checkpoints would sit beside the other run's as -v1 siblings, which ModelCheckpoint never deletes, and --ckpt-path last would resolve to the other run's last.ckpt; the monitor would also find its records file taken.

Raises:

Type Description
ValueError

If the directory is not empty, or is not a directory.

refuse_a_resume_from_another_directory(checkpoint, checkpoint_dir, *, resolved_dirpath=None)

Refuse an exact resume from a checkpoint that ModelCheckpoint saved in another directory (P19).

Lightning 2.5.5 restores ModelCheckpoint's best-k state only into the directory it was saved from: elsewhere it keeps no earlier epoch, so a lower later epoch would be saved as the best, and last.ckpt would become last-v1.ckpt. The directories compare as ModelCheckpoint stores them, as real paths, and the state is the one this run's callback reads.

The keyword-only resolved_dirpath is supplied only by a transport identity probe. It names an absolute normalized instance directory and is compared literally, without resolving it on this host; it is never an arbitrary unverified directory override.

Raises:

Type Description
ValueError

If the checkpoint holds no state of this run's ModelCheckpoint, or holds one saved in another directory.

refuse_a_resume_under_other_settings(checkpoint, settings, *, seed)

Refuse an exact resume whose checkpoint was saved under other run settings or another seed (P19).

train builds the model and the data from the config, and the checkpoint contract's encoder record holds only the fusion, the dimension and the backbone, so a changed setting (the radius bound, a learning rate, the epoch budget) would change the run mid-way. The seed draws every epoch's permutations and names every monitor read, so it must not change either.

Parameters:

Name Type Description Default
checkpoint Mapping[str, Any]

The checkpoint's contents (read_checkpoint).

 required
settings Mapping[str, Any]

This run's run_settings.

 required
seed int

This run's seed.

 required

Raises:

Type Description
ValueError

If the checkpoint records no run settings, or a setting or the seed differs, naming each difference with its saved and current values.

constructor_settings(cfg)

Behavior and shape controls saved in model hyperparameters, supplementary to P21's 21 keys.

LoRA always participates. Expert settings affect construction and the objective only under MoE fusion (R11); changing inactive expert settings cannot change a masked-mean/attention run.

refuse_other_constructor_settings(checkpoint, cfg)

Refuse missing or changed saved constructor controls before exact resume or seed export.

Lightning tensor restoration does not restore controls such as LoRA scaling/dropout or MoE routing. Current checkpoints save these constructor arguments separately from run_settings; checking them preserves the literal 21-key arm identity without rewriting checkpoint files.

refuse_a_resume_of_a_stopped_run(checkpoint, patience)

Refuse an exact resume of a run that early stopping ended (P19, spec 4.4).

Lightning 2.5.5 restores EarlyStopping's state, stopped_epoch with it, but not the trainer's stop, so a resumed run would train the epochs after its stop: it would append more monitor records and could keep another epoch, changing the run's selection after the fact. Early stopping checks at an epoch's end before ModelCheckpoint saves, so the last.ckpt of a stopped run records the stop, and its kept checkpoint and monitor_reads.jsonl are final. A run that spent its epoch budget without an early stop is not refused, deliberately: its resume trains nothing. The state read is the one this run's EarlyStopping (outcome_early_stopping) restores.

Parameters:

Name Type Description Default
checkpoint Mapping[str, Any]

The checkpoint's contents (read_checkpoint).

 required
patience int

This run's training.early_stopping_patience.

 required

Raises:

Type Description
ValueError

If the checkpoint records no state of this run's EarlyStopping, or one whose stopped_epoch is above 0, naming that epoch.

outcome_checkpoint(checkpoint_dir)

The ModelCheckpoint the monitor's MRR drives (P17, spec 4.4).

It keeps one epoch, epoch=<NNN>.ckpt: the earliest with the highest val/outcome_mrr, since a later epoch replaces it only by beating it. It also keeps last.ckpt for exact resume. Both save at each training epoch's end, after the module's hook has logged the MRR. The exact-resume guard reads the saved state of this same callback.

Parameters:

Name Type Description Default
checkpoint_dir Path

The run's checkpoint directory.

 required

outcome_early_stopping(patience)

The EarlyStopping the monitor's MRR drives (P17, spec 4.4).

It ends the run after patience epochs without a higher val/outcome_mrr, checked at each training epoch's end, before ModelCheckpoint saves: Lightning runs checkpoint callbacks last. An equal MRR is no gain (min_delta 0), so a tie counts against the patience. The stopped-run guard reads the saved state of this same callback.

Parameters:

Name Type Description Default
patience int

The run's training.early_stopping_patience.

 required

create_trainer(cfg, hardware, checkpoint_dir, callbacks=None, tb_logger=None)

Create the text stage's Trainer (P17), which train fits.

The monitor's val/outcome_mrr drives the kept checkpoint (outcome_checkpoint) and early stopping (outcome_early_stopping), both checked at each training epoch's end: there is no validation loop, since validation is the outcome monitor (spec 4.4). TrainDatasetEpochCallback hands each epoch to the step dataset. The trainer runs on one device at the precision detect_hardware resolved.

Parameters:

Name Type Description Default
cfg Config

Training configuration.

 required
hardware HardwareInfo

Detected hardware information.

 required
checkpoint_dir Path

The run's checkpoint directory.

 required
callbacks Optional[List[Callback]]

Optional additional callbacks to include.

 None
tb_logger Optional[TensorBoardLogger]

Optional TensorBoard logger (created if not provided).

 None

Returns:

Type Description
Tuple[Trainer, ModelCheckpoint, EarlyStopping]

Tuple of (Trainer, ModelCheckpoint callback, EarlyStopping callback).
Example

hw = detect_hardware() trainer, ckpt_cb, es_cb = create_trainer(cfg, hw, Path('checkpoints')) trainer.fit(model, datamodule)

save_training_summary(result, config, hardware, output_dir, format='both')

Save training summary artifacts for downstream evaluation and documentation.

Creates YAML and/or JSON summary files containing training results, configuration snapshot, and hardware information. These artifacts can be used for evaluation scripts, MkDocs documentation, or CI/CD pipelines.

Parameters:

Name Type Description Default
result TrainingResult

TrainingResult from the completed training run.

 required
config Config

Configuration used for training.

 required
hardware HardwareInfo

Hardware information used during training.

 required
output_dir Path

Directory to save summary files.

 required
format str

Output format - 'yaml', 'json', or 'both'.

 'both'

Returns:

Type Description
Dict[str, str]

Dictionary mapping format to output file path.
Example

paths = save_training_summary(result, cfg, hw, Path('outputs')) print(f'Summary saved to: {paths}')