From 38fa59c8de40f395c9bdb0c114339ce5cac3fd62 Mon Sep 17 00:00:00 2001 From: Xiaolong Huang Date: Tue, 9 Jun 2026 16:50:31 -0400 Subject: [PATCH 1/7] Add CELO2 and ELO-CELO2 learned optimizers (naive PyTorch) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port the inference-time forward pass of the CELO2 / ELO-CELO2 learned optimizers from the original JAX/optax implementation into pure PyTorch, as drop-in pylo optimizers. - pylo/models/CELO2_MLP.py: CELO2MLP, the split-input per-parameter MLP backbone (14 split first-layer weights + dense layers, HF Hub mixin). - pylo/optim/CELO2_naive.py: CELO2_naive optimizer — momentum / RMS / factored Adafactor accumulators, CELO2 feature stack, Newton-Schulz orthogonalization for 2D+ params, AdamW for 1D params over the shared accumulators, and a warmup + cosine LR schedule. Loads its meta-model from HuggingFace (default DiamondXL/celo2); a local converted checkpoint or an explicit network take precedence. - pylo/optim/ELO_CELO2_naive.py: ELO_CELO2_naive — at inference the ELO expert mechanism is disabled, so it reduces to the CELO2 forward with the ELO default hyper-parameters (weight_decay=0.1, clip_grad=True); default meta-model DiamondXL/elo-celo2. - scripts/convert_celo2_checkpoint.py: convert a JAX/Haiku theta checkpoint into a CELO2MLP state_dict. - tests/test_celo2.py: step/update, higher-rank param, state-dict resume, and a JAX numerical-alignment test (2D update matches the reference to ~3e-6; auto-skips when the JAX source is unavailable). - Register the new classes in the pylo / pylo.optim / pylo.models inits. Co-Authored-By: Claude Opus 4.8 (1M context) --- pylo/__init__.py | 5 + pylo/models/CELO2_MLP.py | 123 +++++++++ pylo/models/__init__.py | 2 + pylo/optim/CELO2_naive.py | 392 ++++++++++++++++++++++++++++ pylo/optim/ELO_CELO2_naive.py | 90 +++++++ pylo/optim/__init__.py | 17 +- scripts/convert_celo2_checkpoint.py | 104 ++++++++ tests/test_celo2.py | 179 +++++++++++++ 8 files changed, 911 insertions(+), 1 deletion(-) create mode 100644 pylo/models/CELO2_MLP.py create mode 100644 pylo/optim/CELO2_naive.py create mode 100644 pylo/optim/ELO_CELO2_naive.py create mode 100644 scripts/convert_celo2_checkpoint.py create mode 100644 tests/test_celo2.py diff --git a/pylo/__init__.py b/pylo/__init__.py index 4edaaa8..82bb494 100644 --- a/pylo/__init__.py +++ b/pylo/__init__.py @@ -5,9 +5,14 @@ "MetaMLP", "VeLOMLP", "VeLORNN", + "CELO2MLP", "AdafacLO_naive", "MuLO_naive", "VeLO_naive", + "CELO2_naive", + "ELO_CELO2_naive", + "CELO2", + "ELO_CELO2", # Default aliases (re-exported from pylo.optim). These resolve to the # CUDA implementations when available, falling back to the naive ones # otherwise, so downstream code can simply `from pylo import VeLO`. diff --git a/pylo/models/CELO2_MLP.py b/pylo/models/CELO2_MLP.py new file mode 100644 index 0000000..56d94b3 --- /dev/null +++ b/pylo/models/CELO2_MLP.py @@ -0,0 +1,123 @@ +"""CELO2MLP: the per-parameter MLP backbone of the CELO2 learned optimizer. + +This mirrors the Haiku ``_ff_mod`` network from the original JAX/optax CELO2 +implementation (https://arxiv.org/abs/2602.19142). The defining feature of the +network is a *split* first layer: every input feature group owns its own weight +matrix ``(feature_dim, hidden_size)`` and the first hidden pre-activation is the +sum of ``feature @ weight`` over all groups. The remaining layers are dense. + +The forward pass multiplies as ``x @ w`` (Haiku/optax convention, where the +weight is stored ``(in, out)``) rather than ``x @ w.T`` (``torch.nn.Linear``), +so that a checkpoint converted from the JAX implementation reproduces the JAX +output bit-for-bit. +""" + +import torch +import torch.nn as nn +from huggingface_hub import PyTorchModelHubMixin + +# Per-feature input dimensions, in the exact order CELO2 builds its features for +# a 2D+ (factored) parameter: +# g, clip(g), p, m[3], rms[1], m*rsqrt[3], rsqrt[1], fac_g[3], g*rsqrt[1], +# row_feat[3], col_feat[3], rsqrt(row)[3], rsqrt(col)[3], fac_mom_mult[3] +# A 1D parameter only produces the first 9 groups (no factored features); the +# trailing weights simply go unused for those parameters. +CELO2_FEATURE_DIMS = (1, 1, 1, 3, 1, 3, 1, 3, 1, 3, 3, 3, 3, 3) + + +class CELO2MLP( + nn.Module, + PyTorchModelHubMixin, + license="apache-2.0", + tags=["learned-optimizer"], +): + """Split-input MLP used by the CELO2 / ELO-CELO2 optimizers. + + Args: + feature_dims: Input width of each split first-layer weight. Defaults to + the 14 feature groups produced by a 2D parameter. + hidden_size: Width of the hidden layers. + hidden_layers: Number of hidden weight applications. The architecture is + ``[hidden_size] * hidden_layers + [output_size]`` (matching the + original ``[ff_hidden_size] * ff_hidden_layers + [3]``), so the total + number of weight matrices is ``hidden_layers + 1`` (one split input + layer plus ``hidden_layers`` dense layers). + output_size: Width of the final output (CELO2 uses 3; only the first two + channels, direction and magnitude, are consumed by the optimizer). + activation: ``"relu"`` or ``"tanh"``. + """ + + def __init__( + self, + feature_dims=CELO2_FEATURE_DIMS, + hidden_size=8, + hidden_layers=2, + output_size=3, + activation="relu", + ): + super().__init__() + self.feature_dims = tuple(feature_dims) + self.hidden_size = hidden_size + self.hidden_layers = hidden_layers + self.output_size = output_size + self.activation = activation + + if activation == "relu": + self.act_fn = torch.relu + elif activation == "tanh": + self.act_fn = torch.tanh + else: + raise ValueError(f"Invalid MLP activation: {activation}") + + # Split first layer: one (feature_dim, hidden_size) weight per feature. + self.w0 = nn.ParameterList( + [nn.Parameter(torch.empty(fd, hidden_size)) for fd in self.feature_dims] + ) + self.b0 = nn.Parameter(torch.zeros(hidden_size)) + + # Dense layers: sizes [hidden_size] * hidden_layers + [output_size]. + layer_sizes = [hidden_size] * hidden_layers + [output_size] + self.dense_w = nn.ParameterList() + self.dense_b = nn.ParameterList() + last = hidden_size + for size in layer_sizes[1:]: + self.dense_w.append(nn.Parameter(torch.empty(last, size))) + self.dense_b.append(nn.Parameter(torch.zeros(size))) + last = size + + self.reset_parameters() + + def reset_parameters(self): + # Truncated-normal init mirroring the Haiku stddev = 1 / sqrt(fan_in). + total_in = sum(self.feature_dims) + for w in self.w0: + nn.init.trunc_normal_(w, std=1.0 / (total_in ** 0.5)) + last = self.hidden_size + for w in self.dense_w: + nn.init.trunc_normal_(w, std=1.0 / (last ** 0.5)) + last = w.shape[-1] + + def forward(self, inps): + """Run the split-input MLP. + + Args: + inps: List of per-feature tensors, each of shape ``[..., feature_dim]`` + and already normalized by the caller. Its length may be shorter + than ``feature_dims`` (e.g. 9 groups for a 1D parameter), in which + case only the leading split weights are used. + + Returns: + Output tensor of shape ``[..., output_size]``. + """ + o = inps[0] @ self.w0[0] + for x, w in zip(inps[1:], self.w0[1:]): + o = o + x @ w + o = o + self.b0 + o = self.act_fn(o) + + n = len(self.dense_w) + for i, (w, b) in enumerate(zip(self.dense_w, self.dense_b)): + o = o @ w + b + if i != n - 1: + o = self.act_fn(o) + return o diff --git a/pylo/models/__init__.py b/pylo/models/__init__.py index 90a2c38..6d37c1e 100644 --- a/pylo/models/__init__.py +++ b/pylo/models/__init__.py @@ -1,9 +1,11 @@ from pylo.models.Meta_MLP import MetaMLP from pylo.models.VeLO_MLP import VeLOMLP from pylo.models.VeLO_RNN import VeLORNN +from pylo.models.CELO2_MLP import CELO2MLP __all__ = [ "MetaMLP", "VeLOMLP", "VeLORNN", + "CELO2MLP", ] diff --git a/pylo/optim/CELO2_naive.py b/pylo/optim/CELO2_naive.py new file mode 100644 index 0000000..69d2bf9 --- /dev/null +++ b/pylo/optim/CELO2_naive.py @@ -0,0 +1,392 @@ +"""CELO2_naive: a PyTorch port of the CELO2 learned optimizer. + +This is a faithful, pure-PyTorch reimplementation of the inference-time forward +pass of the original JAX/optax CELO2 optimizer +(``scaling_l2o/src/learned_optimizers/celo2_optax.py``, +https://arxiv.org/abs/2602.19142). + +The optimizer keeps, per parameter, three families of accumulators: + * momentum at several decays (default ``(0.9, 0.99, 0.999)``), + * an RMS / second-moment accumulator (default ``(0.95,)``), + * an Adafactor-style factored second moment (default ``(0.9, 0.99, 0.999)``). + +For each 2D+ parameter it builds a stack of features, feeds them through the +:class:`~pylo.models.CELO2_MLP.CELO2MLP` (split-input MLP), optionally +orthogonalizes the resulting update via Newton-Schulz iteration, normalizes it, +and applies it with a warmup + cosine learning-rate schedule. 1D parameters +(biases, norms) are updated with AdamW using the shared momentum/RMS +accumulators, matching the original implementation. +""" + +from typing import Optional + +import numpy as np +import torch +from torch.optim import Optimizer + +from pylo.models.CELO2_MLP import CELO2MLP + + +# ============================================================================= +# Shared math helpers (mirroring celo2_optax.py) +# ============================================================================= + +def factored_dims(shape): + """The two largest dims used for the factored second moment, or None.""" + if len(shape) < 2: + return None + sorted_dims = np.argsort(shape) + return int(sorted_dims[-2]), int(sorted_dims[-1]) + + +def safe_rsqrt(x): + return torch.rsqrt(torch.clamp(x, min=1e-9)) + + +def second_moment_normalizer(x, axis, eps=1e-9): + mean_squared = torch.mean(x ** 2, dim=axis, keepdim=True) + return x * torch.rsqrt(eps + mean_squared) + + +def orthogonalize_via_newton_schulz(x, ns_coeffs, ns_steps=5, eps=1e-8): + """Newton-Schulz orthogonalization over the last two dims of ``x``.""" + transposed = False + if x.shape[-2] > x.shape[-1]: + x = x.transpose(-2, -1) + transposed = True + x = x / (torch.linalg.norm(x, dim=(-2, -1), keepdim=True) + eps) + c0, c1, c2 = ns_coeffs + for _ in range(ns_steps): + a = x @ x.transpose(-2, -1) + b = c1 * a + c2 * (a @ a) + x = c0 * x + b @ x + if transposed: + x = x.transpose(-2, -1) + return x + + +class CELO2_naive(Optimizer): + """Pure-PyTorch CELO2 learned optimizer. + + Args: + 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. + num_steps: Total number of inner training steps. Required: it defines the + warmup + cosine learning-rate schedule. + init_lr, peak_lr, warmup_steps, warmup_fraction, end_lr: Warmup + cosine + schedule parameters. ``warmup_fraction`` (a fraction of ``num_steps``) + takes priority over ``warmup_steps`` when > 0. + weight_decay: Decoupled weight decay for the CELO2 (2D+) path. + adam_lr_mult, adam_beta1, adam_beta2, adam_weight_decay, use_adamw_for_1d: + AdamW configuration for 1D parameters. ``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, clip_norm: Optional global-norm gradient clipping. + ff_hidden_size, ff_hidden_layers, initial_momentum_decays, + initial_rms_decays, initial_adafactor_decays, exp_mult, rmsmult, + ns_coeffs, ns_iters, ns_eps: CELO2 model / accumulator configuration. + 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 :class:`CELO2MLP` (overrides the above). + """ + + def __init__( + self, + params, + num_steps, + # LR schedule + init_lr=0.0, + peak_lr=1e-3, + warmup_steps=0, + warmup_fraction=0.05, + end_lr=0.0, + # Weight decay + weight_decay=0.0, + # AdamW for 1D params (reuses momentum_decays[0] / rms_decays[-1] as betas, + # matching the shared-accumulator design of the JAX ELO-CELO2) + adam_lr_mult=1.0, + adam_weight_decay=None, + use_adamw_for_1d=True, + # CELO2 backbone + 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.7750, 2.0315), + ns_iters=5, + ns_eps=1e-8, + grad_clip_val=1000.0, + # Weights + hf_key: Optional[str] = "DiamondXL/celo2", + checkpoint_path: Optional[str] = None, + network: Optional[CELO2MLP] = None, + ): + if num_steps is None: + raise ValueError("CELO2_naive requires num_steps for the LR schedule.") + + self.device = "cuda" if torch.cuda.is_available() else "cpu" + + defaults = dict( + num_steps=num_steps, + init_lr=init_lr, + peak_lr=peak_lr, + warmup_steps=warmup_steps, + warmup_fraction=warmup_fraction, + end_lr=end_lr, + weight_decay=weight_decay, + adam_lr_mult=adam_lr_mult, + adam_weight_decay=( + adam_weight_decay if adam_weight_decay is not None else weight_decay + ), + use_adamw_for_1d=use_adamw_for_1d, + orthogonalize=orthogonalize, + clip_grad=clip_grad, + clip_norm=clip_norm, + exp_mult=exp_mult, + rmsmult=rmsmult, + param_scale_mult=param_scale_mult, + grad_clip_val=grad_clip_val, + is_embedding=False, + ) + super().__init__(params, defaults) + + self.initial_momentum_decays = torch.tensor( + initial_momentum_decays, dtype=torch.float32, device=self.device + ) + self.initial_rms_decays = torch.tensor( + initial_rms_decays, dtype=torch.float32, device=self.device + ) + self.initial_adafactor_decays = torch.tensor( + initial_adafactor_decays, dtype=torch.float32, device=self.device + ) + self.ns_coeffs = tuple(float(c) for c in ns_coeffs) + self.ns_iters = ns_iters + self.ns_eps = ns_eps + + # Precedence: explicit network > local checkpoint > HuggingFace Hub. + if network is not None: + self.network = network + elif checkpoint_path is not None: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network.load_state_dict( + torch.load(checkpoint_path, map_location="cpu") + ) + elif hf_key is not None: + self.network = CELO2MLP.from_pretrained(hf_key) + else: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network = self.network.to(self.device) + + # ---------------------------------------------------------------- helpers + def _lr_schedule(self, step, group): + """Warmup + cosine decay learning rate (matches the JAX schedule).""" + num_steps = group["num_steps"] + warmup = group["warmup_steps"] + if group["warmup_fraction"] > 0: + warmup = group["warmup_fraction"] * num_steps + warmup_f = max(float(warmup), 1.0) + decay_f = max(float(num_steps), 1.0) + step = float(step) + peak_lr, init_lr, end_lr = group["peak_lr"], group["init_lr"], group["end_lr"] + if step < warmup_f: + return init_lr + (peak_lr - init_lr) * min(step / warmup_f, 1.0) + progress = min(max((step - warmup_f) / max(decay_f - warmup_f, 1.0), 0.0), 1.0) + return end_lr + (peak_lr - end_lr) * 0.5 * (1.0 + np.cos(np.pi * progress)) + + def _global_grad_scale(self, clip_norm): + """optax.clip_by_global_norm scale factor over all parameter grads.""" + total = torch.zeros((), device=self.device) + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + total = total + p.grad.detach().float().pow(2).sum() + global_norm = torch.sqrt(total) + denom = torch.clamp(global_norm, min=clip_norm) + return (clip_norm / denom).item() if denom > 0 else 1.0 + + def _update_accumulators(self, state, batch_g, p_shape): + """Advance momentum, RMS and factored accumulators in place; return fac_g.""" + beta_m = self.initial_momentum_decays + beta_rms = self.initial_rms_decays + beta_fac = self.initial_adafactor_decays + + # momentum: m = decay * m + (1 - decay) * g + mom = state["mom"] + mom.mul_(beta_m).add_((1 - beta_m) * batch_g) + # rms: rms = clip(decay) * rms + (1 - clip(decay)) * g^2 + crms = torch.clamp(beta_rms, 0.0, 1.0) + rms = state["rms"] + rms.mul_(crms).add_((1 - crms) * (batch_g ** 2)) + + # factored Adafactor accumulator (over the adafactor decays) + cdec = torch.clamp(beta_fac, 0.0, 1.0) + mixing = 1.0 - cdec + g_rep = batch_g.repeat([1] * len(p_shape) + [beta_fac.shape[-1]]) + grad_sqr = g_rep * g_rep + 1e-30 + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + new_v_row = cdec * state["fac_vec_row"] + mixing * grad_sqr.mean(dim=d0) + new_v_col = cdec * state["fac_vec_col"] + mixing * grad_sqr.mean(dim=d1) + reduced_d1 = d1 - 1 if d1 > d0 else d1 + row_col_mean = new_v_row.mean(dim=reduced_d1, keepdim=True) + row_factor = safe_rsqrt(new_v_row / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(new_v_col) + fac_g = g_rep * row_factor.unsqueeze(d0) * col_factor.unsqueeze(d1) + state["fac_vec_row"], state["fac_vec_col"] = new_v_row, new_v_col + else: + new_v = cdec * state["fac_vec_v"] + mixing * grad_sqr + fac_g = g_rep * safe_rsqrt(new_v + 1e-9) + state["fac_vec_v"] = new_v + return mom, rms, fac_g + + def _celo2_step(self, p, grad, state, group): + """The CELO2 MLP update for a single (2D+) parameter, shape == p.shape.""" + p_shape = tuple(p.shape) + batch_p = p.unsqueeze(-1) + batch_g = grad.unsqueeze(-1) + + mom, rms, fac_g = self._update_accumulators(state, batch_g, p_shape) + rsqrt = torch.rsqrt(rms + 1e-8) + + inps = [ + batch_g, + torch.clamp(batch_g, -0.1, 0.1), + batch_p, + mom, + rms, + mom * rsqrt, + rsqrt, + fac_g, + batch_g * rsqrt, + ] + + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + v_row, v_col = state["fac_vec_row"], state["fac_vec_col"] + rp_row = [1] * (1 + len(p_shape)) + rp_col = [1] * (1 + len(p_shape)) + rp_row[d0] = p_shape[d0] + rp_col[d1] = p_shape[d1] + row_feat = v_row.unsqueeze(d0).repeat(rp_row) + col_feat = v_col.unsqueeze(d1).repeat(rp_col) + inps.extend([ + row_feat, + col_feat, + torch.rsqrt(row_feat + 1e-8), + torch.rsqrt(col_feat + 1e-8), + ]) + reduced_d1 = d1 - 1 if d1 > d0 else d1 + row_col_mean = v_row.mean(dim=reduced_d1, keepdim=True) + row_factor = safe_rsqrt(v_row / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(v_col) + fac_mom_mult = mom * row_factor.unsqueeze(d0) * col_factor.unsqueeze(d1) + inps.append(fac_mom_mult) + + axis = list(range(len(p_shape)))[-2:] + inps = [second_moment_normalizer(i, axis=axis) for i in inps] + + out = self.network(inps) + direction = out[..., 0] + magnitude = out[..., 1] + + mag_param = torch.exp(magnitude * group["exp_mult"]) + if group["param_scale_mult"]: + param_scale = torch.sqrt(torch.mean(p ** 2) + 1e-9) + step = direction * (param_scale * mag_param) + else: + step = direction * mag_param + + if group["orthogonalize"] and step.dim() >= 2: + step = orthogonalize_via_newton_schulz( + step, self.ns_coeffs, self.ns_iters, self.ns_eps + ) + step = second_moment_normalizer(step, axis=axis) + step = step * group["rmsmult"] + return step + + # ------------------------------------------------------------------ step + @torch.no_grad() + def step(self, loss=None): + # Global-norm clipping needs all grads up front. + grad_scales = {} + for group in self.param_groups: + if group["clip_grad"]: + grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) + + for group in self.param_groups: + group["step"] = group.get("step", 0) + 1 + t = group["step"] + lr = self._lr_schedule(t, group) + adam_lr = group["adam_lr_mult"] * lr + beta1 = float(self.initial_momentum_decays[0]) + beta2 = float(self.initial_rms_decays[-1]) + scale = grad_scales.get(id(group), 1.0) + clip_val = group["grad_clip_val"] + use_adamw_for_1d = group["use_adamw_for_1d"] + force_adam = group["is_embedding"] + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if scale != 1.0: + grad = grad * scale + grad = torch.clamp(grad, -clip_val, clip_val) + p_shape = tuple(p.shape) + + state = self.state[p] + if len(state) == 0: + n_mom = self.initial_momentum_decays.shape[-1] + n_rms = self.initial_rms_decays.shape[-1] + n_fac = self.initial_adafactor_decays.shape[-1] + state["mom"] = torch.zeros( + p_shape + (n_mom,), device=self.device + ) + state["rms"] = torch.zeros( + p_shape + (n_rms,), device=self.device + ) + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + full = p_shape + (n_fac,) + vr = tuple(d for i, d in enumerate(full) if i != d0) + vc = tuple(d for i, d in enumerate(full) if i != d1) + state["fac_vec_row"] = torch.zeros(vr, device=self.device) + state["fac_vec_col"] = torch.zeros(vc, device=self.device) + state["fac_vec_v"] = torch.empty(0, device=self.device) + else: + state["fac_vec_row"] = torch.empty(0, device=self.device) + state["fac_vec_col"] = torch.empty(0, device=self.device) + state["fac_vec_v"] = torch.zeros( + p_shape + (n_fac,), device=self.device + ) + + is_1d = p.dim() <= 1 or force_adam + if use_adamw_for_1d and is_1d: + batch_g = grad.unsqueeze(-1) + mom, rms, _ = self._update_accumulators(state, batch_g, p_shape) + m_bc = mom[..., 0] / (1.0 - beta1 ** t) + v_bc = rms[..., -1] / (1.0 - beta2 ** t) + adam_step = m_bc / (torch.sqrt(v_bc) + 1e-8) + p.add_(adam_step + group["adam_weight_decay"] * p, alpha=-adam_lr) + else: + step = self._celo2_step(p, grad, state, group) + p.add_(step + group["weight_decay"] * p, alpha=-lr) + return loss diff --git a/pylo/optim/ELO_CELO2_naive.py b/pylo/optim/ELO_CELO2_naive.py new file mode 100644 index 0000000..89b9f7c --- /dev/null +++ b/pylo/optim/ELO_CELO2_naive.py @@ -0,0 +1,90 @@ +"""ELO_CELO2_naive: the ELO-CELO2 learned optimizer (PyTorch). + +ELO-CELO2 fuses a CELO2 MLP backbone with an ELO expert mechanism *during +meta-training*. At inference / meta-test time the expert trajectory and the IMT +losses are disabled (``meta_train=False``), so the parameter update reduces +exactly to the CELO2 forward pass: CELO2 MLP steps for 2D+ parameters and AdamW +(over the shared accumulators) for 1D parameters. + +Consequently this class is a thin wrapper over :class:`CELO2_naive` that only +changes the default hyper-parameters to match the ELO-CELO2 configuration +(``config/learned_optimizer/elo_celo2.py``): nonzero weight decay and enabled +gradient clipping. The distinct learned weights live in the loaded checkpoint. +""" + +from typing import Optional + +from pylo.optim.CELO2_naive import CELO2_naive + + +class ELO_CELO2_naive(CELO2_naive): + """Inference-time ELO-CELO2 optimizer (CELO2 forward with ELO defaults).""" + + def __init__( + self, + params, + num_steps, + # LR schedule + init_lr=0.0, + peak_lr=1e-3, + warmup_steps=0, + warmup_fraction=0.05, + end_lr=0.0, + # ELO-CELO2 defaults differ from CELO2 here: + weight_decay=0.1, + clip_grad=True, + clip_norm=1.0, + # AdamW for 1D params + adam_lr_mult=1.0, + adam_weight_decay=None, + use_adamw_for_1d=True, + # CELO2 backbone + 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.7750, 2.0315), + ns_iters=5, + ns_eps=1e-8, + grad_clip_val=1000.0, + # Weights + hf_key: Optional[str] = "DiamondXL/elo-celo2", + checkpoint_path: Optional[str] = None, + network=None, + ): + super().__init__( + params, + num_steps=num_steps, + init_lr=init_lr, + peak_lr=peak_lr, + warmup_steps=warmup_steps, + warmup_fraction=warmup_fraction, + end_lr=end_lr, + weight_decay=weight_decay, + adam_lr_mult=adam_lr_mult, + adam_weight_decay=adam_weight_decay, + use_adamw_for_1d=use_adamw_for_1d, + orthogonalize=orthogonalize, + clip_grad=clip_grad, + clip_norm=clip_norm, + ff_hidden_size=ff_hidden_size, + ff_hidden_layers=ff_hidden_layers, + initial_momentum_decays=initial_momentum_decays, + initial_rms_decays=initial_rms_decays, + initial_adafactor_decays=initial_adafactor_decays, + exp_mult=exp_mult, + rmsmult=rmsmult, + param_scale_mult=param_scale_mult, + ns_coeffs=ns_coeffs, + ns_iters=ns_iters, + ns_eps=ns_eps, + grad_clip_val=grad_clip_val, + hf_key=hf_key, + checkpoint_path=checkpoint_path, + network=network, + ) diff --git a/pylo/optim/__init__.py b/pylo/optim/__init__.py index f8f2bed..049d35a 100644 --- a/pylo/optim/__init__.py +++ b/pylo/optim/__init__.py @@ -1,9 +1,24 @@ from pylo.optim.AdafacLO_naive import AdafacLO_naive from pylo.optim.MuLO_naive import MuLO_naive from pylo.optim.Velo_naive import VeLO_naive +from pylo.optim.CELO2_naive import CELO2_naive +from pylo.optim.ELO_CELO2_naive import ELO_CELO2_naive # Initialize with optimizers we know we can import -__all__ = ["AdafacLO_naive", "MuLO_naive", "VeLO_naive"] +__all__ = [ + "AdafacLO_naive", + "MuLO_naive", + "VeLO_naive", + "CELO2_naive", + "ELO_CELO2_naive", +] + +# CELO2 / ELO-CELO2 currently ship a naive (pure-PyTorch) implementation only. +# Expose the bare public aliases now; if a CUDA build is added later these can be +# overridden in the try-block below, mirroring VeLO / AdafacLO. +CELO2 = CELO2_naive +ELO_CELO2 = ELO_CELO2_naive +__all__.extend(["CELO2", "ELO_CELO2"]) # Try to import CUDA-based optimizers try: diff --git a/scripts/convert_celo2_checkpoint.py b/scripts/convert_celo2_checkpoint.py new file mode 100644 index 0000000..24c0bfd --- /dev/null +++ b/scripts/convert_celo2_checkpoint.py @@ -0,0 +1,104 @@ +"""Convert a JAX/Haiku CELO2 ``theta`` checkpoint into a PyTorch state_dict. + +The original CELO2 meta-parameters are stored as a nested dict +``{"ff_mod_stack": {"": {"w0__0": ..., "b0": ..., "w1": ..., ...}}}`` +where every leaf is a ``(in, out)`` weight or a bias. This script flattens that +tree, maps each leaf onto the corresponding :class:`pylo.models.CELO2_MLP.CELO2MLP` +parameter, and saves a ``state_dict`` loadable via +``CELO2_naive(..., checkpoint_path=)``. + +Supported input formats (auto-detected by extension / content): + * ``.pickle`` — a pickled meta-train checkpoint (the framework's format). + * otherwise — a flax msgpack checkpoint (requires ``flax`` to be installed). + +Run this in an environment that can unpickle the checkpoint (e.g. the +``scaling_l2o`` JAX environment); only ``numpy`` and ``torch`` are required for +the conversion itself. + +Usage: + python scripts/convert_celo2_checkpoint.py \ + --input /path/to/global_step200000.pickle \ + --output /path/to/celo2_theta.pt +""" + +import argparse +import pickle +import re + +import numpy as np +import torch + +from pylo.models.CELO2_MLP import CELO2MLP + + +def _load_raw(path): + if str(path).endswith(".pickle"): + with open(path, "rb") as f: + return pickle.load(f) + # Flax msgpack: deserialize against a reference tree. + import flax # noqa: local import; only needed for this branch + + with open(path, "rb") as f: + return flax.serialization.msgpack_restore(f.read()) + + +def _flatten(tree, prefix=""): + """Yield (leaf_name, array) pairs from a nested dict of arrays.""" + if isinstance(tree, dict): + for k, v in tree.items(): + yield from _flatten(v, f"{prefix}/{k}" if prefix else str(k)) + else: + yield prefix, np.asarray(tree) + + +def convert(raw, hidden_size=8, hidden_layers=2): + # Collect leaves keyed by their final component (b0, w0__3, w1, ...). + leaves = {} + for name, arr in _flatten(raw): + leaves[name.split("/")[-1]] = arr + + model = CELO2MLP(hidden_size=hidden_size, hidden_layers=hidden_layers) + sd = model.state_dict() + + def assign(dst, src_key): + src = leaves[src_key] + if tuple(sd[dst].shape) != tuple(src.shape): + raise ValueError( + f"shape mismatch for {dst} <- {src_key}: " + f"{tuple(sd[dst].shape)} vs {tuple(src.shape)}" + ) + sd[dst] = torch.tensor(src, dtype=torch.float32) + + # Split first-layer weights: w0__{i} -> w0.{i} + n_split = len(model.w0) + for i in range(n_split): + assign(f"w0.{i}", f"w0__{i}") + assign("b0", "b0") + + # Dense layers: w1, w2, ... -> dense_w.0, dense_w.1, ...; b1, b2 -> dense_b.* + n_dense = len(model.dense_w) + for j in range(n_dense): + assign(f"dense_w.{j}", f"w{j + 1}") + assign(f"dense_b.{j}", f"b{j + 1}") + + model.load_state_dict(sd) + return model + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--input", required=True, help="Path to the JAX theta checkpoint.") + ap.add_argument("--output", required=True, help="Path for the PyTorch state_dict (.pt).") + ap.add_argument("--hidden-size", type=int, default=8) + ap.add_argument("--hidden-layers", type=int, default=2) + args = ap.parse_args() + + raw = _load_raw(args.input) + model = convert(raw, hidden_size=args.hidden_size, hidden_layers=args.hidden_layers) + torch.save(model.state_dict(), args.output) + n = sum(p.numel() for p in model.parameters()) + print(f"Saved CELO2MLP state_dict ({n} params) to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_celo2.py b/tests/test_celo2.py new file mode 100644 index 0000000..50bfbba --- /dev/null +++ b/tests/test_celo2.py @@ -0,0 +1,179 @@ +"""Tests for the CELO2 / ELO-CELO2 learned optimizers.""" + +import importlib.util +import os + +import numpy as np +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from pylo.models.CELO2_MLP import CELO2MLP +from pylo.optim import CELO2_naive, ELO_CELO2_naive +from pylo.optim.CELO2_naive import factored_dims + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +SCALING_L2O = "/home/mila/h/huangx/scaling_l2o" +CELO2_OPTAX = os.path.join(SCALING_L2O, "src/learned_optimizers/celo2_optax.py") + + +class SmallNet(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(10, 8) # 2D weight + 1D bias + self.fc2 = nn.Linear(8, 1) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +@pytest.mark.parametrize("optimizer_class", [CELO2_naive, ELO_CELO2_naive]) +def test_step_runs_and_updates(optimizer_class): + """Optimizer runs end-to-end and produces finite, non-trivial updates.""" + torch.manual_seed(0) + model = SmallNet().to(DEVICE) + before = [p.clone() for p in model.parameters()] + + optimizer = optimizer_class(model.parameters(), num_steps=50, peak_lr=1e-2) + x = torch.randn(16, 10, device=DEVICE) + y = torch.randn(16, 1, device=DEVICE) + + for _ in range(5): + optimizer.zero_grad() + loss = F.mse_loss(model(x), y) + loss.backward() + optimizer.step(loss) + + for p in model.parameters(): + assert torch.isfinite(p).all(), "non-finite parameter after step" + moved = any(not torch.allclose(b, p) for b, p in zip(before, model.parameters())) + assert moved, "optimizer did not update any parameters" + + +def test_higher_rank_parameter(): + """A >2D parameter (conv-like) goes through the orthogonalization path.""" + torch.manual_seed(0) + w = nn.Parameter(torch.randn(3, 4, 5, device=DEVICE)) + optimizer = CELO2_naive([w], num_steps=20, peak_lr=1e-2) + for _ in range(3): + optimizer.zero_grad() + (w.sum() ** 2).backward() + optimizer.step() + assert torch.isfinite(w).all() + + +def test_celo2_resume(tmp_path): + """State-dict save/load reproduces the exact continuation trajectory.""" + torch.manual_seed(42) + np.random.seed(42) + net = CELO2MLP().to(DEVICE) # shared meta-model (random but fixed) + + model = SmallNet().to(DEVICE) + optimizer = CELO2_naive(model.parameters(), num_steps=50, peak_lr=1e-2, network=net) + + x = torch.randn(32, 10, device=DEVICE) + y = torch.randn(32, 1, device=DEVICE) + + n_steps = 5 + for _ in range(n_steps): + optimizer.zero_grad() + F.mse_loss(model(x), y).backward() + optimizer.step() + + ckpt = {"model": model.state_dict(), "optimizer": optimizer.state_dict()} + path = tmp_path / "celo2_ckpt.pt" + torch.save(ckpt, path) + + # Continue with the original. + orig_final = [] + for _ in range(n_steps): + optimizer.zero_grad() + F.mse_loss(model(x), y).backward() + optimizer.step() + orig_final = [p.clone().detach() for p in model.parameters()] + + # Restart from checkpoint and continue. + loaded_model = SmallNet().to(DEVICE) + loaded_opt = CELO2_naive( + loaded_model.parameters(), num_steps=50, peak_lr=1e-2, network=net + ) + ckpt = torch.load(path) + loaded_model.load_state_dict(ckpt["model"]) + loaded_opt.load_state_dict(ckpt["optimizer"]) + for _ in range(n_steps): + loaded_opt.zero_grad() + F.mse_loss(loaded_model(x), y).backward() + loaded_opt.step() + loaded_final = [p.clone().detach() for p in loaded_model.parameters()] + + for a, b in zip(orig_final, loaded_final): + assert torch.max(torch.abs(a - b)).item() < 1e-5 + + +@pytest.mark.skipif( + not os.path.isfile(CELO2_OPTAX), reason="scaling_l2o JAX source not available" +) +@pytest.mark.parametrize("ortho", [True, False]) +@pytest.mark.parametrize("shape", [(7, 5), (4, 9), (6, 6)]) +def test_jax_numerical_alignment(ortho, shape): + """The 2D CELO2 update matches the reference JAX implementation bitwise-ish.""" + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + spec = importlib.util.spec_from_file_location("celo2_optax_ref", CELO2_OPTAX) + ref = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ref) + + r, c = shape + tr = ref.Celo2Transformation(orthogonalize=ortho) + theta = tr.init_meta_params(jax.random.PRNGKey(7)) + tr_t = ref.Celo2Transformation(theta=theta, orthogonalize=ortho) + + params = {"w": jnp.asarray(np.random.randn(r, c).astype(np.float32))} + jstate = tr_t.init(params) + + # Convert theta -> torch CELO2MLP. + leaves = {} + + def flat(t, pre=""): + if isinstance(t, dict): + for k, v in t.items(): + flat(v, f"{pre}/{k}" if pre else k) + else: + leaves[pre.split("/")[-1]] = np.asarray(t) + + flat(theta) + model = CELO2MLP() + sd = model.state_dict() + for i in range(len(model.w0)): + sd[f"w0.{i}"] = torch.tensor(leaves[f"w0__{i}"], dtype=torch.float32) + sd["b0"] = torch.tensor(leaves["b0"], dtype=torch.float32) + sd["dense_w.0"] = torch.tensor(leaves["w1"], dtype=torch.float32) + sd["dense_b.0"] = torch.tensor(leaves["b1"], dtype=torch.float32) + sd["dense_w.1"] = torch.tensor(leaves["w2"], dtype=torch.float32) + sd["dense_b.1"] = torch.tensor(leaves["b2"], dtype=torch.float32) + model.load_state_dict(sd) + + p = torch.nn.Parameter(torch.tensor(np.asarray(params["w"]), dtype=torch.float32)) + opt = CELO2_naive([p], num_steps=100, orthogonalize=ortho, network=model) + opt.device = "cpu" + opt.initial_momentum_decays = opt.initial_momentum_decays.cpu() + opt.initial_rms_decays = opt.initial_rms_decays.cpu() + opt.initial_adafactor_decays = opt.initial_adafactor_decays.cpu() + group = opt.param_groups[0] + + state = opt.state[p] + state["mom"] = torch.zeros((r, c, 3)) + state["rms"] = torch.zeros((r, c, 1)) + d1, d0 = factored_dims((r, c)) + full = (r, c, 3) + state["fac_vec_row"] = torch.zeros(tuple(d for i, d in enumerate(full) if i != d0)) + state["fac_vec_col"] = torch.zeros(tuple(d for i, d in enumerate(full) if i != d1)) + state["fac_vec_v"] = torch.empty(0) + + for _ in range(5): + g = np.random.randn(r, c).astype(np.float32) + jstep, jstate = tr_t.update({"w": jnp.asarray(g)}, jstate, params) + tstep = opt._celo2_step(p, torch.tensor(g), state, group).detach().numpy() + assert np.max(np.abs(np.asarray(jstep["w"]) - tstep)) < 1e-4 From b6a45e4ba1ce674b71acdd26bb6e3b216c119aaa Mon Sep 17 00:00:00 2001 From: Xiaolong Huang Date: Tue, 9 Jun 2026 17:45:45 -0400 Subject: [PATCH 2/7] Add ELO (Adafactor-MLP) learned optimizer (naive PyTorch) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port the inference-time forward pass of the ELO learned optimizer from the original JAX implementation (ELO_AdafacMLPLOpt) into pure PyTorch. At inference the ELO expert mechanism is disabled, so the update reduces to the Adafactor-MLP forward — identical features and meta-model (MetaMLP, 39 inputs / 2 outputs) to AdafacLO. ELO differs only in using raw accumulator decays, a warmup-then-constant (optionally cosine) LR schedule, and the update rule p -= lr * (dir*exp(mag*exp_mult) + wd*p). - pylo/optim/ELO_naive.py: ELO_naive optimizer (reuses MetaMLP and the AdafacLO feature helpers); default meta-model DiamondXL/elo. - scripts/convert_elo_checkpoint.py: convert a JAX/Haiku ELO theta into a MetaMLP state_dict (transposes the dense weights). - tests/test_elo.py: step/update, state-dict resume, and a JAX numerical-alignment test (matches the reference to ~1.5e-8; auto-skips when the JAX source is unavailable). - Register ELO_naive / ELO in the pylo and pylo.optim inits. Co-Authored-By: Claude Opus 4.8 (1M context) --- pylo/__init__.py | 2 + pylo/optim/ELO_naive.py | 277 ++++++++++++++++++++++++++++++ pylo/optim/__init__.py | 5 +- scripts/convert_elo_checkpoint.py | 91 ++++++++++ tests/test_elo.py | 148 ++++++++++++++++ 5 files changed, 522 insertions(+), 1 deletion(-) create mode 100644 pylo/optim/ELO_naive.py create mode 100644 scripts/convert_elo_checkpoint.py create mode 100644 tests/test_elo.py diff --git a/pylo/__init__.py b/pylo/__init__.py index 82bb494..8ec79ee 100644 --- a/pylo/__init__.py +++ b/pylo/__init__.py @@ -11,8 +11,10 @@ "VeLO_naive", "CELO2_naive", "ELO_CELO2_naive", + "ELO_naive", "CELO2", "ELO_CELO2", + "ELO", # Default aliases (re-exported from pylo.optim). These resolve to the # CUDA implementations when available, falling back to the naive ones # otherwise, so downstream code can simply `from pylo import VeLO`. diff --git a/pylo/optim/ELO_naive.py b/pylo/optim/ELO_naive.py new file mode 100644 index 0000000..7d04d7d --- /dev/null +++ b/pylo/optim/ELO_naive.py @@ -0,0 +1,277 @@ +"""ELO_naive: a PyTorch port of the ELO (Adafactor-MLP) learned optimizer. + +ELO trains an Adafactor-style MLP learned optimizer with an auxiliary "expert" +mechanism. The expert trajectory and the imitation (IMT) losses are only active +during meta-training; at inference (``meta_train=False``) the parameter update +reduces to the plain Adafactor-MLP forward pass. This file ports that +inference-time forward pass. + +The per-parameter features and the meta-model are identical to +:class:`pylo.optim.AdafacLO_naive.AdafacLO_naive` / :class:`pylo.models.Meta_MLP.MetaMLP` +(39 input features, 2 outputs). ELO differs only in: + + * raw accumulator decays (no decay reparameterization), + * a warmup-then-constant (optionally cosine) learning-rate schedule, and + * the update rule ``p -= scheduled_lr * (direction * exp(magnitude * exp_mult) + + weight_decay * p)`` (the schedule plays the role of ``step_mult``). + +Reference: ``scaling_l2o/src/learned_optimizers/elo_adfac_mlp_lopt.py``. +""" + +from typing import Optional + +import numpy as np +import torch +from torch.optim import Optimizer + +from pylo.models.Meta_MLP import MetaMLP +from pylo.optim.AdafacLO_naive import ( + factored_dims, + init_factors, + safe_rsqrt, + second_moment_normalizer, + tanh_embedding, + update_factors, +) + + +class ELO_naive(Optimizer): + """Pure-PyTorch ELO (Adafactor-MLP) learned optimizer. + + Args: + params: Iterable of parameters or ``param_groups``. + num_steps: Total number of inner training steps (defines the warmup / + cosine schedule). Required. + exp_mult, step_mult: Magnitude exponent multiplier and base step size + (``step_mult`` is the post-warmup learning rate). + init_lr, warmup_fraction, warmup_steps: Linear warmup from ``init_lr`` to + ``step_mult``. ``warmup_fraction`` (fraction of ``num_steps``) takes + priority over ``warmup_steps`` when > 0. + use_lo_cosine_scheduler, step_mult_min: If enabled, cosine-decay the + post-warmup learning rate from ``step_mult`` down to ``step_mult_min``. + weight_decay: Decoupled weight decay (scaled by the schedule). + hidden_size, 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, initial_rms_decays, initial_adafactor_decays: + Raw accumulator decays (used directly, no reparameterization). + clip_grad, 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 :class:`MetaMLP` (overrides the above). + """ + + def __init__( + self, + params, + num_steps, + exp_mult=0.001, + step_mult=0.001, + init_lr=0.0, + warmup_fraction=0.05, + warmup_steps=0, + use_lo_cosine_scheduler=False, + step_mult_min=1e-4, + 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: Optional[str] = "DiamondXL/elo", + checkpoint_path: Optional[str] = None, + network: Optional[MetaMLP] = None, + ): + if num_steps is None: + raise ValueError("ELO_naive requires num_steps for the LR schedule.") + + self.device = "cuda" if torch.cuda.is_available() else "cpu" + + defaults = dict( + num_steps=num_steps, + exp_mult=exp_mult, + step_mult=step_mult, + init_lr=init_lr, + warmup_fraction=warmup_fraction, + warmup_steps=warmup_steps, + use_lo_cosine_scheduler=use_lo_cosine_scheduler, + step_mult_min=step_mult_min, + weight_decay=weight_decay, + clip_grad=clip_grad, + clip_norm=clip_norm, + ) + super().__init__(params, defaults) + + self.initial_momentum_decays = torch.tensor( + initial_momentum_decays, dtype=torch.float32, device=self.device + ) + self.initial_rms_decays = torch.tensor( + initial_rms_decays, dtype=torch.float32, device=self.device + ) + self.initial_adafactor_decays = torch.tensor( + initial_adafactor_decays, dtype=torch.float32, device=self.device + ) + + # Precedence: explicit network > local checkpoint > HuggingFace Hub. + if network is not None: + self.network = network + elif checkpoint_path is not None: + self.network = MetaMLP( + input_size=39, hidden_size=hidden_size, hidden_layers=hidden_layers + ) + self.network.load_state_dict(torch.load(checkpoint_path, map_location="cpu")) + elif hf_key is not None: + self.network = MetaMLP.from_pretrained(hf_key) + else: + self.network = MetaMLP( + input_size=39, hidden_size=hidden_size, hidden_layers=hidden_layers + ) + self.network = self.network.to(self.device) + + def _scheduled_lr(self, iteration, group): + """Warmup → constant (or cosine) schedule, matching the JAX version.""" + num_steps = group["num_steps"] + step_mult = group["step_mult"] + init_lr = group["init_lr"] + if group["warmup_fraction"] > 0: + warmup_n = group["warmup_fraction"] * num_steps + else: + warmup_n = group["warmup_steps"] + warmup_n = max(float(warmup_n), 1.0) + it = float(iteration) + + if group["use_lo_cosine_scheduler"]: + frac = min(max(it / max(num_steps - 1, 1), 0.0), 1.0) + step_mult_min = group["step_mult_min"] + base = step_mult_min + (step_mult - step_mult_min) * 0.5 * ( + 1.0 + np.cos(np.pi * frac) + ) + else: + base = step_mult + + warmup_lr = init_lr + (step_mult - init_lr) * min(it / warmup_n, 1.0) + return warmup_lr if it < warmup_n else base + + def _global_grad_scale(self, clip_norm): + total = torch.zeros((), device=self.device) + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + total = total + p.grad.detach().float().pow(2).sum() + denom = torch.clamp(torch.sqrt(total), min=clip_norm) + return (clip_norm / denom).item() if denom > 0 else 1.0 + + @torch.no_grad() + def step(self, loss=None): + grad_scales = {} + for group in self.param_groups: + if group["clip_grad"]: + grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) + + for group in self.param_groups: + iteration = group.get("step", 0) # 0-indexed, as in the JAX update + group["step"] = iteration + 1 + exp_mult = group["exp_mult"] + weight_decay = group["weight_decay"] + scheduled_lr = self._scheduled_lr(iteration, group) + scale = grad_scales.get(id(group), 1.0) + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if scale != 1.0: + grad = grad * scale + grad = torch.nan_to_num(grad) + p_shape = p.shape + + beta_m = self.initial_momentum_decays + beta_rms = self.initial_rms_decays + beta_adafactor = self.initial_adafactor_decays + + state = self.state[p] + if len(state) == 0: + state["mom"] = torch.zeros(p_shape + (beta_m.shape[-1],)).to(self.device) + state["rms"] = torch.zeros(p_shape + (beta_rms.shape[-1],)).to(self.device) + fr, fc, fv = init_factors(p) + state["fac_vec_row"] = fr.to(self.device) + state["fac_vec_col"] = fc.to(self.device) + state["fac_vec_v"] = fv.to(self.device) + + batch_p = p.unsqueeze(-1) + batch_g = grad.unsqueeze(-1) + + training_step_feature = tanh_embedding(iteration).to(self.device) + axis = list(range(len(p_shape))) + for _ in axis: + beta_m = beta_m[None, ...] + beta_rms = beta_rms[None, ...] + beta_adafactor = beta_adafactor[None, ...] + training_step_feature = training_step_feature[None, ...] + training_step_feature = training_step_feature.repeat(p_shape + (1,)) + + mom = state["mom"] + rms = state["rms"] + mom.mul_(beta_m).add_((1 - beta_m) * batch_g) + rms.mul_(beta_rms).add_((1 - beta_rms) * (batch_g ** 2)) + ( + state["fac_vec_col"], + state["fac_vec_row"], + state["fac_vec_v"], + fac_g, + ) = update_factors( + state["fac_vec_col"], + state["fac_vec_row"], + state["fac_vec_v"], + batch_g, + p_shape, + beta_adafactor, + ) + fac_vec_col = state["fac_vec_col"] + fac_vec_row = state["fac_vec_row"] + fac_vec_v = state["fac_vec_v"] + + rsqrt = torch.rsqrt(rms + 1e-6) + inps = [batch_g, batch_p, mom, rms, mom * rsqrt, rsqrt, fac_g] + + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + rp_row = [1] * (1 + len(p_shape)) + rp_col = [1] * (1 + len(p_shape)) + rp_row[d0] = p_shape[d0] + rp_col[d1] = p_shape[d1] + row_feat = fac_vec_row.unsqueeze(d0).repeat(rp_row) + col_feat = fac_vec_col.unsqueeze(d1).repeat(rp_col) + inps.extend([ + row_feat, + col_feat, + torch.rsqrt(row_feat + 1e-8), + torch.rsqrt(col_feat + 1e-8), + ]) + reduced_d1 = d1 - 1 if d1 > d0 else d1 + row_col_mean = fac_vec_row.mean(dim=reduced_d1, keepdim=True) + row_factor = safe_rsqrt(fac_vec_row / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(fac_vec_col) + fac_mom_mult = mom * row_factor.unsqueeze(d0) * col_factor.unsqueeze(d1) + inps.append(fac_mom_mult) + else: + inps.extend([ + fac_vec_v, + fac_vec_v, + torch.rsqrt(fac_vec_v + 1e-8), + torch.rsqrt(fac_vec_v + 1e-8), + ]) + fac_mom_mult = mom * torch.pow(fac_vec_v + 1e-6, -0.5) + inps.append(fac_mom_mult) + + inps = torch.cat(inps, dim=-1) + inps = second_moment_normalizer(inps, axis=axis) + inp_stack = torch.cat([inps, training_step_feature], dim=-1) + + direction, magnitude = self.network(inp_stack).split(1, dim=-1) + step = (direction * torch.exp(magnitude * exp_mult)).squeeze(-1) + p.add_(step + weight_decay * p, alpha=-scheduled_lr) + return loss diff --git a/pylo/optim/__init__.py b/pylo/optim/__init__.py index 049d35a..5c95a4b 100644 --- a/pylo/optim/__init__.py +++ b/pylo/optim/__init__.py @@ -3,6 +3,7 @@ from pylo.optim.Velo_naive import VeLO_naive from pylo.optim.CELO2_naive import CELO2_naive from pylo.optim.ELO_CELO2_naive import ELO_CELO2_naive +from pylo.optim.ELO_naive import ELO_naive # Initialize with optimizers we know we can import __all__ = [ @@ -11,6 +12,7 @@ "VeLO_naive", "CELO2_naive", "ELO_CELO2_naive", + "ELO_naive", ] # CELO2 / ELO-CELO2 currently ship a naive (pure-PyTorch) implementation only. @@ -18,7 +20,8 @@ # overridden in the try-block below, mirroring VeLO / AdafacLO. CELO2 = CELO2_naive ELO_CELO2 = ELO_CELO2_naive -__all__.extend(["CELO2", "ELO_CELO2"]) +ELO = ELO_naive +__all__.extend(["CELO2", "ELO_CELO2", "ELO"]) # Try to import CUDA-based optimizers try: diff --git a/scripts/convert_elo_checkpoint.py b/scripts/convert_elo_checkpoint.py new file mode 100644 index 0000000..663ee11 --- /dev/null +++ b/scripts/convert_elo_checkpoint.py @@ -0,0 +1,91 @@ +"""Convert a JAX/Haiku ELO (Adafactor-MLP) ``theta`` checkpoint to a PyTorch +:class:`pylo.models.Meta_MLP.MetaMLP` state_dict. + +The ELO meta-parameters are stored as +``{"nn": {"": {"w0": (39, 32), "b0": (32,), "w1": (32, 32), "b1": ..., +"w2": (32, 2), "b2": (2,)}}}`` — a standard dense MLP. Haiku stores linear +weights as ``(in, out)`` and applies ``x @ w``; ``torch.nn.Linear`` stores +``(out, in)`` and applies ``x @ w.T``, so each weight is transposed. + +Run this in an environment that can unpickle the checkpoint (e.g. the +``scaling_l2o`` JAX environment); only ``numpy`` and ``torch`` are required. + +Usage: + python scripts/convert_elo_checkpoint.py \ + --input /path/to/global_step100000.pickle \ + --output /path/to/elo_theta.pt +""" + +import argparse +import pickle + +import numpy as np +import torch + +from pylo.models.Meta_MLP import MetaMLP + + +def _load_raw(path): + if str(path).endswith(".pickle"): + with open(path, "rb") as f: + return pickle.load(f) + import flax # noqa: local import; only needed for the msgpack branch + + with open(path, "rb") as f: + return flax.serialization.msgpack_restore(f.read()) + + +def _flatten(tree, prefix=""): + if isinstance(tree, dict): + for k, v in tree.items(): + yield from _flatten(v, f"{prefix}/{k}" if prefix else str(k)) + else: + yield prefix, np.asarray(tree) + + +def convert(raw, input_size=39, hidden_size=32, hidden_layers=1): + leaves = {name.split("/")[-1]: arr for name, arr in _flatten(raw)} + + model = MetaMLP( + input_size=input_size, hidden_size=hidden_size, hidden_layers=hidden_layers + ) + sd = model.state_dict() + + # MetaMLP dense layers in forward order: input, linear_0..linear_{n-1}, output. + targets = ["network.input"] + [ + f"network.linear_{i}" for i in range(hidden_layers) + ] + ["network.output"] + + for j, dst in enumerate(targets): + w = leaves[f"w{j}"] # (in, out) -> nn.Linear (out, in) + b = leaves[f"b{j}"] + sd[f"{dst}.weight"] = torch.tensor(w, dtype=torch.float32).t().contiguous() + sd[f"{dst}.bias"] = torch.tensor(b, dtype=torch.float32) + + model.load_state_dict(sd) + return model + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--input", required=True) + ap.add_argument("--output", required=True) + ap.add_argument("--input-size", type=int, default=39) + ap.add_argument("--hidden-size", type=int, default=32) + ap.add_argument("--hidden-layers", type=int, default=1) + args = ap.parse_args() + + raw = _load_raw(args.input) + model = convert( + raw, + input_size=args.input_size, + hidden_size=args.hidden_size, + hidden_layers=args.hidden_layers, + ) + torch.save(model.state_dict(), args.output) + n = sum(p.numel() for p in model.parameters()) + print(f"Saved MetaMLP state_dict ({n} params) to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_elo.py b/tests/test_elo.py new file mode 100644 index 0000000..3f12784 --- /dev/null +++ b/tests/test_elo.py @@ -0,0 +1,148 @@ +"""Tests for the ELO (Adafactor-MLP) learned optimizer.""" + +import importlib.util +import os + +import numpy as np +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from pylo.models.Meta_MLP import MetaMLP +from pylo.optim import ELO_naive + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +SCALING_L2O = "/home/mila/h/huangx/scaling_l2o" +ELO_SRC = os.path.join(SCALING_L2O, "src/learned_optimizers/elo_adfac_mlp_lopt.py") + + +class SmallNet(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(10, 8) + self.fc2 = nn.Linear(8, 1) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +def test_step_runs_and_updates(): + torch.manual_seed(0) + model = SmallNet().to(DEVICE) + before = [p.clone() for p in model.parameters()] + optimizer = ELO_naive(model.parameters(), num_steps=50, step_mult=1e-2) + x = torch.randn(16, 10, device=DEVICE) + y = torch.randn(16, 1, device=DEVICE) + for _ in range(6): + optimizer.zero_grad() + F.mse_loss(model(x), y).backward() + optimizer.step() + for p in model.parameters(): + assert torch.isfinite(p).all() + assert any(not torch.allclose(b, p) for b, p in zip(before, model.parameters())) + + +def test_elo_resume(tmp_path): + torch.manual_seed(42) + net = MetaMLP(input_size=39, hidden_size=32, hidden_layers=1).to(DEVICE) + + model = SmallNet().to(DEVICE) + optimizer = ELO_naive(model.parameters(), num_steps=50, step_mult=1e-2, network=net) + x = torch.randn(32, 10, device=DEVICE) + y = torch.randn(32, 1, device=DEVICE) + + n_steps = 5 + for _ in range(n_steps): + optimizer.zero_grad() + F.mse_loss(model(x), y).backward() + optimizer.step() + + ckpt = {"model": model.state_dict(), "optimizer": optimizer.state_dict()} + path = tmp_path / "elo_ckpt.pt" + torch.save(ckpt, path) + + for _ in range(n_steps): + optimizer.zero_grad() + F.mse_loss(model(x), y).backward() + optimizer.step() + orig_final = [p.clone().detach() for p in model.parameters()] + + loaded_model = SmallNet().to(DEVICE) + loaded_opt = ELO_naive(loaded_model.parameters(), num_steps=50, step_mult=1e-2, network=net) + ckpt = torch.load(path) + loaded_model.load_state_dict(ckpt["model"]) + loaded_opt.load_state_dict(ckpt["optimizer"]) + for _ in range(n_steps): + loaded_opt.zero_grad() + F.mse_loss(loaded_model(x), y).backward() + loaded_opt.step() + loaded_final = [p.clone().detach() for p in loaded_model.parameters()] + + for a, b in zip(orig_final, loaded_final): + assert torch.max(torch.abs(a - b)).item() < 1e-5 + + +@pytest.mark.skipif( + not os.path.isfile(ELO_SRC), reason="scaling_l2o JAX source not available" +) +def test_jax_numerical_alignment(): + """ELO_naive matches the reference JAX ELO_AdafacMLPLOpt (meta_train=False).""" + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + spec = importlib.util.spec_from_file_location("elo_adfac_ref", ELO_SRC) + elo = importlib.util.module_from_spec(spec) + spec.loader.exec_module(elo) + + N = 50 + lopt = elo.ELO_AdafacMLPLOpt( + meta_train=False, exp_mult=0.001, step_mult=0.001, init_lr=0.0, + warmup_fraction=0.05, weight_decay=0.0, hidden_size=32, hidden_layers=2, + 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, + use_lo_cosine_scheduler=False, + ) + theta = lopt.init(jax.random.PRNGKey(0)) + jopt = lopt.opt_fn(theta) + + rng = np.random.RandomState(0) + W = rng.randn(7, 5).astype(np.float32) + B = rng.randn(5).astype(np.float32) + jstate = jopt.init({"w": jnp.asarray(W), "b": jnp.asarray(B)}, num_steps=N) + + # Convert theta -> torch MetaMLP (transpose haiku (in,out) -> nn.Linear (out,in)). + leaves = {} + + def flat(t, pre=""): + if isinstance(t, dict): + for k, v in t.items(): + flat(v, f"{pre}/{k}" if pre else k) + else: + leaves[pre.split("/")[-1]] = np.asarray(t) + + flat(theta) + net = MetaMLP(input_size=39, hidden_size=32, hidden_layers=1) + sd = net.state_dict() + for j, dst in enumerate(["network.input", "network.linear_0", "network.output"]): + sd[f"{dst}.weight"] = torch.tensor(leaves[f"w{j}"]).t().contiguous() + sd[f"{dst}.bias"] = torch.tensor(leaves[f"b{j}"]) + net.load_state_dict(sd) + + pw = torch.nn.Parameter(torch.tensor(W)) + pb = torch.nn.Parameter(torch.tensor(B)) + topt = ELO_naive([pw, pb], num_steps=N, network=net) + topt.device = "cpu" + for a in ("initial_momentum_decays", "initial_rms_decays", "initial_adafactor_decays"): + setattr(topt, a, getattr(topt, a).cpu()) + + gr = np.random.RandomState(1) + for _ in range(8): + gw = gr.randn(7, 5).astype(np.float32) + gb = gr.randn(5).astype(np.float32) + jstate = jopt.update(jstate, {"w": jnp.asarray(gw), "b": jnp.asarray(gb)}, loss=jnp.asarray(0.0)) + pw.grad = torch.tensor(gw) + pb.grad = torch.tensor(gb) + topt.step() + assert np.max(np.abs(np.asarray(jstate.params["w"]) - pw.detach().numpy())) < 1e-5 + assert np.max(np.abs(np.asarray(jstate.params["b"]) - pb.detach().numpy())) < 1e-5 From 725e3ba1cbe230c91adfa0e3ae43aa73a458500a Mon Sep 17 00:00:00 2001 From: Xiaolong Huang Date: Tue, 9 Jun 2026 18:09:36 -0400 Subject: [PATCH 3/7] Align CELO2 LR-schedule indexing with Celo2LOpt; add end-to-end JAX tests End-to-end comparison against the real JAX optimizers revealed that the standalone Celo2LOpt drives its LR schedule through an optax chain whose step count starts at 0 (so the first update uses schedule(0)), whereas ELO_Celo2LOpt evaluates the schedule at iteration+1. CELO2_naive was 1-indexed (matching ELO-CELO2), leaving a ~1.8e-3 warmup-phase discrepancy vs Celo2LOpt. - Add a per-class LR-schedule offset: CELO2_naive is 0-indexed (matches Celo2LOpt), ELO_CELO2_naive overrides to 1-indexed (matches ELO_Celo2LOpt). AdamW bias correction stays 1-indexed in both. - Add test_jax_end_to_end_alignment: drives the real Celo2LOpt / ELO_Celo2LOpt over a multi-step trajectory with a 2D weight + 1D bias, nonzero weight decay and enabled gradient clipping, exercising the full step() (1D AdamW, schedule, weight decay, global-norm clipping). Both match the reference to ~6e-8 (was only the 2D core verified before). Co-Authored-By: Claude Opus 4.8 (1M context) --- pylo/optim/CELO2_naive.py | 11 ++- pylo/optim/ELO_CELO2_naive.py | 3 + tests/test_celo2.py | 132 +++++++++++++++++++++++++++++++++- 3 files changed, 144 insertions(+), 2 deletions(-) diff --git a/pylo/optim/CELO2_naive.py b/pylo/optim/CELO2_naive.py index 69d2bf9..19db029 100644 --- a/pylo/optim/CELO2_naive.py +++ b/pylo/optim/CELO2_naive.py @@ -192,6 +192,13 @@ def __init__( ) self.network = self.network.to(self.device) + # LR-schedule iteration offset. The standalone Celo2LOpt drives its + # schedule through an optax chain whose step count starts at 0, so the + # first update uses schedule(0); ELO-CELO2 instead evaluates the + # schedule at iteration+1 (1-indexed). CELO2 therefore offsets the + # 1-indexed step counter by 1; ELO_CELO2_naive overrides this to 0. + self._lr_offset = 1 + # ---------------------------------------------------------------- helpers def _lr_schedule(self, step, group): """Warmup + cosine decay learning rate (matches the JAX schedule).""" @@ -333,7 +340,9 @@ def step(self, loss=None): for group in self.param_groups: group["step"] = group.get("step", 0) + 1 t = group["step"] - lr = self._lr_schedule(t, group) + # Schedule index follows the reference (0-indexed for CELO2 via + # _lr_offset=1); AdamW bias correction stays 1-indexed (uses t). + lr = self._lr_schedule(t - self._lr_offset, group) adam_lr = group["adam_lr_mult"] * lr beta1 = float(self.initial_momentum_decays[0]) beta2 = float(self.initial_rms_decays[-1]) diff --git a/pylo/optim/ELO_CELO2_naive.py b/pylo/optim/ELO_CELO2_naive.py index 89b9f7c..363f525 100644 --- a/pylo/optim/ELO_CELO2_naive.py +++ b/pylo/optim/ELO_CELO2_naive.py @@ -88,3 +88,6 @@ def __init__( checkpoint_path=checkpoint_path, network=network, ) + # ELO-CELO2 evaluates the LR schedule at iteration+1 (1-indexed), + # unlike the standalone CELO2 which is 0-indexed via the optax chain. + self._lr_offset = 0 diff --git a/tests/test_celo2.py b/tests/test_celo2.py index 50bfbba..13e78f8 100644 --- a/tests/test_celo2.py +++ b/tests/test_celo2.py @@ -1,5 +1,6 @@ """Tests for the CELO2 / ELO-CELO2 learned optimizers.""" +import functools import importlib.util import os @@ -15,7 +16,62 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") SCALING_L2O = "/home/mila/h/huangx/scaling_l2o" -CELO2_OPTAX = os.path.join(SCALING_L2O, "src/learned_optimizers/celo2_optax.py") +LO_DIR = os.path.join(SCALING_L2O, "src/learned_optimizers") +CELO2_OPTAX = os.path.join(LO_DIR, "celo2_optax.py") + + +@functools.lru_cache(maxsize=1) +def _load_jax_lopts(): + """Import the real JAX Celo2LOpt / ELO_Celo2LOpt without their package init. + + The ``learned_optimizers`` package __init__ pulls in heavy modules; here we + register a throwaway package so the modules' relative imports resolve, then + load only the files we need. Cached so the ``@gin.configurable`` classes are + registered exactly once per process (re-loading would raise a gin error). + """ + import sys + import types + + pkg = types.ModuleType("_lo_ref") + pkg.__path__ = [LO_DIR] + sys.modules["_lo_ref"] = pkg + + def load(name, fn): + spec = importlib.util.spec_from_file_location( + f"_lo_ref.{name}", os.path.join(LO_DIR, fn) + ) + m = importlib.util.module_from_spec(spec) + sys.modules[f"_lo_ref.{name}"] = m + spec.loader.exec_module(m) + return m + + load("celo2_optax", "celo2_optax.py") + load("elo_adfac_mlp_lopt", "elo_adfac_mlp_lopt.py") + return load("celo2_lopt", "celo2_lopt.py"), load("elo_celo2", "elo_celo2.py") + + +def _theta_to_celo2mlp(theta): + leaves = {} + + def flat(t, pre=""): + if isinstance(t, dict): + for k, v in t.items(): + flat(v, f"{pre}/{k}" if pre else k) + else: + leaves[pre.split("/")[-1]] = np.asarray(t) + + flat(theta) + m = CELO2MLP() + sd = m.state_dict() + for i in range(len(m.w0)): + sd[f"w0.{i}"] = torch.tensor(leaves[f"w0__{i}"], dtype=torch.float32) + sd["b0"] = torch.tensor(leaves["b0"], dtype=torch.float32) + sd["dense_w.0"] = torch.tensor(leaves["w1"], dtype=torch.float32) + sd["dense_b.0"] = torch.tensor(leaves["b1"], dtype=torch.float32) + sd["dense_w.1"] = torch.tensor(leaves["w2"], dtype=torch.float32) + sd["dense_b.1"] = torch.tensor(leaves["b2"], dtype=torch.float32) + m.load_state_dict(sd) + return m class SmallNet(nn.Module): @@ -177,3 +233,77 @@ def flat(t, pre=""): jstep, jstate = tr_t.update({"w": jnp.asarray(g)}, jstate, params) tstep = opt._celo2_step(p, torch.tensor(g), state, group).detach().numpy() assert np.max(np.abs(np.asarray(jstep["w"]) - tstep)) < 1e-4 + + +@pytest.mark.skipif( + not os.path.isfile(CELO2_OPTAX), reason="scaling_l2o JAX source not available" +) +@pytest.mark.parametrize("which", ["celo2", "elo_celo2"]) +def test_jax_end_to_end_alignment(which): + """The FULL step() (1D AdamW + warmup/cosine LR + weight decay + global-norm + clipping) matches the reference JAX optimizer over a multi-step trajectory. + + Uses a model with a 2D weight and a 1D bias, nonzero weight decay, and + enabled gradient clipping with large grads, so every code path is exercised. + """ + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + celo2_lopt, elo_celo2 = _load_jax_lopts() + + N = 60 + cfg = dict( + init_lr=0.0, peak_lr=1e-3, warmup_fraction=0.05, end_lr=0.0, + weight_decay=0.1, adam_lr_mult=1.0, use_adamw_for_1d=True, + clip_grad=True, clip_norm=1.0, + ) + spec = importlib.util.spec_from_file_location("celo2_optax_ref", CELO2_OPTAX) + ref = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ref) + theta = ref.Celo2Transformation(orthogonalize=True).init_meta_params( + jax.random.PRNGKey(0) + ) + net = _theta_to_celo2mlp(theta) + + if which == "celo2": + jlopt = celo2_lopt.Celo2LOpt( + orthogonalize=True, adam_beta1=0.9, adam_beta2=0.95, **cfg + ) + topt_cls = CELO2_naive + else: + jlopt = elo_celo2.ELO_Celo2LOpt( + 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), + meta_train=False, **cfg, + ) + topt_cls = ELO_CELO2_naive + + rng = np.random.RandomState(0) + W = rng.randn(7, 5).astype(np.float32) + B = rng.randn(5).astype(np.float32) + jopt = jlopt.opt_fn(theta) + jstate = jopt.init({"w": jnp.asarray(W), "b": jnp.asarray(B)}, num_steps=N) + + pw = torch.nn.Parameter(torch.tensor(W)) + pb = torch.nn.Parameter(torch.tensor(B)) + topt = topt_cls( + [pw, pb], num_steps=N, peak_lr=1e-3, warmup_fraction=0.05, end_lr=0.0, + weight_decay=0.1, clip_grad=True, clip_norm=1.0, network=net, + ) + topt.device = "cpu" + for a in ("initial_momentum_decays", "initial_rms_decays", "initial_adafactor_decays"): + setattr(topt, a, getattr(topt, a).cpu()) + + gr = np.random.RandomState(1) + for _ in range(12): + gw = (gr.randn(7, 5) * 3).astype(np.float32) + gb = (gr.randn(5) * 3).astype(np.float32) + out = jopt.update(jstate, {"w": jnp.asarray(gw), "b": jnp.asarray(gb)}, loss=jnp.asarray(0.0)) + jstate = out[0] if isinstance(out, tuple) else out + pw.grad = torch.tensor(gw) + pb.grad = torch.tensor(gb) + topt.step() + assert np.max(np.abs(np.asarray(jstate.params["w"]) - pw.detach().numpy())) < 1e-4 + assert np.max(np.abs(np.asarray(jstate.params["b"]) - pb.detach().numpy())) < 1e-4 From c5dec39be82193165bd389f960fe98a0a16969d9 Mon Sep 17 00:00:00 2001 From: Xiaolong Huang Date: Thu, 18 Jun 2026 17:12:12 -0400 Subject: [PATCH 4/7] README: merge ELO additions onto xiao_elo (non-destructive) Keep the existing VeLO_CUDA Quick Start intact and append an ELO-CELO2 example plus an "ELO series" entry with the arXiv link. Pure additions (no content removed from the xiao_elo README). Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/README.md b/README.md index 513178d..b866170 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,9 @@ PyLO provides efficient PyTorch implementations of cutting-edge learned optimize - **PyTorch-native API** designed for simplicity and familiarity - **Hugging Face integration** for sharing and loading meta-models +Learned optimizers: +* ELO series: [https://arxiv.org/abs/2506.10315](https://arxiv.org/abs/2506.10315) + # Installation ### Via URL (slow, no Kernels) @@ -90,6 +93,26 @@ for epoch in range(10): optimizer.step(loss) # pass the loss ``` +Taking `ELO-CELO2` (the strongest LO) for example. + +```python +import torch +from pylo.optim import ELO_CELO2_CUDA + +model = torch.nn.Linear(10, 2) + +num_steps = 1000 # total optimization steps (used for the LR schedule) + +# Meta-learned weights download automatically from the Hugging Face Hub on first use. +optimizer = ELO_CELO2_CUDA(model.parameters(), num_steps=num_steps, peak_lr=3.16e-4, end_lr=3.16e-5, weight_decay=0.1, adam_lr_mult=20) + +for step in range(num_steps): + optimizer.zero_grad() + loss = loss_fn(model(input), target) + loss.backward() + optimizer.step() +``` + ## Sharing Learned Optimizers PyLO integrates with Hugging Face Hub for sharing meta-trained optimizers: From ea03d85d4f05b454606cd2630f63f8e02257bfaf Mon Sep 17 00:00:00 2001 From: Xiaolong Huang Date: Thu, 18 Jun 2026 17:23:56 -0400 Subject: [PATCH 5/7] Add ELO-CELO2 fused CUDA implementation and wire it in Merge the CUDA path from the ELO-torch line into xiao_elo (additive; the existing naive optimizers and VeLO/AdafacLO kernels are untouched): - pylo/csrc/celo2_kernel.cu: fused feature-construction + split-input MLP forward kernel for CELO2 / ELO-CELO2 - pylo/optim/CELO2_cuda.py, ELO_CELO2_cuda.py: Python wrappers - tests/test_celo2_cuda.py: CUDA-vs-naive numerical alignment tests - setup.py: register the celo2_cuda_kernel CUDAExtension - pylo/optim/__init__.py: override CELO2 / ELO_CELO2 to the CUDA variants when the extension is available, falling back to naive otherwise - .gitignore: ignore *.pickle LO checkpoints Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 4 +- pylo/csrc/celo2_kernel.cu | 357 +++++++++++++++++++++++++++++++++++ pylo/optim/CELO2_cuda.py | 354 ++++++++++++++++++++++++++++++++++ pylo/optim/ELO_CELO2_cuda.py | 89 +++++++++ pylo/optim/__init__.py | 19 +- setup.py | 7 + tests/test_celo2_cuda.py | 129 +++++++++++++ 7 files changed, 955 insertions(+), 4 deletions(-) create mode 100644 pylo/csrc/celo2_kernel.cu create mode 100644 pylo/optim/CELO2_cuda.py create mode 100644 pylo/optim/ELO_CELO2_cuda.py create mode 100644 tests/test_celo2_cuda.py diff --git a/.gitignore b/.gitignore index ede05bf..c870888 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,6 @@ wandb mup _build dist -snippets \ No newline at end of file +snippets +# Learned-optimizer checkpoints (hosted on the Hugging Face Hub) +*.pickle diff --git a/pylo/csrc/celo2_kernel.cu b/pylo/csrc/celo2_kernel.cu new file mode 100644 index 0000000..55f15da --- /dev/null +++ b/pylo/csrc/celo2_kernel.cu @@ -0,0 +1,357 @@ +/** + * CUDA-accelerated per-element forward pass of the CELO2 learned optimizer. + * + * This mirrors the feature construction + split-input MLP of the pure-PyTorch + * `CELO2_naive._celo2_step` (pylo/optim/CELO2_naive.py). The kernel handles, per + * 2D+ parameter element: + * 1. building the 30 CELO2 input features (no time embedding), + * 2. per-channel second-moment normalization, and + * 3. the 30 -> 8 -> 8 -> 3 MLP inference, + * and writes the *raw* per-element step `direction * exp(magnitude * exp_mult)` + * into `step_out`. It intentionally does NOT update the parameter: Newton-Schulz + * orthogonalization, the second second-moment normalization, the rmsmult scale + * and the `p -= lr * (step + wd * p)` apply all happen back in Python, matching + * the naive control flow. + * + * Layout notes (differ from learned_optimizer.cu / AdafacLO): + * - INPUT_DIM = 30 (CELO2_FEATURE_DIMS sums to 30), HIDDEN_DIM = 8, OUTPUT = 3. + * - CELO2MLP stores weights as `(in, out)` and computes `x @ w`, so weights are + * indexed in-major: `w[k * out_dim + j]` (k = input idx, j = output idx). + * - Momentum `m` is leading-decay `(3,)+shape` => `m[idx + k*n_elements]`. + * - `fac_r` (fac_vec_row) and `row_factor` are indexed by the dc-coordinate + * (col_idx); `fac_c` (fac_vec_col) and `col_factor` by the dr-coordinate + * (row_idx) -- identical convention to AdafacLO's populate_vector_inp. + */ + +#include +#include +#include +#include + +#define BLOCK_SIZE 256 +#define NUM_DECAYS 3 +#define INPUT_DIM 30 +#define HIDDEN_DIM 8 +#define OUTPUT_DIM 3 + +// rsqrt(rms + 1e-8) and rsqrt(fac + 1e-8) match the naive feature eps; the +// second-moment normalizer uses 1e-9 (CELO2_naive.second_moment_normalizer). +#define FEAT_EPS 1e-8f +#define NORM_EPS 1e-9f + +__device__ __forceinline__ float relu_(float x) { return fmaxf(x, 0.0f); } + +// Build the 30 CELO2 input features for element `idx` into `f`. +// For 1D (vector_like) params the factored groups (channels 15..29) are zeroed, +// reproducing CELO2MLP.forward being fed only the first 9 feature groups; note +// that CELO2 routes 1D params to AdamW, so this branch is effectively unused. +template +__device__ void populate_celo2_features( + float *f, + const T *g, + const T *p, + const T *m, + const T *v, + const T *fac_r, + const T *fac_c, + const T *row_factor, + const T *col_factor, + const int idx, + const int num_rows, + const int num_cols, + const int row_stride, + const int col_stride, + const int n_elements, + const int vector_like) +{ + const int row_idx = vector_like ? idx : (idx / row_stride) % num_rows; + const int col_idx = vector_like ? idx : (idx / col_stride) % num_cols; + + const float gv = static_cast(g[idx]); + const float pv = static_cast(p[idx]); + const float m0 = static_cast(m[idx]); + const float m1 = static_cast(m[idx + n_elements]); + const float m2 = static_cast(m[idx + 2 * n_elements]); + const float rms = static_cast(v[idx]); + const float rs = __frsqrt_rn(rms + FEAT_EPS); + + // row_factor / fac_r are indexed by col_idx; col_factor / fac_c by row_idx. + const float rf0 = static_cast(row_factor[col_idx]); + const float rf1 = static_cast(row_factor[col_idx + num_cols]); + const float rf2 = static_cast(row_factor[col_idx + 2 * num_cols]); + const float cf0 = vector_like ? 1.0f : static_cast(col_factor[row_idx]); + const float cf1 = vector_like ? 1.0f : static_cast(col_factor[row_idx + num_rows]); + const float cf2 = vector_like ? 1.0f : static_cast(col_factor[row_idx + 2 * num_rows]); + + // Groups 0-8 (15 channels): always present. + f[0] = gv; // g + f[1] = fminf(fmaxf(gv, -0.1f), 0.1f); // clip(g, -0.1, 0.1) + f[2] = pv; // p + f[3] = m0; f[4] = m1; f[5] = m2; // mom[3] + f[6] = rms; // rms + f[7] = m0 * rs; f[8] = m1 * rs; f[9] = m2 * rs; // mom * rsqrt(rms) + f[10] = rs; // rsqrt(rms) + f[11] = gv * rf0 * cf0; // fac_g[3] = g * row_factor * col_factor + f[12] = gv * rf1 * cf1; + f[13] = gv * rf2 * cf2; + f[14] = gv * rs; // g * rsqrt(rms) + + if (vector_like) + { +#pragma unroll + for (int j = 15; j < INPUT_DIM; j++) f[j] = 0.0f; + return; + } + + // Groups 9-13 (channels 15-29): factored features for 2D+ params. + const float vr0 = static_cast(fac_r[col_idx]); + const float vr1 = static_cast(fac_r[col_idx + num_cols]); + const float vr2 = static_cast(fac_r[col_idx + 2 * num_cols]); + const float vc0 = static_cast(fac_c[row_idx]); + const float vc1 = static_cast(fac_c[row_idx + num_rows]); + const float vc2 = static_cast(fac_c[row_idx + 2 * num_rows]); + + f[15] = vr0; f[16] = vr1; f[17] = vr2; // row_feat (v_row) + f[18] = vc0; f[19] = vc1; f[20] = vc2; // col_feat (v_col) + f[21] = __frsqrt_rn(vr0 + FEAT_EPS); // rsqrt(row_feat) + f[22] = __frsqrt_rn(vr1 + FEAT_EPS); + f[23] = __frsqrt_rn(vr2 + FEAT_EPS); + f[24] = __frsqrt_rn(vc0 + FEAT_EPS); // rsqrt(col_feat) + f[25] = __frsqrt_rn(vc1 + FEAT_EPS); + f[26] = __frsqrt_rn(vc2 + FEAT_EPS); + f[27] = m0 * rf0 * cf0; // fac_mom_mult = mom * row_factor * col_factor + f[28] = m1 * rf1 * cf1; + f[29] = m2 * rf2 * cf2; +} + +// Pass 1: accumulate sum of squares of each of the 30 feature channels across +// all elements into second_moment[30]. +template +__global__ void celo2_moment_kernel( + const T *__restrict__ g, + const T *__restrict__ p, + const T *__restrict__ m, + const T *__restrict__ v, + const T *__restrict__ fac_r, + const T *__restrict__ fac_c, + const T *__restrict__ row_factor, + const T *__restrict__ col_factor, + float *__restrict__ second_moment, + const int n_elements, + const int num_rows, + const int num_cols, + const int row_stride, + const int col_stride, + const int vector_like) +{ + const int tid = threadIdx.x; + const int warp_id = tid / warpSize; + const int lane_id = tid % warpSize; + const int num_warps = blockDim.x / warpSize; + __shared__ float s_warp_results[BLOCK_SIZE / 32][INPUT_DIM]; + + if (tid < num_warps * INPUT_DIM) + { + int wid = tid / INPUT_DIM; + int j = tid % INPUT_DIM; + s_warp_results[wid][j] = 0.0f; + } + __syncthreads(); + + float thread_accum[INPUT_DIM] = {0}; + + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_elements; i += blockDim.x * gridDim.x) + { + float f[INPUT_DIM]; + populate_celo2_features(f, g, p, m, v, fac_r, fac_c, row_factor, col_factor, + i, num_rows, num_cols, row_stride, col_stride, + n_elements, vector_like); +#pragma unroll + for (int j = 0; j < INPUT_DIM; j++) + thread_accum[j] += f[j] * f[j]; + } + +#pragma unroll + for (int j = 0; j < INPUT_DIM; j++) + { + float val = thread_accum[j]; +#pragma unroll + for (int offset = warpSize / 2; offset > 0; offset /= 2) + val += __shfl_down_sync(0xffffffff, val, offset); + if (lane_id == 0) + s_warp_results[warp_id][j] += val; + } + __syncthreads(); + + if (warp_id == 0) + { +#pragma unroll + for (int j = 0; j < INPUT_DIM; j++) + { + float sum = (lane_id < num_warps) ? s_warp_results[lane_id][j] : 0.0f; +#pragma unroll + for (int offset = warpSize / 2; offset > 0; offset /= 2) + sum += __shfl_down_sync(0xffffffff, sum, offset); + if (lane_id == 0) + atomicAdd(&second_moment[j], sum); + } + } +} + +// Pass 2: normalize features by rsqrt(mean + 1e-9), run the 30->8->8->3 MLP, +// and write step_out[idx] = out0 * exp(out1 * exp_mult). p is left untouched. +template +__global__ void celo2_apply_kernel( + const T *__restrict__ g, + const T *__restrict__ p, + const T *__restrict__ m, + const T *__restrict__ v, + const T *__restrict__ fac_r, + const T *__restrict__ fac_c, + const T *__restrict__ row_factor, + const T *__restrict__ col_factor, + const float *__restrict__ second_moment, + const T *__restrict__ w_in, // (INPUT_DIM, HIDDEN_DIM), in-major + const T *__restrict__ b_in, // (HIDDEN_DIM,) + const T *__restrict__ w_h, // (HIDDEN_DIM, HIDDEN_DIM) + const T *__restrict__ b_h, // (HIDDEN_DIM,) + const T *__restrict__ w_out, // (HIDDEN_DIM, OUTPUT_DIM) + const T *__restrict__ b_out, // (OUTPUT_DIM,) + T *__restrict__ step_out, + const int n_elements, + const int num_rows, + const int num_cols, + const int row_stride, + const int col_stride, + const float exp_mult, + const int vector_like) +{ + const int tid = threadIdx.x; + __shared__ float s_norm[INPUT_DIM]; + + if (tid < INPUT_DIM) + s_norm[tid] = rsqrtf(second_moment[tid] / n_elements + NORM_EPS); + __syncthreads(); + + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_elements; i += blockDim.x * gridDim.x) + { + float f[INPUT_DIM]; + populate_celo2_features(f, g, p, m, v, fac_r, fac_c, row_factor, col_factor, + i, num_rows, num_cols, row_stride, col_stride, + n_elements, vector_like); + + // Input (split) layer: 30 -> 8, then ReLU. Weights are in-major (in, out). + float h0[HIDDEN_DIM]; +#pragma unroll + for (int j = 0; j < HIDDEN_DIM; j++) + { + float acc = static_cast(b_in[j]); +#pragma unroll + for (int k = 0; k < INPUT_DIM; k++) + acc += static_cast(w_in[k * HIDDEN_DIM + j]) * (f[k] * s_norm[k]); + h0[j] = relu_(acc); + } + + // Hidden dense layer: 8 -> 8, ReLU. + float h1[HIDDEN_DIM]; +#pragma unroll + for (int j = 0; j < HIDDEN_DIM; j++) + { + float acc = static_cast(b_h[j]); +#pragma unroll + for (int k = 0; k < HIDDEN_DIM; k++) + acc += static_cast(w_h[k * HIDDEN_DIM + j]) * h0[k]; + h1[j] = relu_(acc); + } + + // Output layer: 8 -> 3, linear. + float out[OUTPUT_DIM]; +#pragma unroll + for (int j = 0; j < OUTPUT_DIM; j++) + { + float acc = static_cast(b_out[j]); +#pragma unroll + for (int k = 0; k < HIDDEN_DIM; k++) + acc += static_cast(w_out[k * OUTPUT_DIM + j]) * h1[k]; + out[j] = acc; + } + + // direction = out[0], magnitude = out[1] (out[2] unused, as in naive). + step_out[i] = static_cast(out[0] * __expf(out[1] * exp_mult)); + } +} + +void celo2_kernel( + at::Tensor &g, + at::Tensor &p, + at::Tensor &m, + at::Tensor &v, + at::Tensor &fac_r, + at::Tensor &fac_c, + at::Tensor &row_factor, + at::Tensor &col_factor, + at::Tensor &second_moment, + at::Tensor &w_in, + at::Tensor &b_in, + at::Tensor &w_h, + at::Tensor &b_h, + at::Tensor &w_out, + at::Tensor &b_out, + at::Tensor &step_out, + const float exp_mult, + const int dc, + const int dr, + const int vector_like) +{ + const int n_elements = p.numel(); + const int num_rows = vector_like ? n_elements : p.size(dr); + const int num_cols = vector_like ? n_elements : p.size(dc); + const int row_stride = vector_like ? 1 : p.stride(dr); + const int col_stride = vector_like ? 1 : p.stride(dc); + const int blocks_needed = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; + const int blocks = min(blocks_needed, 1728); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(p.scalar_type(), "celo2_moment_kernel", ([&] { + celo2_moment_kernel<<>>( + g.data_ptr(), + p.data_ptr(), + m.data_ptr(), + v.data_ptr(), + fac_r.data_ptr(), + fac_c.data_ptr(), + row_factor.data_ptr(), + col_factor.data_ptr(), + second_moment.data_ptr(), + n_elements, num_rows, num_cols, row_stride, col_stride, vector_like); + })); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(p.scalar_type(), "celo2_apply_kernel", ([&] { + celo2_apply_kernel<<>>( + g.data_ptr(), + p.data_ptr(), + m.data_ptr(), + v.data_ptr(), + fac_r.data_ptr(), + fac_c.data_ptr(), + row_factor.data_ptr(), + col_factor.data_ptr(), + second_moment.data_ptr(), + w_in.data_ptr(), + b_in.data_ptr(), + w_h.data_ptr(), + b_h.data_ptr(), + w_out.data_ptr(), + b_out.data_ptr(), + step_out.data_ptr(), + n_elements, num_rows, num_cols, row_stride, col_stride, exp_mult, vector_like); + })); + + AT_CUDA_CHECK(cudaGetLastError()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("celo2_kernel", &celo2_kernel, + "Fused CUDA kernel for the CELO2 learned optimizer (feature + MLP)", + py::call_guard()); +} diff --git a/pylo/optim/CELO2_cuda.py b/pylo/optim/CELO2_cuda.py new file mode 100644 index 0000000..38b81a5 --- /dev/null +++ b/pylo/optim/CELO2_cuda.py @@ -0,0 +1,354 @@ +"""CELO2_CUDA: CUDA-accelerated CELO2 learned optimizer. + +This is the fused-kernel counterpart of :class:`pylo.optim.CELO2_naive.CELO2_naive`. +The per-element heavy lifting -- CELO2 feature construction, per-channel +second-moment normalization, and the split-input MLP inference -- is done by the +``celo2_cuda_kernel.celo2_kernel`` CUDA extension. Everything that is cheap or +matrix-shaped stays in Python and is shared verbatim with the naive optimizer: + + * the warmup + cosine LR schedule (``CELO2_naive._lr_schedule``), + * global-norm gradient clipping, + * Newton-Schulz orthogonalization of the 2D+ update, + * the post-MLP second-moment normalization + rmsmult scaling, and + * the inline AdamW branch for 1D / embedding parameters. + +The kernel only computes the *raw* per-element step +``direction * exp(magnitude * exp_mult)``; it deliberately does not touch the +parameter, so the Python side can orthogonalize/normalize/apply exactly as the +naive code does. Correctness is checked against the (JAX-aligned) naive optimizer +in ``tests/test_celo2_cuda.py``. + +Layout: accumulators use the *leading* decay axis ``(n_decay,) + shape`` expected +by the kernel (the naive optimizer uses a trailing axis); the factored +second-moment handling mirrors :mod:`pylo.optim.AdafacLO_cuda`, whose +``exp_avg_sq_r/c`` + ``row_factor/col_factor`` layout is exactly what the kernel +indexes. + +Only ``ff_hidden_layers == 2`` (one split input layer + one hidden + one output, +i.e. 30 -> 8 -> 8 -> 3) is supported by the hardcoded kernel. +""" + +from typing import Optional + +import numpy as np +import torch +from torch.optim import Optimizer + +import celo2_cuda_kernel + +from pylo.models.CELO2_MLP import CELO2MLP +from pylo.optim.CELO2_naive import ( + factored_dims, + safe_rsqrt, + second_moment_normalizer, + orthogonalize_via_newton_schulz, +) + +# Must match the kernel's INPUT_DIM / HIDDEN_DIM / OUTPUT_DIM. +_KERNEL_INPUT_DIM = 30 +_KERNEL_HIDDEN_DIM = 8 +_KERNEL_OUTPUT_DIM = 3 + + +class CELO2_CUDA(Optimizer): + """CUDA CELO2 learned optimizer. Drop-in for :class:`CELO2_naive`. + + The constructor signature mirrors :class:`CELO2_naive`. See that class for a + description of the hyper-parameters. + """ + + def __init__( + self, + params, + num_steps, + # LR schedule + init_lr=0.0, + peak_lr=1e-3, + warmup_steps=0, + warmup_fraction=0.05, + end_lr=0.0, + # Weight decay + weight_decay=0.0, + # AdamW for 1D params + adam_lr_mult=1.0, + adam_weight_decay=None, + use_adamw_for_1d=True, + # CELO2 backbone + 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.7750, 2.0315), + ns_iters=5, + ns_eps=1e-8, + grad_clip_val=1000.0, + # Weights + hf_key: Optional[str] = "DiamondXL/celo2", + checkpoint_path: Optional[str] = None, + network: Optional[CELO2MLP] = None, + ): + if num_steps is None: + raise ValueError("CELO2_CUDA requires num_steps for the LR schedule.") + if not torch.cuda.is_available(): + raise RuntimeError("CELO2_CUDA requires a CUDA device.") + + self.device = torch.device("cuda") + + defaults = dict( + num_steps=num_steps, + init_lr=init_lr, + peak_lr=peak_lr, + warmup_steps=warmup_steps, + warmup_fraction=warmup_fraction, + end_lr=end_lr, + weight_decay=weight_decay, + adam_lr_mult=adam_lr_mult, + adam_weight_decay=( + adam_weight_decay if adam_weight_decay is not None else weight_decay + ), + use_adamw_for_1d=use_adamw_for_1d, + orthogonalize=orthogonalize, + clip_grad=clip_grad, + clip_norm=clip_norm, + exp_mult=exp_mult, + rmsmult=rmsmult, + param_scale_mult=param_scale_mult, + grad_clip_val=grad_clip_val, + is_embedding=False, + ) + super().__init__(params, defaults) + + # Fixed CELO2 decays (no learned offset, unlike AdafacLO), clamped to [0, 1]. + self.beta_m = torch.clamp( + torch.tensor(initial_momentum_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.beta_rms = torch.clamp( + torch.tensor(initial_rms_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.beta_fac = torch.clamp( + torch.tensor(initial_adafactor_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.n_mom = self.beta_m.shape[-1] + self.n_rms = self.beta_rms.shape[-1] + self.n_fac = self.beta_fac.shape[-1] + self.ns_coeffs = tuple(float(c) for c in ns_coeffs) + self.ns_iters = ns_iters + self.ns_eps = ns_eps + + # Precedence: explicit network > local checkpoint > HuggingFace Hub. + if network is not None: + self.network = network + elif checkpoint_path is not None: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network.load_state_dict(torch.load(checkpoint_path, map_location="cpu")) + elif hf_key is not None: + self.network = CELO2MLP.from_pretrained(hf_key) + else: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network = self.network.to(self.device) + + # The kernel hardcodes a 30 -> 8 -> 8 -> 3 MLP (one split input layer, one + # hidden, one output): validate the loaded network matches that shape. + if sum(self.network.feature_dims) != _KERNEL_INPUT_DIM: + raise ValueError( + f"CELO2_CUDA kernel expects feature_dims summing to " + f"{_KERNEL_INPUT_DIM}, got {sum(self.network.feature_dims)}." + ) + if self.network.hidden_size != _KERNEL_HIDDEN_DIM or len(self.network.dense_w) != 2: + raise ValueError( + "CELO2_CUDA kernel only supports ff_hidden_size=8, ff_hidden_layers=2." + ) + + # Pre-stack the split first layer into a single (30, 8) matrix: the kernel + # treats it as a plain (in, out) dense layer (CELO2MLP.forward computes + # sum_i feature_i @ w0[i] == concat(features) @ cat(w0)). + self.w_in = torch.cat([w.detach() for w in self.network.w0], dim=0).contiguous() + self.b_in = self.network.b0.detach().contiguous() + self.w_h = self.network.dense_w[0].detach().contiguous() + self.b_h = self.network.dense_b[0].detach().contiguous() + self.w_out = self.network.dense_w[1].detach().contiguous() + self.b_out = self.network.dense_b[1].detach().contiguous() + + # CELO2 offsets the 1-indexed step counter by 1 (schedule(0) on first step); + # ELO_CELO2_CUDA overrides this to 0. + self._lr_offset = 1 + + # ---------------------------------------------------------------- helpers + def _lr_schedule(self, step, group): + """Warmup + cosine decay learning rate (identical to CELO2_naive).""" + num_steps = group["num_steps"] + warmup = group["warmup_steps"] + if group["warmup_fraction"] > 0: + warmup = group["warmup_fraction"] * num_steps + warmup_f = max(float(warmup), 1.0) + decay_f = max(float(num_steps), 1.0) + step = float(step) + peak_lr, init_lr, end_lr = group["peak_lr"], group["init_lr"], group["end_lr"] + if step < warmup_f: + return init_lr + (peak_lr - init_lr) * min(step / warmup_f, 1.0) + progress = min(max((step - warmup_f) / max(decay_f - warmup_f, 1.0), 0.0), 1.0) + return end_lr + (peak_lr - end_lr) * 0.5 * (1.0 + np.cos(np.pi * progress)) + + def _global_grad_scale(self, clip_norm): + total = torch.zeros((), device=self.device) + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + total = total + p.grad.detach().float().pow(2).sum() + global_norm = torch.sqrt(total) + denom = torch.clamp(global_norm, min=clip_norm) + return (clip_norm / denom).item() if denom > 0 else 1.0 + + # ----------------------------------------------------- accumulator update + def _init_state(self, state, p): + shape = tuple(p.shape) + state["mom"] = torch.zeros((self.n_mom,) + shape, device=self.device) + state["rms"] = torch.zeros(shape, device=self.device) + f_dims = factored_dims(shape) + if f_dims is not None: + dc, dr = f_dims + r_shape = list(shape) + r_shape[dr] = 1 + c_shape = list(shape) + c_shape[dc] = 1 + state["fac_r"] = torch.zeros([self.n_fac] + r_shape, device=self.device) + state["fac_c"] = torch.zeros([self.n_fac] + c_shape, device=self.device) + + def _update_mom_rms(self, state, grad): + """EMA of momentum (leading layout) and the single rms accumulator.""" + m = state["mom"] + m.lerp_(grad[None, ...], (1 - self.beta_m).view([-1] + [1] * grad.dim())) + rms = state["rms"] + rms.lerp_(grad * grad, float(1 - self.beta_rms[-1])) + return m, rms + + def _update_factored(self, state, grad): + """Mirror AdafacLO_cuda / CELO2_naive factored second-moment update. + + Returns (fac_r, fac_c, row_factor, col_factor, dc, dr). + """ + grad_sqr = grad * grad + 1e-30 + dc, dr = factored_dims(tuple(grad.shape)) + beta_fac = self.beta_fac.view([-1] + [1] * grad.dim()) + fac_r = state["fac_r"] + fac_c = state["fac_c"] + fac_r.lerp_(grad_sqr.mean(dim=dr, keepdim=True)[None, ...], 1 - beta_fac) + fac_c.lerp_(grad_sqr.mean(dim=dc, keepdim=True)[None, ...], 1 - beta_fac) + # Match CELO2_naive: row_col_mean reduces the dc axis of the (kept-dim) + # row accumulator. In this leading-decay layout the dc parameter axis sits + # at tensor position dc + 1 (the decay axis is prepended, all param dims + # are kept via keepdim above). + row_col_mean = fac_r.mean(dim=dc + 1, keepdim=True) + row_factor = safe_rsqrt(fac_r / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(fac_c) + return fac_r, fac_c, row_factor, col_factor, dc, dr + + # ------------------------------------------------------------------ step + @torch.no_grad() + def step(self, loss=None): + grad_scales = {} + for group in self.param_groups: + if group["clip_grad"]: + grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) + + for group in self.param_groups: + group["step"] = group.get("step", 0) + 1 + t = group["step"] + lr = self._lr_schedule(t - self._lr_offset, group) + adam_lr = group["adam_lr_mult"] * lr + beta1 = float(self.beta_m[0]) + beta2 = float(self.beta_rms[-1]) + scale = grad_scales.get(id(group), 1.0) + clip_val = group["grad_clip_val"] + use_adamw_for_1d = group["use_adamw_for_1d"] + force_adam = group["is_embedding"] + exp_mult = group["exp_mult"] + + w_in = self.w_in + b_in = self.b_in + w_h = self.w_h + b_h = self.b_h + w_out = self.w_out + b_out = self.b_out + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if scale != 1.0: + grad = grad * scale + grad = torch.clamp(grad, -clip_val, clip_val) + + state = self.state[p] + if len(state) == 0: + self._init_state(state, p) + + m, rms = self._update_mom_rms(state, grad) + + is_1d = p.dim() <= 1 or force_adam + if use_adamw_for_1d and is_1d: + # Inline AdamW using the shared momentum[0] / rms accumulators. + m_bc = m[0] / (1.0 - beta1 ** t) + v_bc = rms / (1.0 - beta2 ** t) + adam_step = m_bc / (torch.sqrt(v_bc) + 1e-8) + p.add_(adam_step + group["adam_weight_decay"] * p, alpha=-adam_lr) + continue + + # CELO2 (2D+) path: kernel computes the raw per-element step. + fac_r, fac_c, row_factor, col_factor, dc, dr = self._update_factored( + state, grad + ) + + second_moment = torch.zeros( + _KERNEL_INPUT_DIM, dtype=torch.float32, device=self.device + ) + step_out = torch.empty_like(p) + dt = grad.dtype + celo2_cuda_kernel.celo2_kernel( + grad.contiguous(), + p, + m.to(dt).contiguous(), + rms.to(dt).contiguous(), + fac_r.to(dt).contiguous(), + fac_c.to(dt).contiguous(), + row_factor.to(dt).contiguous(), + col_factor.to(dt).contiguous(), + second_moment, + w_in.to(dt), + b_in.to(dt), + w_h.to(dt), + b_h.to(dt), + w_out.to(dt), + b_out.to(dt), + step_out, + float(exp_mult), + int(dc), + int(dr), + 0, + ) + + step = step_out + if group["orthogonalize"] and step.dim() >= 2: + step = orthogonalize_via_newton_schulz( + step, self.ns_coeffs, self.ns_iters, self.ns_eps + ) + axis = list(range(p.dim()))[-2:] + step = second_moment_normalizer(step, axis=axis) + step = step * group["rmsmult"] + p.add_(step + group["weight_decay"] * p, alpha=-lr) + + return loss diff --git a/pylo/optim/ELO_CELO2_cuda.py b/pylo/optim/ELO_CELO2_cuda.py new file mode 100644 index 0000000..ad169ed --- /dev/null +++ b/pylo/optim/ELO_CELO2_cuda.py @@ -0,0 +1,89 @@ +"""ELO_CELO2_CUDA: the ELO-CELO2 learned optimizer (CUDA-accelerated). + +CUDA counterpart of :class:`pylo.optim.ELO_CELO2_naive.ELO_CELO2_naive`. At +inference time ELO-CELO2 reduces exactly to the CELO2 forward pass, so this is a +thin wrapper over :class:`pylo.optim.CELO2_cuda.CELO2_CUDA` that only changes the +default hyper-parameters (nonzero weight decay, gradient clipping enabled, the +ELO-CELO2 checkpoint) and the LR-schedule offset. The learned weights live in the +loaded checkpoint. +""" + +from typing import Optional + +from pylo.optim.CELO2_cuda import CELO2_CUDA + + +class ELO_CELO2_CUDA(CELO2_CUDA): + """Inference-time ELO-CELO2 optimizer (CUDA CELO2 forward with ELO defaults).""" + + def __init__( + self, + params, + num_steps, + # LR schedule + init_lr=0.0, + peak_lr=1e-3, + warmup_steps=0, + warmup_fraction=0.05, + end_lr=0.0, + # ELO-CELO2 defaults differ from CELO2 here: + weight_decay=0.1, + clip_grad=True, + clip_norm=1.0, + # AdamW for 1D params + adam_lr_mult=20.0, + adam_weight_decay=None, + use_adamw_for_1d=True, + # CELO2 backbone + 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.7750, 2.0315), + ns_iters=5, + ns_eps=1e-8, + grad_clip_val=1000.0, + # Weights + hf_key: Optional[str] = "DiamondXL/elo-celo2", + checkpoint_path: Optional[str] = None, + network=None, + ): + super().__init__( + params, + num_steps=num_steps, + init_lr=init_lr, + peak_lr=peak_lr, + warmup_steps=warmup_steps, + warmup_fraction=warmup_fraction, + end_lr=end_lr, + weight_decay=weight_decay, + adam_lr_mult=adam_lr_mult, + adam_weight_decay=adam_weight_decay, + use_adamw_for_1d=use_adamw_for_1d, + orthogonalize=orthogonalize, + clip_grad=clip_grad, + clip_norm=clip_norm, + ff_hidden_size=ff_hidden_size, + ff_hidden_layers=ff_hidden_layers, + initial_momentum_decays=initial_momentum_decays, + initial_rms_decays=initial_rms_decays, + initial_adafactor_decays=initial_adafactor_decays, + exp_mult=exp_mult, + rmsmult=rmsmult, + param_scale_mult=param_scale_mult, + ns_coeffs=ns_coeffs, + ns_iters=ns_iters, + ns_eps=ns_eps, + grad_clip_val=grad_clip_val, + hf_key=hf_key, + checkpoint_path=checkpoint_path, + network=network, + ) + # ELO-CELO2 evaluates the LR schedule at iteration+1 (1-indexed), + # unlike the standalone CELO2 which is 0-indexed via the optax chain. + self._lr_offset = 0 diff --git a/pylo/optim/__init__.py b/pylo/optim/__init__.py index 5c95a4b..822826a 100644 --- a/pylo/optim/__init__.py +++ b/pylo/optim/__init__.py @@ -15,9 +15,9 @@ "ELO_naive", ] -# CELO2 / ELO-CELO2 currently ship a naive (pure-PyTorch) implementation only. -# Expose the bare public aliases now; if a CUDA build is added later these can be -# overridden in the try-block below, mirroring VeLO / AdafacLO. +# CELO2 / ELO-CELO2 default to the naive (pure-PyTorch) implementation; the +# dedicated try-block below overrides these to the CUDA variants when the +# celo2_cuda_kernel extension is available (mirroring VeLO / AdafacLO). CELO2 = CELO2_naive ELO_CELO2 = ELO_CELO2_naive ELO = ELO_naive @@ -51,3 +51,16 @@ AdafacLO = AdafacLO_naive MuLO = MuLO_naive __all__.extend(["VeLO", "AdafacLO", "MuLO"]) + +# CELO2 / ELO-CELO2 CUDA are wired independently: they only need the +# celo2_cuda_kernel extension, so a missing/failed build here leaves the other +# CUDA optimizers (and the naive CELO2 fallback set above) untouched. +try: + from pylo.optim.CELO2_cuda import CELO2_CUDA + from pylo.optim.ELO_CELO2_cuda import ELO_CELO2_CUDA + + CELO2 = CELO2_CUDA + ELO_CELO2 = ELO_CELO2_CUDA + __all__.extend(["CELO2_CUDA", "ELO_CELO2_CUDA"]) +except ImportError: + pass # keep the naive CELO2 / ELO_CELO2 aliases set above diff --git a/setup.py b/setup.py index 4988286..f4a1035 100644 --- a/setup.py +++ b/setup.py @@ -71,6 +71,13 @@ def get_build_config(): extra_compile_args=extra_compile_args, ) ) + ext_modules.append( + CUDAExtension( + name="celo2_cuda_kernel", + sources=["pylo/csrc/celo2_kernel.cu"], + extra_compile_args=extra_compile_args, + ) + ) cmdclass["build_ext"] = BuildExtension setup( diff --git a/tests/test_celo2_cuda.py b/tests/test_celo2_cuda.py new file mode 100644 index 0000000..10d7400 --- /dev/null +++ b/tests/test_celo2_cuda.py @@ -0,0 +1,129 @@ +"""Numerical-alignment tests for the CUDA CELO2 / ELO-CELO2 optimizers. + +We validate the CUDA implementation against the pure-PyTorch naive one. The naive +optimizer is itself aligned to the original JAX/optax CELO2 (see test_celo2.py), +so CUDA-vs-naive agreement transitively implies CUDA-vs-JAX agreement, without +needing a JAX environment. + +Both optimizers are fed the *same* meta-model weights (``network=net``), the same +initial parameters, and the same gradient sequence; we then compare the parameter +trajectories step by step. The tolerance (~1e-3) is looser than the CPU-vs-JAX +1e-4 because the kernel accumulates features/MLP in fp32 with fast-math rsqrt/exp. + +These tests require a CUDA device and the compiled ``celo2_cuda_kernel`` +extension; they skip otherwise. +""" + +import numpy as np +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CELO2 CUDA optimizer requires a GPU" +) + +# Skip cleanly if the extension hasn't been built. +pytest.importorskip("celo2_cuda_kernel") + +from pylo.models.CELO2_MLP import CELO2MLP +from pylo.optim.CELO2_naive import CELO2_naive +from pylo.optim.ELO_CELO2_naive import ELO_CELO2_naive +from pylo.optim.CELO2_cuda import CELO2_CUDA +from pylo.optim.ELO_CELO2_cuda import ELO_CELO2_CUDA + +DEVICE = torch.device("cuda") + +# Matrix params in both orientations (RC, to exercise both dc/dr index +# paths) plus a 1D param to exercise the shared-accumulator AdamW branch. +SHAPES = [(64, 32), (32, 64), (32,)] + + +def _make_init(seed): + g = torch.Generator(device=DEVICE).manual_seed(seed) + return [torch.randn(s, generator=g, device=DEVICE) for s in SHAPES] + + +def _make_grads(n_steps, seed, scale=3.0): + g = torch.Generator(device=DEVICE).manual_seed(seed) + return [ + [torch.randn(s, generator=g, device=DEVICE) * scale for s in SHAPES] + for _ in range(n_steps) + ] + + +def _run(opt_class, init_vals, grads_seq, net, **cfg): + params = [torch.nn.Parameter(v.clone()) for v in init_vals] + opt = opt_class(params, num_steps=cfg.pop("num_steps", 50), network=net, **cfg) + traj = [] + for grads in grads_seq: + for p, gg in zip(params, grads): + p.grad = gg.clone() + opt.step() + traj.append([p.detach().clone() for p in params]) + return traj + + +def _max_traj_diff(traj_a, traj_b): + worst = 0.0 + for step_a, step_b in zip(traj_a, traj_b): + for a, b in zip(step_a, step_b): + worst = max(worst, torch.max(torch.abs(a - b)).item()) + return worst + + +@pytest.mark.parametrize("orthogonalize", [True, False]) +@pytest.mark.parametrize("exp_mult", [0.0, 0.01]) +def test_celo2_cuda_matches_naive(orthogonalize, exp_mult): + torch.manual_seed(0) + np.random.seed(0) + net = CELO2MLP().to(DEVICE) # shared random-but-fixed meta-model + + init_vals = _make_init(seed=1) + grads_seq = _make_grads(n_steps=50, seed=2) + cfg = dict( + num_steps=100, + peak_lr=1e-2, + weight_decay=0.1, + clip_grad=True, + clip_norm=1.0, + orthogonalize=orthogonalize, + exp_mult=exp_mult, + ) + + naive_traj = _run(CELO2_naive, init_vals, grads_seq, net, **cfg) + cuda_traj = _run(CELO2_CUDA, init_vals, grads_seq, net, **cfg) + + assert _max_traj_diff(naive_traj, cuda_traj) < 1e-5 + + +def test_elo_celo2_cuda_matches_naive(): + """ELO-CELO2 (different defaults + _lr_offset=0) stays aligned.""" + torch.manual_seed(0) + np.random.seed(0) + net = CELO2MLP().to(DEVICE) + + init_vals = _make_init(seed=3) + grads_seq = _make_grads(n_steps=50, seed=4) + cfg = dict(num_steps=100, peak_lr=1e-3) + + naive_traj = _run(ELO_CELO2_naive, init_vals, grads_seq, net, **cfg) + cuda_traj = _run(ELO_CELO2_CUDA, init_vals, grads_seq, net, **cfg) + + assert _max_traj_diff(naive_traj, cuda_traj) < 1e-5 + + +def test_celo2_cuda_single_matrix_alignment(): + """Tight check on a single matrix param (pure CELO2 path, no AdamW).""" + torch.manual_seed(0) + np.random.seed(0) + net = CELO2MLP().to(DEVICE) + + g = torch.Generator(device=DEVICE).manual_seed(7) + init = [torch.randn(48, 24, generator=g, device=DEVICE)] + grads = [[torch.randn(48, 24, generator=g, device=DEVICE) * 2.0] for _ in range(50)] + cfg = dict(num_steps=100, peak_lr=1e-2, orthogonalize=True) + + naive_traj = _run(CELO2_naive, init, grads, net, **cfg) + cuda_traj = _run(CELO2_CUDA, init, grads, net, **cfg) + + assert _max_traj_diff(naive_traj, cuda_traj) < 1e-5 From 6ed73989fe92186f705f5b9e064516107a072ddf Mon Sep 17 00:00:00 2001 From: DiamondH <82591747+xiaol827@users.noreply.github.com> Date: Thu, 18 Jun 2026 17:38:04 -0400 Subject: [PATCH 6/7] Update README.md --- README.md | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/README.md b/README.md index b866170..a6a77ea 100644 --- a/README.md +++ b/README.md @@ -75,24 +75,6 @@ python -m pylo.util.patch_mup ## Quick Start -```python -import torch -from pylo.optim import VeLO_CUDA - -# Initialize a model -model = torch.nn.Linear(10, 2) - -# Create a learned optimizer instance -optimizer = VeLO_CUDA(model.parameters()) - -# Use it like any PyTorch optimizer -for epoch in range(10): - optimizer.zero_grad() - loss = loss_fn(model(input), target) - loss.backward() - optimizer.step(loss) # pass the loss -``` - Taking `ELO-CELO2` (the strongest LO) for example. ```python From dd899e772f936f2048b5062bf6171126dccc40c7 Mon Sep 17 00:00:00 2001 From: Xiaolong Huang Date: Sat, 20 Jun 2026 01:11:02 -0400 Subject: [PATCH 7/7] Remove built-in LR scheduler; CELO2 independent AdamW; add ELO CUDA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Optimizers no longer carry an internal warmup/cosine schedule — drive it with an external torch.optim.lr_scheduler via a plain `lr`. Applies to CELO2 (naive + cuda), ELO-CELO2 (naive + cuda), and ELO (naive). - CELO2_naive/CELO2_cuda: AdamW branch for 1D/embedding params now keeps its own exp_avg/exp_avg_sq moments (adam_betas/adam_eps), decoupled from the learned optimizer's momentum/RMS accumulators (numerically identical at default betas). - ELO_CELO2_naive/ELO_CELO2_cuda: no longer subclass CELO2; standalone classes that keep the original shared-accumulator AdamW design. - ELO_naive: drop num_steps/step_mult/warmup schedule, keep iteration counter for the tanh_embedding meta-model feature. - Add ELO_CUDA: CUDA ELO reusing the cuda_lo kernel (shared with AdafacLO), raw decays, lr-only scaling, exact pre-step weight-decay match; wired in __init__ and covered by tests/test_elo_cuda.py. - Fix device selection in naive optimizers to follow the parameters' device instead of torch.cuda.is_available() (CELO2/ELO/ELO-CELO2/AdafacLO/VeLO). - README quick start + tests updated to the new lr API; JAX end-to-end alignment pinned to a constant schedule (init_lr==peak_lr==end_lr) so it still holds. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 8 +- pylo/optim/AdafacLO_naive.py | 9 +- pylo/optim/CELO2_cuda.py | 101 +++++----- pylo/optim/CELO2_naive.py | 174 ++++++++--------- pylo/optim/ELO_CELO2_cuda.py | 289 +++++++++++++++++++++++---- pylo/optim/ELO_CELO2_naive.py | 293 ++++++++++++++++++++++++---- pylo/optim/ELO_cuda.py | 288 +++++++++++++++++++++++++++ pylo/optim/ELO_naive.py | 81 ++------ pylo/optim/Velo_naive.py | 9 +- pylo/optim/__init__.py | 10 + pylo/optim/velo_cuda.py | 9 +- scripts/convert_celo2_checkpoint.py | 1 - tests/test_celo2.py | 37 ++-- tests/test_celo2_cuda.py | 11 +- tests/test_elo.py | 20 +- tests/test_elo_cuda.py | 109 +++++++++++ 16 files changed, 1123 insertions(+), 326 deletions(-) create mode 100644 pylo/optim/ELO_cuda.py create mode 100644 tests/test_elo_cuda.py diff --git a/README.md b/README.md index a6a77ea..8c5f4ae 100644 --- a/README.md +++ b/README.md @@ -83,16 +83,20 @@ from pylo.optim import ELO_CELO2_CUDA model = torch.nn.Linear(10, 2) -num_steps = 1000 # total optimization steps (used for the LR schedule) +num_steps = 1000 # total optimization steps # Meta-learned weights download automatically from the Hugging Face Hub on first use. -optimizer = ELO_CELO2_CUDA(model.parameters(), num_steps=num_steps, peak_lr=3.16e-4, end_lr=3.16e-5, weight_decay=0.1, adam_lr_mult=20) +# The optimizer has no built-in LR schedule; drive it with a standard +# torch.optim.lr_scheduler (warmup, cosine, etc.). +optimizer = ELO_CELO2_CUDA(model.parameters(), lr=3.16e-4, weight_decay=0.1, adam_lr_mult=20) +scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps, eta_min=3.16e-5) for step in range(num_steps): optimizer.zero_grad() loss = loss_fn(model(input), target) loss.backward() optimizer.step() + scheduler.step() ``` ## Sharing Learned Optimizers diff --git a/pylo/optim/AdafacLO_naive.py b/pylo/optim/AdafacLO_naive.py index 4d32ad3..07f8a90 100644 --- a/pylo/optim/AdafacLO_naive.py +++ b/pylo/optim/AdafacLO_naive.py @@ -133,7 +133,14 @@ def __init__( mup_lrs=None, hf_key: Optional[str] = "btherien/mulo", ): - self.device = "cuda" if torch.cuda.is_available() else "cpu" + # Place state / meta-network on the parameters' device (materialize the + # iterable first so peeking does not consume a generator before super()). + params = list(params) + probe = params[0] + if isinstance(probe, dict): + probe["params"] = list(probe["params"]) + probe = probe["params"][0] + self.device = probe.device.type momentum_decays = torch.tensor(momentum_decays).to(self.device) rms_decays = torch.tensor(rms_decays).to(self.device) adafactor_decays = torch.tensor(adafactor_decays).to(self.device) diff --git a/pylo/optim/CELO2_cuda.py b/pylo/optim/CELO2_cuda.py index 38b81a5..8e56141 100644 --- a/pylo/optim/CELO2_cuda.py +++ b/pylo/optim/CELO2_cuda.py @@ -6,17 +6,19 @@ ``celo2_cuda_kernel.celo2_kernel`` CUDA extension. Everything that is cheap or matrix-shaped stays in Python and is shared verbatim with the naive optimizer: - * the warmup + cosine LR schedule (``CELO2_naive._lr_schedule``), * global-norm gradient clipping, * Newton-Schulz orthogonalization of the 2D+ update, * the post-MLP second-moment normalization + rmsmult scaling, and - * the inline AdamW branch for 1D / embedding parameters. + * the independent AdamW branch for 1D / embedding parameters. The kernel only computes the *raw* per-element step ``direction * exp(magnitude * exp_mult)``; it deliberately does not touch the parameter, so the Python side can orthogonalize/normalize/apply exactly as the -naive code does. Correctness is checked against the (JAX-aligned) naive optimizer -in ``tests/test_celo2_cuda.py``. +naive code does. Correctness is checked against the naive optimizer in +``tests/test_celo2_cuda.py``. + +No learning-rate schedule is built in: ``lr`` is used as-is, so an external +``torch.optim.lr_scheduler`` can drive the schedule. Layout: accumulators use the *leading* decay axis ``(n_decay,) + shape`` expected by the kernel (the naive optimizer uses a trailing axis); the factored @@ -30,7 +32,6 @@ from typing import Optional -import numpy as np import torch from torch.optim import Optimizer @@ -54,24 +55,21 @@ class CELO2_CUDA(Optimizer): """CUDA CELO2 learned optimizer. Drop-in for :class:`CELO2_naive`. The constructor signature mirrors :class:`CELO2_naive`. See that class for a - description of the hyper-parameters. + description of the hyper-parameters. Like the naive optimizer, the AdamW + branch for 1D / embedding parameters maintains independent moments. """ def __init__( self, params, - num_steps, - # LR schedule - init_lr=0.0, - peak_lr=1e-3, - warmup_steps=0, - warmup_fraction=0.05, - end_lr=0.0, + lr=1e-3, # Weight decay weight_decay=0.0, - # AdamW for 1D params + # AdamW for 1D / embedding params (independent accumulators) adam_lr_mult=1.0, adam_weight_decay=None, + adam_betas=(0.9, 0.95), + adam_eps=1e-8, use_adamw_for_1d=True, # CELO2 backbone orthogonalize=True, @@ -94,20 +92,13 @@ def __init__( checkpoint_path: Optional[str] = None, network: Optional[CELO2MLP] = None, ): - if num_steps is None: - raise ValueError("CELO2_CUDA requires num_steps for the LR schedule.") if not torch.cuda.is_available(): raise RuntimeError("CELO2_CUDA requires a CUDA device.") self.device = torch.device("cuda") defaults = dict( - num_steps=num_steps, - init_lr=init_lr, - peak_lr=peak_lr, - warmup_steps=warmup_steps, - warmup_fraction=warmup_fraction, - end_lr=end_lr, + lr=lr, weight_decay=weight_decay, adam_lr_mult=adam_lr_mult, adam_weight_decay=( @@ -141,6 +132,8 @@ def __init__( self.n_mom = self.beta_m.shape[-1] self.n_rms = self.beta_rms.shape[-1] self.n_fac = self.beta_fac.shape[-1] + self.adam_beta1, self.adam_beta2 = (float(b) for b in adam_betas) + self.adam_eps = adam_eps self.ns_coeffs = tuple(float(c) for c in ns_coeffs) self.ns_iters = ns_iters self.ns_eps = ns_eps @@ -183,26 +176,7 @@ def __init__( self.w_out = self.network.dense_w[1].detach().contiguous() self.b_out = self.network.dense_b[1].detach().contiguous() - # CELO2 offsets the 1-indexed step counter by 1 (schedule(0) on first step); - # ELO_CELO2_CUDA overrides this to 0. - self._lr_offset = 1 - # ---------------------------------------------------------------- helpers - def _lr_schedule(self, step, group): - """Warmup + cosine decay learning rate (identical to CELO2_naive).""" - num_steps = group["num_steps"] - warmup = group["warmup_steps"] - if group["warmup_fraction"] > 0: - warmup = group["warmup_fraction"] * num_steps - warmup_f = max(float(warmup), 1.0) - decay_f = max(float(num_steps), 1.0) - step = float(step) - peak_lr, init_lr, end_lr = group["peak_lr"], group["init_lr"], group["end_lr"] - if step < warmup_f: - return init_lr + (peak_lr - init_lr) * min(step / warmup_f, 1.0) - progress = min(max((step - warmup_f) / max(decay_f - warmup_f, 1.0), 0.0), 1.0) - return end_lr + (peak_lr - end_lr) * 0.5 * (1.0 + np.cos(np.pi * progress)) - def _global_grad_scale(self, clip_norm): total = torch.zeros((), device=self.device) for group in self.param_groups: @@ -214,7 +188,7 @@ def _global_grad_scale(self, clip_norm): return (clip_norm / denom).item() if denom > 0 else 1.0 # ----------------------------------------------------- accumulator update - def _init_state(self, state, p): + def _init_celo2_state(self, state, p): shape = tuple(p.shape) state["mom"] = torch.zeros((self.n_mom,) + shape, device=self.device) state["rms"] = torch.zeros(shape, device=self.device) @@ -257,6 +231,27 @@ def _update_factored(self, state, grad): col_factor = safe_rsqrt(fac_c) return fac_r, fac_c, row_factor, col_factor, dc, dr + def _adamw_step(self, p, grad, state, lr, weight_decay): + """Independent AdamW update (own exp_avg / exp_avg_sq moments).""" + if "exp_avg" not in state: + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + state["adam_step"] = 0 + + state["adam_step"] += 1 + t = state["adam_step"] + beta1, beta2 = self.adam_beta1, self.adam_beta2 + + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + m_bc = exp_avg / (1.0 - beta1 ** t) + v_bc = exp_avg_sq / (1.0 - beta2 ** t) + adam_step = m_bc / (torch.sqrt(v_bc) + self.adam_eps) + p.add_(adam_step + weight_decay * p, alpha=-lr) + # ------------------------------------------------------------------ step @torch.no_grad() def step(self, loss=None): @@ -266,12 +261,8 @@ def step(self, loss=None): grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) for group in self.param_groups: - group["step"] = group.get("step", 0) + 1 - t = group["step"] - lr = self._lr_schedule(t - self._lr_offset, group) + lr = group["lr"] adam_lr = group["adam_lr_mult"] * lr - beta1 = float(self.beta_m[0]) - beta2 = float(self.beta_rms[-1]) scale = grad_scales.get(id(group), 1.0) clip_val = group["grad_clip_val"] use_adamw_for_1d = group["use_adamw_for_1d"] @@ -294,21 +285,17 @@ def step(self, loss=None): grad = torch.clamp(grad, -clip_val, clip_val) state = self.state[p] - if len(state) == 0: - self._init_state(state, p) - - m, rms = self._update_mom_rms(state, grad) - is_1d = p.dim() <= 1 or force_adam if use_adamw_for_1d and is_1d: - # Inline AdamW using the shared momentum[0] / rms accumulators. - m_bc = m[0] / (1.0 - beta1 ** t) - v_bc = rms / (1.0 - beta2 ** t) - adam_step = m_bc / (torch.sqrt(v_bc) + 1e-8) - p.add_(adam_step + group["adam_weight_decay"] * p, alpha=-adam_lr) + self._adamw_step( + p, grad, state, adam_lr, group["adam_weight_decay"] + ) continue # CELO2 (2D+) path: kernel computes the raw per-element step. + if "mom" not in state: + self._init_celo2_state(state, p) + m, rms = self._update_mom_rms(state, grad) fac_r, fac_c, row_factor, col_factor, dc, dr = self._update_factored( state, grad ) diff --git a/pylo/optim/CELO2_naive.py b/pylo/optim/CELO2_naive.py index 19db029..a778e15 100644 --- a/pylo/optim/CELO2_naive.py +++ b/pylo/optim/CELO2_naive.py @@ -5,17 +5,25 @@ (``scaling_l2o/src/learned_optimizers/celo2_optax.py``, https://arxiv.org/abs/2602.19142). -The optimizer keeps, per parameter, three families of accumulators: +The optimizer keeps, per 2D+ parameter, three families of accumulators: * momentum at several decays (default ``(0.9, 0.99, 0.999)``), * an RMS / second-moment accumulator (default ``(0.95,)``), * an Adafactor-style factored second moment (default ``(0.9, 0.99, 0.999)``). For each 2D+ parameter it builds a stack of features, feeds them through the :class:`~pylo.models.CELO2_MLP.CELO2MLP` (split-input MLP), optionally -orthogonalizes the resulting update via Newton-Schulz iteration, normalizes it, -and applies it with a warmup + cosine learning-rate schedule. 1D parameters -(biases, norms) are updated with AdamW using the shared momentum/RMS -accumulators, matching the original implementation. +orthogonalizes the resulting update via Newton-Schulz iteration, and normalizes +it. 1D parameters (biases, norms) and embeddings are updated with AdamW. + +Unlike the original shared-accumulator design (still used by +:class:`pylo.optim.ELO_CELO2_naive.ELO_CELO2_naive`), the AdamW branch here keeps +its *own* ``exp_avg`` / ``exp_avg_sq`` moments, fully decoupled from the learned +optimizer's momentum/RMS accumulators. This makes the AdamW betas/eps independent +hyper-parameters and avoids allocating the (unused) CELO2 accumulators for 1D +parameters. + +No learning-rate schedule is built in: ``lr`` is used as-is, so an external +``torch.optim.lr_scheduler`` can drive the schedule (warmup, cosine, etc.). """ from typing import Optional @@ -72,15 +80,14 @@ class CELO2_naive(Optimizer): 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. - num_steps: Total number of inner training steps. Required: it defines the - warmup + cosine learning-rate schedule. - init_lr, peak_lr, warmup_steps, warmup_fraction, end_lr: Warmup + cosine - schedule parameters. ``warmup_fraction`` (a fraction of ``num_steps``) - takes priority over ``warmup_steps`` when > 0. + 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, adam_beta1, adam_beta2, adam_weight_decay, use_adamw_for_1d: - AdamW configuration for 1D parameters. ``adam_weight_decay`` defaults - to ``weight_decay`` when None. + adam_lr_mult, adam_weight_decay, adam_betas, adam_eps, 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, clip_norm: Optional global-norm gradient clipping. @@ -97,19 +104,14 @@ class CELO2_naive(Optimizer): def __init__( self, params, - num_steps, - # LR schedule - init_lr=0.0, - peak_lr=1e-3, - warmup_steps=0, - warmup_fraction=0.05, - end_lr=0.0, + lr=1e-3, # Weight decay weight_decay=0.0, - # AdamW for 1D params (reuses momentum_decays[0] / rms_decays[-1] as betas, - # matching the shared-accumulator design of the JAX ELO-CELO2) + # AdamW for 1D / embedding params (independent accumulators) adam_lr_mult=1.0, adam_weight_decay=None, + adam_betas=(0.9, 0.95), + adam_eps=1e-8, use_adamw_for_1d=True, # CELO2 backbone orthogonalize=True, @@ -132,18 +134,8 @@ def __init__( checkpoint_path: Optional[str] = None, network: Optional[CELO2MLP] = None, ): - if num_steps is None: - raise ValueError("CELO2_naive requires num_steps for the LR schedule.") - - self.device = "cuda" if torch.cuda.is_available() else "cpu" - defaults = dict( - num_steps=num_steps, - init_lr=init_lr, - peak_lr=peak_lr, - warmup_steps=warmup_steps, - warmup_fraction=warmup_fraction, - end_lr=end_lr, + lr=lr, weight_decay=weight_decay, adam_lr_mult=adam_lr_mult, adam_weight_decay=( @@ -161,6 +153,11 @@ def __init__( ) super().__init__(params, defaults) + # Place state / meta-network on the same device as the parameters, so + # CPU params on a GPU machine don't trigger device-mismatch errors. + first_param = next(p for g in self.param_groups for p in g["params"]) + self.device = first_param.device.type + self.initial_momentum_decays = torch.tensor( initial_momentum_decays, dtype=torch.float32, device=self.device ) @@ -170,6 +167,8 @@ def __init__( self.initial_adafactor_decays = torch.tensor( initial_adafactor_decays, dtype=torch.float32, device=self.device ) + self.adam_beta1, self.adam_beta2 = (float(b) for b in adam_betas) + self.adam_eps = adam_eps self.ns_coeffs = tuple(float(c) for c in ns_coeffs) self.ns_iters = ns_iters self.ns_eps = ns_eps @@ -192,29 +191,7 @@ def __init__( ) self.network = self.network.to(self.device) - # LR-schedule iteration offset. The standalone Celo2LOpt drives its - # schedule through an optax chain whose step count starts at 0, so the - # first update uses schedule(0); ELO-CELO2 instead evaluates the - # schedule at iteration+1 (1-indexed). CELO2 therefore offsets the - # 1-indexed step counter by 1; ELO_CELO2_naive overrides this to 0. - self._lr_offset = 1 - # ---------------------------------------------------------------- helpers - def _lr_schedule(self, step, group): - """Warmup + cosine decay learning rate (matches the JAX schedule).""" - num_steps = group["num_steps"] - warmup = group["warmup_steps"] - if group["warmup_fraction"] > 0: - warmup = group["warmup_fraction"] * num_steps - warmup_f = max(float(warmup), 1.0) - decay_f = max(float(num_steps), 1.0) - step = float(step) - peak_lr, init_lr, end_lr = group["peak_lr"], group["init_lr"], group["end_lr"] - if step < warmup_f: - return init_lr + (peak_lr - init_lr) * min(step / warmup_f, 1.0) - progress = min(max((step - warmup_f) / max(decay_f - warmup_f, 1.0), 0.0), 1.0) - return end_lr + (peak_lr - end_lr) * 0.5 * (1.0 + np.cos(np.pi * progress)) - def _global_grad_scale(self, clip_norm): """optax.clip_by_global_norm scale factor over all parameter grads.""" total = torch.zeros((), device=self.device) @@ -226,6 +203,27 @@ def _global_grad_scale(self, clip_norm): denom = torch.clamp(global_norm, min=clip_norm) return (clip_norm / denom).item() if denom > 0 else 1.0 + def _init_celo2_state(self, state, p_shape): + """Allocate the CELO2 momentum / RMS / factored accumulators.""" + n_mom = self.initial_momentum_decays.shape[-1] + n_rms = self.initial_rms_decays.shape[-1] + n_fac = self.initial_adafactor_decays.shape[-1] + state["mom"] = torch.zeros(p_shape + (n_mom,), device=self.device) + state["rms"] = torch.zeros(p_shape + (n_rms,), device=self.device) + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + full = p_shape + (n_fac,) + vr = tuple(d for i, d in enumerate(full) if i != d0) + vc = tuple(d for i, d in enumerate(full) if i != d1) + state["fac_vec_row"] = torch.zeros(vr, device=self.device) + state["fac_vec_col"] = torch.zeros(vc, device=self.device) + state["fac_vec_v"] = torch.empty(0, device=self.device) + else: + state["fac_vec_row"] = torch.empty(0, device=self.device) + state["fac_vec_col"] = torch.empty(0, device=self.device) + state["fac_vec_v"] = torch.zeros(p_shape + (n_fac,), device=self.device) + def _update_accumulators(self, state, batch_g, p_shape): """Advance momentum, RMS and factored accumulators in place; return fac_g.""" beta_m = self.initial_momentum_decays @@ -262,6 +260,27 @@ def _update_accumulators(self, state, batch_g, p_shape): state["fac_vec_v"] = new_v return mom, rms, fac_g + def _adamw_step(self, p, grad, state, lr, weight_decay): + """Independent AdamW update (own exp_avg / exp_avg_sq moments).""" + if "exp_avg" not in state: + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + state["adam_step"] = 0 + + state["adam_step"] += 1 + t = state["adam_step"] + beta1, beta2 = self.adam_beta1, self.adam_beta2 + + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + m_bc = exp_avg / (1.0 - beta1 ** t) + v_bc = exp_avg_sq / (1.0 - beta2 ** t) + adam_step = m_bc / (torch.sqrt(v_bc) + self.adam_eps) + p.add_(adam_step + weight_decay * p, alpha=-lr) + def _celo2_step(self, p, grad, state, group): """The CELO2 MLP update for a single (2D+) parameter, shape == p.shape.""" p_shape = tuple(p.shape) @@ -338,14 +357,8 @@ def step(self, loss=None): grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) for group in self.param_groups: - group["step"] = group.get("step", 0) + 1 - t = group["step"] - # Schedule index follows the reference (0-indexed for CELO2 via - # _lr_offset=1); AdamW bias correction stays 1-indexed (uses t). - lr = self._lr_schedule(t - self._lr_offset, group) + lr = group["lr"] adam_lr = group["adam_lr_mult"] * lr - beta1 = float(self.initial_momentum_decays[0]) - beta2 = float(self.initial_rms_decays[-1]) scale = grad_scales.get(id(group), 1.0) clip_val = group["grad_clip_val"] use_adamw_for_1d = group["use_adamw_for_1d"] @@ -361,41 +374,14 @@ def step(self, loss=None): p_shape = tuple(p.shape) state = self.state[p] - if len(state) == 0: - n_mom = self.initial_momentum_decays.shape[-1] - n_rms = self.initial_rms_decays.shape[-1] - n_fac = self.initial_adafactor_decays.shape[-1] - state["mom"] = torch.zeros( - p_shape + (n_mom,), device=self.device - ) - state["rms"] = torch.zeros( - p_shape + (n_rms,), device=self.device - ) - f_dims = factored_dims(p_shape) - if f_dims is not None: - d1, d0 = f_dims - full = p_shape + (n_fac,) - vr = tuple(d for i, d in enumerate(full) if i != d0) - vc = tuple(d for i, d in enumerate(full) if i != d1) - state["fac_vec_row"] = torch.zeros(vr, device=self.device) - state["fac_vec_col"] = torch.zeros(vc, device=self.device) - state["fac_vec_v"] = torch.empty(0, device=self.device) - else: - state["fac_vec_row"] = torch.empty(0, device=self.device) - state["fac_vec_col"] = torch.empty(0, device=self.device) - state["fac_vec_v"] = torch.zeros( - p_shape + (n_fac,), device=self.device - ) - is_1d = p.dim() <= 1 or force_adam if use_adamw_for_1d and is_1d: - batch_g = grad.unsqueeze(-1) - mom, rms, _ = self._update_accumulators(state, batch_g, p_shape) - m_bc = mom[..., 0] / (1.0 - beta1 ** t) - v_bc = rms[..., -1] / (1.0 - beta2 ** t) - adam_step = m_bc / (torch.sqrt(v_bc) + 1e-8) - p.add_(adam_step + group["adam_weight_decay"] * p, alpha=-adam_lr) + self._adamw_step( + p, grad, state, adam_lr, group["adam_weight_decay"] + ) else: + if "mom" not in state: + self._init_celo2_state(state, p_shape) step = self._celo2_step(p, grad, state, group) p.add_(step + group["weight_decay"] * p, alpha=-lr) return loss diff --git a/pylo/optim/ELO_CELO2_cuda.py b/pylo/optim/ELO_CELO2_cuda.py index ad169ed..7af8e0f 100644 --- a/pylo/optim/ELO_CELO2_cuda.py +++ b/pylo/optim/ELO_CELO2_cuda.py @@ -1,36 +1,53 @@ """ELO_CELO2_CUDA: the ELO-CELO2 learned optimizer (CUDA-accelerated). CUDA counterpart of :class:`pylo.optim.ELO_CELO2_naive.ELO_CELO2_naive`. At -inference time ELO-CELO2 reduces exactly to the CELO2 forward pass, so this is a -thin wrapper over :class:`pylo.optim.CELO2_cuda.CELO2_CUDA` that only changes the -default hyper-parameters (nonzero weight decay, gradient clipping enabled, the -ELO-CELO2 checkpoint) and the LR-schedule offset. The learned weights live in the -loaded checkpoint. +inference time ELO-CELO2 reduces exactly to the CELO2 forward pass, fused into the +``celo2_cuda_kernel`` extension for the 2D+ path. + +Like the naive ELO-CELO2 (and unlike :class:`pylo.optim.CELO2_cuda.CELO2_CUDA`), +the AdamW branch for 1D parameters reuses the *shared* momentum[0] / rms +accumulators that feed the CELO2 MLP. This class is therefore a standalone copy +of the original shared-accumulator CELO2 CUDA optimizer with the ELO-CELO2 +defaults (nonzero weight decay, gradient clipping, ELO-CELO2 checkpoint, larger +AdamW LR multiplier); it no longer inherits ``CELO2_CUDA``. + +No learning-rate schedule is built in: ``lr`` is used as-is, so an external +``torch.optim.lr_scheduler`` can drive the schedule. """ from typing import Optional -from pylo.optim.CELO2_cuda import CELO2_CUDA +import torch +from torch.optim import Optimizer + +import celo2_cuda_kernel + +from pylo.models.CELO2_MLP import CELO2MLP +from pylo.optim.CELO2_naive import ( + factored_dims, + safe_rsqrt, + second_moment_normalizer, + orthogonalize_via_newton_schulz, +) +# Must match the kernel's INPUT_DIM / HIDDEN_DIM / OUTPUT_DIM. +_KERNEL_INPUT_DIM = 30 +_KERNEL_HIDDEN_DIM = 8 +_KERNEL_OUTPUT_DIM = 3 -class ELO_CELO2_CUDA(CELO2_CUDA): - """Inference-time ELO-CELO2 optimizer (CUDA CELO2 forward with ELO defaults).""" + +class ELO_CELO2_CUDA(Optimizer): + """Inference-time ELO-CELO2 optimizer (CUDA CELO2 forward, shared-accumulator AdamW).""" def __init__( self, params, - num_steps, - # LR schedule - init_lr=0.0, - peak_lr=1e-3, - warmup_steps=0, - warmup_fraction=0.05, - end_lr=0.0, + lr=1e-3, # ELO-CELO2 defaults differ from CELO2 here: weight_decay=0.1, clip_grad=True, clip_norm=1.0, - # AdamW for 1D params + # AdamW for 1D params (shares momentum[0] / rms accumulators) adam_lr_mult=20.0, adam_weight_decay=None, use_adamw_for_1d=True, @@ -51,39 +68,231 @@ def __init__( # Weights hf_key: Optional[str] = "DiamondXL/elo-celo2", checkpoint_path: Optional[str] = None, - network=None, + network: Optional[CELO2MLP] = None, ): - super().__init__( - params, - num_steps=num_steps, - init_lr=init_lr, - peak_lr=peak_lr, - warmup_steps=warmup_steps, - warmup_fraction=warmup_fraction, - end_lr=end_lr, + if not torch.cuda.is_available(): + raise RuntimeError("ELO_CELO2_CUDA requires a CUDA device.") + + self.device = torch.device("cuda") + + defaults = dict( + lr=lr, weight_decay=weight_decay, adam_lr_mult=adam_lr_mult, - adam_weight_decay=adam_weight_decay, + adam_weight_decay=( + adam_weight_decay if adam_weight_decay is not None else weight_decay + ), use_adamw_for_1d=use_adamw_for_1d, orthogonalize=orthogonalize, clip_grad=clip_grad, clip_norm=clip_norm, - ff_hidden_size=ff_hidden_size, - ff_hidden_layers=ff_hidden_layers, - initial_momentum_decays=initial_momentum_decays, - initial_rms_decays=initial_rms_decays, - initial_adafactor_decays=initial_adafactor_decays, exp_mult=exp_mult, rmsmult=rmsmult, param_scale_mult=param_scale_mult, - ns_coeffs=ns_coeffs, - ns_iters=ns_iters, - ns_eps=ns_eps, grad_clip_val=grad_clip_val, - hf_key=hf_key, - checkpoint_path=checkpoint_path, - network=network, + is_embedding=False, ) - # ELO-CELO2 evaluates the LR schedule at iteration+1 (1-indexed), - # unlike the standalone CELO2 which is 0-indexed via the optax chain. - self._lr_offset = 0 + super().__init__(params, defaults) + + # Fixed CELO2 decays (no learned offset, unlike AdafacLO), clamped to [0, 1]. + self.beta_m = torch.clamp( + torch.tensor(initial_momentum_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.beta_rms = torch.clamp( + torch.tensor(initial_rms_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.beta_fac = torch.clamp( + torch.tensor(initial_adafactor_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.n_mom = self.beta_m.shape[-1] + self.n_rms = self.beta_rms.shape[-1] + self.n_fac = self.beta_fac.shape[-1] + self.ns_coeffs = tuple(float(c) for c in ns_coeffs) + self.ns_iters = ns_iters + self.ns_eps = ns_eps + + # Precedence: explicit network > local checkpoint > HuggingFace Hub. + if network is not None: + self.network = network + elif checkpoint_path is not None: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network.load_state_dict(torch.load(checkpoint_path, map_location="cpu")) + elif hf_key is not None: + self.network = CELO2MLP.from_pretrained(hf_key) + else: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network = self.network.to(self.device) + + # The kernel hardcodes a 30 -> 8 -> 8 -> 3 MLP (one split input layer, one + # hidden, one output): validate the loaded network matches that shape. + if sum(self.network.feature_dims) != _KERNEL_INPUT_DIM: + raise ValueError( + f"ELO_CELO2_CUDA kernel expects feature_dims summing to " + f"{_KERNEL_INPUT_DIM}, got {sum(self.network.feature_dims)}." + ) + if self.network.hidden_size != _KERNEL_HIDDEN_DIM or len(self.network.dense_w) != 2: + raise ValueError( + "ELO_CELO2_CUDA kernel only supports ff_hidden_size=8, ff_hidden_layers=2." + ) + + # Pre-stack the split first layer into a single (30, 8) matrix. + self.w_in = torch.cat([w.detach() for w in self.network.w0], dim=0).contiguous() + self.b_in = self.network.b0.detach().contiguous() + self.w_h = self.network.dense_w[0].detach().contiguous() + self.b_h = self.network.dense_b[0].detach().contiguous() + self.w_out = self.network.dense_w[1].detach().contiguous() + self.b_out = self.network.dense_b[1].detach().contiguous() + + # ---------------------------------------------------------------- helpers + def _global_grad_scale(self, clip_norm): + total = torch.zeros((), device=self.device) + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + total = total + p.grad.detach().float().pow(2).sum() + global_norm = torch.sqrt(total) + denom = torch.clamp(global_norm, min=clip_norm) + return (clip_norm / denom).item() if denom > 0 else 1.0 + + # ----------------------------------------------------- accumulator update + def _init_state(self, state, p): + shape = tuple(p.shape) + state["mom"] = torch.zeros((self.n_mom,) + shape, device=self.device) + state["rms"] = torch.zeros(shape, device=self.device) + f_dims = factored_dims(shape) + if f_dims is not None: + dc, dr = f_dims + r_shape = list(shape) + r_shape[dr] = 1 + c_shape = list(shape) + c_shape[dc] = 1 + state["fac_r"] = torch.zeros([self.n_fac] + r_shape, device=self.device) + state["fac_c"] = torch.zeros([self.n_fac] + c_shape, device=self.device) + + def _update_mom_rms(self, state, grad): + """EMA of momentum (leading layout) and the single rms accumulator.""" + m = state["mom"] + m.lerp_(grad[None, ...], (1 - self.beta_m).view([-1] + [1] * grad.dim())) + rms = state["rms"] + rms.lerp_(grad * grad, float(1 - self.beta_rms[-1])) + return m, rms + + def _update_factored(self, state, grad): + """Mirror AdafacLO_cuda / CELO2_naive factored second-moment update. + + Returns (fac_r, fac_c, row_factor, col_factor, dc, dr). + """ + grad_sqr = grad * grad + 1e-30 + dc, dr = factored_dims(tuple(grad.shape)) + beta_fac = self.beta_fac.view([-1] + [1] * grad.dim()) + fac_r = state["fac_r"] + fac_c = state["fac_c"] + fac_r.lerp_(grad_sqr.mean(dim=dr, keepdim=True)[None, ...], 1 - beta_fac) + fac_c.lerp_(grad_sqr.mean(dim=dc, keepdim=True)[None, ...], 1 - beta_fac) + row_col_mean = fac_r.mean(dim=dc + 1, keepdim=True) + row_factor = safe_rsqrt(fac_r / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(fac_c) + return fac_r, fac_c, row_factor, col_factor, dc, dr + + # ------------------------------------------------------------------ step + @torch.no_grad() + def step(self, loss=None): + grad_scales = {} + for group in self.param_groups: + if group["clip_grad"]: + grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) + + for group in self.param_groups: + group["step"] = group.get("step", 0) + 1 + t = group["step"] + lr = group["lr"] + adam_lr = group["adam_lr_mult"] * lr + beta1 = float(self.beta_m[0]) + beta2 = float(self.beta_rms[-1]) + scale = grad_scales.get(id(group), 1.0) + clip_val = group["grad_clip_val"] + use_adamw_for_1d = group["use_adamw_for_1d"] + force_adam = group["is_embedding"] + exp_mult = group["exp_mult"] + + w_in = self.w_in + b_in = self.b_in + w_h = self.w_h + b_h = self.b_h + w_out = self.w_out + b_out = self.b_out + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if scale != 1.0: + grad = grad * scale + grad = torch.clamp(grad, -clip_val, clip_val) + + state = self.state[p] + if len(state) == 0: + self._init_state(state, p) + + m, rms = self._update_mom_rms(state, grad) + + is_1d = p.dim() <= 1 or force_adam + if use_adamw_for_1d and is_1d: + # AdamW over the shared momentum[0] / rms accumulators. + m_bc = m[0] / (1.0 - beta1 ** t) + v_bc = rms / (1.0 - beta2 ** t) + adam_step = m_bc / (torch.sqrt(v_bc) + 1e-8) + p.add_(adam_step + group["adam_weight_decay"] * p, alpha=-adam_lr) + continue + + # CELO2 (2D+) path: kernel computes the raw per-element step. + fac_r, fac_c, row_factor, col_factor, dc, dr = self._update_factored( + state, grad + ) + + second_moment = torch.zeros( + _KERNEL_INPUT_DIM, dtype=torch.float32, device=self.device + ) + step_out = torch.empty_like(p) + dt = grad.dtype + celo2_cuda_kernel.celo2_kernel( + grad.contiguous(), + p, + m.to(dt).contiguous(), + rms.to(dt).contiguous(), + fac_r.to(dt).contiguous(), + fac_c.to(dt).contiguous(), + row_factor.to(dt).contiguous(), + col_factor.to(dt).contiguous(), + second_moment, + w_in.to(dt), + b_in.to(dt), + w_h.to(dt), + b_h.to(dt), + w_out.to(dt), + b_out.to(dt), + step_out, + float(exp_mult), + int(dc), + int(dr), + 0, + ) + + step = step_out + if group["orthogonalize"] and step.dim() >= 2: + step = orthogonalize_via_newton_schulz( + step, self.ns_coeffs, self.ns_iters, self.ns_eps + ) + axis = list(range(p.dim()))[-2:] + step = second_moment_normalizer(step, axis=axis) + step = step * group["rmsmult"] + p.add_(step + group["weight_decay"] * p, alpha=-lr) + + return loss diff --git a/pylo/optim/ELO_CELO2_naive.py b/pylo/optim/ELO_CELO2_naive.py index 363f525..ae7ac2d 100644 --- a/pylo/optim/ELO_CELO2_naive.py +++ b/pylo/optim/ELO_CELO2_naive.py @@ -4,37 +4,48 @@ meta-training*. At inference / meta-test time the expert trajectory and the IMT losses are disabled (``meta_train=False``), so the parameter update reduces exactly to the CELO2 forward pass: CELO2 MLP steps for 2D+ parameters and AdamW -(over the shared accumulators) for 1D parameters. +for 1D parameters. -Consequently this class is a thin wrapper over :class:`CELO2_naive` that only -changes the default hyper-parameters to match the ELO-CELO2 configuration -(``config/learned_optimizer/elo_celo2.py``): nonzero weight decay and enabled -gradient clipping. The distinct learned weights live in the loaded checkpoint. +Unlike :class:`pylo.optim.CELO2_naive.CELO2_naive` (which gives AdamW its own +independent moments), ELO-CELO2 keeps the original *shared-accumulator* design: +the AdamW branch for 1D parameters reuses ``momentum[..., 0]`` / ``rms[..., -1]`` +of the same momentum/RMS accumulators that feed the CELO2 MLP, matching the JAX +ELO-CELO2 (``config/learned_optimizer/elo_celo2.py``). Because of that this class +no longer inherits from ``CELO2_naive`` -- it is a standalone copy that only +differs in the AdamW accumulator wiring and in the ELO-CELO2 default +hyper-parameters (nonzero weight decay, gradient clipping enabled, ELO-CELO2 +checkpoint). The distinct learned weights live in the loaded checkpoint. + +No learning-rate schedule is built in: ``lr`` is used as-is, so an external +``torch.optim.lr_scheduler`` can drive the schedule. """ from typing import Optional -from pylo.optim.CELO2_naive import CELO2_naive +import torch +from torch.optim import Optimizer + +from pylo.models.CELO2_MLP import CELO2MLP +from pylo.optim.CELO2_naive import ( + factored_dims, + safe_rsqrt, + second_moment_normalizer, + orthogonalize_via_newton_schulz, +) -class ELO_CELO2_naive(CELO2_naive): - """Inference-time ELO-CELO2 optimizer (CELO2 forward with ELO defaults).""" +class ELO_CELO2_naive(Optimizer): + """Inference-time ELO-CELO2 optimizer (CELO2 forward, shared-accumulator AdamW).""" def __init__( self, params, - num_steps, - # LR schedule - init_lr=0.0, - peak_lr=1e-3, - warmup_steps=0, - warmup_fraction=0.05, - end_lr=0.0, + lr=1e-3, # ELO-CELO2 defaults differ from CELO2 here: weight_decay=0.1, clip_grad=True, clip_norm=1.0, - # AdamW for 1D params + # AdamW for 1D params (shares momentum[0] / rms[-1] accumulators) adam_lr_mult=1.0, adam_weight_decay=None, use_adamw_for_1d=True, @@ -55,39 +66,241 @@ def __init__( # Weights hf_key: Optional[str] = "DiamondXL/elo-celo2", checkpoint_path: Optional[str] = None, - network=None, + network: Optional[CELO2MLP] = None, ): - super().__init__( - params, - num_steps=num_steps, - init_lr=init_lr, - peak_lr=peak_lr, - warmup_steps=warmup_steps, - warmup_fraction=warmup_fraction, - end_lr=end_lr, + defaults = dict( + lr=lr, weight_decay=weight_decay, adam_lr_mult=adam_lr_mult, - adam_weight_decay=adam_weight_decay, + adam_weight_decay=( + adam_weight_decay if adam_weight_decay is not None else weight_decay + ), use_adamw_for_1d=use_adamw_for_1d, orthogonalize=orthogonalize, clip_grad=clip_grad, clip_norm=clip_norm, - ff_hidden_size=ff_hidden_size, - ff_hidden_layers=ff_hidden_layers, - initial_momentum_decays=initial_momentum_decays, - initial_rms_decays=initial_rms_decays, - initial_adafactor_decays=initial_adafactor_decays, exp_mult=exp_mult, rmsmult=rmsmult, param_scale_mult=param_scale_mult, - ns_coeffs=ns_coeffs, - ns_iters=ns_iters, - ns_eps=ns_eps, grad_clip_val=grad_clip_val, - hf_key=hf_key, - checkpoint_path=checkpoint_path, - network=network, + is_embedding=False, + ) + super().__init__(params, defaults) + + # Place state / meta-network on the same device as the parameters. + first_param = next(p for g in self.param_groups for p in g["params"]) + self.device = first_param.device.type + + self.initial_momentum_decays = torch.tensor( + initial_momentum_decays, dtype=torch.float32, device=self.device + ) + self.initial_rms_decays = torch.tensor( + initial_rms_decays, dtype=torch.float32, device=self.device + ) + self.initial_adafactor_decays = torch.tensor( + initial_adafactor_decays, dtype=torch.float32, device=self.device ) - # ELO-CELO2 evaluates the LR schedule at iteration+1 (1-indexed), - # unlike the standalone CELO2 which is 0-indexed via the optax chain. - self._lr_offset = 0 + self.ns_coeffs = tuple(float(c) for c in ns_coeffs) + self.ns_iters = ns_iters + self.ns_eps = ns_eps + + # Precedence: explicit network > local checkpoint > HuggingFace Hub. + if network is not None: + self.network = network + elif checkpoint_path is not None: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network.load_state_dict( + torch.load(checkpoint_path, map_location="cpu") + ) + elif hf_key is not None: + self.network = CELO2MLP.from_pretrained(hf_key) + else: + self.network = CELO2MLP( + hidden_size=ff_hidden_size, hidden_layers=ff_hidden_layers + ) + self.network = self.network.to(self.device) + + # ---------------------------------------------------------------- helpers + def _global_grad_scale(self, clip_norm): + """optax.clip_by_global_norm scale factor over all parameter grads.""" + total = torch.zeros((), device=self.device) + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + total = total + p.grad.detach().float().pow(2).sum() + global_norm = torch.sqrt(total) + denom = torch.clamp(global_norm, min=clip_norm) + return (clip_norm / denom).item() if denom > 0 else 1.0 + + def _update_accumulators(self, state, batch_g, p_shape): + """Advance momentum, RMS and factored accumulators in place; return fac_g.""" + beta_m = self.initial_momentum_decays + beta_rms = self.initial_rms_decays + beta_fac = self.initial_adafactor_decays + + # momentum: m = decay * m + (1 - decay) * g + mom = state["mom"] + mom.mul_(beta_m).add_((1 - beta_m) * batch_g) + # rms: rms = clip(decay) * rms + (1 - clip(decay)) * g^2 + crms = torch.clamp(beta_rms, 0.0, 1.0) + rms = state["rms"] + rms.mul_(crms).add_((1 - crms) * (batch_g ** 2)) + + # factored Adafactor accumulator (over the adafactor decays) + cdec = torch.clamp(beta_fac, 0.0, 1.0) + mixing = 1.0 - cdec + g_rep = batch_g.repeat([1] * len(p_shape) + [beta_fac.shape[-1]]) + grad_sqr = g_rep * g_rep + 1e-30 + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + new_v_row = cdec * state["fac_vec_row"] + mixing * grad_sqr.mean(dim=d0) + new_v_col = cdec * state["fac_vec_col"] + mixing * grad_sqr.mean(dim=d1) + reduced_d1 = d1 - 1 if d1 > d0 else d1 + row_col_mean = new_v_row.mean(dim=reduced_d1, keepdim=True) + row_factor = safe_rsqrt(new_v_row / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(new_v_col) + fac_g = g_rep * row_factor.unsqueeze(d0) * col_factor.unsqueeze(d1) + state["fac_vec_row"], state["fac_vec_col"] = new_v_row, new_v_col + else: + new_v = cdec * state["fac_vec_v"] + mixing * grad_sqr + fac_g = g_rep * safe_rsqrt(new_v + 1e-9) + state["fac_vec_v"] = new_v + return mom, rms, fac_g + + def _celo2_step(self, p, grad, state, group): + """The CELO2 MLP update for a single (2D+) parameter, shape == p.shape.""" + p_shape = tuple(p.shape) + batch_p = p.unsqueeze(-1) + batch_g = grad.unsqueeze(-1) + + mom, rms, fac_g = self._update_accumulators(state, batch_g, p_shape) + rsqrt = torch.rsqrt(rms + 1e-8) + + inps = [ + batch_g, + torch.clamp(batch_g, -0.1, 0.1), + batch_p, + mom, + rms, + mom * rsqrt, + rsqrt, + fac_g, + batch_g * rsqrt, + ] + + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + v_row, v_col = state["fac_vec_row"], state["fac_vec_col"] + rp_row = [1] * (1 + len(p_shape)) + rp_col = [1] * (1 + len(p_shape)) + rp_row[d0] = p_shape[d0] + rp_col[d1] = p_shape[d1] + row_feat = v_row.unsqueeze(d0).repeat(rp_row) + col_feat = v_col.unsqueeze(d1).repeat(rp_col) + inps.extend([ + row_feat, + col_feat, + torch.rsqrt(row_feat + 1e-8), + torch.rsqrt(col_feat + 1e-8), + ]) + reduced_d1 = d1 - 1 if d1 > d0 else d1 + row_col_mean = v_row.mean(dim=reduced_d1, keepdim=True) + row_factor = safe_rsqrt(v_row / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(v_col) + fac_mom_mult = mom * row_factor.unsqueeze(d0) * col_factor.unsqueeze(d1) + inps.append(fac_mom_mult) + + axis = list(range(len(p_shape)))[-2:] + inps = [second_moment_normalizer(i, axis=axis) for i in inps] + + out = self.network(inps) + direction = out[..., 0] + magnitude = out[..., 1] + + mag_param = torch.exp(magnitude * group["exp_mult"]) + if group["param_scale_mult"]: + param_scale = torch.sqrt(torch.mean(p ** 2) + 1e-9) + step = direction * (param_scale * mag_param) + else: + step = direction * mag_param + + if group["orthogonalize"] and step.dim() >= 2: + step = orthogonalize_via_newton_schulz( + step, self.ns_coeffs, self.ns_iters, self.ns_eps + ) + step = second_moment_normalizer(step, axis=axis) + step = step * group["rmsmult"] + return step + + def _init_shared_state(self, state, p_shape): + """Allocate the (shared) CELO2 momentum / RMS / factored accumulators.""" + n_mom = self.initial_momentum_decays.shape[-1] + n_rms = self.initial_rms_decays.shape[-1] + n_fac = self.initial_adafactor_decays.shape[-1] + state["mom"] = torch.zeros(p_shape + (n_mom,), device=self.device) + state["rms"] = torch.zeros(p_shape + (n_rms,), device=self.device) + f_dims = factored_dims(p_shape) + if f_dims is not None: + d1, d0 = f_dims + full = p_shape + (n_fac,) + vr = tuple(d for i, d in enumerate(full) if i != d0) + vc = tuple(d for i, d in enumerate(full) if i != d1) + state["fac_vec_row"] = torch.zeros(vr, device=self.device) + state["fac_vec_col"] = torch.zeros(vc, device=self.device) + state["fac_vec_v"] = torch.empty(0, device=self.device) + else: + state["fac_vec_row"] = torch.empty(0, device=self.device) + state["fac_vec_col"] = torch.empty(0, device=self.device) + state["fac_vec_v"] = torch.zeros(p_shape + (n_fac,), device=self.device) + + # ------------------------------------------------------------------ step + @torch.no_grad() + def step(self, loss=None): + # Global-norm clipping needs all grads up front. + grad_scales = {} + for group in self.param_groups: + if group["clip_grad"]: + grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) + + for group in self.param_groups: + group["step"] = group.get("step", 0) + 1 + t = group["step"] + lr = group["lr"] + adam_lr = group["adam_lr_mult"] * lr + beta1 = float(self.initial_momentum_decays[0]) + beta2 = float(self.initial_rms_decays[-1]) + scale = grad_scales.get(id(group), 1.0) + clip_val = group["grad_clip_val"] + use_adamw_for_1d = group["use_adamw_for_1d"] + force_adam = group["is_embedding"] + + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if scale != 1.0: + grad = grad * scale + grad = torch.clamp(grad, -clip_val, clip_val) + p_shape = tuple(p.shape) + + state = self.state[p] + if len(state) == 0: + self._init_shared_state(state, p_shape) + + is_1d = p.dim() <= 1 or force_adam + if use_adamw_for_1d and is_1d: + # AdamW over the shared momentum[0] / rms[-1] accumulators. + batch_g = grad.unsqueeze(-1) + mom, rms, _ = self._update_accumulators(state, batch_g, p_shape) + m_bc = mom[..., 0] / (1.0 - beta1 ** t) + v_bc = rms[..., -1] / (1.0 - beta2 ** t) + adam_step = m_bc / (torch.sqrt(v_bc) + 1e-8) + p.add_(adam_step + group["adam_weight_decay"] * p, alpha=-adam_lr) + else: + step = self._celo2_step(p, grad, state, group) + p.add_(step + group["weight_decay"] * p, alpha=-lr) + return loss diff --git a/pylo/optim/ELO_cuda.py b/pylo/optim/ELO_cuda.py new file mode 100644 index 0000000..3e3707f --- /dev/null +++ b/pylo/optim/ELO_cuda.py @@ -0,0 +1,288 @@ +"""ELO_CUDA: CUDA-accelerated ELO (Adafactor-MLP) learned optimizer. + +This is the fused-kernel counterpart of :class:`pylo.optim.ELO_naive.ELO_naive`. +ELO and AdafacLO share the *identical* MetaMLP feature set (39 inputs, 2 outputs) +and per-element update, so this optimizer reuses the same fused +``cuda_lo.learned_optimizer_kernel`` extension as +:class:`pylo.optim.AdafacLO_cuda.AdafacLO_CUDA`. The kernel constructs the +features, runs the split MLP, computes ``direction * exp(magnitude * exp_mult)`` +and applies it to the parameter in-place. + +ELO differs from AdafacLO only in: + + * raw accumulator decays (no ``param_to_decay(decay_to_param(.) + offset)`` + reparameterization -- the decays are used directly), and + * the update scale: ELO uses ``lr`` directly (no separate ``step_mult``), and + decoupled weight decay scaled by ``lr``. + +No learning-rate schedule is built in: ``lr`` is used as-is, so an external +``torch.optim.lr_scheduler`` can drive the schedule. Correctness is checked +against the (JAX-aligned) naive optimizer in ``tests/test_elo_cuda.py``. + +The kernel hardcodes a 39 -> 32 -> 32 -> 2 MLP (``input`` / ``linear_0`` / +``output``), i.e. ``hidden_size=32`` and ``hidden_layers=1``. +""" + +from typing import Optional, Tuple + +import torch +from torch.optim import Optimizer + +import cuda_lo + +from pylo.models.Meta_MLP import MetaMLP + + +def _get_scalar_dtype(): + return torch.float64 + + +def safe_rsqrt(x): + return torch.rsqrt( + torch.maximum(x, torch.tensor(1e-9, dtype=x.dtype, device=x.device)) + ) + + +def _factored_dims(shape: Tuple[int, ...]) -> Optional[Tuple[int, int]]: + """Two largest axes (dc, dr) for the factored second moment, or None for <2D.""" + if len(shape) < 2: + return None + sorted_dims = sorted(((x, i) for i, x in enumerate(shape))) + return int(sorted_dims[-2][1]), int(sorted_dims[-1][1]) + + +class ELO_CUDA(Optimizer): + """CUDA ELO (Adafactor-MLP) learned optimizer. Drop-in for :class:`ELO_naive`. + + Args: + 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, hidden_layers: MetaMLP geometry. The fused kernel only + supports ``hidden_size=32`` and ``hidden_layers=1`` (39->32->32->2). + initial_momentum_decays, initial_rms_decays, initial_adafactor_decays: + Raw accumulator decays (used directly, no reparameterization). + clip_grad, 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 :class:`MetaMLP` (overrides the above). + """ + + def __init__( + self, + 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: Optional[str] = "DiamondXL/elo", + checkpoint_path: Optional[str] = None, + network: Optional[MetaMLP] = None, + ): + if not torch.cuda.is_available(): + raise RuntimeError("ELO_CUDA requires a CUDA device.") + + self.device = torch.device("cuda") + + defaults = dict( + lr=lr, + exp_mult=exp_mult, + weight_decay=weight_decay, + clip_grad=clip_grad, + clip_norm=clip_norm, + ) + super().__init__(params, defaults) + + # Raw decays (no reparameterization, unlike AdafacLO), clamped to [0, 1]. + self.beta_m = torch.clamp( + torch.tensor(initial_momentum_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.beta_rms = torch.clamp( + torch.tensor(initial_rms_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + self.beta_adafactor = torch.clamp( + torch.tensor(initial_adafactor_decays, dtype=torch.float32, device=self.device), + 0.0, 1.0, + ) + + # Precedence: explicit network > local checkpoint > HuggingFace Hub. + if network is not None: + self.network = network + elif checkpoint_path is not None: + self.network = MetaMLP( + input_size=39, hidden_size=hidden_size, hidden_layers=hidden_layers + ) + self.network.load_state_dict(torch.load(checkpoint_path, map_location="cpu")) + elif hf_key is not None: + self.network = MetaMLP.from_pretrained(hf_key) + else: + self.network = MetaMLP( + input_size=39, hidden_size=hidden_size, hidden_layers=hidden_layers + ) + self.network = self.network.to(self.device) + for p in self.network.parameters(): + p.requires_grad = False + + # The kernel hardcodes a 39 -> 32 -> 32 -> 2 MLP (input / linear_0 / output). + net = self.network.network + if not (hasattr(net, "input") and hasattr(net, "linear_0") and hasattr(net, "output")): + raise ValueError( + "ELO_CUDA kernel requires a MetaMLP with input/linear_0/output " + "layers (hidden_layers=1)." + ) + if hasattr(net, "linear_1"): + raise ValueError( + "ELO_CUDA kernel only supports hidden_layers=1 (39->32->32->2)." + ) + + # ---------------------------------------------------------------- helpers + def _global_grad_scale(self, clip_norm): + """optax.clip_by_global_norm scale factor over all parameter grads.""" + total = torch.zeros((), device=self.device) + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + total = total + p.grad.detach().float().pow(2).sum() + denom = torch.clamp(torch.sqrt(total), min=clip_norm) + return (clip_norm / denom).item() if denom > 0 else 1.0 + + def _init_state(self, state, grad): + state["step"] = torch.tensor(0.0, dtype=_get_scalar_dtype()) + shape = grad.shape + f_dims = _factored_dims(tuple(shape)) + if f_dims is not None: + dc, dr = f_dims + row_shape = list(shape) + row_shape[dr] = 1 + col_shape = list(shape) + col_shape[dc] = 1 + state["exp_avg_sq_r"] = grad.new_zeros([3] + row_shape) + state["exp_avg_sq_c"] = grad.new_zeros([3] + col_shape) + else: + state["exp_avg_sq_r"] = grad.new_zeros((3,) + tuple(shape)) + state["exp_avg_sq_c"] = grad.new_zeros((3,) + tuple(shape)) + state["exp_avg_sq"] = torch.zeros_like(grad, memory_format=torch.preserve_format) + state["exp_avg"] = grad.new_zeros((3,) + tuple(shape)) + + # ------------------------------------------------------------------ step + @torch.no_grad() + def step(self, loss=None): + grad_scales = {} + for group in self.param_groups: + if group["clip_grad"]: + grad_scales[id(group)] = self._global_grad_scale(group["clip_norm"]) + + for group in self.param_groups: + lr = group["lr"] + exp_mult = group["exp_mult"] + weight_decay = group["weight_decay"] + scale = grad_scales.get(id(group), 1.0) + + w_in = self.network.network.input.weight + b_in = self.network.network.input.bias + w_h = self.network.network.linear_0.weight + b_h = self.network.network.linear_0.bias + w_out = self.network.network.output.weight + b_out = self.network.network.output.bias + + for p in group["params"]: + if p.grad is None: + continue + if p.grad.is_sparse: + raise RuntimeError("Sparse gradients not supported") + + grad = p.grad + if scale != 1.0: + grad = grad * scale + grad = torch.nan_to_num(grad) + + state = self.state[p] + if len(state) == 0: + self._init_state(state, grad) + + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + exp_avg_sq_r = state["exp_avg_sq_r"] + exp_avg_sq_c = state["exp_avg_sq_c"] + step_t = state["step"] + step_t += 1 + + eps = 1e-7 if grad.dtype == torch.float16 else 1e-30 + grad_sqr = torch.square(grad) + eps + one_minus_beta2 = (1 - self.beta_rms).to(grad_sqr.dtype) + + f_dims = _factored_dims(tuple(grad.shape)) + if f_dims is not None: + dc, dr = f_dims + beta_fac = 1 - self.beta_adafactor.view(-1, *[1] * grad.dim()).to(grad_sqr.dtype) + exp_avg_sq_r.lerp_(grad_sqr.mean(dim=dr, keepdim=True)[None, ...], beta_fac) + exp_avg_sq_c.lerp_(grad_sqr.mean(dim=dc, keepdim=True)[None, ...], beta_fac) + exp_avg_sq.lerp_(grad_sqr, one_minus_beta2) + reduce_dc = dc - 1 if dc > dr else dc + row_col_mean = exp_avg_sq_r.mean(dim=reduce_dc, keepdim=True) + row_factor = safe_rsqrt(exp_avg_sq_r / (row_col_mean + 1e-9)) + col_factor = safe_rsqrt(exp_avg_sq_c) + vector_like = 0 + else: + dc = dr = 0 + beta_fac = 1 - self.beta_adafactor.view(-1, 1).to(grad_sqr.dtype) + exp_avg_sq_r.lerp_(grad_sqr[None, ...], beta_fac) + exp_avg_sq_c.lerp_(grad_sqr[None, ...], beta_fac) + exp_avg_sq.lerp_(grad_sqr, one_minus_beta2) + row_factor = safe_rsqrt(exp_avg_sq_r + 1e-9) + col_factor = torch.ones_like(row_factor) + vector_like = 1 + + exp_avg.lerp_( + grad[None, ...], + (1 - self.beta_m.view([-1] + [1] * grad.dim())).to(grad.dtype), + ) + + # ELO_naive applies decoupled weight decay against the *pre-step* + # parameter (p -= lr * (step + wd * p)); the kernel updates p + # in-place, so snapshot p first to reproduce that exactly. + p_before = p.detach().clone() if weight_decay > 0 else None + + second_moment = torch.zeros([28], device=self.device) + cuda_lo.learned_optimizer_kernel( + grad, + p, + exp_avg, + exp_avg_sq, + exp_avg_sq_r, + exp_avg_sq_c, + row_factor, + col_factor, + second_moment, + w_in.to(grad.dtype), + b_in.to(grad.dtype), + w_h.to(grad.dtype), + b_h.to(grad.dtype), + w_out.to(grad.dtype), + b_out.to(grad.dtype), + lr, # ELO update scale (no separate step_mult) + 1.0, # step_mult == 1.0 for ELO + exp_mult, + 1e-6, # rms rsqrt epsilon + step_t - 1, + 0.0, # weight decay handled in Python below + dc, + dr, + vector_like, + ) + + if weight_decay > 0: + p.add_(p_before, alpha=-weight_decay * lr) + + return loss diff --git a/pylo/optim/ELO_naive.py b/pylo/optim/ELO_naive.py index 7d04d7d..88e70b0 100644 --- a/pylo/optim/ELO_naive.py +++ b/pylo/optim/ELO_naive.py @@ -10,17 +10,20 @@ :class:`pylo.optim.AdafacLO_naive.AdafacLO_naive` / :class:`pylo.models.Meta_MLP.MetaMLP` (39 input features, 2 outputs). ELO differs only in: - * raw accumulator decays (no decay reparameterization), - * a warmup-then-constant (optionally cosine) learning-rate schedule, and - * the update rule ``p -= scheduled_lr * (direction * exp(magnitude * exp_mult) - + weight_decay * p)`` (the schedule plays the role of ``step_mult``). + * raw accumulator decays (no decay reparameterization), and + * the update rule ``p -= lr * (direction * exp(magnitude * exp_mult) + + weight_decay * p)``. + +No learning-rate schedule is built in: ``lr`` is used as-is, so an external +``torch.optim.lr_scheduler`` can drive any warmup/cosine schedule. The optimizer +still tracks the iteration counter, which feeds the ``tanh_embedding`` training +step feature consumed by the meta-model. Reference: ``scaling_l2o/src/learned_optimizers/elo_adfac_mlp_lopt.py``. """ from typing import Optional -import numpy as np import torch from torch.optim import Optimizer @@ -40,16 +43,11 @@ class ELO_naive(Optimizer): Args: params: Iterable of parameters or ``param_groups``. - num_steps: Total number of inner training steps (defines the warmup / - cosine schedule). Required. - exp_mult, step_mult: Magnitude exponent multiplier and base step size - (``step_mult`` is the post-warmup learning rate). - init_lr, warmup_fraction, warmup_steps: Linear warmup from ``init_lr`` to - ``step_mult``. ``warmup_fraction`` (fraction of ``num_steps``) takes - priority over ``warmup_steps`` when > 0. - use_lo_cosine_scheduler, step_mult_min: If enabled, cosine-decay the - post-warmup learning rate from ``step_mult`` down to ``step_mult_min``. - weight_decay: Decoupled weight decay (scaled by the schedule). + 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, 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). @@ -64,14 +62,8 @@ class ELO_naive(Optimizer): def __init__( self, params, - num_steps, + lr=0.001, exp_mult=0.001, - step_mult=0.001, - init_lr=0.0, - warmup_fraction=0.05, - warmup_steps=0, - use_lo_cosine_scheduler=False, - step_mult_min=1e-4, weight_decay=0.0, hidden_size=32, hidden_layers=1, @@ -84,26 +76,19 @@ def __init__( checkpoint_path: Optional[str] = None, network: Optional[MetaMLP] = None, ): - if num_steps is None: - raise ValueError("ELO_naive requires num_steps for the LR schedule.") - - self.device = "cuda" if torch.cuda.is_available() else "cpu" - defaults = dict( - num_steps=num_steps, + lr=lr, exp_mult=exp_mult, - step_mult=step_mult, - init_lr=init_lr, - warmup_fraction=warmup_fraction, - warmup_steps=warmup_steps, - use_lo_cosine_scheduler=use_lo_cosine_scheduler, - step_mult_min=step_mult_min, weight_decay=weight_decay, clip_grad=clip_grad, clip_norm=clip_norm, ) super().__init__(params, defaults) + # Place state / meta-network on the same device as the parameters. + first_param = next(p for g in self.param_groups for p in g["params"]) + self.device = first_param.device.type + self.initial_momentum_decays = torch.tensor( initial_momentum_decays, dtype=torch.float32, device=self.device ) @@ -130,30 +115,6 @@ def __init__( ) self.network = self.network.to(self.device) - def _scheduled_lr(self, iteration, group): - """Warmup → constant (or cosine) schedule, matching the JAX version.""" - num_steps = group["num_steps"] - step_mult = group["step_mult"] - init_lr = group["init_lr"] - if group["warmup_fraction"] > 0: - warmup_n = group["warmup_fraction"] * num_steps - else: - warmup_n = group["warmup_steps"] - warmup_n = max(float(warmup_n), 1.0) - it = float(iteration) - - if group["use_lo_cosine_scheduler"]: - frac = min(max(it / max(num_steps - 1, 1), 0.0), 1.0) - step_mult_min = group["step_mult_min"] - base = step_mult_min + (step_mult - step_mult_min) * 0.5 * ( - 1.0 + np.cos(np.pi * frac) - ) - else: - base = step_mult - - warmup_lr = init_lr + (step_mult - init_lr) * min(it / warmup_n, 1.0) - return warmup_lr if it < warmup_n else base - def _global_grad_scale(self, clip_norm): total = torch.zeros((), device=self.device) for group in self.param_groups: @@ -175,7 +136,7 @@ def step(self, loss=None): group["step"] = iteration + 1 exp_mult = group["exp_mult"] weight_decay = group["weight_decay"] - scheduled_lr = self._scheduled_lr(iteration, group) + lr = group["lr"] scale = grad_scales.get(id(group), 1.0) for p in group["params"]: @@ -273,5 +234,5 @@ def step(self, loss=None): direction, magnitude = self.network(inp_stack).split(1, dim=-1) step = (direction * torch.exp(magnitude * exp_mult)).squeeze(-1) - p.add_(step + weight_decay * p, alpha=-scheduled_lr) + p.add_(step + weight_decay * p, alpha=-lr) return loss diff --git a/pylo/optim/Velo_naive.py b/pylo/optim/Velo_naive.py index c107f00..62ed7cf 100644 --- a/pylo/optim/Velo_naive.py +++ b/pylo/optim/Velo_naive.py @@ -235,7 +235,14 @@ def __init__( hf_key_rnn="Pauljanson002/VeLO_RNN", hf_key_mlp="Pauljanson002/VeLO_MLP", ): - self.device = "cuda" if torch.cuda.is_available() else "cpu" + # Place state / meta-network on the parameters' device (materialize the + # iterable first so peeking does not consume a generator before super()). + params = list(params) + probe = params[0] + if isinstance(probe, dict): + probe["params"] = list(probe["params"]) + probe = probe["params"][0] + self.device = probe.device.type momentum_decays = torch.tensor(momentum_decays).to(self.device) rms_decays = torch.tensor(rms_decays).to(self.device) adafactor_decays = torch.tensor(adafactor_decays).to(self.device) diff --git a/pylo/optim/__init__.py b/pylo/optim/__init__.py index 822826a..6044564 100644 --- a/pylo/optim/__init__.py +++ b/pylo/optim/__init__.py @@ -64,3 +64,13 @@ __all__.extend(["CELO2_CUDA", "ELO_CELO2_CUDA"]) except ImportError: pass # keep the naive CELO2 / ELO_CELO2 aliases set above + +# ELO CUDA reuses the cuda_lo kernel (shared with AdafacLO); wire it independently +# so a failed AdafacLO_cuda import (its optional deps) doesn't disable ELO_CUDA. +try: + from pylo.optim.ELO_cuda import ELO_CUDA + + ELO = ELO_CUDA + __all__.append("ELO_CUDA") +except ImportError: + pass # keep the naive ELO alias set above diff --git a/pylo/optim/velo_cuda.py b/pylo/optim/velo_cuda.py index 744300b..88c993d 100644 --- a/pylo/optim/velo_cuda.py +++ b/pylo/optim/velo_cuda.py @@ -282,7 +282,14 @@ def __init__( hf_key_mlp="Pauljanson002/VeLO_MLP", legacy=False, ): - self.device = "cuda" if torch.cuda.is_available() else "cpu" + # Place state / meta-network on the parameters' device (materialize the + # iterable first so peeking does not consume a generator before super()). + params = list(params) + probe = params[0] + if isinstance(probe, dict): + probe["params"] = list(probe["params"]) + probe = probe["params"][0] + self.device = probe.device.type momentum_decays = torch.tensor(momentum_decays).to(self.device) rms_decays = torch.tensor(rms_decays).to(self.device) adafactor_decays = torch.tensor(adafactor_decays).to(self.device) diff --git a/scripts/convert_celo2_checkpoint.py b/scripts/convert_celo2_checkpoint.py index 24c0bfd..eccff59 100644 --- a/scripts/convert_celo2_checkpoint.py +++ b/scripts/convert_celo2_checkpoint.py @@ -23,7 +23,6 @@ import argparse import pickle -import re import numpy as np import torch diff --git a/tests/test_celo2.py b/tests/test_celo2.py index 13e78f8..842f9f8 100644 --- a/tests/test_celo2.py +++ b/tests/test_celo2.py @@ -15,9 +15,9 @@ from pylo.optim.CELO2_naive import factored_dims DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") -SCALING_L2O = "/home/mila/h/huangx/scaling_l2o" -LO_DIR = os.path.join(SCALING_L2O, "src/learned_optimizers") -CELO2_OPTAX = os.path.join(LO_DIR, "celo2_optax.py") +SCALING_L2O = os.environ.get("SCALING_L2O", "") +LO_DIR = os.path.join(SCALING_L2O, "src/learned_optimizers") if SCALING_L2O else "" +CELO2_OPTAX = os.path.join(LO_DIR, "celo2_optax.py") if LO_DIR else "" @functools.lru_cache(maxsize=1) @@ -91,7 +91,7 @@ def test_step_runs_and_updates(optimizer_class): model = SmallNet().to(DEVICE) before = [p.clone() for p in model.parameters()] - optimizer = optimizer_class(model.parameters(), num_steps=50, peak_lr=1e-2) + optimizer = optimizer_class(model.parameters(), lr=1e-2) x = torch.randn(16, 10, device=DEVICE) y = torch.randn(16, 1, device=DEVICE) @@ -111,7 +111,7 @@ def test_higher_rank_parameter(): """A >2D parameter (conv-like) goes through the orthogonalization path.""" torch.manual_seed(0) w = nn.Parameter(torch.randn(3, 4, 5, device=DEVICE)) - optimizer = CELO2_naive([w], num_steps=20, peak_lr=1e-2) + optimizer = CELO2_naive([w], lr=1e-2) for _ in range(3): optimizer.zero_grad() (w.sum() ** 2).backward() @@ -126,7 +126,7 @@ def test_celo2_resume(tmp_path): net = CELO2MLP().to(DEVICE) # shared meta-model (random but fixed) model = SmallNet().to(DEVICE) - optimizer = CELO2_naive(model.parameters(), num_steps=50, peak_lr=1e-2, network=net) + optimizer = CELO2_naive(model.parameters(), lr=1e-2, network=net) x = torch.randn(32, 10, device=DEVICE) y = torch.randn(32, 1, device=DEVICE) @@ -152,7 +152,7 @@ def test_celo2_resume(tmp_path): # Restart from checkpoint and continue. loaded_model = SmallNet().to(DEVICE) loaded_opt = CELO2_naive( - loaded_model.parameters(), num_steps=50, peak_lr=1e-2, network=net + loaded_model.parameters(), lr=1e-2, network=net ) ckpt = torch.load(path) loaded_model.load_state_dict(ckpt["model"]) @@ -212,7 +212,7 @@ def flat(t, pre=""): model.load_state_dict(sd) p = torch.nn.Parameter(torch.tensor(np.asarray(params["w"]), dtype=torch.float32)) - opt = CELO2_naive([p], num_steps=100, orthogonalize=ortho, network=model) + opt = CELO2_naive([p], orthogonalize=ortho, network=model) opt.device = "cpu" opt.initial_momentum_decays = opt.initial_momentum_decays.cpu() opt.initial_rms_decays = opt.initial_rms_decays.cpu() @@ -240,11 +240,15 @@ def flat(t, pre=""): ) @pytest.mark.parametrize("which", ["celo2", "elo_celo2"]) def test_jax_end_to_end_alignment(which): - """The FULL step() (1D AdamW + warmup/cosine LR + weight decay + global-norm - clipping) matches the reference JAX optimizer over a multi-step trajectory. - - Uses a model with a 2D weight and a 1D bias, nonzero weight decay, and - enabled gradient clipping with large grads, so every code path is exercised. + """The FULL step() (1D AdamW + weight decay + global-norm clipping) matches + the reference JAX optimizer over a multi-step trajectory. + + The pylo optimizers no longer carry an internal LR schedule (it is applied + externally now), so we pin the JAX reference to a *constant* schedule + (``init_lr == peak_lr == end_lr``) and run the torch side at that same + constant ``lr``; the per-step math is then identical. Uses a model with a 2D + weight and a 1D bias, nonzero weight decay, and enabled gradient clipping + with large grads, so every code path is exercised. """ jax = pytest.importorskip("jax") import jax.numpy as jnp @@ -252,8 +256,9 @@ def test_jax_end_to_end_alignment(which): celo2_lopt, elo_celo2 = _load_jax_lopts() N = 60 + LR = 1e-3 cfg = dict( - init_lr=0.0, peak_lr=1e-3, warmup_fraction=0.05, end_lr=0.0, + init_lr=LR, peak_lr=LR, warmup_fraction=0.05, end_lr=LR, weight_decay=0.1, adam_lr_mult=1.0, use_adamw_for_1d=True, clip_grad=True, clip_norm=1.0, ) @@ -289,8 +294,8 @@ def test_jax_end_to_end_alignment(which): pw = torch.nn.Parameter(torch.tensor(W)) pb = torch.nn.Parameter(torch.tensor(B)) topt = topt_cls( - [pw, pb], num_steps=N, peak_lr=1e-3, warmup_fraction=0.05, end_lr=0.0, - weight_decay=0.1, clip_grad=True, clip_norm=1.0, network=net, + [pw, pb], lr=LR, weight_decay=0.1, clip_grad=True, clip_norm=1.0, + network=net, ) topt.device = "cpu" for a in ("initial_momentum_decays", "initial_rms_decays", "initial_adafactor_decays"): diff --git a/tests/test_celo2_cuda.py b/tests/test_celo2_cuda.py index 10d7400..2c66304 100644 --- a/tests/test_celo2_cuda.py +++ b/tests/test_celo2_cuda.py @@ -53,7 +53,7 @@ def _make_grads(n_steps, seed, scale=3.0): def _run(opt_class, init_vals, grads_seq, net, **cfg): params = [torch.nn.Parameter(v.clone()) for v in init_vals] - opt = opt_class(params, num_steps=cfg.pop("num_steps", 50), network=net, **cfg) + opt = opt_class(params, network=net, **cfg) traj = [] for grads in grads_seq: for p, gg in zip(params, grads): @@ -81,8 +81,7 @@ def test_celo2_cuda_matches_naive(orthogonalize, exp_mult): init_vals = _make_init(seed=1) grads_seq = _make_grads(n_steps=50, seed=2) cfg = dict( - num_steps=100, - peak_lr=1e-2, + lr=1e-2, weight_decay=0.1, clip_grad=True, clip_norm=1.0, @@ -97,14 +96,14 @@ def test_celo2_cuda_matches_naive(orthogonalize, exp_mult): def test_elo_celo2_cuda_matches_naive(): - """ELO-CELO2 (different defaults + _lr_offset=0) stays aligned.""" + """ELO-CELO2 (different defaults, shared-accumulator AdamW) stays aligned.""" torch.manual_seed(0) np.random.seed(0) net = CELO2MLP().to(DEVICE) init_vals = _make_init(seed=3) grads_seq = _make_grads(n_steps=50, seed=4) - cfg = dict(num_steps=100, peak_lr=1e-3) + cfg = dict(lr=1e-3) naive_traj = _run(ELO_CELO2_naive, init_vals, grads_seq, net, **cfg) cuda_traj = _run(ELO_CELO2_CUDA, init_vals, grads_seq, net, **cfg) @@ -121,7 +120,7 @@ def test_celo2_cuda_single_matrix_alignment(): g = torch.Generator(device=DEVICE).manual_seed(7) init = [torch.randn(48, 24, generator=g, device=DEVICE)] grads = [[torch.randn(48, 24, generator=g, device=DEVICE) * 2.0] for _ in range(50)] - cfg = dict(num_steps=100, peak_lr=1e-2, orthogonalize=True) + cfg = dict(lr=1e-2, orthogonalize=True) naive_traj = _run(CELO2_naive, init, grads, net, **cfg) cuda_traj = _run(CELO2_CUDA, init, grads, net, **cfg) diff --git a/tests/test_elo.py b/tests/test_elo.py index 3f12784..6fe9543 100644 --- a/tests/test_elo.py +++ b/tests/test_elo.py @@ -13,8 +13,12 @@ from pylo.optim import ELO_naive DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") -SCALING_L2O = "/home/mila/h/huangx/scaling_l2o" -ELO_SRC = os.path.join(SCALING_L2O, "src/learned_optimizers/elo_adfac_mlp_lopt.py") +SCALING_L2O = os.environ.get("SCALING_L2O", "") +ELO_SRC = ( + os.path.join(SCALING_L2O, "src/learned_optimizers/elo_adfac_mlp_lopt.py") + if SCALING_L2O + else "" +) class SmallNet(nn.Module): @@ -31,7 +35,7 @@ def test_step_runs_and_updates(): torch.manual_seed(0) model = SmallNet().to(DEVICE) before = [p.clone() for p in model.parameters()] - optimizer = ELO_naive(model.parameters(), num_steps=50, step_mult=1e-2) + optimizer = ELO_naive(model.parameters(), lr=1e-2) x = torch.randn(16, 10, device=DEVICE) y = torch.randn(16, 1, device=DEVICE) for _ in range(6): @@ -48,7 +52,7 @@ def test_elo_resume(tmp_path): net = MetaMLP(input_size=39, hidden_size=32, hidden_layers=1).to(DEVICE) model = SmallNet().to(DEVICE) - optimizer = ELO_naive(model.parameters(), num_steps=50, step_mult=1e-2, network=net) + optimizer = ELO_naive(model.parameters(), lr=1e-2, network=net) x = torch.randn(32, 10, device=DEVICE) y = torch.randn(32, 1, device=DEVICE) @@ -69,7 +73,7 @@ def test_elo_resume(tmp_path): orig_final = [p.clone().detach() for p in model.parameters()] loaded_model = SmallNet().to(DEVICE) - loaded_opt = ELO_naive(loaded_model.parameters(), num_steps=50, step_mult=1e-2, network=net) + loaded_opt = ELO_naive(loaded_model.parameters(), lr=1e-2, network=net) ckpt = torch.load(path) loaded_model.load_state_dict(ckpt["model"]) loaded_opt.load_state_dict(ckpt["optimizer"]) @@ -96,8 +100,10 @@ def test_jax_numerical_alignment(): spec.loader.exec_module(elo) N = 50 + # pylo's ELO_naive has no internal schedule; pin the JAX reference to a + # constant LR (init_lr == step_mult) so the constant-lr torch side matches. lopt = elo.ELO_AdafacMLPLOpt( - meta_train=False, exp_mult=0.001, step_mult=0.001, init_lr=0.0, + meta_train=False, exp_mult=0.001, step_mult=0.001, init_lr=0.001, warmup_fraction=0.05, weight_decay=0.0, hidden_size=32, hidden_layers=2, 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, @@ -131,7 +137,7 @@ def flat(t, pre=""): pw = torch.nn.Parameter(torch.tensor(W)) pb = torch.nn.Parameter(torch.tensor(B)) - topt = ELO_naive([pw, pb], num_steps=N, network=net) + topt = ELO_naive([pw, pb], lr=0.001, network=net) topt.device = "cpu" for a in ("initial_momentum_decays", "initial_rms_decays", "initial_adafactor_decays"): setattr(topt, a, getattr(topt, a).cpu()) diff --git a/tests/test_elo_cuda.py b/tests/test_elo_cuda.py new file mode 100644 index 0000000..f8a7949 --- /dev/null +++ b/tests/test_elo_cuda.py @@ -0,0 +1,109 @@ +"""Numerical-alignment tests for the CUDA ELO (Adafactor-MLP) optimizer. + +We validate the CUDA implementation against the pure-PyTorch naive one. The naive +optimizer is itself aligned to the reference JAX ELO (see test_elo.py), so +CUDA-vs-naive agreement transitively implies CUDA-vs-JAX agreement without needing +a JAX environment. + +Both optimizers are fed the *same* meta-model weights (``network=net``), the same +initial parameters, and the same gradient sequence; we then compare the parameter +trajectories step by step. The tolerance is looser than a pure-PyTorch comparison +because the kernel accumulates features / MLP in fp32 with fast-math rsqrt/exp. + +These tests require a CUDA device and the compiled ``cuda_lo`` extension; they +skip otherwise. +""" + +import numpy as np +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="ELO CUDA optimizer requires a GPU" +) + +# Skip cleanly if the extension hasn't been built. +pytest.importorskip("cuda_lo") + +from pylo.models.Meta_MLP import MetaMLP +from pylo.optim.ELO_naive import ELO_naive +from pylo.optim.ELO_cuda import ELO_CUDA + +DEVICE = torch.device("cuda") + +# Matrix params in both orientations (RC, to exercise both dc/dr index +# paths) plus a 1D param to exercise the non-factored (vector_like) path. +SHAPES = [(64, 32), (32, 64), (32,)] + + +def _make_init(seed): + g = torch.Generator(device=DEVICE).manual_seed(seed) + return [torch.randn(s, generator=g, device=DEVICE) for s in SHAPES] + + +def _make_grads(n_steps, seed, scale=3.0): + g = torch.Generator(device=DEVICE).manual_seed(seed) + return [ + [torch.randn(s, generator=g, device=DEVICE) * scale for s in SHAPES] + for _ in range(n_steps) + ] + + +def _run(opt_class, init_vals, grads_seq, net, **cfg): + params = [torch.nn.Parameter(v.clone()) for v in init_vals] + opt = opt_class(params, network=net, **cfg) + traj = [] + for grads in grads_seq: + for p, gg in zip(params, grads): + p.grad = gg.clone() + opt.step() + traj.append([p.detach().clone() for p in params]) + return traj + + +def _max_traj_diff(traj_a, traj_b): + worst = 0.0 + for step_a, step_b in zip(traj_a, traj_b): + for a, b in zip(step_a, step_b): + worst = max(worst, torch.max(torch.abs(a - b)).item()) + return worst + + +@pytest.mark.parametrize("clip_grad", [False, True]) +@pytest.mark.parametrize("exp_mult", [0.001, 0.01]) +def test_elo_cuda_matches_naive(clip_grad, exp_mult): + torch.manual_seed(0) + np.random.seed(0) + net = MetaMLP(input_size=39, hidden_size=32, hidden_layers=1).to(DEVICE) + + init_vals = _make_init(seed=1) + grads_seq = _make_grads(n_steps=50, seed=2) + cfg = dict( + lr=1e-2, + exp_mult=exp_mult, + weight_decay=0.1, + clip_grad=clip_grad, + clip_norm=1.0, + ) + + naive_traj = _run(ELO_naive, init_vals, grads_seq, net, **cfg) + cuda_traj = _run(ELO_CUDA, init_vals, grads_seq, net, **cfg) + + assert _max_traj_diff(naive_traj, cuda_traj) < 1e-3 + + +def test_elo_cuda_single_matrix_alignment(): + """Tight-ish check on a single matrix param (factored path).""" + torch.manual_seed(0) + np.random.seed(0) + net = MetaMLP(input_size=39, hidden_size=32, hidden_layers=1).to(DEVICE) + + g = torch.Generator(device=DEVICE).manual_seed(7) + init = [torch.randn(48, 24, generator=g, device=DEVICE)] + grads = [[torch.randn(48, 24, generator=g, device=DEVICE) * 2.0] for _ in range(50)] + cfg = dict(lr=1e-2, exp_mult=0.001) + + naive_traj = _run(ELO_naive, init, grads, net, **cfg) + cuda_traj = _run(ELO_CUDA, init, grads, net, **cfg) + + assert _max_traj_diff(naive_traj, cuda_traj) < 1e-3