dew.training

The trainer and what it is built from.

dew.training knows no modality. It imports nothing from dew.diffusion, dew.inputs or dew.sampling, and wandb only when a WandbTracker logs.

Name	Summary
DEFAULT_RULES
Aux	Everything a loss returns besides its statistics. Documented in dew.objectives.base.
Checkpoints	Holds the checkpoints of one run, in one directory. Documented in dew.
EMASpec	Say which leaves the EMA copy tracks, and how fast it follows them. Documented in dew.objectives.base.
Evaluation	Holds one evaluation event, with bounded hosted previews on rank zero.
Layout	Says how a train state is placed on a mesh.
LocalTracker	Writes synchronous reports in a tracking directory.
MLflowTracker	An MLflow run in experiment, opened on the first value logged into it.
MeshSpec	Says how many devices each sharding axis takes; data parallelism fills the rest.
Metric	Reduce a validation pass to one scalar, on the host. Documented in dew.objectives.base.
Objective	Define what is being learned: the parameters, the loss, what evaluation produces. Documented in dew.objectives.base.
Preempted	Trainer.fit stopped at a preemption notice, at step, and wrote that step’s checkpoint and data position. Documented in dew.training.runtime.
ProfileWindow	Asks for one profiler window per fit: steps steps traced into directory after warmup steps have run, so the trace holds the loop and not the compile.
Quantization	Says how a run quantizes its trunk matmuls, for Qwix’s provider. Documented in dew.training.quantization.
Rollout	Produces a batch on the host, before the compiled step reads it.
Step	What the trainer tells an objective about the current step. Documented in dew.objectives.base.
TensorBoardTracker	A TensorBoard event file in directory, opened on the first value logged into it.
Tracker
Trackers	Fan out without dropping reports; attempt every sink, raise the first failure.
TrainState	Hold everything a run must checkpoint to resume where it stopped. Documented in dew.training.state.
Trainer	Runs an Objective: gradients, sharding, EMA, checkpoints, logging.
WandbTracker	A Weights & Biases run, opened on the first value logged into it.
apply_quantization	Wrap model so its trunk matmuls train in spec’s dtype. Documented in dew.training.quantization.
build_mesh	Build the six-axis device mesh spec describes.
build_optimizer	Build the solver a config describes, with its schedule, parameter groups and clipping. Documented in dew.training.optim.
data_partition	The share of every global batch this process reads on mesh.
ema_update	Update selected EMA leaves in their initialized storage dtypes.
evaluate	Evaluate a finite coordinated prefix without an optimizer or tracker.
everything	Documented in dew.objectives.base.
prepare_process	Raise the fd/core limits, set the env vars, join the JAX process pool. Documented in dew.training.runtime.
run_timestamp	Return process 0’s wall clock as %Y-%m-%d_%H:%M:%S, on every process. Documented in dew.training.runtime.
under	Build a filter that accepts the leaves below prefix. Documented in dew.objectives.base.
write_back	Replace nonparameter collections whole; the optimizer owns params.

DEFAULT_RULES

attribute source

DEFAULT_RULES: LogicalAxisRules

Evaluation

dataclass source

class Evaluation(
    step: int,
    split: str,
    scores: dict[str, float],
    coordinated_batches: int,
    records: int,
    uneven_shards: bool,
    event_key: tuple[int, ...],
    elapsed_seconds: float,
    previews: tuple[Artifact, ...],
)

Holds one evaluation event, with bounded hosted previews on rank zero.

Scores, logical row counts and RNG identity agree across ranks. Scores include the split prefix. elapsed_seconds is rank zero’s wall time through iterator cleanup. Metric accumulators and validation batches are not kept.

scalars: dict[str, float]

