dew.config

Typed configuration for a training run.

One dataclass tree describes a run: what to build, what to feed it, how to optimize it, and how the trainer runs. Recipes parse it with tyro, build the objective and the data it names, and hand both to RunConfig.train, so to_dict() is a full record of a run and from_dict() puts it back together.

Model kwargs are an opaque JSON dict; the registry knows which architecture takes which fields. A dataset is the registered spec itself, which tyro turns into a subcommand (data:token-windows --data.path ...).

The resolved config is the run’s spec. A recipe writes it to run.json next to the checkpoints with save, and load reads it back into the same class, so inference rebuilds a run from what training was built from. A field the class does not have, or one the file lacks, raises.

Name	Summary
`JsonDict`	The fields a model is built from, written as a single JSON string on the command line.
`DataSpec`
`ScheduleSpec`
`ModelConfig`	Holds the architecture name and the fields `models.build` receives.
`OptimConfig`	Holds the optimizer, learning-rate schedule and gradient clipping.
`Wandb`	Says where a run reports to.
`TrainerConfig`	Holds the run length, checkpointing, sharding and run tracking.
`RunConfig`	Describes a whole run.

JsonDict

attribute source

JsonDict

The fields a model is built from, written as a single JSON string on the command line. The registry knows which architecture takes which field and narrows each one where it builds it, so the values are read there.

DataSpec

attribute source

ScheduleSpec

attribute source

ModelConfig

dataclass source

class ModelConfig(
    architecture: str = 'simple_dit',
    config: JsonDict = dict(),
    dtype: registry.DtypeName = 'bfloat16',
    param_dtype: registry.DtypeName | None = None,
    matmul_precision: Literal['default', 'high', 'highest'] | None = None,
    attention_impl: AttentionImpl = 'auto',
)

Holds the architecture name and the fields models.build receives.

dtype: registry.DtypeName: Compute dtype; parameter storage is independent.
param_dtype: registry.DtypeName | None: Parameter storage, where the model declares the field. Unset stores float32, which is the model’s own default.
matmul_precision: Literal['default', 'high', 'highest'] | None: What every matmul of the model asks XLA for, where the model declares a precision field: default is the backend’s fastest algorithm, high and highest trade throughput for mantissa bits (on Ampere and later, tf32 and fp32 against bf16x3). Unset leaves the model’s own.
attention_impl: AttentionImpl: Attention kernel; ‘auto’ is cudnn on a GPU for the shapes cudnn supports and xla for the rest, xla on any other backend.

ModelConfig.fields

def fields() -> Mapping[str, object]

Return the model’s fields with the run’s precision settings in them.

ModelConfig.precision_settings

def precision_settings() -> frozenset[str]

Return the names fields() writes that config did not carry: the run’s precision settings, as this architecture takes them. A resolved record leaves them out, since this value writes them again every time it builds.

ModelConfig.from_dict

def from_dict(values: Mapping[str, object]) -> Self

Read back the record RunConfig.to_dict writes for this field.

ModelConfig.build

def build()

OptimConfig

dataclass source

class OptimConfig(
    optimizer: Literal['adam', 'adamw', 'lamb', 'muon', 'muonclip'] = 'adamw',
    optimizer_opts: JsonDict = dict(),
    learning_rate: float = 0.00027,
    schedule: ScheduleSpec | None = None,
    weight_decay: float | None = None,
    param_groups: Annotated[tuple[ParamGroup, ...], json_list_argument(ParamGroup)] = (),
    clip_grads: float = 0.0,
    state_dtype: Literal['float32', 'bfloat16'] = 'float32',
)

Holds the optimizer, learning-rate schedule and gradient clipping.

learning_rate: float: The constant rate, when no schedule is named.
schedule: ScheduleSpec | None: The learning-rate schedule, one typed record per kind (dew.training.optim): cosine, power (lm-engine’s power law with an optional linear tail) or linear; each holds only its own fields.
param_groups: Annotated[tuple[ParamGroup, ...], json_list_argument(ParamGroup)]: Per-group learning-rate multipliers and weight decay, first match wins; empty moves every parameter alike. dew.training.optim.mup_param_groups is lm-engine’s muP split.
state_dtype: Literal['float32', 'bfloat16']: Adam’s moments in memory. bfloat16 stores both stochastically rounded (dew.training.optim.bf16_moments), for adam and adamw only: half the optimizer state and less of the update’s memory traffic.

Wandb

dataclass source

class Wandb(project: str, entity: str | None = None, offline: bool = False)

Says where a run reports to. Setting it turns tracking on; the entity and the offline switch mean nothing without a project.

TrainerConfig

dataclass source

