diff --git a/baseline/MNIST_MUON_MICROBATCH_RG.md b/baseline/MNIST_MUON_MICROBATCH_RG.md new file mode 100644 index 00000000..9aa793b2 --- /dev/null +++ b/baseline/MNIST_MUON_MICROBATCH_RG.md @@ -0,0 +1,124 @@ +# MNIST MLP3 Muon microbatch RG capture + +This opt-in experiment trains the existing `784-512-512-10` MNIST MLP3 with +the baseline's exact Muon-on-hidden-layers plus auxiliary-AdamW recipe. The +baseline has no gradient accumulation, so each DataLoader minibatch is one +optimizer microbatch. + +## Install + +```bash +cd baseline +python -m pip install -e '.[experiment]' +``` + +## Run + +A bounded first test that saves the three weight matrices after every update: + +```bash +rg-mnist-muon-microbatch \ + --data-dir ./data \ + --output-dir ./results/mnist_mlp3_muon_microbatch_500 \ + --max-steps 500 \ + --capture-every 1 \ + --checkpoint-dtype float32 \ + --overwrite +``` + +The full 30-epoch baseline has about 12,900 optimizer microbatches. Saving all +three matrices in float32 at every step is roughly 32 GiB before container +overhead, so the runner refuses captures above 8 GiB unless explicitly enabled: + +```bash +rg-mnist-muon-microbatch \ + --data-dir ./data \ + --output-dir ./results/mnist_mlp3_muon_microbatch_full \ + --capture-every 1 \ + --checkpoint-dtype float32 \ + --allow-large-capture \ + --overwrite +``` + +To reduce storage, use `--checkpoint-dtype float16`, increase +`--capture-every`, or set `--max-capture-step` while allowing training to +continue. + +## Artifacts + +```text +/ + manifest.json + training_metrics.csv + final_state.pt + microbatch_checkpoints/ + manifest.json + checkpoint_index.csv + frames/ + step_0000000.pt + step_0000001.pt + ... +``` + +Each frame stores `fc1.weight`, `fc2.weight`, and `fc3.weight` only. + +## Original pseudoinverse analysis + +Open: + +```text +notebooks/MNIST_MLP3_Muon_Microbatch_RG_ESD.ipynb +``` + +This exploratory notebook computes the ordinary weight ESD and the supported +pseudoinverse relative-flow spectrum. The latter is complete for square +full-rank matrices but mixes core deformation with subspace overlap for +rectangular matrices. + +## Gauge-aligned rectangular analysis + +Open: + +```text +notebooks/MNIST_MLP3_Muon_Rectangular_RG_ESD.ipynb +``` + +or run: + +```bash +rg-mnist-muon-rectangular-analysis \ + --run-dir ./results/mnist_mlp3_muon_microbatch_500 \ + --step-stride 1 +``` + +For a wide full-row-rank matrix such as `fc1.weight`, write + +```text +W_t = B_t V_t^T, +V_t^T V_t = I. +``` + +The row-space bases at successive steps are aligned by orthogonal Procrustes. +The analysis then reports two independent spectra: + +1. Aligned square-core flow: + `abs(log(sigma(B_t_aligned B_{t-1}^{-1})^2))`. +2. Grassmann angular flow: the squared principal angles `theta_i^2` between + successive row spaces. + +For `fc1.weight` (`512 x 784`), two 512-dimensional row spaces must intersect +in at least 240 dimensions, so there are at most 272 nontrivial angular modes. +The implementation removes those dimension-forced zero angles before fitting. + +For square full-rank `fc2.weight`, the angular sector vanishes and the aligned +core operator reduces numerically to `W_t W_{t-1}^{-1}`. This gives a direct +control showing that the rectangular construction agrees with the original +square relative Jacobian. + +The analysis writes power-law fits, tail sizes, tail fractions, KS distances, +condition numbers, principal-angle diagnostics, ESD archives, and alpha-versus- +step plots. + +`powerlaw` 2.0 uses a built-in upper bound of `alpha = 3` for its power-law +model. Both notebooks explicitly expand the fitting range to +`1.01 <= alpha <= 10` and mark fits that reach the expanded boundary. diff --git a/baseline/notebooks/MNIST_MLP3_Muon_Microbatch_RG_ESD.ipynb b/baseline/notebooks/MNIST_MLP3_Muon_Microbatch_RG_ESD.ipynb new file mode 100644 index 00000000..38a2dc96 --- /dev/null +++ b/baseline/notebooks/MNIST_MLP3_Muon_Microbatch_RG_ESD.ipynb @@ -0,0 +1,470 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "02b664f6", + "metadata": {}, + "source": [ + "# MNIST MLP3 Muon microbatch RG spectra\n", + "\n", + "This notebook analyzes the matrix-only checkpoints written by\n", + "`rg-mnist-muon-microbatch`.\n", + "\n", + "For every layer and checkpoint it computes the ordinary weight ESD\n", + "\n", + "\\[\n", + "\\lambda_i(W_t)=\\sigma_i(W_t)^2.\n", + "\\]\n", + "\n", + "For every successive pair it also constructs the supported relative-flow map\n", + "\n", + "\\[\n", + "J_t=W_tW_{t-1}^{+}\n", + "\\]\n", + "\n", + "in output space for wide matrices, or\n", + "\n", + "\\[\n", + "J_t=W_{t-1}^{+}W_t\n", + "\\]\n", + "\n", + "in input space for tall matrices. Its ESD is \\(\\sigma_i(J_t)^2\\).\n", + "\n", + "The third spectrum,\n", + "\n", + "\\[\n", + "|\\log \\lambda_i(J_t)|,\n", + "\\]\n", + "\n", + "drops the trivial identity/orthogonal mode at \\(\\lambda=1\\). Each positive\n", + "spectrum is fitted with `powerlaw`, and the fitted exponent \\(\\alpha\\) is\n", + "plotted against optimizer step. The horizontal line at \\(\\alpha=2\\) is the\n", + "RG power-counting hypothesis, not a fitted constraint.\n", + "\n", + "**Fit convention.** `powerlaw` 2.0 defaults to an upper bound\n", + "\\(\\alpha\\leq 3\\). This notebook explicitly expands the admissible range to\n", + "\\(1.01\\leq\\alpha\\leq10\\), and records whether a result lands on that expanded\n", + "boundary.\n", + "\n", + "**Caveats.** The relative-flow operator is basis dependent; this is a\n", + "numerical test of the proposal, not a proof of basis invariance. The\n", + "classifier `fc3.weight` has only ten singular values, so its power-law fits are\n", + "low-sample and should be treated as qualitative.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "54ac3aaf", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pandas as pd\n", + "import powerlaw\n", + "\n", + "from rg_baselines.muon_microbatch_capture import (\n", + " load_microbatch_checkpoint,\n", + " load_microbatch_index,\n", + " log_flow_deviation,\n", + " matrix_esd_eigenvalues,\n", + " relative_flow_esd_eigenvalues,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1dc9548c", + "metadata": {}, + "outputs": [], + "source": [ + "# Point this at the output directory created by rg-mnist-muon-microbatch.\n", + "RUN_DIR = Path(\"../results/mnist_mlp3_muon_microbatch\")\n", + "\n", + "# Optional limits for exploratory analysis.\n", + "MAX_CHECKPOINTS = None # e.g. 500; None analyzes all saved checkpoints\n", + "CHECKPOINT_STRIDE = 1 # 1 means truly successive saved checkpoints\n", + "PINV_RTOL = 1e-6\n", + "POWERLAW_MIN_POINTS = 8\n", + "POWERLAW_ALPHA_RANGE = [1.01, 10.0]\n", + "ALPHA_BOUNDARY_ATOL = 1e-3\n", + "LOG_DEVIATION_ZERO_TOL = 1e-12\n", + "\n", + "index = load_microbatch_index(RUN_DIR)\n", + "if MAX_CHECKPOINTS is not None:\n", + " index = index.iloc[: int(MAX_CHECKPOINTS)].copy()\n", + "index = index.iloc[:: int(CHECKPOINT_STRIDE)].reset_index(drop=True)\n", + "index.head(), index.tail(), len(index)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "44a59e92", + "metadata": {}, + "outputs": [], + "source": [ + "def fit_powerlaw(values, *, min_points=POWERLAW_MIN_POINTS):\n", + " values = np.asarray(values, dtype=float)\n", + " values = values[np.isfinite(values) & (values > 0.0)]\n", + " base = {\n", + " \"alpha\": np.nan,\n", + " \"xmin\": np.nan,\n", + " \"ks_distance\": np.nan,\n", + " \"n_tail\": 0,\n", + " \"n_values\": int(values.size),\n", + " \"tail_fraction\": 0.0,\n", + " \"alpha_lower_bound\": float(POWERLAW_ALPHA_RANGE[0]),\n", + " \"alpha_upper_bound\": float(POWERLAW_ALPHA_RANGE[1]),\n", + " \"alpha_at_boundary\": False,\n", + " \"fit_error\": \"\",\n", + " }\n", + " if values.size < int(min_points):\n", + " return {**base, \"fit_error\": \"too_few_values\"}\n", + "\n", + " try:\n", + " # powerlaw 2.x uses parameter_ranges (plural). Its built-in default\n", + " # range for Power_Law is alpha in [0, 3], so expand it explicitly.\n", + " fit = powerlaw.Fit(\n", + " values,\n", + " discrete=False,\n", + " verbose=False,\n", + " parameter_ranges={\"alpha\": list(POWERLAW_ALPHA_RANGE)},\n", + " )\n", + " alpha = float(fit.power_law.alpha)\n", + " xmin = float(fit.power_law.xmin)\n", + " n_tail = int(np.count_nonzero(values >= xmin))\n", + " at_boundary = bool(\n", + " np.isclose(\n", + " alpha,\n", + " POWERLAW_ALPHA_RANGE[0],\n", + " atol=ALPHA_BOUNDARY_ATOL,\n", + " rtol=0.0,\n", + " )\n", + " or np.isclose(\n", + " alpha,\n", + " POWERLAW_ALPHA_RANGE[1],\n", + " atol=ALPHA_BOUNDARY_ATOL,\n", + " rtol=0.0,\n", + " )\n", + " )\n", + " return {\n", + " **base,\n", + " \"alpha\": alpha,\n", + " \"xmin\": xmin,\n", + " \"ks_distance\": float(fit.power_law.D),\n", + " \"n_tail\": n_tail,\n", + " \"tail_fraction\": n_tail / max(values.size, 1),\n", + " \"alpha_at_boundary\": at_boundary,\n", + " }\n", + " except Exception as exc:\n", + " return {\n", + " **base,\n", + " \"fit_error\": f\"{type(exc).__name__}: {exc}\",\n", + " }\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ef36ab38", + "metadata": {}, + "outputs": [], + "source": [ + "fit_rows = []\n", + "esd_records = {}\n", + "previous_matrices = None\n", + "previous_step = None\n", + "\n", + "for row in index.itertuples(index=False):\n", + " payload = load_microbatch_checkpoint(row.checkpoint_path)\n", + " step = int(payload[\"global_step\"])\n", + " matrices = payload[\"matrices\"]\n", + "\n", + " for layer_name, matrix in matrices.items():\n", + " weight_eigs = matrix_esd_eigenvalues(matrix)\n", + " esd_records[(\"weight\", step, layer_name)] = weight_eigs\n", + " fit_rows.append({\n", + " \"spectrum\": \"weight\",\n", + " \"global_step\": step,\n", + " \"previous_step\": np.nan,\n", + " \"step_delta\": np.nan,\n", + " \"epoch\": int(payload[\"epoch\"]),\n", + " \"layer\": layer_name,\n", + " \"operator_side\": \"weight\",\n", + " **fit_powerlaw(weight_eigs),\n", + " })\n", + "\n", + " if previous_matrices is None:\n", + " continue\n", + "\n", + " flow_eigs, side = relative_flow_esd_eigenvalues(\n", + " previous_matrices[layer_name],\n", + " matrix,\n", + " pinv_rtol=PINV_RTOL,\n", + " )\n", + " esd_records[(\"relative_flow\", step, layer_name)] = flow_eigs\n", + " fit_rows.append({\n", + " \"spectrum\": \"relative_flow\",\n", + " \"global_step\": step,\n", + " \"previous_step\": int(previous_step),\n", + " \"step_delta\": step - int(previous_step),\n", + " \"epoch\": int(payload[\"epoch\"]),\n", + " \"layer\": layer_name,\n", + " \"operator_side\": side,\n", + " **fit_powerlaw(flow_eigs),\n", + " })\n", + "\n", + " log_modes = log_flow_deviation(\n", + " flow_eigs,\n", + " zero_tol=LOG_DEVIATION_ZERO_TOL,\n", + " )\n", + " esd_records[(\"log_flow_deviation\", step, layer_name)] = log_modes\n", + " fit_rows.append({\n", + " \"spectrum\": \"log_flow_deviation\",\n", + " \"global_step\": step,\n", + " \"previous_step\": int(previous_step),\n", + " \"step_delta\": step - int(previous_step),\n", + " \"epoch\": int(payload[\"epoch\"]),\n", + " \"layer\": layer_name,\n", + " \"operator_side\": side,\n", + " **fit_powerlaw(log_modes),\n", + " })\n", + "\n", + " previous_matrices = {\n", + " name: value.detach().clone() for name, value in matrices.items()\n", + " }\n", + " previous_step = step\n", + "\n", + "fits = pd.DataFrame(fit_rows).sort_values(\n", + " [\"spectrum\", \"layer\", \"global_step\"]\n", + ").reset_index(drop=True)\n", + "\n", + "fits.to_csv(RUN_DIR / \"microbatch_powerlaw_fits.csv\", index=False)\n", + "np.savez_compressed(\n", + " RUN_DIR / \"microbatch_esd_spectra.npz\",\n", + " **{\n", + " f\"{kind}__step_{step:07d}__{layer.replace('.', '_')}\": values\n", + " for (kind, step, layer), values in esd_records.items()\n", + " },\n", + ")\n", + "\n", + "fits.head(), fits.shape\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e4438912", + "metadata": {}, + "outputs": [], + "source": [ + "boundary_hits = fits[fits[\"alpha_at_boundary\"].astype(bool)]\n", + "print(f\"Power-law fit rows: {len(fits):,}\")\n", + "print(f\"Expanded-range boundary hits: {len(boundary_hits):,}\")\n", + "if len(boundary_hits):\n", + " display(\n", + " boundary_hits[\n", + " [\n", + " \"spectrum\",\n", + " \"global_step\",\n", + " \"layer\",\n", + " \"alpha\",\n", + " \"n_tail\",\n", + " \"n_values\",\n", + " \"ks_distance\",\n", + " ]\n", + " ].head(30)\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b516e625", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_alpha_vs_step(frame, spectrum):\n", + " selected = frame[frame[\"spectrum\"].eq(spectrum)].copy()\n", + " fig, ax = plt.subplots(figsize=(9, 5.5))\n", + " for layer, group in selected.groupby(\"layer\", sort=True):\n", + " group = group[np.isfinite(group[\"alpha\"])].sort_values(\"global_step\")\n", + " ax.plot(\n", + " group[\"global_step\"],\n", + " group[\"alpha\"],\n", + " marker=\".\",\n", + " linewidth=1.2,\n", + " label=layer,\n", + " )\n", + "\n", + " boundary = group[group[\"alpha_at_boundary\"].astype(bool)]\n", + " if not boundary.empty:\n", + " ax.scatter(\n", + " boundary[\"global_step\"],\n", + " boundary[\"alpha\"],\n", + " marker=\"x\",\n", + " s=28,\n", + " )\n", + "\n", + " ax.axhline(\n", + " 2.0,\n", + " linestyle=\"--\",\n", + " linewidth=1.4,\n", + " label=\"RG hypothesis α=2\",\n", + " )\n", + " ax.set_xlabel(\"optimizer step\")\n", + " ax.set_ylabel(\"power-law exponent α\")\n", + " ax.set_title(f\"MNIST MLP3 Muon: {spectrum} α versus step\")\n", + " ax.grid(True, alpha=0.25)\n", + " ax.legend()\n", + " fig.tight_layout()\n", + " return fig, ax\n", + "\n", + "plot_alpha_vs_step(fits, \"weight\")\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ba37cfb9", + "metadata": {}, + "outputs": [], + "source": [ + "plot_alpha_vs_step(fits, \"relative_flow\")\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "93b898ef", + "metadata": {}, + "outputs": [], + "source": [ + "plot_alpha_vs_step(fits, \"log_flow_deviation\")\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "94a09166", + "metadata": {}, + "outputs": [], + "source": [ + "# Inspect ESDs at the first, middle, and last available steps.\n", + "def selected_steps_for(spectrum, layer):\n", + " steps = sorted(\n", + " step\n", + " for kind, step, name in esd_records\n", + " if kind == spectrum and name == layer\n", + " )\n", + " if not steps:\n", + " return []\n", + " return sorted(set([steps[0], steps[len(steps) // 2], steps[-1]]))\n", + "\n", + "\n", + "def plot_selected_esds(spectrum, layer, bins=30):\n", + " fig, ax = plt.subplots(figsize=(8, 5.5))\n", + " for step in selected_steps_for(spectrum, layer):\n", + " values = np.asarray(esd_records[(spectrum, step, layer)], dtype=float)\n", + " values = values[np.isfinite(values) & (values > 0.0)]\n", + " if values.size < 2 or values.min() == values.max():\n", + " continue\n", + " edges = np.logspace(\n", + " np.log10(values.min()),\n", + " np.log10(values.max()),\n", + " bins,\n", + " )\n", + " density, edges = np.histogram(values, bins=edges, density=True)\n", + " centers = np.sqrt(edges[:-1] * edges[1:])\n", + " mask = density > 0.0\n", + " ax.loglog(\n", + " centers[mask],\n", + " density[mask],\n", + " marker=\"o\",\n", + " label=f\"step {step}\",\n", + " )\n", + " ax.set_xlabel(\"eigenvalue / mode magnitude\")\n", + " ax.set_ylabel(\"ESD density\")\n", + " ax.set_title(f\"{spectrum}: {layer}\")\n", + " ax.grid(True, which=\"both\", alpha=0.25)\n", + " ax.legend()\n", + " fig.tight_layout()\n", + " return fig, ax\n", + "\n", + "plot_selected_esds(\"weight\", \"fc1.weight\")\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "39ffd991", + "metadata": {}, + "outputs": [], + "source": [ + "plot_selected_esds(\"relative_flow\", \"fc1.weight\")\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc76b01b", + "metadata": {}, + "outputs": [], + "source": [ + "plot_selected_esds(\"log_flow_deviation\", \"fc1.weight\")\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "861de036", + "metadata": {}, + "outputs": [], + "source": [ + "# Compact summary over the final 25 fitted checkpoints.\n", + "summary = (\n", + " fits[np.isfinite(fits[\"alpha\"])]\n", + " .sort_values(\"global_step\")\n", + " .groupby([\"spectrum\", \"layer\"], as_index=False)\n", + " .tail(25)\n", + " .groupby([\"spectrum\", \"layer\"], as_index=False)\n", + " .agg(\n", + " alpha_median=(\"alpha\", \"median\"),\n", + " alpha_mean=(\"alpha\", \"mean\"),\n", + " alpha_std=(\"alpha\", \"std\"),\n", + " last_step=(\"global_step\", \"max\"),\n", + " median_tail_size=(\"n_tail\", \"median\"),\n", + " median_tail_fraction=(\"tail_fraction\", \"median\"),\n", + " median_ks_distance=(\"ks_distance\", \"median\"),\n", + " boundary_hits=(\"alpha_at_boundary\", \"sum\"),\n", + " )\n", + ")\n", + "summary\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/baseline/notebooks/MNIST_MLP3_Muon_Rectangular_RG_ESD.ipynb b/baseline/notebooks/MNIST_MLP3_Muon_Rectangular_RG_ESD.ipynb new file mode 100644 index 00000000..b5e46379 --- /dev/null +++ b/baseline/notebooks/MNIST_MLP3_Muon_Rectangular_RG_ESD.ipynb @@ -0,0 +1,151 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro", + "metadata": {}, + "source": [ + "# MNIST MLP3 Muon: rectangular RG spectra for FC1 and FC2\n", + "\n", + "This notebook replaces the pseudoinverse-only flow for rectangular matrices with a gauge-aligned decomposition.\n", + "\n", + "For a wide full-row-rank matrix such as\n", + "\n", + "\\[\n", + "W_t\\in\\mathbb{R}^{512\\times 784},\n", + "\\]\n", + "\n", + "write\n", + "\n", + "\\[\n", + "W_t=B_tV_t^\\top,\n", + "\\qquad V_t^\\top V_t=I.\n", + "\\]\n", + "\n", + "The row-space bases at successive steps are aligned by orthogonal Procrustes. The square core flow is then\n", + "\n", + "\\[\n", + "J_t^{\\mathrm{core}}=\\widetilde B_t B_{t-1}^{-1},\n", + "\\]\n", + "\n", + "with spectrum\n", + "\n", + "\\[\n", + "x_i^{\\mathrm{core}}=\\left|\\log\\sigma_i^2(J_t^{\\mathrm{core}})\\right|.\n", + "\\]\n", + "\n", + "The quotient/angular sector is measured independently by the principal angles between the two row spaces:\n", + "\n", + "\\[\n", + "x_i^{\\mathrm{angular}}=\\theta_i^2.\n", + "\\]\n", + "\n", + "For square full-rank `fc2.weight`, the angular sector vanishes and the core operator reduces to\n", + "\n", + "\\[\n", + "J_t^{\\mathrm{core}}=W_tW_{t-1}^{-1}.\n", + "\\]\n", + "\n", + "Both spectra are fitted with `powerlaw.Fit`, using the explicit exponent range \\(1.01\\leq\\alpha\\leq10\\)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "from IPython.display import Image, display\n", + "\n", + "from rg_baselines.mnist_muon_rectangular_analysis import (\n", + " analyze_rectangular_muon_run,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "parameters", + "metadata": {}, + "outputs": [], + "source": [ + "RUN_DIR = Path(\"../results/mnist_mlp3_muon_microbatch\")\n", + "OUTPUT_DIR = RUN_DIR / \"rectangular_rg_analysis\"\n", + "\n", + "# Use 1 for every successive microbatch. Larger values reduce fitting cost,\n", + "# while each selected step is still compared with the immediately prior step.\n", + "STEP_STRIDE = 1\n", + "MAX_STEP = None\n", + "POWERLAW_ALPHA_RANGE = (1.01, 10.0)\n", + "ANGLE_ZERO_TOL = 1e-12\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "run", + "metadata": {}, + "outputs": [], + "source": [ + "RESULT = analyze_rectangular_muon_run(\n", + " RUN_DIR,\n", + " output_dir=OUTPUT_DIR,\n", + " layers=(\"fc1.weight\", \"fc2.weight\"),\n", + " step_stride=STEP_STRIDE,\n", + " max_step=MAX_STEP,\n", + " angle_zero_tol=ANGLE_ZERO_TOL,\n", + " powerlaw_alpha_range=POWERLAW_ALPHA_RANGE,\n", + ")\n", + "\n", + "RESULT[\"summary\"]\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "tables", + "metadata": {}, + "outputs": [], + "source": [ + "display(RESULT[\"summary\"])\n", + "display(RESULT[\"diagnostics\"].groupby(\"layer\").tail(5))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "plots", + "metadata": {}, + "outputs": [], + "source": [ + "for filename in (\n", + " \"alpha_core_log_deviation_vs_step.png\",\n", + " \"alpha_angular_theta_squared_vs_step.png\",\n", + " \"maximum_principal_angle_vs_step.png\",\n", + " \"esd_core_log_deviation_fc1_weight.png\",\n", + " \"esd_angular_theta_squared_fc1_weight.png\",\n", + " \"esd_core_log_deviation_fc2_weight.png\",\n", + "):\n", + " path = OUTPUT_DIR / filename\n", + " print(path)\n", + " display(Image(filename=str(path)))\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/baseline/pyproject.toml b/baseline/pyproject.toml index d6a30fc5..655dd866 100644 --- a/baseline/pyproject.toml +++ b/baseline/pyproject.toml @@ -18,8 +18,13 @@ dependencies = [ experiment = [ "matplotlib>=3.7", "weightwatcher==0.7.7", + "powerlaw>=2.0.0,<3", "jupyter>=1.0", ] +[project.scripts] +rg-mnist-muon-microbatch = "rg_baselines.mnist_muon_microbatch:main" +rg-mnist-muon-rectangular-analysis = "rg_baselines.mnist_muon_rectangular_analysis:main" + [tool.setuptools] packages = ["rg_baselines"] diff --git a/baseline/requirements.txt b/baseline/requirements.txt index 0452bfbd..e02a303f 100644 --- a/baseline/requirements.txt +++ b/baseline/requirements.txt @@ -4,4 +4,5 @@ numpy>=1.24 pandas>=2.0 matplotlib>=3.7 weightwatcher==0.7.7 +powerlaw>=2.0.0,<3 jupyter>=1.0 diff --git a/baseline/rg_baselines/mnist_muon_microbatch.py b/baseline/rg_baselines/mnist_muon_microbatch.py new file mode 100644 index 00000000..130bd8c7 --- /dev/null +++ b/baseline/rg_baselines/mnist_muon_microbatch.py @@ -0,0 +1,393 @@ +"""Train the reference MNIST MLP3 with Muon and capture every microbatch. + +This is an opt-in diagnostic runner. It reuses the baseline's exact MLP3, +MNIST split, Muon-with-auxiliary-AdamW optimizer, gradient clipping, and +update-level warmup/cosine learning-rate schedule. Only the three weight +matrices are persisted at microbatch cadence. +""" + +from __future__ import annotations + +import argparse +import csv +from dataclasses import asdict +import json +import math +from pathlib import Path +import shutil +from typing import Any, Sequence + +import torch +import torch.nn.functional as F + +from .config import BaselineConfig +from .engine import choose_device, evaluate, set_seed +from .model import MLP3 +from .muon_microbatch_capture import ( + DEFAULT_MATRIX_NAMES, + MuonMicrobatchCheckpointRecorder, + estimated_capture_bytes, +) +from .optimizers import build_optimizer, set_scheduled_learning_rates +from .runner import _make_datasets_and_loaders + + +def _atomic_json(payload: dict[str, Any], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(payload, indent=2, sort_keys=True, default=str), + encoding="utf-8", + ) + temporary.replace(path) + + +def _append_metrics(path: Path, row: dict[str, Any]) -> None: + fields = [ + "epoch", + "global_step", + "examples_seen", + "online_train_loss", + "online_train_accuracy", + "validation_loss", + "validation_accuracy", + "test_loss", + "test_accuracy", + "primary_lr", + "auxiliary_lr", + "partial_epoch", + ] + rows: list[dict[str, Any]] = [] + if path.is_file() and path.stat().st_size: + with path.open("r", newline="", encoding="utf-8") as handle: + rows.extend(csv.DictReader(handle)) + epoch = int(row["epoch"]) + rows = [item for item in rows if int(item["epoch"]) != epoch] + rows.append({key: row.get(key, "") for key in fields}) + rows.sort(key=lambda item: int(item["epoch"])) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + temporary.replace(path) + + +def _resolve_device(value: str) -> torch.device: + name = str(value).lower() + if name == "auto": + return choose_device() + if name not in {"cpu", "cuda", "mps"}: + raise ValueError("device must be auto, cpu, cuda, or mps") + device = torch.device(name) + if name == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is unavailable") + if name == "mps" and not ( + hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + ): + raise RuntimeError("MPS was requested but is unavailable") + return device + + +def _capture_count(total_steps: int, capture_every: int, max_capture_step: int) -> int: + upper = min(total_steps, max_capture_step) if max_capture_step else total_steps + return 1 + upper // capture_every + + +def run_muon_microbatch_capture( + *, + data_dir: str | Path, + output_dir: str | Path, + epochs: int = 30, + batch_size: int = 128, + seed: int = 1337, + device: str | torch.device = "auto", + capture_every: int = 1, + max_steps: int = 0, + max_capture_step: int = 0, + checkpoint_dtype: str = "float32", + allow_large_capture: bool = False, + large_capture_gib: float = 8.0, + overwrite: bool = False, + progress: bool = True, +) -> Path: + """Run the exact Muon baseline recipe with matrix-only step checkpoints.""" + + config = BaselineConfig( + optimizer="sgd_momentum_muon", + epochs=int(epochs), + batch_size=int(batch_size), + seed=int(seed), + strict_metrics=False, + ) + config.validate() + if int(max_steps) < 0 or int(max_capture_step) < 0: + raise ValueError("max_steps and max_capture_step must be nonnegative") + if int(capture_every) < 1: + raise ValueError("capture_every must be positive") + if not math.isfinite(float(large_capture_gib)) or large_capture_gib <= 0: + raise ValueError("large_capture_gib must be positive and finite") + + resolved_device = ( + device if isinstance(device, torch.device) else _resolve_device(device) + ) + set_seed(config.seed) + ( + train_loader, + _, + validation_loader, + test_loader, + _, + train_indices, + validation_indices, + ) = _make_datasets_and_loaders( + config, data_dir=data_dir, device=resolved_device + ) + + model = MLP3().to(resolved_device) + optimizer = build_optimizer(model, config) + steps_per_epoch = len(train_loader) + total_steps = int(config.epochs) * steps_per_epoch + training_limit = min(total_steps, int(max_steps)) if max_steps else total_steps + capture_count = _capture_count( + training_limit, int(capture_every), int(max_capture_step) + ) + raw_bytes = estimated_capture_bytes( + model, + matrix_names=DEFAULT_MATRIX_NAMES, + dtype=checkpoint_dtype, + checkpoint_count=capture_count, + ) + estimated_gib = raw_bytes / float(1024**3) + if estimated_gib > float(large_capture_gib) and not allow_large_capture: + raise RuntimeError( + "requested microbatch capture is estimated at " + f"{estimated_gib:.2f} GiB of raw tensors. Re-run with " + "allow_large_capture=True/--allow-large-capture, increase " + "--capture-every, lower --max-steps, or use float16 checkpoints." + ) + + run_dir = Path(output_dir) + if run_dir.exists() and overwrite: + shutil.rmtree(run_dir) + if run_dir.exists() and any(run_dir.iterdir()): + raise FileExistsError( + f"output directory is not empty: {run_dir}; use --overwrite" + ) + run_dir.mkdir(parents=True, exist_ok=True) + + _atomic_json( + { + "schema_version": 1, + "purpose": "mnist_mlp3_muon_microbatch_training", + "config": asdict(config), + "device": str(resolved_device), + "train_examples": len(train_indices), + "validation_examples": len(validation_indices), + "test_examples": 10_000, + "steps_per_epoch": steps_per_epoch, + "baseline_total_steps": total_steps, + "training_step_limit": training_limit, + "capture_every": int(capture_every), + "max_capture_step": int(max_capture_step), + "checkpoint_dtype": checkpoint_dtype, + "estimated_checkpoint_count": capture_count, + "estimated_raw_capture_gib": estimated_gib, + "completed": False, + }, + run_dir / "manifest.json", + ) + + recorder = MuonMicrobatchCheckpointRecorder( + run_dir=run_dir, + model=model, + capture_every=int(capture_every), + max_capture_step=int(max_capture_step), + dtype=checkpoint_dtype, + ) + initial_lrs = set_scheduled_learning_rates( + optimizer, + config, + update_index=0, + total_steps=total_steps, + steps_per_epoch=steps_per_epoch, + ) + recorder.capture( + global_step=0, + epoch=0, + batch_index=0, + examples_seen=0, + learning_rates=initial_lrs, + ) + + global_step = 0 + examples_seen = 0 + last_lrs = dict(initial_lrs) + metrics_path = run_dir / "training_metrics.csv" + stop = False + + for epoch in range(1, config.epochs + 1): + model.train() + loss_sum = 0.0 + correct = 0 + seen = 0 + batches = 0 + for batch_index, (inputs, targets) in enumerate(train_loader, start=1): + last_lrs = set_scheduled_learning_rates( + optimizer, + config, + update_index=global_step, + total_steps=total_steps, + steps_per_epoch=steps_per_epoch, + ) + inputs = inputs.to(resolved_device) + targets = targets.to(resolved_device) + optimizer.zero_grad(set_to_none=True) + logits = model(inputs) + loss = F.cross_entropy(logits, targets) + loss.backward() + torch.nn.utils.clip_grad_norm_( + model.parameters(), float(config.grad_clip_norm) + ) + optimizer.step() + + global_step += 1 + batch_examples = int(targets.numel()) + examples_seen += batch_examples + seen += batch_examples + batches += 1 + loss_value = float(loss.detach().cpu()) + loss_sum += loss_value * batch_examples + correct += int((logits.argmax(1) == targets).sum().detach().cpu()) + recorder.capture( + global_step=global_step, + epoch=epoch, + batch_index=batch_index, + examples_seen=examples_seen, + training_loss=loss_value, + learning_rates=last_lrs, + ) + + if progress and (global_step == 1 or global_step % 100 == 0): + print( + "[mnist-muon-microbatch] " + f"step={global_step}/{training_limit} epoch={epoch} " + f"loss={loss_value:.5f} capture_gib_est={estimated_gib:.2f}", + flush=True, + ) + if global_step >= training_limit: + stop = True + break + + validation = evaluate(model, validation_loader, device=resolved_device) + test = ( + evaluate(model, test_loader, device=resolved_device) + if stop or epoch == config.epochs + else {"loss": float("nan"), "accuracy": float("nan")} + ) + _append_metrics( + metrics_path, + { + "epoch": epoch, + "global_step": global_step, + "examples_seen": examples_seen, + "online_train_loss": loss_sum / max(seen, 1), + "online_train_accuracy": correct / max(seen, 1), + "validation_loss": float(validation["loss"]), + "validation_accuracy": float(validation["accuracy"]), + "test_loss": float(test["loss"]), + "test_accuracy": float(test["accuracy"]), + "primary_lr": float(last_lrs.get("primary", float("nan"))), + "auxiliary_lr": float(last_lrs.get("auxiliary", float("nan"))), + "partial_epoch": int(batches < steps_per_epoch), + }, + ) + if stop: + break + + torch.save( + { + "schema_version": 1, + "purpose": "mnist_mlp3_muon_microbatch_final_state", + "global_step": global_step, + "examples_seen": examples_seen, + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "config": asdict(config), + }, + run_dir / "final_state.pt", + ) + manifest = json.loads((run_dir / "manifest.json").read_text(encoding="utf-8")) + manifest.update( + { + "completed": bool(global_step >= training_limit), + "global_step": global_step, + "examples_seen": examples_seen, + } + ) + _atomic_json(manifest, run_dir / "manifest.json") + return run_dir + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Train the reference MNIST MLP3 Muon baseline and save the three " + "weight matrices at microbatch cadence." + ) + ) + parser.add_argument("--data-dir", default="./data") + parser.add_argument( + "--output-dir", default="./results/mnist_mlp3_muon_microbatch" + ) + parser.add_argument("--epochs", type=int, default=30) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--seed", type=int, default=1337) + parser.add_argument("--device", default="auto") + parser.add_argument("--capture-every", type=int, default=1) + parser.add_argument( + "--max-steps", + type=int, + default=0, + help="Stop training after this many optimizer steps; zero means full run.", + ) + parser.add_argument( + "--max-capture-step", + type=int, + default=0, + help="Stop writing matrix checkpoints after this step; zero means no limit.", + ) + parser.add_argument( + "--checkpoint-dtype", + choices=("float32", "float16", "bfloat16"), + default="float32", + ) + parser.add_argument("--allow-large-capture", action="store_true") + parser.add_argument("--large-capture-gib", type=float, default=8.0) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--quiet", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> None: + args = build_parser().parse_args(argv) + run_muon_microbatch_capture( + data_dir=args.data_dir, + output_dir=args.output_dir, + epochs=args.epochs, + batch_size=args.batch_size, + seed=args.seed, + device=args.device, + capture_every=args.capture_every, + max_steps=args.max_steps, + max_capture_step=args.max_capture_step, + checkpoint_dtype=args.checkpoint_dtype, + allow_large_capture=args.allow_large_capture, + large_capture_gib=args.large_capture_gib, + overwrite=args.overwrite, + progress=not args.quiet, + ) + + +if __name__ == "__main__": + main() diff --git a/baseline/rg_baselines/mnist_muon_rectangular_analysis.py b/baseline/rg_baselines/mnist_muon_rectangular_analysis.py new file mode 100644 index 00000000..0dd74013 --- /dev/null +++ b/baseline/rg_baselines/mnist_muon_rectangular_analysis.py @@ -0,0 +1,488 @@ +"""Analyze aligned core and Grassmann spectra from MNIST Muon checkpoints.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Iterable +import warnings + +import numpy as np +import pandas as pd + +from .muon_microbatch_capture import ( + load_microbatch_checkpoint, + load_microbatch_index, +) +from .rectangular_rg import rectangular_flow_spectra + +DEFAULT_LAYERS = ("fc1.weight", "fc2.weight") +DEFAULT_ALPHA_RANGE = (1.01, 10.0) + + +def fit_powerlaw_spectrum( + values: np.ndarray, + *, + min_points: int = 8, + alpha_range: tuple[float, float] = DEFAULT_ALPHA_RANGE, + boundary_atol: float = 1e-3, +) -> dict[str, Any]: + """Fit one positive spectrum with powerlaw and retain quality diagnostics.""" + + import powerlaw + + data = np.asarray(values, dtype=float) + data = data[np.isfinite(data) & (data > 0.0)] + lower, upper = (float(alpha_range[0]), float(alpha_range[1])) + if not 1.0 < lower < upper: + raise ValueError("alpha_range must satisfy 1 < lower < upper") + if int(min_points) < 2: + raise ValueError("min_points must be at least two") + base: dict[str, Any] = { + "alpha": np.nan, + "xmin": np.nan, + "ks_distance": np.nan, + "n_tail": 0, + "n_values": int(data.size), + "tail_fraction": 0.0, + "alpha_lower_bound": lower, + "alpha_upper_bound": upper, + "alpha_at_boundary": False, + "fit_warning": "", + "fit_error": "", + } + if data.size < int(min_points): + return {**base, "fit_error": "too_few_values"} + + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fit = powerlaw.Fit( + data, + discrete=False, + verbose=False, + parameter_ranges={"alpha": [lower, upper]}, + ) + alpha = float(fit.power_law.alpha) + xmin = float(fit.power_law.xmin) + n_tail = int(np.count_nonzero(data >= xmin)) + at_boundary = bool( + np.isclose(alpha, lower, atol=float(boundary_atol), rtol=0.0) + or np.isclose(alpha, upper, atol=float(boundary_atol), rtol=0.0) + ) + warning_text = " | ".join( + sorted({str(item.message) for item in caught}) + ) + return { + **base, + "alpha": alpha, + "xmin": xmin, + "ks_distance": float(fit.power_law.D), + "n_tail": n_tail, + "tail_fraction": n_tail / max(data.size, 1), + "alpha_at_boundary": at_boundary, + "fit_warning": warning_text, + } + except Exception as exc: + return { + **base, + "fit_error": f"{type(exc).__name__}: {exc}", + } + + +def _selected_steps( + available_steps: Iterable[int], + *, + step_stride: int, + max_step: int | None, + include_final: bool, +) -> list[int]: + steps = sorted({int(step) for step in available_steps if int(step) > 0}) + if max_step is not None: + steps = [step for step in steps if step <= int(max_step)] + selected = [step for step in steps if step % int(step_stride) == 0] + if include_final and steps and steps[-1] not in selected: + selected.append(steps[-1]) + return sorted(set(selected)) + + +def _late_step_summary(fits: pd.DataFrame, *, tail_points: int = 10) -> pd.DataFrame: + valid = fits[np.isfinite(fits["alpha"])].copy() + if valid.empty: + return pd.DataFrame() + return ( + valid.sort_values("global_step") + .groupby(["spectrum", "layer"], as_index=False) + .tail(int(tail_points)) + .groupby(["spectrum", "layer"], as_index=False) + .agg( + alpha_median=("alpha", "median"), + alpha_mean=("alpha", "mean"), + alpha_std=("alpha", "std"), + last_step=("global_step", "max"), + median_tail_size=("n_tail", "median"), + median_tail_fraction=("tail_fraction", "median"), + median_ks_distance=("ks_distance", "median"), + median_mode_count=("n_values", "median"), + boundary_hits=("alpha_at_boundary", "sum"), + ) + ) + + +def _plot_alpha(fits: pd.DataFrame, spectrum: str, output_path: Path) -> None: + import matplotlib.pyplot as plt + + selected = fits[fits["spectrum"].eq(spectrum)].copy() + fig, ax = plt.subplots(figsize=(9, 5.5)) + for layer, group in selected.groupby("layer", sort=True): + group = group[np.isfinite(group["alpha"])].sort_values("global_step") + ax.plot( + group["global_step"], + group["alpha"], + marker=".", + linewidth=1.2, + label=layer, + ) + boundary = group[group["alpha_at_boundary"].astype(bool)] + if not boundary.empty: + ax.scatter( + boundary["global_step"], + boundary["alpha"], + marker="x", + s=28, + ) + ax.axhline( + 2.0, + linestyle="--", + linewidth=1.4, + label="RG hypothesis alpha=2", + ) + ax.set_xlabel("optimizer step") + ax.set_ylabel("power-law exponent alpha") + ax.set_title(f"MNIST MLP3 Muon: {spectrum} alpha versus step") + ax.grid(True, alpha=0.25) + ax.legend() + fig.tight_layout() + fig.savefig(output_path, dpi=180) + plt.close(fig) + + +def _plot_diagnostic( + diagnostics: pd.DataFrame, + column: str, + ylabel: str, + title: str, + output_path: Path, +) -> None: + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(figsize=(9, 5.5)) + for layer, group in diagnostics.groupby("layer", sort=True): + group = group.sort_values("global_step") + ax.plot( + group["global_step"], + group[column], + marker=".", + linewidth=1.2, + label=layer, + ) + ax.set_xlabel("optimizer step") + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.grid(True, alpha=0.25) + ax.legend() + fig.tight_layout() + fig.savefig(output_path, dpi=180) + plt.close(fig) + + +def _plot_selected_esds( + spectra: dict[tuple[str, int, str], np.ndarray], + *, + spectrum: str, + layer: str, + steps: list[int], + output_path: Path, +) -> None: + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(figsize=(8.5, 5.5)) + for step in steps: + values = np.asarray(spectra.get((spectrum, step, layer), []), dtype=float) + values = values[np.isfinite(values) & (values > 0.0)] + if values.size < 2 or values.min() == values.max(): + continue + edges = np.logspace(np.log10(values.min()), np.log10(values.max()), 30) + density, edges = np.histogram(values, bins=edges, density=True) + centers = np.sqrt(edges[:-1] * edges[1:]) + mask = density > 0.0 + ax.loglog( + centers[mask], + density[mask], + marker="o", + linewidth=1.1, + label=f"step {step}", + ) + ax.set_xlabel("mode magnitude") + ax.set_ylabel("ESD density") + ax.set_title(f"{spectrum}: {layer}") + ax.grid(True, which="both", alpha=0.25) + ax.legend() + fig.tight_layout() + fig.savefig(output_path, dpi=180) + plt.close(fig) + + +def analyze_rectangular_muon_run( + run_dir: str | Path, + *, + output_dir: str | Path | None = None, + layers: Iterable[str] = DEFAULT_LAYERS, + step_stride: int = 1, + max_step: int | None = None, + include_final: bool = True, + rank_rtol: float = 1e-10, + log_zero_tol: float = 1e-12, + angle_zero_tol: float = 1e-12, + powerlaw_min_points: int = 8, + powerlaw_alpha_range: tuple[float, float] = DEFAULT_ALPHA_RANGE, +) -> dict[str, Any]: + """Analyze FC1/FC2 aligned-core and angular spectra over a checkpoint run.""" + + if int(step_stride) < 1: + raise ValueError("step_stride must be positive") + root = Path(run_dir) + destination = ( + Path(output_dir) + if output_dir is not None + else root / "rectangular_rg_analysis" + ) + destination.mkdir(parents=True, exist_ok=True) + selected_layers = tuple(str(layer) for layer in layers) + index = load_microbatch_index(root) + by_step = { + int(row.global_step): Path(row.checkpoint_path) + for row in index.itertuples(index=False) + } + selected_steps = _selected_steps( + by_step, + step_stride=int(step_stride), + max_step=max_step, + include_final=bool(include_final), + ) + selected_steps = [step for step in selected_steps if step - 1 in by_step] + if not selected_steps: + raise ValueError("no successive checkpoint pairs matched the selection") + + fit_rows: list[dict[str, Any]] = [] + diagnostic_rows: list[dict[str, Any]] = [] + spectra: dict[tuple[str, int, str], np.ndarray] = {} + + for step in selected_steps: + previous = load_microbatch_checkpoint(by_step[step - 1]) + current = load_microbatch_checkpoint(by_step[step]) + for layer in selected_layers: + if layer not in previous["matrices"] or layer not in current["matrices"]: + raise KeyError(f"checkpoint pair does not contain {layer}") + result = rectangular_flow_spectra( + previous["matrices"][layer], + current["matrices"][layer], + rank_rtol=rank_rtol, + log_zero_tol=log_zero_tol, + angle_zero_tol=angle_zero_tol, + ) + core_values = np.asarray(result["core_log_deviation"], dtype=float) + spectra[("core_log_deviation", step, layer)] = core_values + fit_rows.append( + { + "spectrum": "core_log_deviation", + "global_step": step, + "previous_step": step - 1, + "layer": layer, + "subspace": result["subspace"], + **fit_powerlaw_spectrum( + core_values, + min_points=powerlaw_min_points, + alpha_range=powerlaw_alpha_range, + ), + } + ) + + angular_values = np.asarray(result["angular_eigenvalues"], dtype=float) + spectra[("angular_theta_squared", step, layer)] = angular_values + fit_rows.append( + { + "spectrum": "angular_theta_squared", + "global_step": step, + "previous_step": step - 1, + "layer": layer, + "subspace": result["subspace"], + **fit_powerlaw_spectrum( + angular_values, + min_points=powerlaw_min_points, + alpha_range=powerlaw_alpha_range, + ), + } + ) + + angles = np.asarray(result["principal_angles"], dtype=float) + positive_angles = np.sqrt(angular_values) + diagnostic_rows.append( + { + "global_step": step, + "previous_step": step - 1, + "layer": layer, + "shape": "x".join(str(value) for value in result["shape"]), + "subspace": result["subspace"], + "rank": int(result["rank"]), + "ambient_dimension": int(result["ambient_dimension"]), + "forced_intersection_dimension": int( + result["forced_intersection_dimension"] + ), + "maximum_angular_modes": int(result["maximum_angular_modes"]), + "observed_angular_modes": int(angular_values.size), + "maximum_principal_angle": float( + angles.max() if angles.size else 0.0 + ), + "median_positive_principal_angle": float( + np.median(positive_angles) if positive_angles.size else 0.0 + ), + "previous_condition_number": float( + result["previous_condition_number"] + ), + "current_condition_number": float( + result["current_condition_number"] + ), + } + ) + + fits = pd.DataFrame(fit_rows).sort_values( + ["spectrum", "layer", "global_step"] + ).reset_index(drop=True) + diagnostics = pd.DataFrame(diagnostic_rows).sort_values( + ["layer", "global_step"] + ).reset_index(drop=True) + summary = _late_step_summary(fits) + + fits.to_csv(destination / "rectangular_powerlaw_fits.csv", index=False) + diagnostics.to_csv(destination / "rectangular_flow_diagnostics.csv", index=False) + summary.to_csv(destination / "late_step_alpha_summary.csv", index=False) + np.savez_compressed( + destination / "rectangular_spectra.npz", + **{ + f"{kind}__step_{step:07d}__{layer.replace('.', '_')}": values + for (kind, step, layer), values in spectra.items() + }, + ) + + _plot_alpha( + fits, + "core_log_deviation", + destination / "alpha_core_log_deviation_vs_step.png", + ) + _plot_alpha( + fits, + "angular_theta_squared", + destination / "alpha_angular_theta_squared_vs_step.png", + ) + _plot_diagnostic( + diagnostics, + "maximum_principal_angle", + "maximum principal angle (radians)", + "MNIST MLP3 Muon: maximum subspace angle versus step", + destination / "maximum_principal_angle_vs_step.png", + ) + + display_steps = sorted( + { + selected_steps[0], + selected_steps[len(selected_steps) // 2], + selected_steps[-1], + } + ) + for layer in selected_layers: + _plot_selected_esds( + spectra, + spectrum="core_log_deviation", + layer=layer, + steps=display_steps, + output_path=destination + / f"esd_core_log_deviation_{layer.replace('.', '_')}.png", + ) + _plot_selected_esds( + spectra, + spectrum="angular_theta_squared", + layer=layer, + steps=display_steps, + output_path=destination + / f"esd_angular_theta_squared_{layer.replace('.', '_')}.png", + ) + + manifest = { + "run_dir": str(root), + "output_dir": str(destination), + "layers": list(selected_layers), + "selected_steps": selected_steps, + "step_stride": int(step_stride), + "pair_lag": 1, + "rank_rtol": float(rank_rtol), + "log_zero_tol": float(log_zero_tol), + "angle_zero_tol": float(angle_zero_tol), + "powerlaw_min_points": int(powerlaw_min_points), + "powerlaw_alpha_range": list(powerlaw_alpha_range), + "fit_rows": int(len(fits)), + } + (destination / "analysis_manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8" + ) + return { + "fits": fits, + "diagnostics": diagnostics, + "summary": summary, + "spectra": spectra, + "manifest": manifest, + "output_dir": destination, + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Analyze FC1/FC2 rectangular RG spectra from MNIST Muon checkpoints" + ) + ) + parser.add_argument("--run-dir", required=True) + parser.add_argument("--output-dir", default=None) + parser.add_argument("--step-stride", type=int, default=1) + parser.add_argument("--max-step", type=int, default=None) + parser.add_argument("--rank-rtol", type=float, default=1e-10) + parser.add_argument("--log-zero-tol", type=float, default=1e-12) + parser.add_argument("--angle-zero-tol", type=float, default=1e-12) + parser.add_argument("--alpha-min", type=float, default=DEFAULT_ALPHA_RANGE[0]) + parser.add_argument("--alpha-max", type=float, default=DEFAULT_ALPHA_RANGE[1]) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + result = analyze_rectangular_muon_run( + args.run_dir, + output_dir=args.output_dir, + step_stride=args.step_stride, + max_step=args.max_step, + rank_rtol=args.rank_rtol, + log_zero_tol=args.log_zero_tol, + angle_zero_tol=args.angle_zero_tol, + powerlaw_alpha_range=(args.alpha_min, args.alpha_max), + ) + summary = result["summary"] + if summary.empty: + print("No valid power-law fits") + else: + print(summary.to_string(index=False)) + print(f"Outputs: {result['output_dir']}") + + +if __name__ == "__main__": + main() diff --git a/baseline/rg_baselines/muon_microbatch_capture.py b/baseline/rg_baselines/muon_microbatch_capture.py new file mode 100644 index 00000000..3dcf89dc --- /dev/null +++ b/baseline/rg_baselines/muon_microbatch_capture.py @@ -0,0 +1,347 @@ +"""Microbatch weight capture and relative-flow spectra for MNIST MLP3 Muon. + +The MNIST baseline has no gradient accumulation, so each DataLoader minibatch is +also one optimizer microbatch. This module saves only the three 2-D weight +matrices, keeping the per-step artifact substantially smaller than a full model +and optimizer checkpoint. +""" + +from __future__ import annotations + +import csv +import json +import math +from pathlib import Path +from typing import Any, Iterable, Mapping + +import numpy as np +import torch + +DEFAULT_MATRIX_NAMES = ("fc1.weight", "fc2.weight", "fc3.weight") +CAPTURE_DIRNAME = "microbatch_checkpoints" +INDEX_FILENAME = "checkpoint_index.csv" +MANIFEST_FILENAME = "manifest.json" + +_DTYPE_BY_NAME = { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, +} + + +def _atomic_torch_save(payload: dict[str, Any], path: Path) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + torch.save(payload, temporary) + temporary.replace(path) + return path + + +def _atomic_json(payload: Mapping[str, Any], path: Path) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(dict(payload), indent=2, sort_keys=True, default=str), + encoding="utf-8", + ) + temporary.replace(path) + return path + + +def _append_index(path: Path, row: Mapping[str, Any]) -> None: + fields = [ + "global_step", + "epoch", + "batch_index", + "examples_seen", + "training_loss", + "primary_lr", + "auxiliary_lr", + "checkpoint_path", + ] + rows: list[dict[str, Any]] = [] + if path.is_file() and path.stat().st_size: + with path.open("r", newline="", encoding="utf-8") as handle: + rows.extend(csv.DictReader(handle)) + step = int(row["global_step"]) + rows = [item for item in rows if int(item["global_step"]) != step] + rows.append({key: row.get(key, "") for key in fields}) + rows.sort(key=lambda item: int(item["global_step"])) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + temporary.replace(path) + + +def _named_matrix_parameters( + model: torch.nn.Module, + matrix_names: Iterable[str], +) -> dict[str, torch.nn.Parameter]: + requested = tuple(str(name) for name in matrix_names) + available = dict(model.named_parameters()) + missing = [name for name in requested if name not in available] + if missing: + raise ValueError(f"matrix parameters not found: {missing}") + selected = {name: available[name] for name in requested} + bad = [name for name, value in selected.items() if value.ndim != 2] + if bad: + raise ValueError(f"captured parameters must be matrices: {bad}") + return selected + + +def estimated_capture_bytes( + model: torch.nn.Module, + *, + matrix_names: Iterable[str] = DEFAULT_MATRIX_NAMES, + dtype: str = "float32", + checkpoint_count: int = 1, +) -> int: + """Estimate raw tensor bytes, excluding zip/container overhead.""" + + if dtype not in _DTYPE_BY_NAME: + raise ValueError(f"unsupported checkpoint dtype: {dtype}") + if checkpoint_count < 0: + raise ValueError("checkpoint_count must be nonnegative") + matrices = _named_matrix_parameters(model, matrix_names) + element_size = torch.empty((), dtype=_DTYPE_BY_NAME[dtype]).element_size() + values = sum(int(value.numel()) for value in matrices.values()) + return int(values * element_size * checkpoint_count) + + +class MuonMicrobatchCheckpointRecorder: + """Append-safe recorder for MLP3 weight matrices after optimizer updates.""" + + def __init__( + self, + *, + run_dir: str | Path, + model: torch.nn.Module, + matrix_names: Iterable[str] = DEFAULT_MATRIX_NAMES, + capture_every: int = 1, + max_capture_step: int = 0, + dtype: str = "float32", + ) -> None: + if int(capture_every) < 1: + raise ValueError("capture_every must be positive") + if int(max_capture_step) < 0: + raise ValueError("max_capture_step must be nonnegative") + if dtype not in _DTYPE_BY_NAME: + raise ValueError( + f"dtype must be one of {tuple(_DTYPE_BY_NAME)}, got {dtype!r}" + ) + self.run_dir = Path(run_dir) + self.model = model + self.matrix_names = tuple(str(name) for name in matrix_names) + self.capture_every = int(capture_every) + self.max_capture_step = int(max_capture_step) + self.dtype_name = str(dtype) + self.dtype = _DTYPE_BY_NAME[self.dtype_name] + self.parameters = _named_matrix_parameters(model, self.matrix_names) + self.capture_dir = self.run_dir / CAPTURE_DIRNAME + self.frame_dir = self.capture_dir / "frames" + self.index_path = self.capture_dir / INDEX_FILENAME + self.frame_dir.mkdir(parents=True, exist_ok=True) + self._write_manifest() + + def _write_manifest(self) -> None: + shapes = { + name: list(parameter.shape) + for name, parameter in self.parameters.items() + } + per_checkpoint = estimated_capture_bytes( + self.model, + matrix_names=self.matrix_names, + dtype=self.dtype_name, + ) + _atomic_json( + { + "schema_version": 1, + "purpose": "mnist_mlp3_muon_microbatch_weight_capture", + "microbatch_semantics": ( + "one MNIST DataLoader minibatch and one optimizer update; " + "the baseline does not use gradient accumulation" + ), + "matrix_names": list(self.matrix_names), + "matrix_shapes": shapes, + "checkpoint_dtype": self.dtype_name, + "capture_every": self.capture_every, + "max_capture_step": self.max_capture_step, + "max_capture_step_semantics": "0 means no capture limit", + "estimated_raw_bytes_per_checkpoint": per_checkpoint, + }, + self.capture_dir / MANIFEST_FILENAME, + ) + + def should_capture(self, global_step: int) -> bool: + step = int(global_step) + if step < 0: + raise ValueError("global_step must be nonnegative") + if self.max_capture_step and step > self.max_capture_step: + return False + if step != 0 and step % self.capture_every != 0: + return False + return not self.checkpoint_path(step).is_file() + + def checkpoint_path(self, global_step: int) -> Path: + return self.frame_dir / f"step_{int(global_step):07d}.pt" + + @torch.no_grad() + def capture( + self, + *, + global_step: int, + epoch: int, + batch_index: int, + examples_seen: int, + training_loss: float = float("nan"), + learning_rates: Mapping[str, float] | None = None, + ) -> Path | None: + step = int(global_step) + if not self.should_capture(step): + return None + rates = dict(learning_rates or {}) + matrices = { + name: parameter.detach().to(device="cpu", dtype=self.dtype).clone() + for name, parameter in self.parameters.items() + } + path = self.checkpoint_path(step) + _atomic_torch_save( + { + "schema_version": 1, + "purpose": "mnist_mlp3_muon_microbatch_weights", + "global_step": step, + "epoch": int(epoch), + "batch_index": int(batch_index), + "examples_seen": int(examples_seen), + "training_loss": float(training_loss), + "learning_rates": { + "primary": float(rates.get("primary", float("nan"))), + "auxiliary": float(rates.get("auxiliary", float("nan"))), + }, + "matrix_names": list(self.matrix_names), + "checkpoint_dtype": self.dtype_name, + "matrices": matrices, + }, + path, + ) + _append_index( + self.index_path, + { + "global_step": step, + "epoch": int(epoch), + "batch_index": int(batch_index), + "examples_seen": int(examples_seen), + "training_loss": float(training_loss), + "primary_lr": float(rates.get("primary", float("nan"))), + "auxiliary_lr": float(rates.get("auxiliary", float("nan"))), + "checkpoint_path": str(path.relative_to(self.run_dir)), + }, + ) + return path + + +def load_microbatch_checkpoint(path: str | Path) -> dict[str, Any]: + payload = torch.load(Path(path), map_location="cpu", weights_only=False) + if payload.get("purpose") != "mnist_mlp3_muon_microbatch_weights": + raise ValueError(f"not an MLP3 Muon microbatch checkpoint: {path}") + matrices = payload.get("matrices") + if not isinstance(matrices, dict) or not matrices: + raise ValueError(f"checkpoint contains no matrices: {path}") + return payload + + +def load_microbatch_index(run_dir: str | Path): + """Load the checkpoint index as a DataFrame without importing pandas early.""" + + import pandas as pd + + path = Path(run_dir) / CAPTURE_DIRNAME / INDEX_FILENAME + if not path.is_file(): + raise FileNotFoundError(path) + frame = pd.read_csv(path) + if frame.empty: + raise ValueError(f"microbatch checkpoint index is empty: {path}") + root = Path(run_dir) + frame["checkpoint_path"] = [ + str(value if Path(str(value)).is_absolute() else root / str(value)) + for value in frame["checkpoint_path"] + ] + return frame.sort_values("global_step").reset_index(drop=True) + + +def matrix_esd_eigenvalues(matrix: torch.Tensor) -> np.ndarray: + """Return positive eigenvalues of W W^T (equivalently sigma(W)^2).""" + + value = matrix.detach().to(device="cpu", dtype=torch.float64) + if value.ndim != 2: + raise ValueError(f"matrix must be 2-D, got shape={tuple(value.shape)}") + singular_values = torch.linalg.svdvals(value) + eigenvalues = singular_values.square().numpy() + eigenvalues = eigenvalues[np.isfinite(eigenvalues) & (eigenvalues > 0.0)] + return np.sort(eigenvalues) + + +def relative_flow_operator( + previous: torch.Tensor, + current: torch.Tensor, + *, + pinv_rtol: float = 1e-6, +) -> tuple[torch.Tensor, str]: + """Construct the supported square map between successive weight matrices. + + For a wide/full-row-rank matrix W, use J = W_t pinv(W_{t-1}) in output + space. For a tall matrix, use J = pinv(W_{t-1}) W_t in input space. This + chooses the smaller min(m, n)-dimensional supported space. + """ + + left = previous.detach().to(device="cpu", dtype=torch.float64) + right = current.detach().to(device="cpu", dtype=torch.float64) + if left.ndim != 2 or right.ndim != 2 or left.shape != right.shape: + raise ValueError( + "successive matrices must be 2-D with the same shape, got " + f"{tuple(left.shape)} and {tuple(right.shape)}" + ) + if not math.isfinite(float(pinv_rtol)) or pinv_rtol <= 0.0: + raise ValueError("pinv_rtol must be positive and finite") + pseudo = torch.linalg.pinv(left, rtol=float(pinv_rtol)) + rows, columns = left.shape + if rows <= columns: + return right @ pseudo, "output" + return pseudo @ right, "input" + + +def relative_flow_esd_eigenvalues( + previous: torch.Tensor, + current: torch.Tensor, + *, + pinv_rtol: float = 1e-6, +) -> tuple[np.ndarray, str]: + operator, side = relative_flow_operator( + previous, current, pinv_rtol=pinv_rtol + ) + return matrix_esd_eigenvalues(operator), side + + +def log_flow_deviation(eigenvalues: np.ndarray, *, zero_tol: float = 1e-12) -> np.ndarray: + """Return |log(lambda)|, dropping the trivial identity/orthogonal mode.""" + + values = np.asarray(eigenvalues, dtype=float) + values = values[np.isfinite(values) & (values > 0.0)] + deviations = np.abs(np.log(values)) + return np.sort(deviations[deviations > float(zero_tol)]) + + +__all__ = [ + "CAPTURE_DIRNAME", + "DEFAULT_MATRIX_NAMES", + "MuonMicrobatchCheckpointRecorder", + "estimated_capture_bytes", + "load_microbatch_checkpoint", + "load_microbatch_index", + "log_flow_deviation", + "matrix_esd_eigenvalues", + "relative_flow_esd_eigenvalues", + "relative_flow_operator", +] diff --git a/baseline/rg_baselines/rectangular_rg.py b/baseline/rg_baselines/rectangular_rg.py new file mode 100644 index 00000000..1cf4de15 --- /dev/null +++ b/baseline/rg_baselines/rectangular_rg.py @@ -0,0 +1,250 @@ +"""Gauge-aligned core and Grassmann spectra for rectangular matrix flows. + +A full-rank matrix W in R^{m x n} is decomposed into a square invertible core +and an orthonormal basis for its row space (m <= n) or column space (m > n). +Successive bases are aligned with the orthogonal Procrustes solution before the +square relative-flow operator is formed. The remaining subspace motion is +reported through principal angles. +""" + +from __future__ import annotations + +import math +from typing import Any + +import numpy as np +import torch + + +def _as_matrix(value: torch.Tensor, *, name: str) -> torch.Tensor: + matrix = value.detach().to(device="cpu", dtype=torch.float64) + if matrix.ndim != 2: + raise ValueError(f"{name} must be 2-D, got shape={tuple(matrix.shape)}") + if not torch.isfinite(matrix).all(): + raise ValueError(f"{name} contains non-finite values") + return matrix + + +def _validate_pair( + previous: torch.Tensor, + current: torch.Tensor, + *, + rank_rtol: float, +) -> tuple[torch.Tensor, torch.Tensor]: + left = _as_matrix(previous, name="previous") + right = _as_matrix(current, name="current") + if left.shape != right.shape: + raise ValueError( + "successive matrices must have the same shape, got " + f"{tuple(left.shape)} and {tuple(right.shape)}" + ) + if not math.isfinite(float(rank_rtol)) or not 0.0 < rank_rtol < 1.0: + raise ValueError("rank_rtol must be finite and lie in (0, 1)") + return left, right + + +def _full_rank_svd( + matrix: torch.Tensor, + *, + rank_rtol: float, + name: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]: + left, singular_values, right_h = torch.linalg.svd( + matrix, full_matrices=False + ) + largest = float(singular_values[0]) + smallest = float(singular_values[-1]) + threshold = float(rank_rtol) * largest + if smallest <= threshold: + numerical_rank = int( + torch.count_nonzero(singular_values > threshold).item() + ) + raise ValueError( + f"{name} is not numerically full rank: rank={numerical_rank}, " + f"expected={min(matrix.shape)}, smallest={smallest:.3e}, " + f"threshold={threshold:.3e}" + ) + condition = largest / smallest + return left, singular_values, right_h, condition + + +def _align_subspace_basis( + previous_basis: torch.Tensor, + current_basis: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Align current_basis to previous_basis by orthogonal Procrustes. + + The bases have shape ambient x rank and orthonormal columns. If + previous_basis.T @ current_basis = L diag(cos(theta)) R.T, the minimizing + right action is R L.T. + """ + + overlap = previous_basis.T @ current_basis + left, cosines, right_h = torch.linalg.svd(overlap, full_matrices=False) + alignment = right_h.T @ left.T + aligned = current_basis @ alignment + cosines = cosines.clamp(0.0, 1.0) + principal_angles = torch.acos(cosines) + return aligned, principal_angles, cosines, alignment + + +def principal_angles_sorted(angles: torch.Tensor) -> torch.Tensor: + values = angles.detach().to(device="cpu", dtype=torch.float64).flatten() + return torch.sort(values).values + + +def aligned_core_flow_operator( + previous: torch.Tensor, + current: torch.Tensor, + *, + rank_rtol: float = 1e-10, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Return the square core flow and principal angles for a full-rank pair. + + For a wide matrix, W = B V.T with an invertible m x m core B and a row-space + basis V in St(n, m). The current V is Procrustes-aligned to the previous V, + and J_core = B_current_aligned B_previous^{-1}. + + For a tall matrix, W = U B with an invertible n x n core B and a column-space + basis U in St(m, n). After alignment, + J_core = B_previous^{-1} B_current_aligned. + + For a square full-rank matrix, this reduces (up to roundoff) to + J_core = W_current W_previous^{-1}, with no nontrivial angular sector. + """ + + old, new = _validate_pair(previous, current, rank_rtol=rank_rtol) + rows, columns = old.shape + rank = min(rows, columns) + + old_u, _, old_vh, old_condition = _full_rank_svd( + old, rank_rtol=rank_rtol, name="previous" + ) + new_u, _, new_vh, new_condition = _full_rank_svd( + new, rank_rtol=rank_rtol, name="current" + ) + + if rows <= columns: + old_basis = old_vh.T + new_basis = new_vh.T + aligned_basis, angles, cosines, alignment = _align_subspace_basis( + old_basis, new_basis + ) + old_core = old @ old_basis + new_core = new @ aligned_basis + operator = torch.linalg.solve(old_core.T, new_core.T).T + subspace = "row" + ambient_dimension = columns + else: + old_basis = old_u + new_basis = new_u + aligned_basis, angles, cosines, alignment = _align_subspace_basis( + old_basis, new_basis + ) + old_core = old_basis.T @ old + new_core = aligned_basis.T @ new + operator = torch.linalg.solve(old_core, new_core) + subspace = "column" + ambient_dimension = rows + + forced_intersection = max(0, 2 * rank - ambient_dimension) + metadata: dict[str, Any] = { + "shape": (rows, columns), + "rank": rank, + "subspace": subspace, + "ambient_dimension": ambient_dimension, + "forced_intersection_dimension": forced_intersection, + "maximum_angular_modes": rank - forced_intersection, + "previous_condition_number": old_condition, + "current_condition_number": new_condition, + "cosines": cosines, + "alignment": alignment, + } + return operator, principal_angles_sorted(angles), metadata + + +def squared_singular_value_spectrum(matrix: torch.Tensor) -> np.ndarray: + value = _as_matrix(matrix, name="matrix") + eigenvalues = torch.linalg.svdvals(value).square().numpy() + eigenvalues = eigenvalues[np.isfinite(eigenvalues) & (eigenvalues > 0.0)] + return np.sort(eigenvalues) + + +def core_log_flow_spectrum( + operator: torch.Tensor, + *, + zero_tol: float = 1e-12, +) -> np.ndarray: + """Return |log sigma(J_core)^2| after removing identity modes.""" + + if not math.isfinite(float(zero_tol)) or zero_tol < 0.0: + raise ValueError("zero_tol must be finite and nonnegative") + eigenvalues = squared_singular_value_spectrum(operator) + deviations = np.abs(np.log(eigenvalues)) + return np.sort(deviations[deviations > float(zero_tol)]) + + +def grassmann_angular_spectrum( + principal_angles: torch.Tensor | np.ndarray, + *, + forced_intersection_dimension: int = 0, + zero_tol: float = 1e-12, +) -> np.ndarray: + """Return theta^2 after removing dimension-forced and numerical zeros.""" + + angles = np.asarray( + torch.as_tensor(principal_angles, dtype=torch.float64).cpu(), dtype=float + ).reshape(-1) + angles = np.sort(angles[np.isfinite(angles) & (angles >= 0.0)]) + forced = int(forced_intersection_dimension) + if forced < 0 or forced > angles.size: + raise ValueError( + "forced_intersection_dimension must lie between zero and the " + "number of principal angles" + ) + if not math.isfinite(float(zero_tol)) or zero_tol < 0.0: + raise ValueError("zero_tol must be finite and nonnegative") + values = np.square(angles[forced:]) + return np.sort(values[values > float(zero_tol)]) + + +def rectangular_flow_spectra( + previous: torch.Tensor, + current: torch.Tensor, + *, + rank_rtol: float = 1e-10, + log_zero_tol: float = 1e-12, + angle_zero_tol: float = 1e-12, +) -> dict[str, Any]: + """Compute aligned core and Grassmann spectra for one matrix step.""" + + operator, angles, metadata = aligned_core_flow_operator( + previous, current, rank_rtol=rank_rtol + ) + core_eigenvalues = squared_singular_value_spectrum(operator) + core_log = core_log_flow_spectrum(operator, zero_tol=log_zero_tol) + angular = grassmann_angular_spectrum( + angles, + forced_intersection_dimension=int( + metadata["forced_intersection_dimension"] + ), + zero_tol=angle_zero_tol, + ) + return { + "core_operator": operator, + "core_eigenvalues": core_eigenvalues, + "core_log_deviation": core_log, + "principal_angles": angles.numpy(), + "angular_eigenvalues": angular, + **metadata, + } + + +__all__ = [ + "aligned_core_flow_operator", + "core_log_flow_spectrum", + "grassmann_angular_spectrum", + "principal_angles_sorted", + "rectangular_flow_spectra", + "squared_singular_value_spectrum", +] diff --git a/baseline/tests/test_muon_microbatch_capture.py b/baseline/tests/test_muon_microbatch_capture.py new file mode 100644 index 00000000..1ef6696f --- /dev/null +++ b/baseline/tests/test_muon_microbatch_capture.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +import numpy as np +import torch +import torch.nn as nn + +from rg_baselines.muon_microbatch_capture import ( + MuonMicrobatchCheckpointRecorder, + load_microbatch_checkpoint, + log_flow_deviation, + matrix_esd_eigenvalues, + relative_flow_esd_eigenvalues, + relative_flow_operator, +) + + +class TinyMLP3(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(7, 5) + self.fc2 = nn.Linear(5, 5) + self.fc3 = nn.Linear(5, 3) + + +class MuonMicrobatchCaptureTests(unittest.TestCase): + def test_checkpoint_roundtrip_and_index(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + model = TinyMLP3() + recorder = MuonMicrobatchCheckpointRecorder( + run_dir=root, + model=model, + capture_every=1, + dtype="float32", + ) + first = recorder.capture( + global_step=0, + epoch=0, + batch_index=0, + examples_seen=0, + learning_rates={"primary": 0.1, "auxiliary": 0.01}, + ) + self.assertIsNotNone(first) + assert first is not None + self.assertTrue(first.is_file()) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.add_(0.01) + second = recorder.capture( + global_step=1, + epoch=1, + batch_index=1, + examples_seen=8, + training_loss=1.25, + learning_rates={"primary": 0.1, "auxiliary": 0.01}, + ) + self.assertIsNotNone(second) + assert second is not None + payload = load_microbatch_checkpoint(second) + self.assertEqual(payload["global_step"], 1) + self.assertEqual( + set(payload["matrices"]), + {"fc1.weight", "fc2.weight", "fc3.weight"}, + ) + index = root / "microbatch_checkpoints" / "checkpoint_index.csv" + self.assertTrue(index.is_file()) + self.assertEqual( + len(index.read_text(encoding="utf-8").splitlines()), 3 + ) + + def test_relative_operator_is_identity_for_unchanged_full_rank_matrices( + self, + ) -> None: + for shape in ((5, 7), (5, 5), (7, 5)): + matrix = torch.zeros(*shape, dtype=torch.float64) + rank = min(shape) + matrix[:rank, :rank] = torch.eye(rank, dtype=torch.float64) + operator, side = relative_flow_operator( + matrix, matrix, pinv_rtol=1e-10 + ) + expected = torch.eye(min(shape), dtype=torch.float64) + self.assertEqual(operator.shape, expected.shape) + self.assertIn(side, {"input", "output"}) + self.assertTrue( + torch.allclose(operator, expected, atol=1e-8, rtol=1e-8) + ) + eigenvalues, _ = relative_flow_esd_eigenvalues( + matrix, matrix, pinv_rtol=1e-10 + ) + self.assertTrue( + np.allclose(eigenvalues, np.ones(min(shape)), atol=1e-8) + ) + self.assertEqual(log_flow_deviation(eigenvalues).size, 0) + + def test_weight_esd_is_squared_singular_values(self) -> None: + matrix = torch.diag(torch.tensor([3.0, 2.0, 0.5])) + eigenvalues = matrix_esd_eigenvalues(matrix) + self.assertTrue( + np.allclose(eigenvalues, np.array([0.25, 4.0, 9.0])) + ) + + def test_training_runner_smoke_with_synthetic_mnist(self) -> None: + from torch.utils.data import DataLoader, TensorDataset + + from rg_baselines import mnist_muon_microbatch as experiment + + generator = torch.Generator().manual_seed(123) + inputs = torch.randn(16, 1, 28, 28, generator=generator) + targets = torch.randint(0, 10, (16,), generator=generator) + dataset = TensorDataset(inputs, targets) + + def fake_loaders(config, *, data_dir, device): + del data_dir, device + loader = DataLoader( + dataset, batch_size=config.batch_size, shuffle=False + ) + return ( + loader, + loader, + loader, + loader, + torch.Generator().manual_seed(config.seed + 101), + list(range(12)), + list(range(12, 16)), + ) + + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + with patch.object( + experiment, "_make_datasets_and_loaders", fake_loaders + ): + run_dir = experiment.run_muon_microbatch_capture( + data_dir=root / "data", + output_dir=root / "run", + epochs=3, + batch_size=8, + max_steps=2, + device="cpu", + capture_every=1, + overwrite=True, + progress=False, + ) + index = ( + Path(run_dir) + / "microbatch_checkpoints" + / "checkpoint_index.csv" + ) + self.assertTrue(index.is_file()) + self.assertEqual( + len(index.read_text(encoding="utf-8").splitlines()), 4 + ) + final = torch.load( + Path(run_dir) / "final_state.pt", + map_location="cpu", + weights_only=False, + ) + self.assertEqual(final["global_step"], 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/baseline/tests/test_muon_microbatch_notebook.py b/baseline/tests/test_muon_microbatch_notebook.py new file mode 100644 index 00000000..ac0b7813 --- /dev/null +++ b/baseline/tests/test_muon_microbatch_notebook.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import json +from pathlib import Path +import unittest + + +class MuonMicrobatchNotebookTests(unittest.TestCase): + def test_notebook_contains_requested_analysis(self) -> None: + path = ( + Path(__file__).resolve().parents[1] + / "notebooks" + / "MNIST_MLP3_Muon_Microbatch_RG_ESD.ipynb" + ) + payload = json.loads(path.read_text(encoding="utf-8")) + source = "\n".join( + "".join(cell.get("source", [])) + for cell in payload.get("cells", []) + ) + self.assertIn("powerlaw.Fit", source) + self.assertIn("parameter_ranges", source) + self.assertIn("POWERLAW_ALPHA_RANGE = [1.01, 10.0]", source) + self.assertIn("alpha_at_boundary", source) + self.assertIn("matrix_esd_eigenvalues", source) + self.assertIn("relative_flow_esd_eigenvalues", source) + self.assertIn("log_flow_deviation", source) + self.assertIn("alpha", source) + self.assertIn("global_step", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/baseline/tests/test_muon_rectangular_notebook.py b/baseline/tests/test_muon_rectangular_notebook.py new file mode 100644 index 00000000..564eabcd --- /dev/null +++ b/baseline/tests/test_muon_rectangular_notebook.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import ast +import json +from pathlib import Path +import unittest + + +class MuonRectangularNotebookTests(unittest.TestCase): + def test_notebook_contains_fc1_fc2_rectangular_analysis(self) -> None: + path = ( + Path(__file__).resolve().parents[1] + / "notebooks" + / "MNIST_MLP3_Muon_Rectangular_RG_ESD.ipynb" + ) + payload = json.loads(path.read_text(encoding="utf-8")) + source = "\n".join( + "".join(cell.get("source", [])) + for cell in payload.get("cells", []) + ) + self.assertIn("analyze_rectangular_muon_run", source) + self.assertIn("fc1.weight", source) + self.assertIn("fc2.weight", source) + self.assertIn("core_log_deviation", source) + self.assertIn("angular_theta_squared", source) + self.assertIn("POWERLAW_ALPHA_RANGE", source) + for index, cell in enumerate(payload.get("cells", [])): + if cell.get("cell_type") != "code": + continue + ast.parse( + "".join(cell.get("source", [])), + filename=f"{path}:cell-{index}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/baseline/tests/test_rectangular_rg.py b/baseline/tests/test_rectangular_rg.py new file mode 100644 index 00000000..7b6250f8 --- /dev/null +++ b/baseline/tests/test_rectangular_rg.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import unittest + +import numpy as np +import torch + +from rg_baselines.rectangular_rg import ( + aligned_core_flow_operator, + grassmann_angular_spectrum, + rectangular_flow_spectra, +) + + +class RectangularRGTests(unittest.TestCase): + @staticmethod + def _orthonormal( + ambient: int, + rank: int, + generator: torch.Generator, + ) -> torch.Tensor: + value = torch.randn( + ambient, + rank, + generator=generator, + dtype=torch.float64, + ) + basis, _ = torch.linalg.qr(value, mode="reduced") + return basis + + def test_square_case_reduces_to_exact_relative_jacobian(self) -> None: + generator = torch.Generator().manual_seed(7) + previous = torch.randn(8, 8, generator=generator, dtype=torch.float64) + current = torch.randn(8, 8, generator=generator, dtype=torch.float64) + previous = previous + 3.0 * torch.eye(8, dtype=torch.float64) + current = current + 3.0 * torch.eye(8, dtype=torch.float64) + + operator, angles, metadata = aligned_core_flow_operator( + previous, current + ) + expected = torch.linalg.solve(previous.T, current.T).T + + self.assertTrue(torch.allclose(operator, expected, atol=1e-10, rtol=1e-10)) + self.assertEqual(metadata["forced_intersection_dimension"], 8) + angular = grassmann_angular_spectrum( + angles, + forced_intersection_dimension=metadata[ + "forced_intersection_dimension" + ], + ) + self.assertEqual(angular.size, 0) + + def test_wide_same_subspace_recovers_known_core_flow(self) -> None: + generator = torch.Generator().manual_seed(11) + basis = self._orthonormal(11, 7, generator) + previous_core = torch.randn( + 7, 7, generator=generator, dtype=torch.float64 + ) + 3.0 * torch.eye(7, dtype=torch.float64) + current_core = torch.randn( + 7, 7, generator=generator, dtype=torch.float64 + ) + 3.0 * torch.eye(7, dtype=torch.float64) + previous = previous_core @ basis.T + current = current_core @ basis.T + + result = rectangular_flow_spectra(previous, current) + expected = torch.linalg.solve(previous_core.T, current_core.T).T + + self.assertTrue( + torch.allclose( + result["core_operator"], expected, atol=1e-10, rtol=1e-10 + ) + ) + self.assertEqual(result["forced_intersection_dimension"], 3) + self.assertEqual(result["maximum_angular_modes"], 4) + self.assertEqual(result["angular_eigenvalues"].size, 0) + + def test_pure_subspace_motion_has_identity_aligned_core(self) -> None: + generator = torch.Generator().manual_seed(13) + previous_basis = self._orthonormal(11, 7, generator) + current_basis = self._orthonormal(11, 7, generator) + overlap = previous_basis.T @ current_basis + left, _, right_h = torch.linalg.svd(overlap, full_matrices=False) + alignment = right_h.T @ left.T + + core = torch.randn( + 7, 7, generator=generator, dtype=torch.float64 + ) + 3.0 * torch.eye(7, dtype=torch.float64) + previous = core @ previous_basis.T + current = (core @ alignment.T) @ current_basis.T + + result = rectangular_flow_spectra(previous, current) + identity = torch.eye(7, dtype=torch.float64) + + self.assertTrue( + torch.allclose( + result["core_operator"], identity, atol=1e-9, rtol=1e-9 + ) + ) + self.assertEqual(result["angular_eigenvalues"].size, 4) + self.assertTrue(np.all(result["angular_eigenvalues"] > 0.0)) + + def test_rank_deficient_pair_is_rejected(self) -> None: + previous = torch.zeros(5, 8, dtype=torch.float64) + current = torch.randn(5, 8, dtype=torch.float64) + with self.assertRaisesRegex(ValueError, "not numerically full rank"): + aligned_core_flow_operator(previous, current) + + +if __name__ == "__main__": + unittest.main()