flashcart.training

Train and evaluate potentials using configurable loss functions, optimizers, and learning-rate schedules. TrainingTask runs training, validation, and checkpointing, while EvaluationTask evaluates a model over a loader and combines batch statistics into the requested metrics.

Loss objects specify both the target property and the predictions needed to compute it. The task uses these requirements to request energies, forces, or stress and to retain the derivative graph when needed for training. Optimizer and scheduler functions read the configuration to create the update rules.

See Quickstart for the training workflow and Multi-GPU training for distributed execution.

flashcart.training.tasks

Run training, validation, checkpointing, and model evaluation.

The best checkpoint is selected by the validation metric. The log checkpoint stores the state needed to resume training. Despite the argument name early_stopping_loss, training continues to max_epochs even if the validation metric stops improving.

flashcart.training.tasks.preserve_grad_state(model)

Freeze all parameters inside the block, restoring their flags on exit.

Used around evaluation so force/stress autograd tracks only the inputs, not the parameters.

Parameters:

model (torch.nn.Module) – Model whose parameters are frozen.

flashcart.training.tasks.predict_kwargs(losses)

Derive the predict() flags from what the given losses require.

Parameters:

losses (list[LossFunction]) – Losses that will consume the predictions.

Returns:

compute_forces and compute_stress flags required by the

losses.

Return type:

dict[str, bool]

flashcart.training.tasks.unwrap_model(model)

Strip Fabric/DDP wrappers down to the underlying model.

Training forwards must go through the wrapper (gradient synchronization). Use this only for evaluation, checkpointing, and attribute access.

Parameters:

model (torch.nn.Module) – Possibly wrapped model.

Returns:

Underlying module after removing wrappers exposing module.

Return type:

torch.nn.Module

flashcart.training.tasks.make_predict_padder(model, predict_compile, pad_kwargs)

Build a batch padder for compiled prediction.

Padding reduces tensor-shape changes while batches fit within the allocated capacities. Eager prediction runs without padding.

Parameters:

• model (torch.nn.Module) – Model providing r_max for fake edges.

• predict_compile (bool) – Whether predictions run compiled.

• pad_kwargs (dict) – PadAtomicData rounding multiples.

Returns:

Padder for compiled prediction, or None for eager

prediction.

Return type:

PadAtomicData | None

flashcart.training.tasks.seed_padder_from_loader(padder, loader)

Initialize padding capacities from the batch sampler’s configured limits.

These initial capacities reduce growth as batch sizes vary. Batches that exceed them may still increase the capacities and trigger recompilation. When sampler limits are unavailable, capacities are determined from observed batches.

Parameters:

• padder (PadAtomicData, optional) – Padder to initialize. Nothing is changed when None.

• loader (DataLoader) – Loader whose batch sampler carries the maxima.

flashcart.training.tasks.padded_predict(model, batch, padder, predict_kwargs, compile_kwargs, create_graph)

Evaluate a batch with optional padding and remove padded output entries.

Parameters:

• model (torch.nn.Module) – Model exposing predict.

• batch (Any) – Batched atomic graph.

• padder (PadAtomicData, optional) – None predicts unpadded.

• predict_kwargs (dict) – Property flags (see predict_kwargs).

• compile_kwargs (dict) – use_compile/mode kwargs, empty for eager.

• create_graph (bool) – Keep derivative graphs when training on energy derivatives.

Returns:

Predictions for the real atoms and structures.

Return type:

dict[str, torch.Tensor]

class flashcart.training.tasks.TrainingTask(model, train_loss, early_stopping_loss, eval_losses, optimizer, scheduler=None, aux_optimizers=None, max_epochs=1000, validate_every=1, save_every=100, grad_clip_norm=None, log_significant_digits=4, log_weight_norms=False, ema_decay=None, predict_compile=False, predict_compile_mode='reduce-overhead', predict_compile_fullgraph=True, predict_compile_dynamic=False, predict_pad_atom_multiple=128, predict_pad_edge_multiple=128, predict_pad_graph_multiple=8)

Bases: object

Train a potential with validation and checkpointing.