Return the metric values, plus the evaluation/* counts if a batch was scored.

Layout

dataclass source

class Layout(
    rules: LogicalAxisRules | Mapping[str, MeshAxes] = DEFAULT_RULES,
    min_shard: int = 2 ** 16,
    tolerance: float = 0.02,
    host: tuple[str, ...] = (),
    host_parameters: tuple[str, ...] = (),
)

Says how a train state is placed on a mesh.

rules map the logical axes the modules declare (dew.nn.sharding) onto the mesh, in precedence order, for the parameters and for the activations a compiled step constrains (the trainer puts them in context). An axis of size 1 shards nothing, so the same table serves every topology.

A parameter the rules place on the data, sequence or stage axis is refused. The first two split the batch, and a parameter placed on either would be gathered on every use. The stage axis holds the layer stack’s pipeline stages, which the decoder places itself from the stored tree.

Below min_shard elements a parameter costs more in collectives than it saves in memory, so it stays replicated. tolerance is the fraction of shardable parameter elements a layout may leave replicated before check refuses it.

host names the fields of HOST_RESIDENT kept in pinned host memory between steps. The step fetches them to the device, updates them as it would have, and writes them back, so what they hold is the same and only where changes. Naming params instead selects canonical CPU ownership of the entire TrainState, including optimizer, EMA and accumulation. The full logical optimizer transaction then runs on a CPU companion of this mesh, and accelerator parameter banks are immutable execution snapshots, never another master. The runtime CPU device count must match the accelerator count on every process before JAX initializes; neither this layout nor the trainer changes it.

host_parameters names the variables an inference placement keeps in pinned host memory, as globs over their logical paths (params/layers_*). Where a parameter sits is independent of how it splits: a selected leaf keeps the spec the rules give it and changes only its memory kind. Only offloaded reads the patterns, because only the stack fetches a layer’s parameters as it reaches it. check refuses a layout that names them to any other placement, rather than place the weights somewhere nothing brings them back from.

axis_rules: LogicalAxisRules

The rules as flax reads them, pairs in precedence order.

Layout.shardings

def shardings(mesh: Mesh, tree: TreeT) -> Placement[TreeT]

Derive a NamedSharding per leaf of tree from the declared axes.

A leaf whose path no module declares takes the largest-divisible-axis heuristic, so a model family can be declared at a time. Flax metadata, if a caller’s own module attached any, is removed here, because the state the trainer materialises against this tree carries plain arrays.

Layout.offloaded

def offloaded(mesh: Mesh, tree: TreeT) -> Placement[TreeT]

Return shardings, with the memory kind each leaf’s path asks for.

The spec is the one the rules give a leaf either way, so a selected parameter is the same shard in another memory space and the collectives its layer issues are the ones a resident run issues.

Every pattern has to name something. One that matches nothing is a typo or a stale path, and a placement that quietly kept those weights on the device would run, with only the memory it did not save to say so. So each pattern is checked on its own, not the table as a whole.

Layout.check

def check(params: Variables, shardings: Placement[Variables], mesh: Mesh) -> None

Reject a layout that left too much of the model replicated, or that asked for host-resident parameters where nothing fetches them.

This is MaxText’s guardrail (base.yml sharding_tolerance) against a mesh whose parameter axes divide none of the model’s dimensions, which the shape heuristic otherwise absorbs by replicating everything.

MaxText measures excess per-chip memory over perfect sharding across every parameter. Here the same ratio is taken over the parameters the threshold policy meant to shard. Anything below min_shard is replicated on purpose, so counting it would fire on models that are merely small.

LocalTracker

class source

class LocalTracker(directory: str | Path, *, plots: bool = False)

Writes synchronous reports in a tracking directory.

Those are JSONL journals, preview files and optional plots. Nonfinite metrics are journaled as the strings NaN, +Inf and -Inf. Plotting runs only on plot() or close and reads the journal rather than retaining history.

LocalTracker.log

def log(scalars: Mapping[str, float], step: int) -> None

LocalTracker.artifact

def artifact(value: Reported, step: int) -> None

LocalTracker.plot

def plot() -> list[Path]

Render journal metrics with matplotlib Agg; never changes the journal.

LocalTracker.close

def close() -> None

MLflowTracker

class source

class MLflowTracker(experiment: str, name: str | None = None, *, uri: str | None = None)

An MLflow run in experiment, opened on the first value logged into it.

Scalars are the run’s metrics and a preview is uploaded as the files the local renderers write. A record becomes a JSON artifact rather than a parameter: MLflow refuses a second value for a parameter, and a run reports several records under one type.

run: tuple[MlflowClient, str]

The client and the run id, creating the experiment and the run once.

MLflowTracker.log

def log(scalars: Mapping[str, float], step: int) -> None

MLflowTracker.artifact

def artifact(value: Reported, step: int) -> None

MLflowTracker.close

def close() -> None

MeshSpec

dataclass source

class MeshSpec(
    fsdp: int = 1,
    expert: int = 1,
    tensor: int = 1,
    sequence: int = 1,
    stage: int = 1,
    microbatches: int | None = None,
    replicas: int = 1,
)

Says how many devices each sharding axis takes; data parallelism fills the rest.

expert: int

Devices the expert dimension of an MoE layer is split over.

tensor: int

Devices the mlp, head and vocabulary widths are split over, beside the fsdp axis they also take; 1 keeps every width on fsdp alone.

sequence: int

Devices the batch’s sequence dimension is split over; 1 keeps whole sequences.

stage: int

Pipeline stages the layer stack is split into, each on its own devices; 1 runs the stack whole on every device.

microbatches: int | None

Microbatches a step feeds through the stages, a multiple of stage; None is one per stage, the smallest schedule. A stage runs one microbatch while the next runs the one before it. So more microbatches shrink the idle time at either end of the step, and make each iteration’s matmuls smaller.

replicas: int

Groups of hosts the data axis spans, for hybrid sharded data parallelism: every other axis, fsdp included, stays inside one group, so the parameter gathers and gradient reduce-scatters run over the fast links and only the gradient all-reduce between replicas crosses the slow one. A group is a whole number of granules: whatever the devices’ slice_index groups (a TPU slice, and on multi-host GPU a host or an NVLink domain), or the process where every device shares one slice. 1 lets jax.make_mesh place every device.

ProfileWindow

dataclass source

class ProfileWindow(directory: str, steps: int, warmup: int = 2)

Asks for one profiler window per fit: steps steps traced into directory after warmup steps have run, so the trace holds the loop and not the compile.

dew.profile is the other way to capture one, a context manager around any code at all; a fit refuses to schedule a window inside one. The window the loop wrote is reported as the ProfileWindow record of dew.telemetry.records.

Rollout

class source

class Rollout(Protocol)

Produces a batch on the host, before the compiled step reads it.

Sampling is effectful and untraceable, so it lives outside jit. The trainer calls the rollout with the state, the prefetched batch and a key folded from the run key and the step, then reshards what comes back with shard_batch. The returned batch must hold arrays in fixed shapes, so the step still compiles once per run.

TensorBoardTracker

class source

class TensorBoardTracker(directory: str | Path)

A TensorBoard event file in directory, opened on the first value logged into it.

Scalars are scalar summaries and a record is the JSON its journal row holds, as a text summary under reporting/<type>. Previews render with _summarize.

writer: EventFileWriter

The event-file writer, opening directory once.

TensorBoardTracker.log

def log(scalars: Mapping[str, float], step: int) -> None

TensorBoardTracker.artifact

def artifact(value: Reported, step: int) -> None

TensorBoardTracker.close

def close() -> None

Tracker

class source

class Tracker(Protocol)

Tracker.log

def log(scalars: Mapping[str, float], step: int) -> None

Tracker.artifact

def artifact(value: Reported, step: int) -> None

Tracker.close

def close() -> None

Trackers

class source

class Trackers(*trackers: Tracker = ())

Fan out without dropping reports; attempt every sink, raise the first failure.

Trackers.log

def log(scalars: Mapping[str, float], step: int) -> None

Trackers.artifact

def artifact(value: Reported, step: int) -> None

Trackers.close

def close() -> None

Trainer

class source

class Trainer(
    objective: Objective[Loss, Effects],
    optimizer: optax.GradientTransformation,
    *,
    key: jax.Array,
    mesh: MeshSpec = MeshSpec(),
    layout: Layout = Layout(),
    accumulation: int = 1,
    dynamic_scale: bool = False,
    checkpoints: Checkpoints | None = None,
    tracker: Tracker | None = None,
    step: Callable[[Objective[Loss, Effects], optax.GradientTransformation], StepFn] | None = None,
    rollout: Rollout | None = None,
    profile: ProfileWindow | None = None,
)

Runs an Objective: gradients, sharding, EMA, checkpoints, logging.

device_mesh: Mesh

Build the mesh MeshSpec describes over this process pool’s devices, on first use.

bank_sites

List the objective’s declared layer stacks, which a host layout streams as banks.

Trainer.from_config

def from_config(
    config: TrainerConfig,
    objective: Objective[ObjectiveLoss, ObjectiveEffects],
    optimizer: optax.GradientTransformation,
    *,
    key: jax.Array,
    checkpoints: Checkpoints | None = None,
    tracker: Tracker | None = None,
    step: Callable[[Objective[ObjectiveLoss, ObjectiveEffects], optax.GradientTransformation], StepFn] | None = None,
    rollout: Rollout | None = None,
) -> Trainer[ObjectiveLoss, ObjectiveEffects]

Build the trainer a TrainerConfig describes.

The mapping from the config’s field names to this constructor’s is written once, here. mesh, layout, accumulation, dynamic_scale and profile are the config fields a trainer holds. key is the run key, which RunConfig.train draws from config.seed.

The rest of the config belongs to the capabilities and to the loop, and reaches them from their own owners. checkpoint_dir and keep build the Checkpoints passed in here, and wandb the tracker. xla_flags, multi_host and compilation_cache_dir are read by prepare_process before JAX opens a backend. batch_ramp wraps the dataset with dew.data.ramped. steps, epochs, log_every, eval_every and checkpoint_every are arguments of fit. step and rollout are not configurable: they are code a caller hands over.

It builds a Trainer, whatever it is called on. The objective’s two parameters are the factory’s own, so a subclass that wants one of itself constructs it.

Trainer.initial_state

def initial_state(
    initializer: Initializer | None = None,
    key: jax.Array | None = None,
) -> TrainState

Build the state a fresh run starts from.

It is pure, so fit traces it once for its shapes and once, sharded, for its values.

Both inputs are the run’s own by default, and place passes them explicitly so that what it compiles takes them as arguments. A held checkpoint then reaches the device as an argument instead of as a constant embedded in the executable. Passing None means resolve the configured input, which is what a no-argument call does. This is the one state implementation, so a subclass overrides it here and every path sees the override.

Trainer.shardings

def shardings(state: TrainState) -> Placement[TrainState]

Place every field of state, each on the axes its own kind takes.

Parameter gradients follow parameters, replay records follow batches, and the layout’s host-resident fields sit in pinned host memory. Under a CPU-owned state the frozen collection is the exception: it sits where the realization reads it (execution.resident) for the whole run.

Trainer.place

def place() -> tuple[TrainState, Placement[TrainState], bytes | None]

Put the state on the mesh, fresh or restored.

Returns it with its shardings and the data position a resume continues from.

Trainer.compile

def compile(state: TrainState, batch: Batch) -> CompiledStep

Compile a transaction over state and one already-produced global batch.

The step consumes the state it is given. The returned state takes over its buffers, so the update runs in place and peak memory holds one copy of the parameters and optimizer state, not two. Keep no reference to a state after stepping it; new = step(old, batch) is the whole contract.

A checkpoint saved before the step is safe. Orbax copies every array to the host before save returns, as long as Checkpoints names no prioritized keys and no concurrent transfer limit. The batch is not donated; the loader owns it.

Trainer.fit

def fit(
    dataset: Dataset,
    *,
    steps: int,
    log_every: int = 100,
    eval_every: int | None = None,
    checkpoint_every: int | None = None,
    metrics: Sequence[Metric] = (),
    preview: bool = False,
) -> TrainState

Train to steps total steps, resuming from the checkpoints’ latest step when the directory holds one.

Every log_every steps the tracker receives the loss, the objective’s metrics and the throughput.

Every eval_every steps, and at the end, the validation split is scored. The objective’s artifacts go to the tracker and to metrics, whose reductions are logged as val/<name>.

Every checkpoint_every steps, and at the end, the state and the data position are written. Every checkpoints.local_every steps they are written to the local directory as well.

A preemption notice (a scheduler’s SIGTERM; PreemptionNotice) stops the run at the next step every process agrees on: that step’s state and data position are written, the final validation is skipped, and fit raises Preempted, which ends the program with 143 unless caught. Run again, fit resumes there.

Previews are generated only when preview=True and a tracker receives them; scalar reporting never triggers preview work.

WandbTracker

class source

class WandbTracker(
    project: str,
    name: str | None = None,
    *,
    entity: str | None = None,
    config: Mapping[str, object] | None = None,
    offline: bool = False,
    id: str | None = None,
)

A Weights & Biases run, opened on the first value logged into it.

WandbTracker.log

def log(scalars: Mapping[str, float], step: int) -> None

WandbTracker.artifact

def artifact(value: Reported, step: int) -> None

WandbTracker.close

def close() -> None

build_mesh

function source

def build_mesh(spec: MeshSpec = MeshSpec(), devices: list | None = None) -> Mesh

Build the six-axis device mesh spec describes.

Parameters shard over ‘fsdp’, ‘expert’ and ‘tensor’, batches over the first four with their sequence dimension on ‘sequence’, and the layer stack over ‘stage’.

An MoE layer’s expert dimension is the one dimension no dense model has, and splitting it is what expert parallelism is. So it gets its own axis and leaves ‘fsdp’ to the model’s widths. The tensor axis is where the mlp, head and vocabulary widths split beside fsdp, as DEFAULT_RULES places them. The sequence axis is where long sequences split. The stage axis is where a decoder’s layers split into pipeline stages, each stage on its own devices.

Sizes of 1 degenerate to plain data parallelism, so the same code path serves every topology without a flag. Axes are Auto so GSPMD infers the collectives.

devices names the devices and their order on one slice: the mesh takes them row-major over MESH_AXES, so the last axes hold neighbouring entries. Unset is every device, in the order jax.make_mesh gives the platform’s topology.

spec.replicas above 1 builds the mesh the way MaxText builds a multislice one, through mesh_utils.create_hybrid_device_mesh: the data axis takes replicas groups of granules as its outer factor, and each group lays out the rest of the mesh over its own devices. A group of more than one granule splits fsdp across them, the one axis whose traffic, a gather and a reduce-scatter per layer, tolerates it.

data_partition

function source

def data_partition(mesh: Mesh) -> DataPartition

The share of every global batch this process reads on mesh.

A batch’s rows split over the batch axes and no others (BATCH_SPEC): the sequence axis splits positions, the tensor axis widths, and the stage axis holds a pipeline’s stages. So the processes whose devices hold the same row shards need the same rows, and the processes fall into groups by the rows they hold. Each group reads one share, numbered by the first row shard it holds, and every process of the group reads it (readers); reader is this process’s place among them, in process order.

Groups whose rows overlap without being the same rows, which a device order built by hand can produce, leave no share each could read whole, so they are refused.

ema_update

function source

def ema_update(
    ema: Variables,
    params: Variables,
    decay: jax.typing.ArrayLike,
) -> Variables

Update selected EMA leaves in their initialized storage dtypes.

Arithmetic uses at least fp32 and preserves explicit fp64. Unit decay selects the original leaf, including nonfinite frozen-reference values.

evaluate

function source

def evaluate(
    objective: Objective[Loss, Effects],
    variables: Variables,
    batches: Reader | None,
    *,
    key: jax.Array,
    metrics: Sequence[Metric] = (),
    step: int | jax.Array = 0,
    averaged: Variables | None = None,
    preview: bool = False,
    mesh: Mesh | None = None,
    split: str = 'val',
    schedule_step: int | jax.Array | None = None,
) -> Evaluation

Evaluate a finite coordinated prefix without an optimizer or tracker.

batches opens a fresh iterator over this process’s share of the split (data_partition of the mesh), owned and closed by this call; Dataset.val can be passed directly. Every rank calls evaluate with the same objective, metrics and numerical settings. Only root’s preview flag controls the once-per-event display. No consumers means no iterator or objective work; preview alone consumes at most the first coordinated batch.

variables is the complete Flax variables tree. averaged, when supplied, is the complete overlay seen as Step.ema, not an optimizer state. Passing state.averaged as variables evaluates those weights directly. step tags the event RNG and report; schedule_step defaults to step and preserves an objective’s accepted-work schedule when training attempts were rejected.

Scalars are broadcast to every rank. Previews remain hosted on root and are bounded by one objective preview, independent of validation length. Reporting is the caller’s job; pool callers must agree reporting failures before entering their next collective. In-flight device failures still require distributed runtime termination.

write_back

function source

def write_back(params: Variables, variables: Variables | None) -> Variables

Replace nonparameter collections whole; the optimizer owns params.