Optimizers

class pylo.optim.AdafacLO_naive(params, momentum_decays=[0.15216392, 0.14245212, 0.06812963], rms_decays=[0.01079706], adafactor_decays=[0.18621896, -0.10864615, -0.06185547], lr=1.0, exp_mult=0.001, step_mult=0.01, input_size=39, hidden_size=32, hidden_layers=1, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.999,), initial_adafactor_decays=(0.9, 0.99, 0.999), max_grad_norm=None, concat_weights=True, make_separate_weights=False, split_weights=False, clip_grad=False, weight_decay=0.0, mup_lrs=None, hf_key: str | None = 'btherien/mulo')

__init__(params, momentum_decays=[0.15216392, 0.14245212, 0.06812963], rms_decays=[0.01079706], adafactor_decays=[0.18621896, -0.10864615, -0.06185547], lr=1.0, exp_mult=0.001, step_mult=0.01, input_size=39, hidden_size=32, hidden_layers=1, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.999,), initial_adafactor_decays=(0.9, 0.99, 0.999), max_grad_norm=None, concat_weights=True, make_separate_weights=False, split_weights=False, clip_grad=False, weight_decay=0.0, mup_lrs=None, hf_key: str | None = 'btherien/mulo')

step(loss=None)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

pylo.optim.MuLO_naive(params, impl=<class 'pylo.optim.AdafacLO_naive.AdafacLO_naive'>, **kwargs)

μP (Maximal Update Parameterization) wrapper for the PyTorch native implementation of the Adafac learned optimizer.

This function applies the μP parameterization to the Adafac learned optimizer, scaling learning rates for matrix-like parameters according to their width multipliers. Parameters are organized into groups based on their infinite-width shape properties.

Note

This implementation requires that all parameters have been processed with mup.set_base_shapes() to establish their infinite-width behavior.

Example

>>> model = MyModel()
>>> mup.set_base_shapes(model, base_model)
>>> optimizer = MuLO_naive(model.parameters())

class pylo.optim.VeLO_naive(params, momentum_decays=[0.0, 0.0, 0.0], rms_decays=[0.0], adafactor_decays=[0.0, 0.0, 0.0], lr=0.001, exp_mult=0.001, step_mult=0.001, input_size=30, hidden_size=4, hidden_layers=1, initial_momentum_decays=(0.9, 0.99, 0.999), lstm_input_size=30, lstm_hidden_size=512, param_inits=256, num_steps=10000, initial_rms_decays=(0.999,), initial_adafactor_decays=(0.9, 0.99, 0.999), concat_weights=True, make_separate_weights=False, split_weights=False, weight_decay=0.0, clip_grad=False, mup_lrs=None, hf_key_rnn='Pauljanson002/VeLO_RNN', hf_key_mlp='Pauljanson002/VeLO_MLP')

__init__(params, momentum_decays=[0.0, 0.0, 0.0], rms_decays=[0.0], adafactor_decays=[0.0, 0.0, 0.0], lr=0.001, exp_mult=0.001, step_mult=0.001, input_size=30, hidden_size=4, hidden_layers=1, initial_momentum_decays=(0.9, 0.99, 0.999), lstm_input_size=30, lstm_hidden_size=512, param_inits=256, num_steps=10000, initial_rms_decays=(0.999,), initial_adafactor_decays=(0.9, 0.99, 0.999), concat_weights=True, make_separate_weights=False, split_weights=False, weight_decay=0.0, clip_grad=False, mup_lrs=None, hf_key_rnn='Pauljanson002/VeLO_RNN', hf_key_mlp='Pauljanson002/VeLO_MLP')

init_state()

collect_rnn_outputs(to_lstm_from_loss)

state_dict()

Return the state of the optimizer as a dict.

It contains two entries:

• state: a Dict holding current optimization state. Its content

  differs between optimizer classes, but some common characteristics hold. For example, state is saved per parameter, and the parameter itself is NOT saved. state is a Dictionary mapping parameter ids to a Dict with state corresponding to each parameter.

