diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 508df64..336fe66 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -48,13 +48,15 @@ jobs: run: | test ! -f setup.py python - <<'PY' - import glob, zipfile + import glob, tomllib, zipfile wheels = glob.glob('dist/*.whl') assert len(wheels) == 1, wheels with zipfile.ZipFile(wheels[0]) as z: meta = [n for n in z.namelist() if n.endswith('METADATA')][0] text = z.read(meta).decode() - assert 'Version: 0.3.0' in text + with open('pyproject.toml', 'rb') as f: + version = tomllib.load(f)['project']['version'] + assert f'Version: {version}' in text, version assert 'Requires-Dist: torch>=2.0.0' in text - print('wheel metadata OK') + print(f'wheel metadata OK ({version})') PY diff --git a/CHANGELOG.md b/CHANGELOG.md index 524af7e..358823c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,102 @@ # Changelog +## 0.4.0 — 2026-11-07 + +The full-verification pass: the optimizer is now checked against two +independent external oracles and an end-to-end objective, and the second +moment was brought to the paper's semantics. + +### Changed (mathematics) + +- **The second moment is the scalar squared Riemannian norm + `lambda(x)^2 * ||g_riem||^2` per row** (Algorithm 1 of Becigneul & Ganea; + geoopt's `component_inner` — the Poincare ball is a single manifold + component), stored broadcast over coordinates. In 0.3.x it was a + per-coordinate accumulator, which is (a) not intrinsic — it depends on the + chart's rotation — and (b) unbounded near the boundary: without the metric + weight the Riemannian length of a first step grows as `lambda(x) * lr` + instead of the bounded `lr * sqrt(dim)`. A scalar needs no parallel + transport, so the sqrt-transport heuristic of 0.3.x is gone: only the + first moment is transported (the paper's `tau` update). The Euclidean + path (k=0) keeps per-coordinate moments (R^d is a product of `d` + components) and remains exactly `torch.optim.Adam`/`AdamW`. +- **State format v3, checkpoints v3.** v1/v2 moments are semantically + incompatible with the scalar second moment and are reset on migration + (weights preserved, `legacy_migrated_from` recorded); v2 states written + by 0.3.0 migrate with the same rule. + +### Verified against independent oracles + +- **geoopt 0.5.1** (`RiemannianAdam` on `PoincareBallExact`, the exact + exponential-map variant): the full multi-step trajectory — `w`, `m` and + `v` over 10 steps, k in {0.7, 1.0}, weight decay in {0, 0.02} — matches + to float32 rounding (~2e-7). `tests/test_geoopt_step.py`; runs whenever + geoopt is importable (install it into a separate directory and set + `PYTHONPATH`; skipped in CI). Near-boundary rows are not compared + directly: geoopt computes entirely in float32 and loses precision + catastrophically there; DSA's near-boundary behaviour is pinned by the + intrinsic invariant below and the stress test. +- **torch.optim** (`tests/test_torch_equivalence.py`): k=0 dense 12-step + trajectories equal `torch.optim.Adam` ("l2") and `torch.optim.AdamW` + ("decoupled") within 2e-6. +- **End-to-end descent** (`tests/test_optimization.py`): 300 steps minimise + the geodesic distance to fixed targets on the ball from ~1.5 to below + 0.05 per row (gradients from torch.autograd through the verified + distance formula); a near-boundary stress test (rows at 0.9999 of the + radius, gradient scale 50, 60 steps) stays finite and inside the ball. +- **Intrinsic invariant** (`test_h2_first_step_riemannian_length_is_bounded`): + the Riemannian length of a first step is at most `lr * sqrt(dim)` at any + point of the ball — the 0.3.x per-coordinate accumulator violated this + near the boundary by a factor of ~lambda. + + +## 0.3.1 — 2026-11-07 + +Audit follow-up: a line-by-line review of the 0.3.0 math and logic against +the 0.2.0 baseline, Algorithm 1 of Becigneul & Ganea (ICLR 2019) and +geoopt's `RiemannianAdam` implementation. Three behavior fixes (each with a +regression test) and a documentation-accuracy fix. + +### Fixed + +- **R2 regression — a crash in the middle of a `resume=False` + reinitialisation left a mixed state that the next `resume=True` accepted + silently.** 0.2.0 marked the transition in meta; the 0.3.0 constructor + path dropped the marker (only `import_checkpoint` kept it), and since all + state files keep their correct size through a partial rewrite, size + verification cannot detect the mixture. The marker is now written before + the data files are touched whenever a reinitialisation runs over an + existing directory; `resume=True` refuses a marked state, `resume=False` + is the documented recovery and clears the marker. Covered by a + crash-simulation test (the write fails right after `w`). +- **A corrupted negative per-row counter crashed `step()` with a misleading + error.** `t < 0` on disk (torn write) survived validation: bias correction + `1 - beta**t` turned negative, `sqrt(v/bc2)` produced NaN and the whole + step aborted with "non-finite weights; reduce lr" instead of the + documented A4 row reset. The counter is now part of the row-corruption + check: the row is fully reset (`w=m=v=0, t=0`) and the step proceeds as a + canonical first Adam step from the origin (pinned analytically). +- **Silent weight-decay semantics switch.** `weight_decay_mode` is stored + in per-param meta and in checkpoints, but a mismatch with the constructor + argument was never reported: a state trained with `"l2"` could silently + continue with `"decoupled"`. `resume=True` and `import_checkpoint` now + log a warning naming both modes. + +### Documentation + +- The hyperbolic algorithm provenance is now stated precisely: + per-coordinate moments (geoopt style — the paper keeps `v` scalar), + first-moment transport (the paper's `tau` update), and a + positivity-preserving sqrt-transport of the second moment (the paper + needs no transport for scalar `v`; geoopt does not transport `v`). + The previous wording ("implements the Riemannian Adam of Becigneul & + Ganea", "the paper's operation") overstated the fidelity; the 0.3.0 + known-limitation note was corrected accordingly. +- The writer now documents that per-row write ordering is guaranteed + through `step()`; direct concurrent `put_update` calls with overlapping + rows can reorder writes. + + ## 0.3.0 — 2026-11-04 Correctness release for hyperbolic training. The headline change is that @@ -121,9 +218,9 @@ CUDA_VISIBLE_DEVICES="" python -m pytest -q chart itself is ill-conditioned there. - `v` is a per-coordinate diagonal moment. Transporting it directly as a tangent vector can produce negative components and NaNs, so its - non-negative square root is transported and squared back. This is the - positivity-preserving version of the paper's operation, not a - component-wise `sqrt` of a rotated vector. + non-negative square root is transported and squared back — a heuristic + that preserves the Riemannian scale of `sqrt(v)`. (Superseded in 0.4.0: + `v` is now the paper's scalar `||g||_x^2` and is not transported at all.) - CUDA-specific memory behavior is tested only through object lifetime on CPU in CI; no GPU is required for the test suite. diff --git a/README.md b/README.md index 8ca65c1..a68ca86 100644 --- a/README.md +++ b/README.md @@ -249,13 +249,15 @@ What `step()` does when `k > 0`: 2. Converts the Euclidean gradient to the Riemannian one, $\operatorname{grad}_R f = \lambda(w)^{-2} g_E$, with $\lambda(w) = 2 / (1 - k\|w\|^2)$. -3. Updates Adam's moments: $m$ is a tangent vector, $v$ is a non-negative - diagonal scale. The active rows are computed in `float64` for numerical +3. Updates Adam's moments: $m$ is a tangent vector, $v$ is the scalar + squared Riemannian norm $\|g\|_x^2$ per row (stored broadcast over + coordinates — the ball is one manifold component, as in the paper and + geoopt). The active rows are computed in `float64` for numerical stability near the boundary. 4. Takes a geodesic step with the exponential map $w \leftarrow \exp_w(-\eta\, \hat m / (\sqrt{\hat v}+\epsilon))$. -5. Transports $m$ (and the square root of $v$) from $T_w$ to $T_{w_{\text{new}}}$ - by parallel transport, so the moments follow the new tangent space. +5. Transports $m$ from $T_w$ to $T_{w_{\text{new}}}$ by parallel transport + (the paper's $\tau$ update); the scalar $v$ needs no transport. 6. Projects the result into the ball as a numerical guard and queues the consistent row update. diff --git a/dsa/__init__.py b/dsa/__init__.py index 4142c40..3740153 100644 --- a/dsa/__init__.py +++ b/dsa/__init__.py @@ -1,6 +1,6 @@ from .optimizer import DiskSparseRiemannianAdam from .async_writer import AsyncDiskWriter, WriterState -__version__ = "0.3.0" +__version__ = "0.4.0" __all__ = ["DiskSparseRiemannianAdam", "AsyncDiskWriter", "WriterState"] diff --git a/dsa/async_writer.py b/dsa/async_writer.py index 6f943f8..dec02f9 100644 --- a/dsa/async_writer.py +++ b/dsa/async_writer.py @@ -10,6 +10,12 @@ * F6: ``shutdown`` is a state machine (``RUNNING -> DRAINING -> STOPPED``, or ``FAILED``) and is retryable: a transient failure (stuck disk, slow msync) does not turn a second call into a false success. + +Per-row write ordering is guaranteed through the optimizer's ``step()``, +which checks the in-flight set before reading. Direct concurrent +``put_update`` calls with overlapping rows can reorder writes: the +synchronous fallback writes at once while older batches may still be +sitting in the queue. """ from __future__ import annotations diff --git a/dsa/optimizer.py b/dsa/optimizer.py index 87c8c53..c5e42bb 100644 --- a/dsa/optimizer.py +++ b/dsa/optimizer.py @@ -4,18 +4,29 @@ moment) and a per-row step counter ``t`` in ``np.memmap`` files, so only the rows touched by a batch are read into RAM/VRAM. -Hyperbolic mode (``k > 0``) implements the Riemannian Adam of Becigneul & -Ganea (ICLR 2019) on the Poincare ball of curvature ``-k``: +Hyperbolic mode (``k > 0``) is the Riemannian Adam of Becigneul & Ganea +(ICLR 2019) on the Poincare ball of curvature ``-k``, matching geoopt's +``RiemannianAdam`` on ``PoincareBallExact`` (verified step-by-step by +tests): * the Euclidean gradient is converted to a Riemannian one at the current point (``lambda(x)^-2 g``); -* Adam moments live in the tangent space at the current point; -* the step is an exponential map, not a Euclidean subtraction plus clip; -* after the step the moments are parallel-transported to the new tangent - space. - -``k = 0`` is plain Euclidean Adam. See ``README.md`` and ``CHANGELOG.md`` -for the exact semantics and the 0.3.0 migration notes. +* the first moment lives in the tangent space at the current point; the + second moment is the scalar squared Riemannian norm ``||g||_x^2`` + (the paper's Algorithm 1; geoopt's ``component_inner`` - the ball is a + single manifold component), stored broadcast over coordinates; +* the step is the exact exponential map, not a Euclidean subtraction plus + clip (geoopt's default ``PoincareBall`` uses a first-order retraction; + the ``Exact`` variant, like DSA, uses the exponential map); +* after the step the first moment is parallel-transported to the new + tangent space (the paper's ``tau`` update); the scalar second moment + needs no transport. + +``k = 0`` is plain per-coordinate Euclidean Adam: ``R^d`` is the product +of ``d`` manifold components, so the second moment stays per-coordinate, +and the update equals ``torch.optim.Adam`` (``"l2"``) / ``torch.optim.AdamW`` +(``"decoupled"``) exactly. See ``README.md`` and ``CHANGELOG.md`` for the +exact semantics and migration notes. """ from __future__ import annotations @@ -36,21 +47,26 @@ from .manifold import ( BALL_EPS, expmap, + lambda_x, parallel_transport, project_to_ball, project_to_ball_np, riemannian_gradient, + sq_norm, ) logger = logging.getLogger("DSA") # Version of the per-parameter and global metadata formats. -STATE_FORMAT_VERSION = 2 -LEGACY_STATE_FORMAT_VERSIONS = (1,) +# v3: the second moment is the scalar ||g||_x^2 per row (paper Algorithm 1 / +# geoopt component_inner); v2/v1 moments are incompatible and are reset on +# migration (weights are preserved). +STATE_FORMAT_VERSION = 3 +LEGACY_STATE_FORMAT_VERSIONS = (1, 2) ALGORITHM = "riemannian_adam" CHECKPOINT_FORMAT = "DSA_checkpoint" -CHECKPOINT_VERSION = 2 +CHECKPOINT_VERSION = 3 CHECKPOINT_MANIFEST_NAME = "checkpoint.json" GLOBAL_META_NAME = "dsa_global_meta.json" @@ -315,10 +331,11 @@ def _read_global_meta_optional(self) -> Optional[dict]: except (OSError, ValueError) as exc: raise RuntimeError(f"глобальный meta-файл {path} не читается: {exc!r}") from exc version = int(meta.get("format_version", -1)) - if version != STATE_FORMAT_VERSION: + if version != STATE_FORMAT_VERSION and version not in LEGACY_STATE_FORMAT_VERSIONS: raise RuntimeError( f"глобальный meta-файл {path}: версия формата {version!r} " - f"не поддерживается (ожидается {STATE_FORMAT_VERSION})") + f"не поддерживается (ожидаются " + f"{list(LEGACY_STATE_FORMAT_VERSIONS) + [STATE_FORMAT_VERSION]})") if "global_step" in meta: try: step = int(meta["global_step"]) @@ -409,6 +426,15 @@ def _check_meta_compat(self, meta: dict, name: str, shape, meta_path: str, raise RuntimeError( f"meta-файл {meta_path}: кривизна k={disk_k} не совпадает " f"с k={self._current_k()}; смешивать геометрии нельзя.") + stored_mode = meta.get("weight_decay_mode") + if stored_mode is not None and stored_mode != self.param_groups[0]["weight_decay_mode"]: + logger.warning( + "meta-файл %s: состояние обучалось с weight_decay_mode=%r, " + "оптимизатор открыт с %r — семантика затухания весов меняется " + "с этого шага (l2 — часть градиента, decoupled — слагаемое " + "направления).", + meta_path, stored_mode, + self.param_groups[0]["weight_decay_mode"]) files = meta.get("files") or {} for suffix, path in paths.items(): expected = os.path.basename(path) @@ -531,6 +557,16 @@ def _init_param_storage(self, name, shape, paths, meta_path, param, f"несовпадение формы состояния DSA для '{name}': на " f"диске {list(disk_shape)}, в параметрах {list(shape)}. " f"Удалите {self.disk_dir} или укажите новый disk_dir.") + if any_data or meta is not None: + # R2: mark the transition BEFORE touching the data files. + # _write_fresh_state rewrites w/m/v/t one by one; a crash in + # the middle would leave a mixed state (some files new, some + # old) that the next resume=True cannot detect by size checks + # alone, because every file keeps its correct size. With the + # marker, resume=True refuses and resume=False is the + # documented recovery (same rule as import_checkpoint). + self._write_meta(meta_path, name, shape, paths, + reinit_in_progress=True) self._write_fresh_state(name, shape, paths, param) self._write_meta(meta_path, name, shape, paths) return @@ -577,15 +613,18 @@ def _init_param_storage(self, name, shape, paths, meta_path, param, f"и снимет маркер.") self._verify_state_file_sizes(name, shape, paths) - if version == 1: + if version != STATE_FORMAT_VERSION: + # v1 (0.2.0): tangent moments without transport; v2 (0.3.0): + # per-coordinate second moment. Both are incompatible with the + # current math - reset the moments/counters, keep the weights. logger.warning( - "DSA '%s': состояние 0.2.0 (format_version=1). Веса сохранены, " - "моменты m/v и счётчики t сброшены: 0.3.0 хранит моменты в " - "касательном пространстве и переносит их параллельно, старые " - "значения к новой математике неприменимы.", name) + "DSA '%s': состояние format_version=%s. Веса сохранены, " + "моменты m/v и счётчики t сброшены: семантика второго " + "момента изменилась, старые значения неприменимы.", name, + version) self._reset_moments(name, shape, paths) self._write_meta(meta_path, name, shape, paths, - legacy_from=1, ball_validated=False) + legacy_from=version, ball_validated=False) else: logger.info( "DSA '%s': продолжаем состояние %s из %s (resume=True)", @@ -813,20 +852,23 @@ def step(self, updates: Dict[str, Tuple[torch.Tensor, torch.Tensor]], # 4. A non-finite row cannot be repaired by projection: reset the # whole row (weights, moments, counter), so this step is a fresh - # Adam step (A4/R1). + # Adam step (A4/R1). A negative counter is corruption too: bias + # correction 1 - beta**t turns negative, sqrt(v/bc2) becomes NaN + # and the step would die with a misleading "non-finite weights". finite = (torch.isfinite(w).all(dim=-1) & torch.isfinite(m).all(dim=-1) & torch.isfinite(v).all(dim=-1) - & (v >= 0).all(dim=-1)) + & (v >= 0).all(dim=-1) + & (t_prev >= 0)) bad = ~finite if bool(bad.any()): count = int(bad.sum()) self._nan_rows_replaced += count logger.warning( - "DSA '%s': %d строк с NaN/inf или отрицательной " - "дисперсией в файлах состояния (строки диска %s); строка " - "полностью сброшена (w=m=v=0, t=0).", name, count, - idx[bad][:10].tolist()) + "DSA '%s': %d строк с NaN/inf, отрицательной дисперсией " + "или битым счётчиком в файлах состояния (строки диска " + "%s); строка полностью сброшена (w=m=v=0, t=0).", name, + count, idx[bad][:10].tolist()) w[bad] = 0.0 m[bad] = 0.0 v[bad] = 0.0 @@ -850,13 +892,29 @@ def step(self, updates: Dict[str, Tuple[torch.Tensor, torch.Tensor]], g_euc = g_euc + weight_decay * w g_riem = riemannian_gradient(g_euc, w, k) - # 7. Adam moments in the tangent space at w. + # 7. Adam moments in the tangent space at w. The second moment + # follows the manifold's component structure (geoopt's + # component_inner, the paper's Algorithm 1): the Poincare ball is + # ONE manifold component, so v accumulates the scalar squared + # Riemannian norm ||g||_x^2 = lambda^2 ||g_riem||^2 broadcast over + # coordinates; the Euclidean R^d of the k=0 path is the product of + # d components, so v stays per-coordinate classic Adam. The + # metric weighting keeps the Riemannian length of a first step + # bounded by lr*sqrt(dim) anywhere in the ball (plain chart + # components g^2 would grow as lambda(x)*lr near the boundary), + # and the scalar form is intrinsic - invariant under chart + # rotations, unlike a per-coordinate norm. t_new = t_prev + 1 bc1 = (1.0 - beta1 ** t_new.double()).unsqueeze(-1) bc2 = (1.0 - beta2 ** t_new.double()).unsqueeze(-1) m.mul_(beta1).add_(g_riem, alpha=1.0 - beta1) - v.mul_(beta2).addcmul_(g_riem, g_riem, value=1.0 - beta2) + if k > 0.0: + lam = lambda_x(w, k) + v_update = (lam * lam) * sq_norm(g_riem) # (B, 1), broadcast + else: + v_update = g_riem * g_riem + v.mul_(beta2).add_(v_update, alpha=1.0 - beta2) direction = (m / bc1) / (torch.sqrt(v / bc2) + eps) # 8. Decoupled decay is a tangent-space term, so in the Euclidean @@ -865,18 +923,15 @@ def step(self, updates: Dict[str, Tuple[torch.Tensor, torch.Tensor]], if weight_decay > 0.0 and mode == "decoupled": direction = direction + weight_decay * w - # 9. Exponential map + parallel transport of the moments. The - # first moment is a tangent vector and is transported directly. - # The second moment is a diagonal (per-coordinate) accumulator; - # transporting it directly can rotate it into negative components, - # so its non-negative square root is transported and squared back - # (same scale, no NaNs). + # 9. Exponential map + parallel transport of the first moment (the + # paper's tau update). The second moment is a scalar per row and + # needs no transport - scalars are invariant under tangent-space + # isometries (geoopt does not transport it either). if k > 0.0: w_new = expmap(w, -lr * direction, k) w_new = project_to_ball(w_new, k) m_new = parallel_transport(w, w_new, m, k) - sqrt_v_new = parallel_transport(w, w_new, torch.sqrt(v), k) - v_new = sqrt_v_new.square() + v_new = v else: w_new = w - lr * direction m_new, v_new = m, v @@ -1139,6 +1194,14 @@ def import_checkpoint(self, path: str, *, verify_hashes: bool = True) -> dict: f"чекпоинт содержит параметры {sorted(saved_names)}, " f"оптимизатор — {sorted(self.param_names)}") k = _validate_k(manifest.get("k"), where="k чекпоинта") + saved_mode = manifest.get("weight_decay_mode") + if saved_mode is not None and saved_mode != self.param_groups[0]["weight_decay_mode"]: + logger.warning( + "чекпоинт %s: состояние обучалось с weight_decay_mode=%r, " + "оптимизатор создан с %r — семантика затухания весов меняется " + "с этого шага.", + manifest_path, saved_mode, + self.param_groups[0]["weight_decay_mode"]) self.writer.wait_all_flushed() staged: Dict[Tuple[str, str], str] = {} diff --git a/pyproject.toml b/pyproject.toml index 92b58cf..590beef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "dsa" -version = "0.3.0" +version = "0.4.0" description = "Disk Sparse Adam (DSA): an out-of-core sparse Riemannian Adam optimizer for PyTorch" readme = "README.md" license = "MIT" diff --git a/tests/reference.py b/tests/reference.py index 1d85a44..351db55 100644 --- a/tests/reference.py +++ b/tests/reference.py @@ -1,12 +1,15 @@ """Independent float64 reference for Riemannian Adam on the Poincare ball. This module deliberately does not import :mod:`dsa.manifold`. It implements -the formulas from Becigneul & Ganea, "Riemannian Adaptive Optimization -Methods" (ICLR 2019) and the standard Poincare-ball gyrovector formulas (the -same conventions as geoopt with ``k_geoopt = -k``). The tests first verify -that this reference satisfies the defining geometric properties, then compare -the optimizer against it, so the optimizer is not merely checked against a -copy of itself. +the standard Poincare-ball gyrovector formulas (the same conventions as +geoopt with ``k_geoopt = -k``) and the optimizer variant used by DSA: +per-coordinate moments (geoopt style, unlike the paper's scalar second +moment), first-moment parallel transport (the paper's ``tau`` update) and a +positivity-preserving sqrt-transport of the second moment (the paper needs +no transport for its scalar ``v``; geoopt does not transport ``v`` at all). +The tests first verify that this reference satisfies the defining geometric +properties, then compare the optimizer against it, so the optimizer is not +merely checked against a copy of itself. """ from __future__ import annotations @@ -109,11 +112,20 @@ def step(self, indices, grads): g = g + self.wd * x if self.k > 0.0: g = g * (1.0 - self.k * _sq(x)) ** 2 / 4.0 + # the ball is one manifold component: scalar squared Riemannian + # norm ||g||_x^2 = lambda^2 ||g_riem||^2, broadcast (the paper's + # Algorithm 1; geoopt's component_inner with keepdim) + g_sq = np.broadcast_to( + (_lambda(x, self.k) ** 2 * _sq(g)).reshape(len(x), 1), + g.shape, + ) + else: + g_sq = g * g # R^d is a product of d components: per-coordinate self.t[idx] += 1 tt = self.t[idx].astype(np.float64)[:, None] m = self.b1 * self.m[idx] + (1.0 - self.b1) * g - v = self.b2 * self.v[idx] + (1.0 - self.b2) * g * g + v = self.b2 * self.v[idx] + (1.0 - self.b2) * g_sq direction = (m / (1.0 - self.b1 ** tt)) / ( np.sqrt(v / (1.0 - self.b2 ** tt)) + self.eps) if self.wd > 0.0 and self.wd_mode == "decoupled": @@ -122,8 +134,7 @@ def step(self, indices, grads): if self.k > 0.0: x_new = expmap(x, -self.lr * direction, self.k) m = parallel_transport(x, x_new, m, self.k) - # Positive-preserving transport of the diagonal second moment. - v = parallel_transport(x, x_new, np.sqrt(v), self.k) ** 2 + # the scalar second moment needs no transport else: x_new = x - self.lr * direction diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 5678e3c..21bba8c 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -35,7 +35,8 @@ import torch from dsa import manifold -from dsa.optimizer import DiskSparseRiemannianAdam +from dsa.optimizer import (CHECKPOINT_VERSION, STATE_FORMAT_VERSION, + DiskSparseRiemannianAdam) from tests.reference import RiemannianAdamReference, expmap BETA1, BETA2, EPS = 0.9, 0.999, 1e-8 @@ -221,7 +222,7 @@ def test_f3_meta_written_on_fresh_init(tmp_path): meta_path = os.path.join(str(tmp_path), "ent_meta.json") with open(meta_path) as f: meta = json.load(f) - assert meta["format_version"] == 2 + assert meta["format_version"] == STATE_FORMAT_VERSION assert meta["algorithm"] == "riemannian_adam" assert meta["shape"] == [n, d] assert meta["dtype"] == "float32" and meta["counter_dtype"] == "int32" @@ -358,6 +359,55 @@ def test_f3_truncated_file_resume_false_reinits(tmp_path): opt2.shutdown() +def test_r2_crash_mid_reinit_is_refused_on_resume(tmp_path, monkeypatch): + """kill -9 посреди resume=False переинициализации (дописан новый w, m/v/t + нет, meta старый валидный): размеры всех файлов совпадают, поэтому без + маркера resume=True принимает смешанное состояние молча. Маркер обязан + стоять ДО перезаписи файлов; resume=True отказывает, resume=False лечит.""" + from dsa import optimizer as dsa_opt + + n, d = 6, 3 + opt = make_opt(str(tmp_path), np.full((n, d), 0.2), lr=0.05, k=0.0) + do_step(opt, [0, 1], np.ones((2, d))) + opt.shutdown() + meta_path = os.path.join(str(tmp_path), "ent_meta.json") + + # имитация смерти процесса после записи w, но до m/v/t + def crash_after_w(self, name, shape, paths, param): + mm = np.memmap(paths["w"], dtype="float32", mode="w+", shape=shape) + mm[:] = np.full(shape, 0.9, dtype=np.float32) + mm.flush() + del mm + raise RuntimeError("имитация kill -9 посреди переинициализации") + + monkeypatch.setattr( + dsa_opt.DiskSparseRiemannianAdam, "_write_fresh_state", crash_after_w) + with pytest.raises(RuntimeError, match="kill -9"): + make_opt(str(tmp_path), np.full((n, d), 0.3), lr=0.05, k=0.0, + resume=False) + monkeypatch.undo() + + # маркер в meta — свидетельство незавершённой переинициализации + with open(meta_path) as f: + meta = json.load(f) + assert meta.get("reinit_in_progress") is True, \ + "переинициализация не помечена: смешанное состояние примут молча" + + with pytest.raises(RuntimeError, match="reinit_in_progress"): + make_opt(str(tmp_path), np.full((n, d), 0.3), lr=0.05, k=0.0, + resume=True) + + # resume=False — документированный путь восстановления, маркер снимается + opt2 = make_opt(str(tmp_path), np.full((n, d), 0.25), lr=0.05, k=0.0, + resume=False) + w, m, v, t = read_state(opt2) + assert np.allclose(w, 0.25, atol=1e-7) + assert np.all(m == 0) and np.all(v == 0) and np.all(t == 0) + with open(meta_path) as f: + assert "reinit_in_progress" not in json.load(f) + opt2.shutdown() + + # ───────────────────────── C1: legacy 0.1.0 / 0.2.0 ───────────────────────── def test_c1_legacy_010_resume_true_migrates(tmp_path): @@ -372,7 +422,7 @@ def test_c1_legacy_010_resume_true_migrates(tmp_path): assert np.all(st[1] == 0) and np.all(st[2] == 0) and np.all(st[3] == 0) with open(os.path.join(str(tmp_path), "ent_meta.json")) as f: meta = json.load(f) - assert meta["format_version"] == 2 + assert meta["format_version"] == STATE_FORMAT_VERSION assert meta["legacy_migrated_from"] == "0.1.0" opt.shutdown() @@ -411,6 +461,47 @@ def test_c1_legacy_020_resets_incompatible_moments(tmp_path): opt.shutdown() +def test_c1_legacy_030_resets_per_coordinate_moments(tmp_path, caplog): + """Состояние 0.3.0 (format_version=2, поэлементный второй момент): + семантика v изменилась на скалярную — моменты обязаны сброситься, + веса сохраниться, meta подняться до текущей версии.""" + n, d = 4, 2 + w = np.full((n, d), 0.03) + _write_legacy_010(str(tmp_path), w, np.full((n, d), 0.4), + np.full((n, d), 0.5)) + mm = np.memmap(os.path.join(str(tmp_path), "ent_t.dat"), dtype="int32", + mode="w+", shape=(n,)) + mm[:] = np.full(n, 9) + mm.flush() + del mm + meta = { + "format_version": 2, + "algorithm": "riemannian_adam", + "param_name": "ent", + "shape": list(w.shape), + "dtype": "float32", + "counter_dtype": "int32", + "k": 0.0, + "weight_decay_mode": "l2", + "files": {s: f"ent_{s}.dat" for s in ("w", "m", "v", "t")}, + "ball_validated": True, + } + with open(os.path.join(str(tmp_path), "ent_meta.json"), "w") as f: + json.dump(meta, f) + + with caplog.at_level(logging.WARNING, logger="DSA"): + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0, resume=True) + st = read_state(opt) + assert np.allclose(st[0], w, atol=1e-7), "веса 0.3.0 должны сохраниться" + assert np.all(st[1] == 0) and np.all(st[2] == 0) and np.all(st[3] == 0) + with open(os.path.join(str(tmp_path), "ent_meta.json")) as f: + migrated = json.load(f) + assert migrated["format_version"] == STATE_FORMAT_VERSION + assert migrated["legacy_migrated_from"] == 2 + assert any("format_version=2" in r.message for r in caplog.records) + opt.shutdown() + + # ───────────────────────── F4 / C4: Adam and weight decay ───────────────────────── def test_f4_first_step_of_fresh_row_is_lr_sign(tmp_path): @@ -527,6 +618,38 @@ def test_c4_invalid_weight_decay_mode(tmp_path): make_opt(str(tmp_path), np.zeros((3, 2)), weight_decay_mode="adamw") +def test_c4_mode_switch_on_existing_state_warns(tmp_path, caplog): + """Режим затухания хранится в meta и в чекпоинте, но смена режима на живом + состоянии раньше проходила молча: состояние, обучавшееся с "l2", могло + тихо продолжиться с "decoupled". Оба пути обязаны предупредить.""" + n, d = 6, 2 + opt = make_opt(str(tmp_path), np.zeros((n, d)), lr=0.05, k=0.0, + weight_decay=0.1, weight_decay_mode="l2") + do_step(opt, [0], np.ones((1, d))) + opt.shutdown() + + # resume=True с другим режимом — предупреждение + with caplog.at_level(logging.WARNING, logger="DSA"): + opt2 = make_opt(str(tmp_path), np.zeros((n, d)), lr=0.05, k=0.0, + weight_decay=0.1, weight_decay_mode="decoupled", + resume=True) + assert any("weight_decay_mode" in r.message for r in caplog.records) + opt2.shutdown() + + # импорт чекпоинта в оптимизатор с другим режимом — тоже предупреждение + caplog.clear() + opt3 = make_opt(str(tmp_path / "b"), np.zeros((n, d)), lr=0.05, k=0.0, + weight_decay=0.1, weight_decay_mode="l2") + ckpt = opt3.export_checkpoint(str(tmp_path / "ck")) + opt3.shutdown() + opt4 = make_opt(str(tmp_path / "c"), np.zeros((n, d)), lr=0.05, k=0.0, + weight_decay=0.1, weight_decay_mode="decoupled") + with caplog.at_level(logging.WARNING, logger="DSA"): + opt4.import_checkpoint(ckpt) + assert any("weight_decay_mode" in r.message for r in caplog.records) + opt4.shutdown() + + # ───────────────────────── F7 / H1: ball invariants ───────────────────────── def test_h1_init_is_scaled_into_ball(tmp_path): @@ -662,9 +785,10 @@ def test_h2_moments_are_transported_isometrically(tmp_path): lam1 = 2.0 / (1 - k * np.dot(w[0], w[0])) assert np.isclose(lam1 * np.linalg.norm(m[0]), lam0 * np.linalg.norm(m_before), rtol=1e-5, atol=1e-7) - v_before = (1 - BETA2) * g_riem ** 2 - assert np.isclose(lam1 ** 2 * np.sum(v[0]), - lam0 ** 2 * np.sum(v_before), rtol=1e-5, atol=1e-9) + # v — скалярная норма на строку: (1-beta2)*lambda^2*||g_riem||^2, без + # переноса (скаляр инвариантен), размноженная по координатам + v_expected = (1 - BETA2) * lam0 ** 2 * np.dot(g_riem, g_riem) + assert np.allclose(v[0], v_expected, rtol=1e-5, atol=1e-9) def test_h2_first_step_from_origin_closed_form(tmp_path): @@ -676,7 +800,11 @@ def test_h2_first_step_from_origin_closed_form(tmp_path): opt.shutdown() w, *_ = read_state(opt) - direction = g[0] / (np.abs(g[0]) + EPS) # first Adam step at t=1 + # Первый шаг из начала координат: g_riem = g/4, lambda(0) = 2; второй + # момент — скалярная норма ||g||_x^2, поэтому знаменатель — норма вектора: + # direction = g_riem / (lambda(0)*||g_riem|| + eps) + g_riem = g[0] / 4.0 + direction = g_riem / (2.0 * np.linalg.norm(g_riem) + EPS) # exp_0(-lr * direction) = -tanh(sqrt(k)*lr*||direction||) / # (sqrt(k)*||direction||) * direction norm = np.linalg.norm(direction) @@ -696,6 +824,43 @@ def test_h2_zero_gradient_is_identity(tmp_path): assert np.all(m == 0) and np.all(v == 0) and np.all(t == 1) +def test_h2_first_step_riemannian_length_is_bounded(tmp_path): + """Семантика статьи/geoopt: риманова длина ПЕРВОГО шага ограничена + lr*sqrt(dim) в любой точке шара. Без метрического веса второго момента + длина растёт как lambda(x)*lr*sqrt(dim) — у границы шара это в тысячи раз + больше задуманного (строка «улетает» на противоположную сторону).""" + d, k, lr = 6, 1.0, 0.05 + radius = math.sqrt((1 - manifold.BALL_EPS) / k) + # строки на разной глубине: центр, середина, вплотную к границе + norms = [0.0, 0.5 * radius, 0.9999 * radius] + w0 = [] + for nrm in norms: + row = np.zeros(d) + row[0] = nrm + w0.append(row) + w0 = np.array(w0) + g = np.tile(np.array([1.0, -2.0, 0.5, 3.0, -1.0, 0.25]), (len(w0), 1)) + + opt = make_opt(str(tmp_path), w0, lr=lr, k=k) + w_before, *_ = read_state(opt) # хранимые float32 значения + do_step(opt, np.arange(len(w0)), g) + opt.shutdown() + w, *_ = read_state(opt) + + # допуск: 1e-4 относительный — запас на округление float32-хранения; + # старый (покоординатный, без метрики) вариант давал длину ~lambda*lr, + # у границы это в тысячи раз больше границы + bound = lr * math.sqrt(d) * (1.0 + 1e-4) + for i in range(len(w0)): + dist = manifold.distance( + torch.tensor(w_before[i], dtype=torch.float64), + torch.tensor(w[i], dtype=torch.float64), k).item() + assert dist <= bound, ( + f"строка {i} (|w0|={norms[i]:.4f}): риманова длина шага {dist:.6f} " + f"превышает lr*sqrt(d)={bound:.6f} — второй момент без " + f"метрического веса?") + + # ───────────────────────── H3: curvature changes ───────────────────────── def test_h3_set_curvature_reprojects_whole_table(tmp_path): @@ -904,6 +1069,39 @@ def test_a4_zero_counter_discards_stale_moments(tmp_path): assert t[1] == 1 +def test_a4_negative_counter_is_treated_as_corruption(tmp_path): + """Битый отрицательный счётчик (torn write) — это порча строки, а не + «древняя» строка: 1 - beta**t < 0 даёт sqrt(отрицательного) = NaN, и шаг + падал с вводящим в заблуждение «не-конечные веса, уменьшите lr». Строка + обязана быть сброшена полностью, а шаг — дать канонический первый шаг.""" + n, d = 5, 3 + lr = 0.1 + opt = make_opt(str(tmp_path), np.full((n, d), 0.1), lr=lr, k=0.0) + do_step(opt, [2], np.ones((1, d))) # валидный шаг, чтобы состояние было живым + opt.shutdown() + + # порча: t[0] = -7 (битый сектор / оборванная запись) + with open(os.path.join(str(tmp_path), "ent_t.dat"), "r+b") as f: + f.write(np.int32(-7).tobytes()) + + opt2 = make_opt(str(tmp_path), np.full((n, d), 0.1), lr=lr, k=0.0, + resume=True) + g = np.full((1, d), 0.5) + do_step(opt2, [0], g) + opt2.shutdown() + + w, m, v, t = read_state(opt2) + # строка полностью сброшена (w=0), поэтому первый шаг идёт от нуля: + # w = 0 - lr*g/(|g|+eps) — канонический первый шаг Adam + expected = -lr * g[0] / (np.abs(g[0]) + EPS) + assert np.allclose(w[0], expected, atol=1e-6), \ + f"шаг по битой строке должен быть каноническим первым: {w[0]}" + assert t[0] == 1, f"счётчик должен перезапуститься с 1, получено {t[0]}" + assert np.isfinite(m[0]).all() and np.isfinite(v[0]).all() + assert opt2._nan_rows_replaced >= 1 + assert t[2] == 1, "неповреждённая строка не должна быть тронута" + + def test_a4_nan_gradient_rejected(tmp_path): n, d = 4, 2 opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) @@ -973,8 +1171,10 @@ def test_riemannian_gradient_scaling_formula(tmp_path): assert np.isclose(lam1 * np.linalg.norm(m[0]), lam0 * np.linalg.norm((1 - BETA1) * g_riem), rtol=1e-5, atol=1e-7) - expected = expmap(w0[0][None, :], - -lr * np.sign(g_riem)[None, :], k)[0] + # скалярный второй момент: знаменатель — норма вектора, поэтому первый + # шаг идёт по -g_riem/||g_riem|| (единичное направление), не по sign + direction = g_riem / (lam0 * np.linalg.norm(g_riem) + EPS) + expected = expmap(w0[0][None, :], -lr * direction[None, :], k)[0] assert np.allclose(w[0], expected, atol=1e-6) @@ -1046,7 +1246,7 @@ def test_c2_checkpoint_is_portable_and_self_contained(tmp_path): with open(manifest_path) as f: manifest = json.load(f) assert manifest["format"] == "DSA_checkpoint" - assert manifest["version"] == 2 + assert manifest["version"] == CHECKPOINT_VERSION assert manifest["global_step"] == 2 for info in manifest["params"].values(): for rec in info["files"].values(): diff --git a/tests/test_geoopt_step.py b/tests/test_geoopt_step.py new file mode 100644 index 0000000..d1a644c --- /dev/null +++ b/tests/test_geoopt_step.py @@ -0,0 +1,60 @@ +"""External oracle: the hyperbolic optimizer vs geoopt's RiemannianAdam. + +Requires the real ``geoopt`` package (skipped otherwise; in CI only the +in-repo oracles run). Install geoopt into a separate directory and run: + + PYTHONPATH=/path/to/geoopt_pkg python -m pytest tests/test_geoopt_step.py + +DSA matches ``geoopt.optim.RiemannianAdam`` on ``PoincareBallExact`` (the +exact-exponential-map variant; the default ``PoincareBall`` uses a first-order +retraction, which DSA deliberately does not): scalar second moment +``lambda^2 ||g_riem||^2`` broadcast over coordinates, first moment updated in +the tangent space and parallel-transported after the exponential-map step, +second moment not transported. The test pins the FULL multi-step trajectory +(weights and both moments), not just one step. +""" + +import numpy as np +import pytest +import torch + +geoopt = pytest.importorskip("geoopt") + +from dsa.optimizer import DiskSparseRiemannianAdam # noqa: E402 + + +@pytest.mark.parametrize("k", [0.7, 1.0]) +@pytest.mark.parametrize("wd", [0.0, 0.02]) +def test_trajectory_matches_geoopt_exact(tmp_path, k, wd): + # Строки на умеренных нормах (0.6 радиуса): здесь совпадение обязано быть + # на уровне округления float32. Приочковые строки НЕ сравниваются с geoopt + # напрямую: geoopt считает целиком в float32 и у границы теряет точность + # катастрофически (лямбда ~1e4); поведение DSA там закреплено отдельными + # инвариантами (ограниченная длина шага) и стресс-тестом. + rng = np.random.default_rng(7) + n, d, lr, steps = 4, 5, 0.03, 10 + w0 = rng.normal(0, 0.3, (n, d)).astype(np.float32) + w0 /= (np.linalg.norm(w0, axis=1, keepdims=True) / 0.6) + + opt = DiskSparseRiemannianAdam( + {"e": torch.from_numpy(w0.copy())}, lr=lr, k=k, disk_dir=str(tmp_path), + resume=False, init_max_norm_fraction=1.0, weight_decay=wd, + weight_decay_mode="l2") + ball = geoopt.PoincareBallExact(c=k) + p = geoopt.ManifoldParameter(torch.from_numpy(w0.copy()), manifold=ball) + gopt = geoopt.optim.RiemannianAdam([p], lr=lr, weight_decay=wd) + + for _ in range(steps): + g = rng.normal(0, 0.7, (n, d)).astype(np.float32) + opt.step(updates={"e": (torch.arange(n), torch.from_numpy(g.copy()))}) + p.grad = torch.from_numpy(g.copy()) + gopt.step() + opt.shutdown() + + st = gopt.state[p] + assert np.allclose(np.array(opt.state_files["e"]["w"]), p.data.numpy(), + atol=1e-5), "траектория w разошлась с geoopt" + assert np.allclose(np.array(opt.state_files["e"]["m"]), + st["exp_avg"].numpy(), atol=1e-5) + assert np.allclose(np.array(opt.state_files["e"]["v"]), + st["exp_avg_sq"].numpy(), atol=1e-5) diff --git a/tests/test_optimization.py b/tests/test_optimization.py new file mode 100644 index 0000000..b5dd202 --- /dev/null +++ b/tests/test_optimization.py @@ -0,0 +1,87 @@ +"""End-to-end optimization sanity: DSA must actually minimize on the ball. + +The objective is the geodesic distance to fixed target points; Euclidean +gradients come from torch.autograd through the (verified) distance formula, +which is exactly the input contract of ``step()``. These tests catch a +class of errors the equality oracles cannot: an update that is internally +consistent but does not descend. +""" + +import math + +import numpy as np +import torch + +from dsa import manifold +from dsa.optimizer import DiskSparseRiemannianAdam + + +def _make(w0, lr, k, tmp_path): + return DiskSparseRiemannianAdam( + {"e": torch.from_numpy(w0.astype(np.float32))}, lr=lr, k=k, + disk_dir=str(tmp_path), resume=False, init_max_norm_fraction=1.0) + + +def test_converges_to_targets_on_the_ball(tmp_path): + rng = np.random.default_rng(3) + n, d, k, lr = 5, 8, 1.0, 0.05 + w0 = rng.normal(0, 0.2, (n, d)) + targets = rng.normal(0, 0.3, (n, d)) + targets = (torch.from_numpy(targets) + * (0.5 * math.sqrt(1.0 / k) + / torch.from_numpy(np.linalg.norm(targets, axis=1, + keepdims=True)))) + + opt = _make(w0, lr, k, tmp_path) + idx = torch.arange(n) + tgt = targets.clone().requires_grad_(False) + + def cur_distances(): + w = opt.read_rows(idx).double() + return manifold.distance(w, tgt.double(), k).squeeze(-1) + + d0 = cur_distances() + for _ in range(300): + w = opt.read_rows(idx).double().requires_grad_(True) + loss = manifold.distance(w, tgt.double(), k).pow(2).sum() + loss.backward() + g = w.grad.detach().float() + if not torch.isfinite(g).all(): + raise AssertionError("градиент расстояления стал не-конечным") + opt.step(updates={"e": (idx, g)}) + opt.shutdown() + d1 = cur_distances() + + assert d0.mean().item() > 1.0, "старт должен быть далеко от целей" + assert d1.mean().item() < 0.05, ( + f"оптимизатор не сошёлся: средняя дистанция {d0.mean():.3f} -> " + f"{d1.mean():.3f} за 300 шагов") + assert bool((d1 < d0).all()), "какая-то строка отдалилась от цели" + + +def test_boundary_stress_stays_finite_and_inside(tmp_path): + """Строки вплотную к границе + большие градиенты: конечно, внутри шара, + без «перелёта» на противоположную сторону (метрическая нормировка + ограничивает риманову длину первого шага lr*sqrt(d)).""" + d, k, lr = 6, 1.0, 0.05 + radius = math.sqrt((1 - manifold.BALL_EPS) / k) + rng = np.random.default_rng(5) + w0 = rng.normal(0, 1, (4, d)) + w0 = (w0 / np.linalg.norm(w0, axis=1, keepdims=True)) * (0.9999 * radius) + + opt = _make(w0.astype(np.float32), lr, k, tmp_path) + idx = torch.arange(4) + for s in range(60): + g = torch.from_numpy( + (rng.normal(0, 50.0, (4, d))).astype(np.float32)) + opt.step(updates={"e": (idx, g)}) + opt.shutdown() + + w = np.array(opt.state_files["e"]["w"]) + assert np.isfinite(w).all(), "веса стали не-конечными у границы" + norms = np.linalg.norm(w, axis=1) + assert np.all(norms <= radius + 1e-5), norms + # не улетели к противоположной границе: средняя норма не выродилась + # к радиусу со всех сторон — строки должны оставаться в разумном диапазоне + assert np.all(norms > 0.05), ( + "строки схлопнулись в начало координата под стрессом") diff --git a/tests/test_packaging.py b/tests/test_packaging.py index d6f5c02..70c87ad 100644 --- a/tests/test_packaging.py +++ b/tests/test_packaging.py @@ -30,7 +30,7 @@ def test_version_matches_package(): assert data["project"]["version"] == dsa.__version__ else: # pragma: no cover text = open(os.path.join(ROOT, "pyproject.toml")).read() - assert re.search(r'version\s*=\s*"0\.3\.0"', text) + assert re.search(r'version\s*=\s*"0\.4\.0"', text) def test_torch_floor_matches_readme_badge(): diff --git a/tests/test_torch_equivalence.py b/tests/test_torch_equivalence.py new file mode 100644 index 0000000..fe1d0d4 --- /dev/null +++ b/tests/test_torch_equivalence.py @@ -0,0 +1,48 @@ +"""External oracle: the Euclidean path (k=0) must equal torch.optim exactly. + +DSA at ``k = 0`` treats ``R^d`` as the product of ``d`` one-dimensional +manifold components, so the moments stay per-coordinate and the update is +classic Adam. This file pins the whole multi-step trajectory (dense batches, +random gradients, weight decay) against ``torch.optim.Adam`` ("l2") and +``torch.optim.AdamW`` ("decoupled") — an oracle independent of this +repository. +""" + +import numpy as np +import torch + +from dsa.optimizer import DiskSparseRiemannianAdam + + +def _run_pair(tmp_path, mode, torch_cls, steps=12): + rng = np.random.default_rng(11) + n, d, lr, wd = 20, 6, 0.05, 0.01 + w0 = rng.normal(0, 0.4, (n, d)) + + opt = DiskSparseRiemannianAdam( + {"e": torch.from_numpy(w0.astype(np.float32))}, lr=lr, k=0.0, + disk_dir=str(tmp_path), resume=False, init_max_norm_fraction=1.0, + weight_decay=wd, weight_decay_mode=mode) + ref_param = torch.from_numpy(w0.astype(np.float32)).requires_grad_(True) + ref = torch_cls([ref_param], lr=lr, weight_decay=wd) + + idx = torch.arange(n) + for _ in range(steps): + g = torch.from_numpy(rng.normal(0, 0.8, (n, d)).astype(np.float32)) + opt.step(updates={"e": (idx, g)}) + ref_param.grad = g.clone() + ref.step() + opt.shutdown() + return np.array(opt.state_files["e"]["w"]), ref_param.detach().numpy() + + +def test_euclidean_matches_torch_adam_l2(tmp_path): + w_dsa, w_torch = _run_pair(tmp_path, "l2", torch.optim.Adam) + assert np.allclose(w_dsa, w_torch, atol=2e-6), \ + np.abs(w_dsa - w_torch).max() + + +def test_euclidean_matches_torch_adamw_decoupled(tmp_path): + w_dsa, w_torch = _run_pair(tmp_path, "decoupled", torch.optim.AdamW) + assert np.allclose(w_dsa, w_torch, atol=2e-6), \ + np.abs(w_dsa - w_torch).max()