class TrainerConfig(
    name: str | None = None,
    checkpoint_dir: str = './checkpoints',
    keep: int = 2,
    batch_size: int = 32,
    seed: int = 0,
    steps: int | None = None,
    epochs: int | None = None,
    log_every: int = 100,
    eval_every: int | Literal['epoch'] | None = 'epoch',
    checkpoint_every: int | Literal['epoch'] | None = 'epoch',
    accumulation: int = 1,
    batch_ramp: Ramp | None = None,
    dynamic_scale: bool = False,
    mesh: MeshSpec = MeshSpec(),
    layout: Layout = Layout(),
    profile: ProfileWindow | None = None,
    compilation_cache_dir: str | None = default_compilation_cache_dir(),
    wandb: Wandb | None = None,
    multi_host: bool | None = None,
    xla_flags: str | None = None,
    quantization: Quantization | None = None,
)

Holds the run length, checkpointing, sharding and run tracking.

keep: int: Latest checkpoints kept, besides the best one.
batch_size: int: Global batch, over every process.
seed: int: Seed of the run key: parameter init and every per-step draw.
epochs: int | None: Run length as passes over the data; steps names it directly instead.
eval_every: int | Literal['epoch'] | None: Steps between validation passes: a number of steps, “epoch” for one pass over the data, None to never validate. “epoch” over a stream that reports no record count raises a ValueError, since it has no pass.
checkpoint_every: int | Literal['epoch'] | None: Steps between checkpoints, the same three answers. None is what a stream whose iterator cannot report a read position trains with; the trainer refuses any other answer for one.
accumulation: int: Micro-batches per optimizer update.
batch_ramp: Ramp | None: Grow batch_size over the run’s first records instead of starting there: the global batch the run starts at, what a stage adds and the records the whole ramp spans. Unset trains at batch_size throughout. One optimizer update a step either way, and the compiled step is traced once per stage.
profile: ProfileWindow | None: One profiler window: the steps to trace, the warmup before it and the directory it is written to. Unset traces nothing.
compilation_cache_dir: str | None: Persisted XLA cache, so a restart skips recompiling the step. None compiles from scratch every run.
wandb: Wandb | None: Optional W&B sink in addition to the local tracking journal.
multi_host: bool | None: Join the JAX process pool. None asks and continues alone only when no cluster is configured; True requires the pool; False never asks.
xla_flags: str | None: Extra XLA_FLAGS for this run, appended to the environment by prepare_process before JAX opens a backend. Library users set XLA_FLAGS themselves; see docs/performance.md for what was measured.
quantization: Quantization | None: Quantized-training spec, wrapped around the module the objective trains before the run initialises it; unset trains in the compute dtype. dew.training.quantization says what the wrap does and what it keeps.

TrainerConfig.total_steps

def total_steps(dataset: Dataset) -> int

Return the run’s length in steps, from steps or from epochs over data.

TrainerConfig.eval_interval

def eval_interval(dataset: Dataset) -> int | None

Return the steps between validation passes over data, or None for never.

TrainerConfig.checkpoint_interval

def checkpoint_interval(dataset: Dataset) -> int | None

Return the steps between checkpoints over data, or None for never.

RunConfig

dataclass source

class RunConfig(
    model: ModelConfig = ModelConfig(),
    data: DataSpec = (lambda: datasets['oxford_flowers102']())(),
    optim: OptimConfig = OptimConfig(),
    trainer: TrainerConfig = TrainerConfig(),
    objective: str | None = None,
    lora: LoRA | None = None,
)

Describes a whole run. Recipes add their objective’s knobs by subclassing this.

lora: LoRA | None: The low-rank adapter the run trains instead of the whole model. The targets are the module paths under params a delta sits on; train adapts the objective’s module and freezes every other leaf.

RunConfig.to_dict

def to_dict() -> dict[str, JSON]

Return a JSON-safe record of the run.

A registered member is written as its name and its fields.

RunConfig.from_dict

def from_dict(values: Mapping[str, object]) -> Self

Read back what to_dict wrote, for subclasses too. An unknown or a missing field raises.

RunConfig.save

def save(directory: str) -> str

Write this config as run.json in directory and return the path.

The path goes through epath, the same filesystem layer orbax writes the checkpoints with, so a gs:// run directory takes the record too.

RunConfig.load

def load(directory: str) -> Self

Read the config a run in directory was built from, as this class.

RunConfig.train

def train(
    objective: Objective[Loss, Effects],
    dataset: Dataset,
    *,
    name: str,
    metrics: Sequence[Metric] = (),
    summary: Mapping[str, object] | None = None,
) -> TrainState

Train objective on data as this run says; every recipe calls this once it has built both.

The run lives under name in trainer.checkpoint_dir, and process zero writes the record there before anything trains. A trainer.wandb opens a tracker under the same name, with the record, summary (the recipe’s own view of the run) and the step count as its config, and the checkpoint the run ends on is published to the registry under the name with slashes and spaces replaced, since an artifact name allows neither. Local tracking journals live in a separate tracking directory.

dataset is the one a recipe loaded at trainer.batch_size, and a dataset at any other batch is refused here: the record, the ramp and the run’s reported throughput all name the configured number, while every step reads the dataset’s, so the two disagreeing is a run that trains at a batch it does not report.

A trainer.quantization wraps the module objective trains before anything initialises it, so the quantized forward is what the run learns through, and a lora adapts the same module the same way, so the run traces the adapted forward and moves only its factors.