• param_groups: a List containing all parameter groups where each

  parameter group is a Dict. Each parameter group contains metadata specific to the optimizer, such as learning rate and weight decay, as well as a List of parameter IDs of the parameters in the group. If a param group was initialized with named_parameters() the names content will also be saved in the state dict.

NOTE: The parameter IDs may look like indices but they are just IDs associating state with param_group. When loading from a state_dict, the optimizer will zip the param_group params (int IDs) and the optimizer param_groups (actual nn.Parameter s) in order to match state WITHOUT additional verification.

A returned state dict might look something like:

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
            'param_names' ['param0']  (optional)
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
            'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional)
        }
    ]
}

load_state_dict(state_dict)

Load the optimizer state.

Parameters:

state_dict (dict) – optimizer state. Should be an object returned from a call to state_dict().

Warning

Make sure this method is called after initializing torch.optim.lr_scheduler.LRScheduler, as calling it beforehand will overwrite the loaded learning rates.

Note

The names of the parameters (if they exist under the “param_names” key of each param group in state_dict()) will not affect the loading process. To use the parameters’ names for custom cases (such as when the parameters in the loaded state dict differ from those initialized in the optimizer), a custom register_load_state_dict_pre_hook should be implemented to adapt the loaded dict accordingly. If param_names exist in loaded state dict param_groups they will be saved and override the current names, if present, in the optimizer state. If they do not exist in loaded state dict, the optimizer param_names will remain unchanged.

Example

>>> # xdoctest: +SKIP
>>> model = torch.nn.Linear(10, 10)
>>> optim = torch.optim.SGD(model.parameters(), lr=3e-4)
>>> scheduler1 = torch.optim.lr_scheduler.LinearLR(
...     optim,
...     start_factor=0.1,
...     end_factor=1,
...     total_iters=20,
... )
>>> scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(
...     optim,
...     T_max=80,
...     eta_min=3e-5,
... )
>>> lr = torch.optim.lr_scheduler.SequentialLR(
...     optim,
...     schedulers=[scheduler1, scheduler2],
...     milestones=[20],
... )
>>> lr.load_state_dict(torch.load("./save_seq.pt"))
>>> # now load the optimizer checkpoint after loading the LRScheduler
>>> optim.load_state_dict(torch.load("./save_optim.pt"))

step(loss)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

class pylo.optim.CELO2_naive(params, lr=0.001, weight_decay=0.0, adam_lr_mult=1.0, adam_weight_decay=None, adam_betas=(0.9, 0.95), adam_eps=1e-08, use_adamw_for_1d=True, orthogonalize=True, clip_grad=False, clip_norm=1.0, ff_hidden_size=8, ff_hidden_layers=2, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.95,), initial_adafactor_decays=(0.9, 0.99, 0.999), exp_mult=0.0, rmsmult=1.0, param_scale_mult=False, ns_coeffs=(3.4445, -4.775, 2.0315), ns_iters=5, ns_eps=1e-08, grad_clip_val=1000.0, hf_key: str | None = 'DiamondXL/celo2', checkpoint_path: str | None = None, network: CELO2MLP | None = None)

Pure-PyTorch CELO2 learned optimizer.

Parameters:

• params – Iterable of parameters or param_groups. A group may carry an is_embedding=True flag to force its (2D) parameters onto the AdamW path, mirroring the 'embed' routing of the JAX version.

• lr – Base learning rate for the CELO2 (2D+) path. No schedule is applied internally; drive any warmup/cosine schedule with an external torch.optim.lr_scheduler.

• weight_decay – Decoupled weight decay for the CELO2 (2D+) path.

• adam_lr_mult – AdamW configuration for 1D / embedding parameters. The AdamW moments are maintained independently of the CELO2 accumulators. adam_weight_decay defaults to weight_decay when None.

