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/README.md b/README.md index 513178d..8c5f4ae 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) @@ -72,22 +75,28 @@ python -m pylo.util.patch_mup ## Quick Start +Taking `ELO-CELO2` (the strongest LO) for example. + ```python import torch -from pylo.optim import VeLO_CUDA +from pylo.optim import ELO_CELO2_CUDA -# Initialize a model model = torch.nn.Linear(10, 2) -# Create a learned optimizer instance -optimizer = VeLO_CUDA(model.parameters()) +num_steps = 1000 # total optimization steps + +# Meta-learned weights download automatically from the Hugging Face Hub on first use. +# 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) -# Use it like any PyTorch optimizer -for epoch in range(10): +for step in range(num_steps): optimizer.zero_grad() loss = loss_fn(model(input), target) loss.backward() - optimizer.step(loss) # pass the loss + optimizer.step() + scheduler.step() ``` ## Sharing Learned Optimizers diff --git a/pylo/__init__.py b/pylo/__init__.py index 4edaaa8..8ec79ee 100644 --- a/pylo/__init__.py +++ b/pylo/__init__.py @@ -5,9 +5,16 @@ "MetaMLP", "VeLOMLP", "VeLORNN", + "CELO2MLP", "AdafacLO_naive", "MuLO_naive", "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/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/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/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 new file mode 100644 index 0000000..8e56141 --- /dev/null +++ b/pylo/optim/CELO2_cuda.py @@ -0,0 +1,341 @@ +"""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: + + * global-norm gradient clipping, + * Newton-Schulz orthogonalization of the 2D+ update, + * the post-MLP second-moment normalization + rmsmult scaling, and + * 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 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 +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 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. Like the naive optimizer, the AdamW + branch for 1D / embedding parameters maintains independent moments. + """ + + def __init__( + self, + params, + lr=1e-3, + # Weight decay + weight_decay=0.0, + # 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, + 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 not torch.cuda.is_available(): + raise RuntimeError("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 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.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 + + # 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() + + # ---------------------------------------------------------------- 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_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) + 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 + + 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): + 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"] + adam_lr = group["adam_lr_mult"] * lr + 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] + is_1d = p.dim() <= 1 or force_adam + if use_adamw_for_1d and is_1d: + 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 + ) + + 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/CELO2_naive.py b/pylo/optim/CELO2_naive.py new file mode 100644 index 0000000..a778e15 --- /dev/null +++ b/pylo/optim/CELO2_naive.py @@ -0,0 +1,387 @@ +"""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 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, 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 + +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. + 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_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. + 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, + lr=1e-3, + # Weight decay + weight_decay=0.0, + # 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, + 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, + ): + defaults = dict( + lr=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) + + # 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 + ) + 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.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 + + # 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 _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 + 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 _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) + 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: + lr = group["lr"] + adam_lr = group["adam_lr_mult"] * lr + 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] + is_1d = p.dim() <= 1 or force_adam + if use_adamw_for_1d and is_1d: + 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 new file mode 100644 index 0000000..7af8e0f --- /dev/null +++ b/pylo/optim/ELO_CELO2_cuda.py @@ -0,0 +1,298 @@ +"""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, 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 + +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(Optimizer): + """Inference-time ELO-CELO2 optimizer (CUDA CELO2 forward, shared-accumulator AdamW).""" + + def __init__( + self, + params, + 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 (shares momentum[0] / rms accumulators) + 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: Optional[CELO2MLP] = None, + ): + 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 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"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 new file mode 100644 index 0000000..ae7ac2d --- /dev/null +++ b/pylo/optim/ELO_CELO2_naive.py @@ -0,0 +1,306 @@ +"""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 +for 1D parameters. + +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 + +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(Optimizer): + """Inference-time ELO-CELO2 optimizer (CELO2 forward, shared-accumulator AdamW).""" + + def __init__( + self, + params, + 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 (shares momentum[0] / rms[-1] accumulators) + 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: Optional[CELO2MLP] = None, + ): + defaults = dict( + lr=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) + + # 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 + ) + 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 new file mode 100644 index 0000000..88e70b0 --- /dev/null +++ b/pylo/optim/ELO_naive.py @@ -0,0 +1,238 @@ +"""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), 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 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``. + 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). + 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, + ): + defaults = dict( + lr=lr, + exp_mult=exp_mult, + 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 + ) + 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 _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"] + lr = group["lr"] + 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=-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 f8f2bed..6044564 100644 --- a/pylo/optim/__init__.py +++ b/pylo/optim/__init__.py @@ -1,9 +1,27 @@ 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 +from pylo.optim.ELO_naive import ELO_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", + "ELO_naive", +] + +# 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 +__all__.extend(["CELO2", "ELO_CELO2", "ELO"]) # Try to import CUDA-based optimizers try: @@ -33,3 +51,26 @@ 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 + +# 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 new file mode 100644 index 0000000..eccff59 --- /dev/null +++ b/scripts/convert_celo2_checkpoint.py @@ -0,0 +1,103 @@ +"""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 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/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/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.py b/tests/test_celo2.py new file mode 100644 index 0000000..842f9f8 --- /dev/null +++ b/tests/test_celo2.py @@ -0,0 +1,314 @@ +"""Tests for the CELO2 / ELO-CELO2 learned optimizers.""" + +import functools +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 = 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) +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): + 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(), 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], 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(), 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(), 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], 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 + + +@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 + 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 + + celo2_lopt, elo_celo2 = _load_jax_lopts() + + N = 60 + LR = 1e-3 + cfg = dict( + 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, + ) + 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], 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"): + 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 diff --git a/tests/test_celo2_cuda.py b/tests/test_celo2_cuda.py new file mode 100644 index 0000000..2c66304 --- /dev/null +++ b/tests/test_celo2_cuda.py @@ -0,0 +1,128 @@ +"""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, 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( + 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, 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(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(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 diff --git a/tests/test_elo.py b/tests/test_elo.py new file mode 100644 index 0000000..6fe9543 --- /dev/null +++ b/tests/test_elo.py @@ -0,0 +1,154 @@ +"""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 = 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): + 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(), lr=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(), 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 / "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(), 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(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 + # 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.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, + 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], 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()) + + 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 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