Optionally maintain exponential moving averages of the model parameters.

The log checkpoint stores the current weights and training state for resumption. The best checkpoint stores the weights that minimize the selected validation metric, using averaged weights when enabled. Training continues to max_epochs. The validation metric selects checkpoints without applying a patience-based stopping rule.

Parameters:

• model (torch.nn.Module) – Model to train.

• train_loss (LossFunction) – Loss backpropagated per step.

• early_stopping_loss (LossFunction) – Validation metric selecting the best checkpoint.

• eval_losses (list[LossFunction]) – Metrics logged for train and valid.

• optimizer (torch.optim.Optimizer) – Main optimizer.

• scheduler (LRScheduler, optional) – Learning-rate scheduler stepped after each epoch. Its first parameter group’s learning rate is copied to the auxiliary optimizers. Default: None.

• aux_optimizers (list, optional) – Additional optimizers stepped after the main optimizer. Default: None.

• max_epochs (int, optional) – Total epochs. Default: 1000.

• validate_every (int, optional) – Epochs between validations. Default: 1.

• save_every (int, optional) – Epochs between rolling checkpoints. Default: 100.

• grad_clip_norm (float, optional) – Maximum global gradient norm, applied before all optimizer steps. This limits gradients rather than directly bounding the resulting parameter updates. Default: None.

• log_significant_digits (int, optional) – Round CSV values to this many significant digits. None logs full precision. Default: 4.

• log_weight_norms (bool, optional) – Include Muon weight and momentum diagnostics in metrics.csv. Default: False.

• ema_decay (float, optional) – Decay factor for the exponential moving average used during validation and best-checkpoint selection. Default: None, disabling averaging.

• predict_compile (bool, optional) – Run predictions through torch.compile (enables batch padding). Default: False.

• predict_compile_mode (str, optional) – Compile mode. Default: “reduce-overhead”.

• predict_compile_fullgraph (bool, optional) – Compile without graph breaks. Default: True.

• predict_compile_dynamic (bool, optional) – Dynamic-shape compile. Default: False.

• predict_pad_atom_multiple (int, optional) – Atom-budget rounding of the padder. Default: 128.

• predict_pad_edge_multiple (int, optional) – Edge-budget rounding. Default: 128.

• predict_pad_graph_multiple (int, optional) – Graph-budget rounding. Default: 8.

run(train_loader, valid_loader, output_path, fabric=None)

Train through max_epochs, restoring a saved training state when present.

If output_path/log/training_state.pt exists, the model weights, optimizer states, scheduler, and averaging state are restored before training continues.

Parameters:

• train_loader (DataLoader) – Training batches.

• valid_loader (DataLoader) – Validation batches.

• output_path (str | Path) – Run directory (checkpoints, metrics.csv).

• fabric (Fabric, optional) – Launched Fabric. A single-process instance is created when omitted.

class flashcart.training.tasks.EvaluationTask(model, eval_losses, predict_compile=False, predict_compile_mode='reduce-overhead', predict_compile_fullgraph=True, predict_compile_dynamic=False, predict_pad_atom_multiple=128, predict_pad_edge_multiple=128, predict_pad_graph_multiple=8)

Bases: object

Evaluate a model over a loader and aggregate the requested metrics.

Parameters:

• model (torch.nn.Module) – Model to evaluate.

• eval_losses (list[LossFunction]) – Metrics to compute.

• predict_compile (bool, optional) – Compiled predictions (enables batch padding). Default: False.

• predict_compile_mode (str, optional) – Compile mode. Default: “reduce-overhead”.

• predict_compile_fullgraph (bool, optional) – Compile without graph breaks. Default: True.

• predict_compile_dynamic (bool, optional) – Dynamic-shape compile. Default: False.

• predict_pad_atom_multiple (int, optional) – Atom-budget rounding of the padder. Default: 128.

• predict_pad_edge_multiple (int, optional) – Edge-budget rounding. Default: 128.

• predict_pad_graph_multiple (int, optional) – Graph-budget rounding. Default: 8.

run(data_loader, fabric=None)

