Skip to content

Configuration

About

The d9d.loop.config package defines the job configuration as Pydantic models. TrainerConfig configures a training job, and InferenceConfig configures an inference job. Pydantic validates the whole configuration when you load it.

Usage

Load the configuration from JSON and pass it to the configurator as parameters:

1
2
3
4
5
from pathlib import Path

from d9d.loop.config import TrainerConfig

config = TrainerConfig.model_validate_json(Path("trainer.json").read_text(encoding="utf-8"))

Periodic actions, such as checkpointing.period_steps, take a StepActionPeriod. It is a step interval (e.g. 100), "last_step" or "disable".

You can leave out the fields and sections that have defaults. An unknown key fails the validation, so a misspelled key stops the job before it starts. The trainer section of the Quickstart config holds only the required fields.

API Reference

Main Config

d9d.loop.config.TrainerConfig

Bases: BaseModel

Configuration for a complete training job.

Attributes:

Name Type Description
run RunConfig

Meta-information about the run (name, ID, tags).
schedule JobScheduleConfig

Job duration settings.
data_prefetch DataPrefetchConfig

Settings for copying data to the device ahead of time.
logging JobLoggerConfig

Experiment tracking settings.
pipelining PipeliningConfig

Pipeline parallelism schedule and settings.
model_stage_factory ModelStageFactoryConfig

Model initialization and additional checkpointing logic.
determinism DeterminismConfig

Random seed settings.
gc GarbageCollectionConfig

Garbage collection settings.
checkpointing CheckpointingConfig

Checkpoint saving settings.
gradient_clipping GradientClippingConfig

Gradient clipping settings.
profiling ProfilingConfig | None

Profiler settings. If None, profiling is disabled.
gradient_manager GradientManagerConfig

Gradient synchronization settings.
timeout TimeoutConfig

Distributed timeout settings.

d9d.loop.config.InferenceConfig

Bases: BaseModel

Configuration for a complete inference or evaluation job.

Attributes:

Name Type Description
schedule JobScheduleConfig

Job duration settings.
data_prefetch DataPrefetchConfig

Settings for copying data to the device ahead of time.
model_stage_factory ModelStageFactoryConfig

Model initialization logic.
determinism DeterminismConfig

Random seed settings.
gc GarbageCollectionConfig

Garbage collection settings.
checkpointing CheckpointingConfig

Checkpointing settings.
profiling ProfilingConfig | None

Profiler settings. If None, profiling is disabled.
timeout TimeoutConfig

Distributed timeout settings.

Diagnostics and Reproducibility

d9d.tracker.RunConfig

Bases: BaseModel

Configuration for a tracked run.

Attributes:

Name Type Description
name str

The display name of the experiment.
description str | None

An optional description of the experiment.
hparams dict[str, Any]

The hyperparameters to log at the start of the run.

d9d.loop.config.JobLoggerConfig

Bases: BaseModel

Configuration for experiment tracking and logging.

Attributes:

Name Type Description
period_steps StepActionPeriod

How often the metrics are logged. The metrics are accumulated over the steps between two logging steps. The loss and the total gradient norm are logged every step.
tracker AnyTrackerConfig

The experiment tracker backend configuration, e.g. Aim.
max_pending_steps int

How many steps the logged values can lag behind the current step, waiting for the GPU to compute them. When they lag more, logging waits for the GPU, which also limits how far the host runs ahead of the GPU.

d9d.loop.config.ProfilingConfig

Bases: BaseModel

Configuration for the PyTorch Profiler.

Attributes:

Name Type Description
traces_dir Path

The directory for trace files.
period_steps int

The total length of a profiling cycle (wait + warmup + active), in steps.
warmup_steps int

The number of profiler warmup steps before recording.
active_steps int

The number of steps to record.
record_shapes bool

Whether to record the input shapes of operators.
with_stack bool

Whether to record the Python call stacks of operators. They make up most of a trace.

d9d.loop.config.DeterminismConfig

Bases: BaseModel

Configuration for reproducibility and random number generation.

Attributes:

Name Type Description
base_seed int

The base seed for the random number generators (Python, NumPy, PyTorch) on all ranks.

Experiment Trackers

d9d.tracker.AnyTrackerConfig = Annotated[AimConfig | LogTrackerConfig | MLflowConfig | NullTrackerConfig | TensorBoardConfig | WandbConfig, Field(discriminator='provider')] module-attribute

Union of all tracker configurations, discriminated by the provider field.

d9d.tracker.provider.log.LogTrackerConfig

Bases: BaseModel

Configuration for the log tracker, which writes scalars to the Python logger of d9d.

Attributes:

Name Type Description
provider Literal['log']

Discriminator field. Always "log".

d9d.tracker.provider.null.NullTrackerConfig

Bases: BaseModel

Configuration for the null tracker, which logs nothing.

Attributes:

Name Type Description
provider Literal['null']

Discriminator field. Always "null".

d9d.tracker.provider.wandb.config.WandbConfig

Bases: BaseModel

Configuration for the W&B tracker backend.

