Utilities and metadata for each parameter block. Attributes: param (Tensor): The original parameter that contains the block. comp

Utilities and metadata for each parameter block specific using DTensor. Attributes: param (Tensor): The original parameter that contains

Implements distributed Shampoo algorithm. -------- Features -------- 1. Layerwise Grafting: In order to tune Shampoo, we can "graft"

Configuration for AdaGrad preconditioner computation. Attributes: epsilon (float): Epsilon term for regularizing square-root of the aggre

Configuration for HybridShard (per-parameter FSDP) distributed computation. Enables distributed computation and optimizer states (like ZeRO-1) vi

Configuration for FullyShard (per-parameter FSDP) distributed computation. Attributes: target_parameter_dimensionality (int | float): Th

Configuration for generalized primal averaging in Shampoo. Generalized Primal Averaging (GPA) maintains two sequences of iterates: - The eval

Configuration for DDP distributed computation. Enables distributed computation and optimizer states (like ZeRO-1) via DTensor for Shampoo. N

Configuration for eigendecomposition with torch.linalg.eigh. The tolerance hyperparameter is used for a criterion that enables an adaptive eigend

Configuration for Adam preconditioner computation. Attributes: beta2 (float): Exponential moving average factor for second moment. (Defau

Configuration for HSDP distributed computation. Enables distributed computation and optimizer states (like ZeRO-1) via DTensor for Shampoo across

Result of wrapping a model for distributed training.

Configuration for Shampoo without any parallelism. Attributes: target_parameter_dimensionality (int | float): The idealized parameter dim

Simple two-layer convolutional network for image classification. Takes in image represented by an order-3 tensor. Used for testing optimizers.

DDP strategy.

Iterator that yields dictionaries by zipping values from iterators in a dictionary. Given a dictionary mapping from strings to iterators,

Configuration for FSDP distributed computation. Passes in additional metadata necessary to run FSDP Shampoo. Attributes: target_para

FSDP Metadata for a parameter. Attributes: fqn (str): Fully qualified name of the parameter. shape (torch.Size): Shape of the par

HSDP v2 (hybrid_shard) strategy.

Configuration for perturbing/damping/regularizing matrix eigenvalues by a small value epsilon (provided in DistributedShampoo arguments) to guara

Configuration for filtering zero/near-zero singular values (i.e., determining rank) to return a pseudo-inverse when the matrix is non-invertible.

Configuration for schedule-free optimization in Shampoo. Schedule-Free is an iterate averaging method that eliminates the need for learning rate

r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b` Compared to torch.nn.Linear, uses a combined parameter for both th

Configuration for matrix root inverse via coupled higher-order method. Attributes: rel_epsilon (float): Relative epsilon for coupled high

Configuration for matrix root inverse via an eigendecomposition. Attributes: rank_deficient_stability_config (RankDeficientStabilityConfi

Configuration for Shampoo preconditioner computation with caching of the eigendecomposed factor matrices. Note: When using custom amortized_compu

FSDP v1 strategy.

HSDP v1 strategy.

NorMuonH: NS-orthogonalised gradient + Adafactor-style row/column variance preconditioning (NorMuon, https://arxiv.org/pdf/2510.05491) + hyperball

Configuration for eigenvalue decomposition via QR algorithm. Determines whether the QR algorithm has converged based on the estimated eigenvalues

Configuration for eigenvalue-corrected Shampoo/SOAP preconditioner computation. Recall that in eigenvalue-corrected Shampoo, the eigenvectors and

FSDP v2 (fully_shard) strategy.

Load balancing configuration for distributing workloads across ranks. The `cost_model` defines how the cost of a tensor is computed, and the dist

Formatter that adds rank to the log message.

Configuration for Shampoo PT2 compilation. Enables Shampoo pytorch compilation with configure to speed up model training. For more details: h

Runtime configuration for Shampoo. Only non-checkpointed options here. Attributes: eager_nan_check (bool): Flag for checking for NaN

Single GPU training (no wrapping).

Configuration for spectral descent computation in DistributedShampoo. NOTE: This config can only be used for 2D parameters, or parameters that ha

Preconditioner list for spectral descent. NOTE: This algorithm can only be used for 2D parameters, or parameters that have been reshaped to 2D.

Adagrad / RMSprop / Adam preconditioners for a list of parameters. Operations are performed in-place with foreach operators. NOTE: Does not