Evaluate over the loader and return {metric name: value}.

Parameters:

• data_loader (DataLoader) – Batches to evaluate.

• fabric (Fabric, optional) – Fabric instance used for batch placement and distributed metric reduction. Default: None, retaining any instance provided by an earlier call.

Returns:

Metric values keyed by loss name.

Return type:

dict[str, float]

flashcart.training.loss

Define property loss functions and combine their batch statistics.

flashcart.training.loss.merge_compute_keys(losses)

Combine the prediction requirements of several losses.

Parameters:

losses (list[LossFunction]) – Losses to merge.

Returns:

Combined force, stress, and virial requirements.

Return type:

dict[str, bool]

class flashcart.training.loss.LossFunction

Bases: object

Interface for training losses and evaluation metrics.

reduce_batch computes the statistics needed from one batch, and reduce_overall combines them using the total atom and structure counts. Means and square roots are applied after aggregation so that their normalization does not depend on batch size.

property name

Column name in logs and metrics.csv, e.g. "forces_mae".

property compute_keys

Prediction requirements used to select energy derivatives.

The mapping contains compute_forces, compute_stress, and compute_virials. Training tasks combine stress and virial requests into the model’s compute_stress flag.

Returns:

Required force, stress, and virial calculations.

Return type:

dict[str, bool]

property is_extensive

True for sum-reduced losses, which scale with the batch size.

__call__(results, batch)

Evaluate the loss for one batch.

Parameters:

• results (dict[str, torch.Tensor]) – Model predictions.

• batch (Data) – Batch containing reference targets.

Returns:

Scalar loss or metric.

Return type:

torch.Tensor

training_loss(results, batch, world_size=1)

Evaluate the per-rank loss for distributed gradient computation.

Extensive (sum-reduced) losses are multiplied by the world size: DDP averages gradients across ranks, and the rescale turns that average back into the global-batch sum.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

• world_size (int, optional) – Number of DDP ranks. Default: 1.

Returns:

Scalar loss, scaled for distributed gradient averaging when

the reduction is extensive.

Return type:

torch.Tensor

reduce_batch(results, batch)

Compute batch statistics to be combined by reduce_overall.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

Returns:

Batch statistics keyed by metric name.

Return type:

dict[str, torch.Tensor]

reduce_overall(losses, n_atoms, n_structures)

Compute the metric from accumulated batch statistics.

Parameters:

• losses (list[dict]) – Batch statistics from reduce_batch. Distributed accumulators combine maxima with max and other statistics with sum before calling this method.

• n_atoms (int) – Total number of atoms seen.

• n_structures (int) – Total number of structures seen.

Returns:

Scalar metric over the accumulated batches.

Return type:

torch.Tensor

class flashcart.training.loss.PropertyLossFunction(key, reduce='mae', normalization='per_atom', delta=None)

Bases: LossFunction

Evaluate a loss function or metric for one predicted property.

The reduction modes sae, sse, huber_sum, and norm_sum are extensive sums. mae, mse, huber_mean, and norm_mean are per-component or per-vector means. rmse is the root mean square error, and maxe is the maximum absolute error. For vector-valued residuals, norm_sum and norm_mean reduce sqrt(sum(residual**2) + 1e-12) along the final dimension. For scalar energies, they reduce absolute residuals.

Parameters:

• key (str) – Property: “energy”, “forces”, “stress”, or “virials”.

• reduce (str, optional) – Reduction mode (see above). Default: “mae”.

• normalization (str, optional) – Residual normalization, “per_atom”, “per_sqrt_atom”, or “none”. Only “none” is valid for the per-atom/already-normalized properties (forces, stress). Default: “per_atom”.

• delta (float, optional) – Positive Huber transition point in the units of the normalized residual. Required for Huber reductions and rejected for other reductions. Default: None.

property name

Column name in logs and metrics.csv, e.g. "forces_mae".

get_residuals(results, batch)

Normalized prediction-minus-target residuals.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

Returns:

Normalized prediction-minus-target residuals with the

predicted property’s shape.

Return type:

torch.Tensor