• adam_weight_decay – AdamW configuration for 1D / embedding parameters. The AdamW moments are maintained independently of the CELO2 accumulators. adam_weight_decay defaults to weight_decay when None.

• adam_betas – AdamW configuration for 1D / embedding parameters. The AdamW moments are maintained independently of the CELO2 accumulators. adam_weight_decay defaults to weight_decay when None.

• adam_eps – AdamW configuration for 1D / embedding parameters. The AdamW moments are maintained independently of the CELO2 accumulators. adam_weight_decay defaults to weight_decay when None.

• use_adamw_for_1d – AdamW configuration for 1D / embedding parameters. The AdamW moments are maintained independently of the CELO2 accumulators. adam_weight_decay defaults to weight_decay when None.

• orthogonalize – Apply Newton-Schulz orthogonalization to 2D+ updates (set False for the “celo2-base” variant).

• clip_grad – Optional global-norm gradient clipping.

• clip_norm – Optional global-norm gradient clipping.

• ff_hidden_size – initial_rms_decays, initial_adafactor_decays, exp_mult, rmsmult, ns_coeffs, ns_iters, ns_eps: CELO2 model / accumulator configuration.

• ff_hidden_layers – initial_rms_decays, initial_adafactor_decays, exp_mult, rmsmult, ns_coeffs, ns_iters, ns_eps: CELO2 model / accumulator configuration.

• initial_momentum_decays – initial_rms_decays, initial_adafactor_decays, exp_mult, rmsmult, ns_coeffs, ns_iters, ns_eps: CELO2 model / accumulator configuration.

:paraminitial_rms_decays, initial_adafactor_decays, exp_mult, rmsmult,

ns_coeffs, ns_iters, ns_eps: CELO2 model / accumulator configuration.

Parameters:

• grad_clip_val – Element-wise gradient clamp applied before preprocessing (matches celo2_optax which clamps to [-1000, 1000]).

• hf_key – HuggingFace Hub id to load the CELO2MLP weights from.

• checkpoint_path – Local path to a converted CELO2MLP state_dict (.pt).

• network – An already-constructed CELO2MLP (overrides the above).

__init__(params, lr=0.001, weight_decay=0.0, adam_lr_mult=1.0, adam_weight_decay=None, adam_betas=(0.9, 0.95), adam_eps=1e-08, use_adamw_for_1d=True, orthogonalize=True, clip_grad=False, clip_norm=1.0, ff_hidden_size=8, ff_hidden_layers=2, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.95,), initial_adafactor_decays=(0.9, 0.99, 0.999), exp_mult=0.0, rmsmult=1.0, param_scale_mult=False, ns_coeffs=(3.4445, -4.775, 2.0315), ns_iters=5, ns_eps=1e-08, grad_clip_val=1000.0, hf_key: str | None = 'DiamondXL/celo2', checkpoint_path: str | None = None, network: CELO2MLP | None = None)

step(loss=None)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

class pylo.optim.ELO_CELO2_naive(params, lr=0.001, weight_decay=0.1, clip_grad=True, clip_norm=1.0, adam_lr_mult=1.0, adam_weight_decay=None, use_adamw_for_1d=True, orthogonalize=True, ff_hidden_size=8, ff_hidden_layers=2, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.95,), initial_adafactor_decays=(0.9, 0.99, 0.999), exp_mult=0.0, rmsmult=1.0, param_scale_mult=False, ns_coeffs=(3.4445, -4.775, 2.0315), ns_iters=5, ns_eps=1e-08, grad_clip_val=1000.0, hf_key: str | None = 'DiamondXL/elo-celo2', checkpoint_path: str | None = None, network: CELO2MLP | None = None)

Inference-time ELO-CELO2 optimizer (CELO2 forward, shared-accumulator AdamW).