Attributes:

Name Type Description
provider Literal['wandb']

Discriminator field. Always "wandb".
project str

The W&B project.
entity str | None

The team or user that owns the project. None selects the default entity of the API key.
group str | None

The group of the run, e.g. to show the runs of one experiment together.
tags list[str]

The tags of the run.
mode Literal['online', 'offline']

"online" sends the values to the server during training. "offline" writes them to local files, which wandb sync uploads later.
dir str | None

The directory for the local files of W&B. None selects the working directory.
capture_terminal_logs bool

Whether to capture stdout and stderr.

d9d.tracker.provider.tensorboard.config.TensorBoardConfig

Bases: BaseModel

Configuration for the TensorBoard tracker backend.

Attributes:

Name Type Description
provider Literal['tensorboard']

Discriminator field. Always "tensorboard".
log_dir str

The directory of the event files of the run.

d9d.tracker.provider.mlflow.config.MLflowConfig

Bases: BaseModel

Configuration for the MLflow tracker backend.

Attributes:

Name Type Description
provider Literal['mlflow']

Discriminator field. Always "mlflow".
experiment str

The name of the MLflow experiment. The tracker creates it if it does not exist.
tracking_uri str | None

The URI of the tracking server or store, e.g. "http://mlflow:5000". None selects the MLFLOW_TRACKING_URI environment variable or the default of MLflow.

d9d.tracker.provider.aim.config.AimConfig

Bases: BaseModel

Configuration for the Aim tracker backend.

Attributes:

Name Type Description
provider Literal['aim']

Discriminator field. Always "aim".
repo str

The path or URL of the Aim repository.
log_system_params bool

Whether to log system resource usage (CPU, GPU, memory).
capture_terminal_logs bool

Whether to capture stdout and stderr.
system_tracking_interval int

The interval of system monitoring in seconds.

Scheduling

d9d.loop.config.JobScheduleConfig

Bases: BaseModel

Configuration for the job's duration.

Attributes:

Name Type Description
total_steps int | None

The total number of steps to run. If None, the length of the microbatch pack stream sets the duration.

Data Prefetching

d9d.loop.config.DataPrefetchConfig

Bases: BaseModel

Configuration for copying microbatch packs to the device ahead of the steps that consume them.

Attributes:

Name Type Description
prefetch_factor int

The number of packs copied to the device ahead of the current step. The copies run on a separate CUDA stream and overlap with compute. Each prefetched pack occupies device memory. 0 disables prefetching, so every pack is copied on the current stream when its step starts.

Checkpointing

d9d.loop.config.CheckpointingConfig

Bases: BaseModel

Configuration for saving model snapshots.

Attributes:

Name Type Description
save_dir Path

The root directory for checkpoints.
period_steps StepActionPeriod

How often to save a checkpoint.
num_to_keep int | None

The maximum number of recent checkpoints to keep. If None, all checkpoints are kept.

Model Initialization

d9d.loop.config.ModelStageFactoryConfig

Bases: BaseModel

Configuration for initializing model weights.

Attributes:

Name Type Description
source_checkpoint Path | None

The path to a checkpoint to load into the model before the job starts. If None, the model is initialized randomly.
checkpoint_only_trainable_parameters bool

If True, checkpoints store only parameters with requires_grad=True. Useful for PEFT, e.g. LoRA.

Optimization

d9d.loop.config.GradientClippingConfig

Bases: BaseModel

Configuration for gradient norm clipping.

Attributes:

Name Type Description
max_norm float | None

The maximum gradient norm. If None, gradients are not clipped.

d9d.loop.config.GradientManagerConfig

Bases: BaseModel

Configuration for gradient synchronization.

Attributes:

Name Type Description
grad_dtype str | None

The name of the torch dtype for gradients, e.g. "float32". If None, gradients use the parameter dtype.
bucket_size_mb int

The maximum size of a gradient communication bucket, in MiB.

Infrastructure

d9d.loop.config.PipeliningConfig

Bases: BaseModel

Configuration for pipeline parallelism orchestration.

Attributes:

Name Type Description
schedule AnyPipelineScheduleConfig

The pipeline schedule configuration.

d9d.loop.config.GarbageCollectionConfig

Bases: BaseModel

Configuration for manual Python garbage collection control.

Attributes:

Name Type Description
period_steps StepActionPeriod

How often to run the Python garbage collector.

d9d.loop.config.TimeoutConfig

Bases: BaseModel

Configuration for distributed process group timeouts.

Attributes:

Name Type Description
init_timeout int

The timeout in seconds for the job setup and the first step.
step_timeout int

The timeout in seconds for communication in later steps.

Types

d9d.loop.config.StepActionPeriod = int | StepActionSpecial module-attribute

Period of a periodic action.

An int is the period in steps. A StepActionSpecial runs the action only at the last step or disables it.

d9d.loop.config.StepActionSpecial

Bases: StrEnum

Special flag values for configuring periodic actions.

Attributes:

Name Type Description
last_step

The action runs once, at the last step of the run.
disable

The action never runs.