property compute_keys

Prediction requirements used to select energy derivatives.

The mapping contains compute_forces, compute_stress, and compute_virials. Training tasks combine stress and virial requests into the model’s compute_stress flag.

Returns:

Required force, stress, and virial calculations.

Return type:

dict[str, bool]

property is_extensive

True for sum-reduced losses, which scale with the batch size.

__call__(results, batch)

Evaluate the loss for one batch.

Parameters:

• results (dict[str, torch.Tensor]) – Model predictions.

• batch (Data) – Batch containing reference targets.

Returns:

Scalar loss or metric.

Return type:

torch.Tensor

reduce_batch(results, batch)

Compute the batch statistic for this reduction, keyed by name.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

Returns:

A scalar sum or maximum keyed by name.

Return type:

dict[str, torch.Tensor]

get_denominator(n_atoms, n_structures)

Count the observations used to normalize mean reductions.

Parameters:

• n_atoms (int) – Total atoms over the accumulated batches.

• n_structures (int) – Total structures over the accumulated batches.

Returns:

Number of scalar components or vectors used to normalize the selected

reduction.

Return type:

int

reduce_overall(losses, n_atoms, n_structures)

Combine accumulated batch statistics into the final metric.

Parameters:

• losses (list[dict]) – Batch sums or maxima from reduce_batch.

• n_atoms (int) – Total atoms represented by the statistics.

• n_structures (int) – Total structures represented by the statistics.

Returns:

Scalar metric over the accumulated batches.

Return type:

torch.Tensor

class flashcart.training.loss.WeightedSumLoss(losses, weights=None)

Bases: LossFunction

Combine property loss functions with specified weights.

Component losses must have distinct names when batch statistics are accumulated. The name identifies the property, normalization, and reduction, but does not include the Huber threshold.

Parameters:

• losses (list[PropertyLossFunction]) – Component losses.

• weights (list[float], optional) – One weight per component. Default: all 1.0.

property name

Column name in logs and metrics.csv, e.g. "forces_mae".

property compute_keys

Prediction requirements used to select energy derivatives.

The mapping contains compute_forces, compute_stress, and compute_virials. Training tasks combine stress and virial requests into the model’s compute_stress flag.

Returns:

Required force, stress, and virial calculations.

Return type:

dict[str, bool]

property is_extensive

True for sum-reduced losses, which scale with the batch size.

__call__(results, batch)

Evaluate the loss for one batch.

Parameters:

• results (dict[str, torch.Tensor]) – Model predictions.

• batch (Data) – Batch containing reference targets.

Returns:

Scalar loss or metric.

Return type:

torch.Tensor

training_loss(results, batch, world_size=1)

Evaluate the per-rank loss for distributed gradient computation.

Extensive (sum-reduced) losses are multiplied by the world size: DDP averages gradients across ranks, and the rescale turns that average back into the global-batch sum.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

• world_size (int, optional) – Number of DDP ranks. Default: 1.

Returns:

Scalar loss, scaled for distributed gradient averaging when

the reduction is extensive.

Return type:

torch.Tensor

reduce_batch(results, batch)

Compute batch statistics to be combined by reduce_overall.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

Returns:

Batch statistics keyed by metric name.

Return type:

dict[str, torch.Tensor]

reduce_overall(losses, n_atoms, n_structures)

Compute the metric from accumulated batch statistics.

Parameters:

• losses (list[dict]) – Batch statistics from reduce_batch. Distributed accumulators combine maxima with max and other statistics with sum before calling this method.

• n_atoms (int) – Total number of atoms seen.

• n_structures (int) – Total number of structures seen.

Returns:

Scalar metric over the accumulated batches.

Return type:

torch.Tensor

class flashcart.training.loss.LossAccumulator(loss)

Bases: object

Accumulate batch statistics and counts for a loss or metric.

update stores detached statistics on CPU together with atom and structure counts. compute combines these statistics and, when a distributed Fabric instance is provided, communicates them across ranks.

Parameters:

loss (LossFunction) – Loss defining how batch statistics are combined.

update(results, batch)