__init__(params, lr=0.001, weight_decay=0.1, clip_grad=True, clip_norm=1.0, adam_lr_mult=1.0, adam_weight_decay=None, use_adamw_for_1d=True, orthogonalize=True, ff_hidden_size=8, ff_hidden_layers=2, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.95,), initial_adafactor_decays=(0.9, 0.99, 0.999), exp_mult=0.0, rmsmult=1.0, param_scale_mult=False, ns_coeffs=(3.4445, -4.775, 2.0315), ns_iters=5, ns_eps=1e-08, grad_clip_val=1000.0, hf_key: str | None = 'DiamondXL/elo-celo2', checkpoint_path: str | None = None, network: CELO2MLP | None = None)

step(loss=None)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

class pylo.optim.ELO_naive(params, lr=0.001, exp_mult=0.001, weight_decay=0.0, hidden_size=32, hidden_layers=1, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.999,), initial_adafactor_decays=(0.9, 0.99, 0.999), clip_grad=False, clip_norm=1.0, hf_key: str | None = 'DiamondXL/elo', checkpoint_path: str | None = None, network: MetaMLP | None = None)

Pure-PyTorch ELO (Adafactor-MLP) learned optimizer.

Parameters:

• params – Iterable of parameters or param_groups.

• lr – Base learning rate (the ELO step_mult). No schedule is applied internally; drive any warmup/cosine schedule with an external torch.optim.lr_scheduler.

• exp_mult – Magnitude exponent multiplier for the MLP output.

• weight_decay – Decoupled weight decay (scaled by lr).

• hidden_size – MetaMLP geometry. Note ELO counts hidden weight layers, so the original hidden_layers=2 maps to MetaMLP(hidden_layers=1) (input + one hidden + output).

• hidden_layers – MetaMLP geometry. Note ELO counts hidden weight layers, so the original hidden_layers=2 maps to MetaMLP(hidden_layers=1) (input + one hidden + output).

• initial_momentum_decays – Raw accumulator decays (used directly, no reparameterization).

• initial_rms_decays – Raw accumulator decays (used directly, no reparameterization).

• initial_adafactor_decays – Raw accumulator decays (used directly, no reparameterization).

• clip_grad – Optional global-norm gradient clipping.

• clip_norm – Optional global-norm gradient clipping.

• hf_key – HuggingFace Hub id for the MetaMLP weights.

• checkpoint_path – Local path to a converted MetaMLP state_dict (.pt).

• network – An already-constructed MetaMLP (overrides the above).

__init__(params, lr=0.001, exp_mult=0.001, weight_decay=0.0, hidden_size=32, hidden_layers=1, initial_momentum_decays=(0.9, 0.99, 0.999), initial_rms_decays=(0.999,), initial_adafactor_decays=(0.9, 0.99, 0.999), clip_grad=False, clip_norm=1.0, hf_key: str | None = 'DiamondXL/elo', checkpoint_path: str | None = None, network: MetaMLP | None = None)

step(loss=None)

Perform a single optimization step to update parameter.

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

pylo.optim.CELO2

alias of CELO2_naive

pylo.optim.ELO_CELO2

alias of ELO_CELO2_naive

pylo.optim.ELO

alias of ELO_naive

pylo.optim.VeLO

alias of VeLO_naive

pylo.optim.AdafacLO

alias of AdafacLO_naive

pylo.optim.MuLO(params, impl=<class 'pylo.optim.AdafacLO_naive.AdafacLO_naive'>, **kwargs)

μP (Maximal Update Parameterization) wrapper for the PyTorch native implementation of the Adafac learned optimizer.

This function applies the μP parameterization to the Adafac learned optimizer, scaling learning rates for matrix-like parameters according to their width multipliers. Parameters are organized into groups based on their infinite-width shape properties.

Note

This implementation requires that all parameters have been processed with mup.set_base_shapes() to establish their infinite-width behavior.

Example

>>> model = MyModel()
>>> mup.set_base_shapes(model, base_model)
>>> optimizer = MuLO_naive(model.parameters())