Resolve a string to WeightDecayType.

Recursively compares two optimizer state dicts for bitwise equivalence, skipping 'param_groups' keys. Args: ref_optim_state (dict[st

Converts nested dictionary with objects with state dict functionality. Args: input_dict (dict[str, Any]): Nested dictionary containing ob

Load model and optimizer state from checkpoint if available.

Allocates a persistent shadow copy of updated parameters.

Redistributes updated parameters to each parameter's rank.

Configure per-rank logging with rank prefix and optional debug level.

Tiled transpose copy: dst = src.T where src is (M, N) and dst is (N, M). Uses a 64x128 tiled Triton kernel with coalesced reads AND writes, a

Cautious weight decay + parameter update. wd_tensor and lr_tensor are 0-D CPU tensors. Mantissa is tracked to enable higher precision

Momentum -> SOAP -> muon_update; returns the raw (pre-radial) update. Mirrors the per-param body in step() — used by _circuit_muon_step so the

Radial scale -> u/w floor -> apply -> radius correction -> SOAP refresh.

Resolve a string to FSDPParamAssignmentStrategy.

Creates a dictionary containing shallow copies of this dataclass's fields. This method creates a dictionary where keys are field nam

Retrieves the inverse exponent override based on the dimension and order. Args: dimension (int): The dimension for which