Store one batch’s detached statistics and update the observation counts.

Parameters:

• results (dict) – Model predictions.

• batch (Data) – Batch with reference targets.

compute(fabric=None)

Compute the metric from accumulated batch statistics.

Parameters:

fabric (Fabric, optional) – Fabric instance for distributed reduction. Default: None, computing locally.

Returns:

Scalar metric computed from the accumulated statistics.

Return type:

torch.Tensor

reset()

Clear the accumulated partials and counts.

flashcart.training.loss.loss_from_config(spec)

Build a property loss or weighted sum from a configuration mapping.

Parameters:

spec (dict) – A property specification with property and optional reduce, normalization, and delta entries, or a weighted_sum list of property specifications with optional weight entries. Defaults are reduce="mae", normalization="none", and weight=1.0.

Returns:

Configured property loss or weighted sum.

Return type:

LossFunction

flashcart.training.optimizers

Construct optimizers from model parameters and configuration settings.

flashcart.training.optimizers.make_optimizer(model, optimizer='muon', lr=0.01, weight_decay=0.001, muon_alignment_damping=False, muon_damping_tau=2.0, muon_damping_ema_decay=0.9, muon_damping_min_scale=0.1, muon_ns_steps=None, muon_ns_polish_steps=None, muon_ns_rms_scale=None)

Construct optimizers from the model’s parameter-selection hooks.

With optimizer="muon", trainable matrix parameters not excluded by non_muon_parameters are assigned to HybridMuon. Remaining parameters are assigned to AdamW, with weight decay disabled for those returned by non_decayable_parameters. With optimizer="adamw", AdamW handles all parameter groups.

The returned AdamW optimizer is the main optimizer. TrainingTask applies its learning-rate schedule to any auxiliary Muon optimizer.

Parameters:

• model (nn.Module) – Model providing the parameter-group hooks.

• optimizer (str, optional) – “muon” or “adamw”. Default: “muon”.

• lr (float, optional) – Learning rate for all groups. Default: 1e-2.

• weight_decay (float, optional) – Decay for the decayable groups. Default: 1e-3.

• muon_alignment_damping (bool, optional) – Magma alignment damping in HybridMuon. Default: False.

• muon_damping_tau (float, optional) – Damping sigmoid temperature. Default: 2.0.

• muon_damping_ema_decay (float, optional) – Damping score EMA decay. Default: 0.9.

• muon_damping_min_scale (float, optional) – Damping floor. Default: 0.1.

• muon_ns_steps (int, optional) – Number of Newton-Schulz iterations. Default: None, using NS_STEPS (5).

• muon_ns_polish_steps (int, optional) – Number of polishing iterations after the Newton-Schulz steps. Default: None, using NS_POLISH_STEPS (0).

• muon_ns_rms_scale (float, optional) – Rescaling factor for the normalized update. Default: None, using NS_RMS_SCALE (0.2).

Returns:

Main AdamW optimizer and a list of auxiliary Muon optimizers, or None for

the auxiliary entry when optimizer="adamw".

Return type:

tuple

flashcart.training.optimizers.optimizer_from_config(model, cfg)

Construct optimizers from a run configuration.

This factory reads the supplied mapping without loading packaged defaults. Omitted settings use the defaults of make_optimizer.

Parameters:

• model (nn.Module) – Model to optimize.

• cfg (dict) – Run config (keys as in make_optimizer).

Returns:

Main AdamW optimizer and a list of auxiliary Muon optimizers, or None for

the auxiliary entry when optimizer="adamw".

Return type:

tuple

flashcart.training.schedulers

Construct learning-rate schedules with warmup and decay.

flashcart.training.schedulers.make_scheduler(optimizer, start_factor=0.1, end_factor=0.001, warmup_epochs=0, max_epochs=100, decay_epochs=0, warmup_shape='linear', decay_shape='cosine', **kwargs)

Construct an epoch-based warmup, constant, and decay schedule.

The learning-rate factor rises from start_factor to one during warmup. It remains constant until the final decay_epochs, when it decreases to end_factor. A zero warmup or decay length disables that phase.

When decay is enabled, end_factor is reached at scheduler epoch max_epochs. If the scheduler is stepped after each training epoch, this occurs after the final optimizer step.

Parameters:

• optimizer (torch.optim.Optimizer) – Optimizer to schedule.

• start_factor (float, optional) – Initial LR fraction, in (0, 1]. Default: 0.1.

• end_factor (float, optional) – Final LR fraction, in [0, 1]. Default: 1e-3.

• warmup_epochs (int, optional) – Warmup length. Default: 0.

• max_epochs (int, optional) – Total epochs. Default: 100.

• decay_epochs (int, optional) – Decay length (counted from the end). Default: 0.

• warmup_shape (str, optional) – “linear”, “cosine”, or “sqrt”. Default: “linear”.

• decay_shape (str, optional) – “linear”, “cosine”, or “sqrt”. Default: “cosine”.

• **kwargs – Additional keyword arguments, accepted but unused.

Returns:

Epoch-based learning-rate schedule.

Return type:

torch.optim.lr_scheduler.LambdaLR

flashcart.training.schedulers.scheduler_from_config(optimizer, cfg)

Construct a learning-rate schedule from a run configuration.

This factory reads the supplied mapping without loading packaged defaults. Omitted warmup and decay lengths are zero.

Parameters:

• optimizer (torch.optim.Optimizer) – Optimizer to schedule.

• cfg (dict) – Run config (keys as in make_scheduler).

Returns:

Configured schedule, or None when both warmup and decay have

zero length.

Return type:

LRScheduler | None

flashcart.training.muon

Hybrid Muon optimizer with batched Newton-Schulz orthogonalisation.

Parts of this module (the Newton-Schulz iteration and the learning-rate adjustment) are adapted from torch.optim._muon in PyTorch, Copyright (c) 2016- Facebook, Inc. and other contributors, BSD 3-Clause License (see NOTICE).

class flashcart.training.muon.HybridMuon(params, alignment_damping=False, damping_tau=2.0, damping_ema_decay=0.9, damping_min_scale=0.1, ns_polish_steps=0, ns_polish_coefficients=(2.0, -1.5, 0.5), ns_rms_scale=0.2, **muon_kwargs)

Bases: Muon

Batched Muon with optional Magma alignment damping.

Follows the Muon optimizer used to train DPA4 (arXiv:2606.02419).

Parameters with the same shape, device, and dtype are grouped for batched Newton-Schulz updates. Optional alignment damping rescales each update using a moving average of the alignment between the gradient and momentum.

Parameters:

• params (iterable) – Iterable of parameters to optimize or dicts defining parameter groups.

• alignment_damping (bool, optional) – Whether to apply alignment damping. Default: False.

• damping_tau (float, optional) – Tau parameter for the alignment damping sigmoid. Default: 2.0.

• damping_ema_decay (float, optional) – EMA decay for the alignment damping score. Default: 0.9.

• damping_min_scale (float, optional) – Minimum scale for the alignment damping. Default: 0.1.

• ns_polish_steps (int, optional) – Exact-Newton polish iterations after the fast ones. Default: 0 (torch behaviour).

• ns_polish_coefficients (tuple, optional) – Coefficients for those iterations. Default: (2.0, -1.5, 0.5).

• ns_rms_scale (float, optional) – Match-RMS rescale constant. Default: 0.2.

• **muon_kwargs – Additional keyword arguments for the base Muon optimizer.

load_state_dict(state_dict)

Restore the optimizer state and reset the internal parameter grouping.

Missing parameter-group settings are filled from the optimizer defaults. Parameters are matched by their order within each group, not by name.

Parameters:

state_dict (dict) – Optimizer state returned by state_dict().

add_param_group(param_group)

Add a param group to the Optimizer s param_groups.

This can be useful when fine tuning a pre-trained network as frozen layers can be made trainable and added to the Optimizer as training progresses.

Parameters:

param_group (dict) – Specifies what Tensors should be optimized along with group specific optimization options.

nonfinite_report()

Report and clear recorded non-finite gradient flags.

Returns:

Parameter names mapped to True when a non-finite gradient

was recorded in the current parameter routing.

Return type:

dict[str, bool]

step(closure=None)

Update parameters using the current gradients.

Parameters:

closure (callable, optional) – Function that reevaluates the loss with gradient recording enabled. Default: None.

Returns:

Loss returned by the closure, or None when no closure

is provided.

Return type:

torch.Tensor | None

flashcart.training.muon_norm

Estimate and measure the RMS magnitude of weights updated by Muon.

The forecast models the squared weight RMS, u, by u_next = (1 - lr*wd)**2 * u + (A*lr)**2, starting from u=1. For constant learning rate and small lr*wd, the stationary RMS is approximately A * sqrt(lr / (2*wd)). The accumulated quantity 2 * wd * sum(lr) estimates the relaxation of the squared RMS toward this stationary value.

The model is an empirical approximation. Its fitted coefficient does not account explicitly for gradient correlations, alignment damping, or changes to the Newton-Schulz update. The reports compare this forecast with observed weight and momentum norms.

flashcart.training.muon_norm.epoch_learning_rates(cfg)

Replay the configured lr schedule, one entry per epoch.

Parameters:

cfg (Mapping) – Run config (lr, max_epochs, scheduler keys).

Returns:

Learning rate used at the start of each training epoch.

Return type:

list[float]

flashcart.training.muon_norm.forecast_weight_norm(epoch_lrs, steps_per_epoch, weight_decay, constant=0.869)

Integrate the norm recursion over a schedule.

Parameters:

• epoch_lrs (iterable of float) – Learning rate for each epoch.

• steps_per_epoch (int) – Optimizer steps per epoch.

• weight_decay (float) – Decoupled weight decay.

• constant (float, optional) – A. Default: MUON_NORM_CONSTANT.

Returns:

ceiling is the approximate stationary RMS at the largest

learning rate. efoldings is the accumulated decay measure. peak_rms and final_rms are estimated from the recurrence. effective_end_factor is the final lr / rms(W) divided by its largest value at the sampled epoch endpoints. The ceiling is infinite when weight decay is nonpositive.

Return type:

dict[str, float]

flashcart.training.muon_norm.format_forecast(forecast)

One line summarizing a forecast, with any warnings appended.

Parameters:

forecast (Mapping[str, float]) – Output of forecast_weight_norm.

Returns:

One-line forecast with warnings for the configured diagnostic thresholds.

Return type:

str

flashcart.training.muon_norm.weight_rms_report(model, muon_optimizers=())

Report aggregate Muon weight RMS and the largest individual matrix RMS.

The aggregate includes trainable floating-point matrices owned by the provided optimizers. The maximum also considers other trainable floating-point matrices unless they are excluded from weight decay. The maximum helps identify individual matrices that grow while the aggregate RMS changes little.

Parameters:

• model (torch.nn.Module) – Unwrapped model.

• muon_optimizers (Iterable[torch.optim.Optimizer], optional) – Optimizers whose parameters define the aggregate. Default: an empty iterable, giving an aggregate RMS of zero.

Returns:

Aggregate weight_rms, largest weight_rms_max, and its

parameter name in weight_rms_max_tensor.

Return type:

dict[str, Any]

flashcart.training.muon_norm.momentum_norm_report(optimizers)

Report the smallest momentum norm relative to the first available epsilon.

The epsilon is taken from the first optimizer group that defines it. Comparisons therefore assume that the provided optimizers use the same Newton-Schulz normalization epsilon.

Parameters:

optimizers (iterable) – The Muon optimizers (aux_optimizers).

Returns:

momentum_norm_min_over_eps, or an empty mapping when no

usable momentum buffers or epsilon are available.

Return type:

dict[str, float]

flashcart.training.muon_norm.muon_weight_decay(optimizers)

Read weight decay from the first available optimizer group.

Parameters:

optimizers (Iterable[torch.optim.Optimizer]) – Optimizers searched in order.

Returns:

First group’s weight decay, defaulting to zero when absent, or

None when no groups are available.

Return type:

float | None