diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 71160a2..508df64 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -9,14 +9,18 @@ on: jobs: test: runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.9', '3.11', '3.12'] steps: - uses: actions/checkout@v4 - - name: Set up Python + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: - python-version: '3.12' + python-version: ${{ matrix.python-version }} - name: Install dependencies run: | @@ -27,4 +31,30 @@ jobs: - name: Run tests run: | - python -m pytest tests/ -v -s + python -m pytest tests/ -v + + package: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - name: Build the wheel + run: | + python -m pip install --upgrade pip build + python -m build + - name: Check single metadata source + run: | + test ! -f setup.py + python - <<'PY' + import glob, 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 + assert 'Requires-Dist: torch>=2.0.0' in text + print('wheel metadata OK') + PY diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1616879 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +build/ +dist/ +.pytest_cache/ +.venv/ +*.dat diff --git a/CHANGELOG.md b/CHANGELOG.md index 57aec16..524af7e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,132 @@ # Changelog +## 0.3.0 — 2026-11-04 + +Correctness release for hyperbolic training. The headline change is that +`k > 0` is now a real Riemannian Adam (exponential map, tangent-space first +moment, parallel transport) instead of Euclidean Adam with gradient scaling +and a clip. Several APIs change incompatibly; see **Breaking changes**. + +### Breaking changes + +- **`resume` defaults to `None`.** An existing state directory is never + silently reused or silently reinitialized: `None` raises, `True` continues, + `False` reinitializes from `params` with a warning. +- **Checkpoints are directories.** `save_checkpoint`/`load_checkpoint` are + directory-only aliases; the primary API is + `export_checkpoint(dir)`/`import_checkpoint(dir)`. 0.2.0 manifests with + absolute paths are rejected. +- **State format v2.** 0.2.0 (`format_version=1`) and 0.1.0 (no meta, no `t`) + states are migrated on `resume=True`: weights are preserved, moments and + per-row counters are reset because the old math is incompatible with + tangent-space moments. +- **`weight_decay_mode` added.** `"l2"` (default) is part of the Euclidean + gradient before the Riemannian conversion (classic Adam L2 / Riemannian + gradient of `f + wd/2 ||w||^2`); `"decoupled"` is a tangent-space term and + equals AdamW when `k = 0`. The old coordinate shrink is gone. +- **Source tensors are released by default** (`retain_params=False`): + `param_groups[...]["params"]` holds zero-length CPU placeholders after + initialization, so a full embedding table is not kept alive in VRAM. Use + `read_rows()` for weights, or `retain_params=True` for the old behavior. +- **Global step is persisted** in `dsa_global_meta.json` and in checkpoints; + `resume=True` restores it. +- **Stricter validation:** negative/out-of-range indices, wrong gradient + shapes, non-finite gradients, invalid `k`, and unknown update keys now + produce explicit errors or warnings. + +### Fixed — audit findings + +- **C1 — legacy 0.1.0 could not be reinitialized.** `resume=False` now + reinitializes even a partial/legacy directory (missing `t`/meta); weights + from `params` are written, moments/counters zeroed. `resume=True` migrates + a 0.1.0 directory (weights preserved, `m/v/t` reset, meta v2 written) and a + 0.2.0 directory (`format_version=1`, same reset rule). +- **C2 — checkpoints were not portable.** `export_checkpoint(dir)` copies + `w/m/v/t` into a directory under relative names and writes + `checkpoint.json` with sizes and SHA-256 checksums. `import_checkpoint(dir)` + verifies shape, size and hash, stages all files, writes a reinitialization + marker, then replaces the live state atomically per file. A moved/copied + checkpoint imports successfully; a corrupted one is rejected before the + live state changes. +- **C3 — the optimizer retained the full source tensor.** With + `retain_params=False` (default), source tensors are replaced by zero-length + placeholders after the on-disk state is written; a weak-reference test + confirms the tensor can be collected. `retain_params=True` restores the old + PyTorch-style retention. +- **C4 — weight decay was coordinate shrink.** `weight_decay_mode="l2"` now + adds `wd * w` to the Euclidean gradient (Riemannian gradient of the L2 + penalty), and `"decoupled"` adds `wd * w` to the tangent direction. The + Euclidean limit matches Adam L2 and AdamW respectively, and the hyperbolic + L2 penalty cancels exactly with `g = -wd * w`. +- **C5 — two packaging metadata sources.** `setup.py` is removed; + `pyproject.toml` is the single source (`torch>=2.0.0`, `numpy>=1.21.0`), + matching the README badge. `dsa.__version__` is `0.3.0`. +- **A1 — a failed `shutdown()` could report success on retry.** The writer is + now a `RUNNING → DRAINING → STOPPED` / `FAILED` state machine. A timeout + leaves the worker alive and the queue unfinished; the next call retries and + only a clean drain/join/flush returns success. Persistent write errors are + re-raised, never hidden. +- **A2 — `global_step` was lost on load.** It is stored in + `dsa_global_meta.json`, exported in checkpoints, restored on `resume=True` + and `import_checkpoint()`. +- **A3 — negative indices silently updated the last row.** Indices are + validated: 1-D integer dtype, `0 <= i < N`; negative, out-of-range, float + and 2-D indices raise before any state is touched. +- **A4 — a NaN row left a corrupted counter/moment.** A non-finite (or + negative-variance) row is fully reset (`w=m=v=0, t=0`), so the current step + is a genuine first Adam step, not a continuation of a broken counter. +- **A9 — gradient shape/finiteness was unchecked.** Wrong shapes, non-float + dtypes and NaN/inf gradients raise `ValueError` before moments are updated. +- **H1 — initialization and resume could leave rows outside the ball.** + Fresh hyperbolic weights are globally scaled so the largest row norm is at + most `init_max_norm_fraction` (default 0.5) of the radius; rows outside are + still projected. On resume, finite out-of-ball rows are projected and + counted; `read_rows()` always projects. `export_checkpoint()` immediately + after initialization is valid. +- **H2 — the step was not Riemannian Adam.** The step now uses + `grad_R = lambda^-2 g_E`, tangent-space moments, the exponential map + `exp_w(-lr * m_hat/(sqrt(v_hat)+eps))` and parallel transport of `m` (and + of `sqrt(v)`, squared back to keep the diagonal second moment + non-negative). Active rows are computed in `float64` (storage stays + `float32`) to avoid catastrophic cancellation in the gyration formula near + the boundary. Tests compare against an independent float64 reference and + check the defining geometric properties (exp map distance, transport + identity/isometry/round-trip, log-exp inverse). +- **H3 — changing `k` left untouched rows outside the new ball.** + `set_curvature(k)` and `step(current_k=...)` reproject the whole table; + `reset_moments=True` optionally restarts the adaptive state. +- **H4 — `k` was not validated.** `k = NaN` no longer falls through to the + Euclidean path, `k = inf` no longer collapses the ball, and `k < 0` raises + at construction instead of a later `math domain error`. +- **F1–F8, R1–R3** from 0.2.0 remain covered: lossless queue fallback, + in-flight read consistency, strict state lifecycle, per-row bias + correction, sorted indices, honest shutdown, projection on read, NaN-row + handling, reinitialization marker, and unknown-key diagnostics. + +### Tests + +The suite now has explicit acceptance tests for C1–C5, A1–A4, A9, H1–H4, +F1–F8 and R1–R3, plus geometric property tests and an independent float64 +reference. Run with a loaded `.env` if your environment requires it: + +```bash +CUDA_VISIBLE_DEVICES="" python -m pytest -q +``` + +### Known limitations + +- The Poincaré chart is numerically singular at the boundary. Points are kept + at least `BALL_EPS = 1e-5` inside, but very long training with a large `lr` + can reach that margin; the optimizer stays finite and projects, while the + 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. +- CUDA-specific memory behavior is tested only through object lifetime on CPU + in CI; no GPU is required for the test suite. + ## 0.2.0 — 2026-11-04 Correctness release. The public API is preserved: the constructor signature diff --git a/README.md b/README.md index a32b2e2..8ca65c1 100644 --- a/README.md +++ b/README.md @@ -17,16 +17,19 @@ ## 🚀 Key Features -* **Near-Zero VRAM Overhead:** Optimizer momentum ($m$), variance ($v$), weights ($w$) and per-row step counters ($t$) reside on NVMe/SSD storage via `np.memmap`. Only active batch parameters $\mathcal{O}(B)$ are loaded into GPU memory. +* **Near-Zero VRAM Overhead:** Optimizer momentum ($m$), variance ($v$), weights ($w$) and per-row step counters ($t$) reside on NVMe/SSD storage via `np.memmap`. Only active batch rows are loaded into RAM/VRAM. By default the optimizer **does not retain the source tensors** after initialization (`retain_params=False`): they are replaced by zero-length CPU placeholders, so a full embedding table cannot stay alive on the GPU through `param_groups`. +* **True Riemannian Adam ($k > 0$):** The Euclidean gradient is converted to a Riemannian one, Adam's first moment lives in the tangent space and is parallel-transported after every step, the step is an **exponential map** (not a Euclidean subtraction followed by a clip), and the second moment is carried as a non-negative diagonal scale. The collection of $O(B)$ rows is computed in `float64`, which keeps the Möbius/gyration formulas stable near the boundary; storage remains `float32`. `k = 0` is standard Euclidean Adam. * **Lossless Async I/O:** A background worker thread with double-buffering (`AsyncDiskWriter`) flushes state updates to disk asynchronously without stalling GPU computations. A full write queue never drops batches: `put_update` waits up to `put_timeout` seconds, then writes the batch synchronously. * **Read Consistency:** The writer tracks in-flight rows; `step()` drains the queue before reading rows the writer has not finished writing, so no step ever computes from stale state. * **Per-Row Adam Bias Correction:** Each row carries its own step counter (`_t.dat`), so the first update of a freshly touched row is a correct Adam step regardless of the global step count. -* **Explicit State Lifecycle:** State directories are guarded by `_meta.json` (shape, dtype, format version). `resume=True` continues existing state; `resume=False` (default) fully reinitializes with a warning; shape/dtype mismatches, partial or truncated state files raise explicit errors instead of being silently reused. -* **Safe Duplicate Aggregation:** Automatically aggregates duplicate node/entity indices (hub nodes) appearing within the same batch via `torch.unique` and `index_add_`. -* **Sorted Sequential I/O:** Batch indices are sorted after deduplication, so reads and writes hit the memmaps sequentially. -* **Riemannian & Euclidean Support:** Supports Riemannian gradient scaling and Poincaré Ball retraction ($k > 0$) for Hyperbolic embeddings, as well as standard Euclidean Adam ($k = 0.0$). Rows read from disk are projected into the ball, so corrupted files cannot break the geometry. -* **Honest Shutdown & Checkpoints:** `shutdown()` drains the queue and raises `RuntimeError` if data could not be fully written; `save_checkpoint`/`load_checkpoint` provide a consistent on-disk checkpoint API (`state_dict()` intentionally raises `NotImplementedError`). -* **Cross-Platform:** Runs seamlessly on local Linux/Windows GPU machines, workstations, and cloud notebooks. +* **Explicit State Lifecycle:** State directories are guarded by `_meta.json` (shape, dtype, format version, curvature, algorithm) and a global `dsa_global_meta.json` (global step, curvature). `resume=None` (default) refuses to guess when a directory already exists: pass `resume=True` to continue or `resume=False` to reinitialize. Shape/dtype mismatches, partial or truncated state files raise explicit errors. 0.1.0/0.2.0 directories are migrated on `resume=True` (weights preserved; incompatible moments and counters reset). +* **Ball Invariants:** Fresh hyperbolic weights are globally scaled to at most `init_max_norm_fraction` (default 0.5) of the ball radius, so `torch.randn` initialization does not start on the boundary. Rows read from disk are projected, resumed tables are revalidated, and `set_curvature()` reprojects the whole table when the ball changes. +* **Safe Duplicate Aggregation & Sorted I/O:** Duplicate node/entity indices are summed; batch indices are sorted after deduplication so reads and writes hit the memmaps sequentially. +* **Explicit Weight Decay Semantics:** `weight_decay_mode="l2"` (default) adds `wd * w` to the Euclidean gradient and is therefore the Riemannian gradient of $f + \tfrac{wd}{2}\|w\|^2$; in the Euclidean case it is classic Adam L2. `"decoupled"` adds `wd * w` to the tangent-space descent direction and reproduces AdamW exactly when `k = 0`. +* **Portable Checkpoints:** `export_checkpoint(dir)` writes a self-contained directory (`checkpoint.json` + relative data files + SHA-256 checksums); `import_checkpoint(dir)` verifies and restores it, including the global step and curvature. Checkpoints can be moved between machines. `save_checkpoint`/`load_checkpoint` remain as directory-only aliases. +* **Honest Shutdown:** `shutdown()` is a retryable state machine (`RUNNING → DRAINING → STOPPED`, or `FAILED`); it never reports success while the worker is alive or the queue is unfinished, and a failed write is re-raised instead of being hidden. +* **Validated Inputs:** Negative/out-of-range indices, wrong gradient shapes, non-finite gradients, NaN/negative moment rows, `k = NaN/inf/negative`, and unknown update keys all have explicit, test-covered behavior. +* **Cross-Platform:** Runs on local Linux/Windows GPU machines, workstations, and cloud notebooks. --- @@ -34,20 +37,25 @@ There is no fixed throughput number: DSA trades compute for memory, and its speed is bounded by your disk — expect optimizer-state I/O to be the -bottleneck on large tables. The correctness guarantees added in 0.2.0 -(blocking writes instead of dropped batches, in-flight row tracking, a fourth -per-row counter file, sorted indices) add roughly ~10% step-time overhead -compared to 0.1.0. Measure on your own hardware before drawing conclusions. +bottleneck on large tables. The 0.3.0 correctness work (blocking writes, +in-flight row tracking, a fourth per-row counter file, sorted indices, +float64 batch math) adds step-time overhead compared to 0.1.0; the float64 +part only touches the active batch, so on large tables it is usually small +relative to disk I/O. Measure on your own hardware before drawing conclusions. What stays true at any scale: | Property | Since | | :--- | :--- | | Optimizer states reside on disk; VRAM footprint scales with batch size $\mathcal{O}(B)$ | 0.1.0 | +| Source tensors are not retained by the optimizer | 0.3.0 | | Duplicate node aggregation within a batch | 0.1.0 | | No silently dropped update batches | 0.2.0 | | Step reads never race the background writer | 0.2.0 | | Correct bias correction for sparsely updated rows | 0.2.0 | +| Full Riemannian Adam (exp map + parallel transport) | 0.3.0 | +| Portable checkpoints with checksums | 0.3.0 | +| Retryable, honest shutdown | 0.3.0 | --- @@ -109,14 +117,10 @@ print("🚀 Starting local training loop...") # 4. Training Loop for epoch in range(1, 101): indices = torch.randint(0, num_entities, (batch_size,)) - idx_np = indices.numpy() - - # Read current weights directly from DSA disk memmap - # (drain the write queue first so the read cannot race a pending write) - optimizer.writer.wait_all_flushed() - vec_np = optimizer.state_files["entity_emb"]["w"][idx_np].copy() - vec = torch.from_numpy(vec_np).to(device).requires_grad_(True) - + + # Public read API: waits for in-flight writes and projects into the ball. + vec = optimizer.read_rows(indices).to(device).requires_grad_(True) + # Compute Loss (MSE against target) target_vec = target_embeddings[indices].to(device) loss = torch.mean((vec - target_vec) ** 2) @@ -186,14 +190,10 @@ print("🚀 Starting training on Kaggle...") # 4. Training Loop for epoch in range(1, 11): indices = torch.randint(0, num_entities, (batch_size,)) - idx_np = indices.numpy() - - # Fetch weights from DSA disk cache - # (drain the write queue first so the read cannot race a pending write) - optimizer.writer.wait_all_flushed() - vec_np = optimizer.state_files["entity_emb"]["w"][idx_np].copy() - vec = torch.from_numpy(vec_np).to(device).requires_grad_(True) - + + # Public read API (in-flight writes are drained, rows are projected). + vec = optimizer.read_rows(indices).to(device).requires_grad_(True) + # Compute Loss loss = (vec ** 2).sum() loss.backward() @@ -213,21 +213,65 @@ if os.path.exists(KAGGLE_CACHE_DIR): ## 🌌 3. Hyperbolic Poincaré Ball Workflow ($k > 0$) -For Hyperbolic Knowledge Graph Embeddings (e.g., Poincaré embeddings, RotatE in hyperbolic space), set the curvature parameter `k > 0.0`: +For hyperbolic knowledge-graph embeddings, set `k > 0` (the ball has curvature +$-k$ and radius $1/\sqrt{k}$): ```python -# Set k=1.0 for Poincaré Ball Conformal Gradient Scaling & Retraction +from dsa.optimizer import DiskSparseRiemannianAdam + optimizer = DiskSparseRiemannianAdam( params={"hyperbolic_nodes": initial_weights}, - lr=0.001, - k=1.0, # Hyperbolic mode - disk_dir="./hyperbolic_cache" + lr=1e-3, + k=1.0, # hyperbolic mode + disk_dir="./hyperbolic_cache", + resume=False, # explicit: new run, reinitialize this directory + weight_decay=0.0, + weight_decay_mode="l2", # classic Adam L2 in the Riemannian gradient ) -# The optimizer automatically applies Riemannian scaling and retraction during step() -optimizer.step(updates={"hyperbolic_nodes": (indices, gradients)}) +for indices, gradients in batches: + optimizer.step(updates={"hyperbolic_nodes": (indices, gradients)}) + +# Read embeddings for inference: projected into the ball, safe public API. +emb = optimizer.read_rows(indices, name="hyperbolic_nodes") + +# Curvature annealing: the ball changes, so every row is reprojected. +optimizer.set_curvature(0.5) + +# Portable, checksummed checkpoint (a directory, not a JSON file). +optimizer.export_checkpoint("./checkpoints/run-001") +optimizer.shutdown() ``` +What `step()` does when `k > 0`: + +1. Reads only the batch rows `(w, m, v, t)` from the memmaps. +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 + 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. +6. Projects the result into the ball as a numerical guard and queues the + consistent row update. + +Notes: + +* `init_max_norm_fraction=0.5` (default) globally scales fresh hyperbolic + weights so that no row starts closer than half the radius to the boundary. + Pass `init_max_norm_fraction=1.0` if you want only strict projection. +* `resume=True` restores the saved curvature; a differing constructor `k` is + ignored with a warning. Use `set_curvature()` to change it deliberately. +* `current_k=...` on `step()` is a shortcut for `set_curvature()`; it + reprojects the whole table because the ball itself changes. +* `weight_decay_mode="decoupled"` applies decay as a tangent-space term + (AdamW semantics when `k = 0`). `"l2"` is part of the gradient and is + therefore scaled by the Poincaré metric. + --- ## 📖 API Reference @@ -244,56 +288,73 @@ optimizer = DiskSparseRiemannianAdam( beta2=0.999, eps=1e-8, weight_decay=0.0, + weight_decay_mode="l2", # "l2" | "decoupled" k=1.0, disk_dir="./disk_cache", max_queue_size=150, - resume=False, - put_timeout=5.0 + resume=None, # None | True | False + put_timeout=5.0, + retain_params=False, + init_max_norm_fraction=0.5, ) ``` #### Parameters: -* **`params`** (*dict* or *list*): Dictionary of `{"param_name": torch.Tensor}` or a list of tensors to optimize. -* **`lr`** (*float*, optional): Learning rate (default: `1e-3`). -* **`beta1`, `beta2`** (*float*, optional): Coefficients for computing running averages of the gradient and its square (default: `0.9`, `0.999`). -* **`eps`** (*float*, optional): Term added to the denominator for numerical stability (default: `1e-8`). -* **`weight_decay`** (*float*, optional): Weight decay coefficient (default: `0.0`). -* **`k`** (*float*, optional): Poincaré Ball curvature parameter ($k > 0$ for Hyperbolic, $k = 0$ for Euclidean) (default: `1.0`). -* **`disk_dir`** (*str*, optional): Directory path where `np.memmap` files (`_w.dat`, `_m.dat`, `_v.dat`, `_t.dat`) and the `_meta.json` descriptors will be created (default: `"./disk_cache"`). -* **`max_queue_size`** (*int*, optional): Maximum capacity of the background I/O queue (default: `150`). -* **`resume`** (*bool*, optional): `False` (default) — if state files already exist in `disk_dir`, they are fully reinitialized (weights from `params`, zeroed moments and counters) with a warning; `True` — continue from the existing state as is. Shape/dtype mismatch raises `RuntimeError` either way. -* **`put_timeout`** (*float*, optional): Seconds to wait for queue space before a write batch falls back to a synchronous write (no data is ever dropped; default: `5.0`). +* **`params`** (*dict* or *list*): `{name: tensor}` or a list of 2-D tensors. Values are read once to initialize the on-disk state. +* **`lr`** (*float*): Learning rate (default `1e-3`). +* **`beta1`, `beta2`** (*float*): Adam moment coefficients, each in `[0, 1)` (defaults `0.9`, `0.999`). +* **`eps`** (*float*): Denominator floor, must be `> 0` (default `1e-8`). +* **`weight_decay`** (*float*): Non-negative decay coefficient (default `0.0`). +* **`weight_decay_mode`** (*str*): `"l2"` (default) puts `wd * w` into the Euclidean gradient before the Riemannian conversion; `"decoupled"` adds `wd * w` to the tangent descent direction (AdamW semantics when `k = 0`). +* **`k`** (*float*): Curvature magnitude. `k > 0` is the Poincaré ball of curvature `-k`; `k = 0` is Euclidean Adam. Must be finite and non-negative; NaN/inf/negative raise immediately. +* **`disk_dir`** (*str*): Directory for `_w/_m/_v/_t.dat`, `_meta.json`, and `dsa_global_meta.json`. +* **`max_queue_size`** (*int*): Background write-queue capacity (default `150`). +* **`resume`** (*None | bool*): `None` (default) raises if `disk_dir` already contains state, forcing an explicit choice; `True` continues it; `False` reinitializes from `params` with a warning. Legacy 0.1.0/0.2.0 state is migrated on `resume=True`. +* **`put_timeout`** (*float*): Seconds to wait for queue space before a synchronous write (default `5.0`; writes are never dropped). +* **`retain_params`** (*bool*): If `False` (default), the source tensors are replaced by zero-length CPU placeholders after initialization, so the optimizer cannot retain an $O(N)$ table in VRAM. Set `True` only if other code inspects `param_groups[...]["params"]`. +* **`init_max_norm_fraction`** (*float*, `(0, 1]`): For `k > 0`, fresh weights whose largest row norm exceeds this fraction of the ball radius are globally scaled down (default `0.5`). `1.0` means strict projection only. #### Methods: -* **`step(updates, current_k=None)`**: Performs a single optimization step. - * `updates` (*dict*): Dictionary mapping `param_name` to a tuple `(indices_tensor, grads_tensor)`. - * `current_k` (*float* or *Tensor*, optional): Dynamic curvature $k$ override. -* **`shutdown()`**: Drains pending writes, flushes all state files and terminates the background writer thread. Raises `RuntimeError` if the data could not be fully written. Must be called at the end of the training process. -* **`save_checkpoint(path)`**: Writes a consistent checkpoint: drains the queue, flushes `w/m/v/t` and atomically stores a manifest with file paths and shapes. -* **`load_checkpoint(path)`**: Verifies the manifest (shapes, paths, file sizes) and reopens the memmaps. Checkpoint files must live in the optimizer's `disk_dir`. -* **`state_dict()` / `load_state_dict()`**: Raise `NotImplementedError` — DSA state lives on disk and cannot be serialized through these methods. +* **`step(updates, current_k=None)`**: One optimizer step. + * `updates` (*dict*): `{param_name: (indices, grads)}`. Indices must be 1-D integer tensors/arrays in `[0, N)`; gradients must have shape `(len(indices), dim)` and be finite. Duplicate indices are summed. + * `current_k` (*float* or *Tensor*, optional): Change curvature for this and subsequent steps; the whole table is reprojected. +* **`read_rows(indices, name=None)`**: Public read API for inference; waits for in-flight writes, validates indices, rejects corrupted rows, and projects into the current ball. +* **`set_curvature(k, reset_moments=False)`**: Change curvature and reproject every row. Moments are kept by default (documented annealing approximation); `reset_moments=True` gives a mathematically clean restart of the adaptive state. +* **`reproject_all(name=None)`**: Scan the whole table, project finite rows into the current ball, and fully reset non-finite/negative-variance rows. Returns the number of projected rows. +* **`export_checkpoint(directory)`**: Write a portable checkpoint directory (`checkpoint.json`, relative data files, SHA-256 checksums) and return the manifest path. +* **`import_checkpoint(directory, verify_hashes=True)`**: Verify and restore a portable checkpoint, including global step and curvature. +* **`save_checkpoint(directory)` / `load_checkpoint(directory)`**: Directory-only aliases for the two methods above. Passing a `.json` path raises a clear migration error. +* **`shutdown()`**: Drain the queue, flush state, and persist the global step. Retryable after a transient failure; raises while data is unfinished or a write failed. Call it at the end of training. +* **`state_dict()` / `load_state_dict()`**: Raise `NotImplementedError` — use the checkpoint API or reopen the directory with `resume=True`. + +State files are an implementation detail; use `read_rows()` and checkpoints instead of reading `state_files` directly. --- -## ⬆️ Upgrading from 0.1.0 +## ⬆️ Upgrading from 0.2.0 / 0.1.0 -0.2.0 keeps the public API (constructor, `step(updates, current_k)`, `shutdown()`, `state_files`) compatible; the new constructor arguments (`resume`, `put_timeout`) are optional. Behavior changes to be aware of: +0.3.0 is a correctness release with several intentional breaking changes: -* **`resume=False` is the default.** Re-running with default arguments fully reinitializes the state in `disk_dir` (weights from `params`, zeroed moments and per-row counters) and logs a warning. In 0.1.0, existing `_w/_m/_v.dat` files were silently reused. Pass `resume=True` to continue a previous run. -* **Four state files per parameter** instead of three: a per-row step counter `_t.dat` (int32) was added, plus a `_meta.json` descriptor. Old 0.1.0 state directories are not compatible and must be removed (or reopened with the default `resume=False`, which reinitializes them). -* **No dropped batches.** A full write queue blocks (up to `put_timeout` seconds) and then writes synchronously instead of discarding the batch. -* **`shutdown()` can raise.** If pending data cannot be fully written, you get a `RuntimeError` with details instead of a false success message. -* **`state_dict()`/`load_state_dict()` raise `NotImplementedError`.** Use `save_checkpoint()`/`load_checkpoint()` or reopen the directory with `resume=True`. -* **Bias correction is per-row.** Rows that are updated rarely (or for the first time late in training) now get correct Adam bias correction; early-training dynamics differ from 0.1.0, which scaled by the global step counter. +* **`resume` is now `None` by default.** An existing state directory is never reused or silently reinitialized: pass `resume=True` to continue or `resume=False` to start over. +* **Checkpoints are directories.** `save_checkpoint`/`load_checkpoint` accept a directory; `export_checkpoint`/`import_checkpoint` are the primary names. 0.2.0 manifests with absolute paths are rejected — create a new checkpoint. +* **State format v2.** 0.2.0 (`format_version=1`) and 0.1.0 (no meta, no `t`) directories are migrated on `resume=True`: weights are preserved, but moments and per-row counters are reset because 0.3.0 stores tangent-space moments and transports them. Re-running with `resume=False` reinitializes from `params` (this also repairs a partial 0.1.0 directory). +* **The Riemannian step changed.** The old code was Adam in ambient coordinates with gradient scaling and a clip. 0.3.0 uses a Riemannian gradient, tangent-space moments, an exponential map and parallel transport. Hyperbolic training dynamics will differ from 0.2.0 even with the same hyperparameters. +* **`weight_decay` semantics changed for `k > 0`.** The default `"l2"` is now part of the Riemannian gradient instead of a coordinate shrink; use `"decoupled"` if you need AdamW-style behavior. +* **Source tensors are released.** `param_groups[...]["params"]` contains zero-length placeholders by default. Use `read_rows()` for weights. Pass `retain_params=True` for the old behavior. +* **Global step is persisted.** `shutdown()` and checkpoints store it; reopening restores `_global_step` (bias correction remains per-row). +* **Validation is strict.** Negative/out-of-range indices, wrong gradient shapes, non-finite gradients, invalid curvature, and unknown update keys now fail or warn explicitly instead of being ignored. --- ## 🛠️ Best Practices & Troubleshooting -1. **Always call `optimizer.shutdown()`**: `shutdown()` drains the write queue and flushes all memory-mapped buffers. Since 0.2.0 it raises `RuntimeError` instead of silently pretending everything was saved. -2. **Disk Space Limit**: Ensure your NVMe drive has enough space for 4 arrays per parameter (`_w.dat`, `_m.dat`, `_v.dat` float32 and `_t.dat` int32) plus a small `_meta.json`. For 1M entities with 128 dimensions, this takes ~1.5 GB + 4 MB. -3. **Queue Overflow Warning**: A full queue no longer drops batches: the write waits `put_timeout` seconds and then completes synchronously (a warning and the `writer._sync_fallbacks` counter are recorded). If this happens regularly, your disk is the bottleneck: increase `max_queue_size` or reduce the step rate. -4. **Reopening a state directory** is controlled by `resume`: leftover files are never silently reused — you either continue explicitly (`resume=True`), reinitialize with a warning (`resume=False`, the default), or get a clear error (shape mismatch, partial or truncated files, missing meta). +1. **Always call `optimizer.shutdown()`**: it drains the write queue, flushes all memmaps, persists the global step and stops the worker. It raises `RuntimeError` instead of pretending that unfinished data was saved; a second call retries after a transient failure. +2. **Disk space.** Four arrays per parameter (`_w/_m/_v.dat` float32, `_t.dat` int32) plus two small JSON descriptors. For 1M entities with 128 dimensions this is ~1.5 GB + 4 MB. +3. **Hyperbolic initialization.** Keep initial norms small. The default `init_max_norm_fraction=0.5` rescales a too-large table globally; a warning tells you when it happens. Starting at the boundary makes $\lambda^2$ huge and learning extremely slow. +4. **Queue overflow.** A full queue waits `put_timeout` seconds and then writes synchronously (warning + `writer._sync_fallbacks`). If this happens regularly, your disk is the bottleneck: increase `max_queue_size` or reduce the step rate. +5. **Reopening a state directory.** Use `resume=True` to continue, or `resume=False` to reinitialize. Shape mismatch, partial/truncated files, a missing meta file, a leftover reinitialization marker, or a curvature mismatch produce explicit errors. +6. **Checkpoints.** Use a dedicated directory per checkpoint. Import verifies sizes and SHA-256 checksums; corrupted or foreign checkpoints are rejected before the live state is touched. +7. **Boundary behavior.** Points are kept strictly inside the ball (`BALL_EPS = 1e-5`). Very long training with a large `lr` can push embeddings to that margin; the optimizer stays finite and projects, but the chart is ill-conditioned there. Prefer a smaller `lr` and early stopping rather than training on the boundary. --- diff --git a/dsa/__init__.py b/dsa/__init__.py index 90a0ba4..4142c40 100644 --- a/dsa/__init__.py +++ b/dsa/__init__.py @@ -1,4 +1,6 @@ from .optimizer import DiskSparseRiemannianAdam -from .async_writer import AsyncDiskWriter +from .async_writer import AsyncDiskWriter, WriterState -__all__ = ["DiskSparseRiemannianAdam", "AsyncDiskWriter"] +__version__ = "0.3.0" + +__all__ = ["DiskSparseRiemannianAdam", "AsyncDiskWriter", "WriterState"] diff --git a/dsa/async_writer.py b/dsa/async_writer.py index 0cce347..6f943f8 100644 --- a/dsa/async_writer.py +++ b/dsa/async_writer.py @@ -1,57 +1,83 @@ -import threading +"""Background writer for DSA state files. + +Guarantees (see CHANGELOG.md): + +* F1: a batch is never dropped. If the queue is full, ``put_update`` waits + ``put_timeout`` seconds and then writes the batch synchronously. +* F2: rows in flight are tracked per parameter, so ``step`` never reads a + row the writer has not finished writing. +* F5: indices are sorted before writing (sequential IO). +* 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. +""" + +from __future__ import annotations + +import logging import queue +import threading import time +from enum import Enum +from typing import Dict, Optional, Set + import numpy as np -import logging -from typing import Dict, Set logger = logging.getLogger("DSA_AsyncWriter") +class WriterState(str, Enum): + RUNNING = "RUNNING" + DRAINING = "DRAINING" + STOPPED = "STOPPED" + FAILED = "FAILED" + + class AsyncDiskWriter: - """ - Фоновый воркер для асинхронной записи обновлений на диск. - Реализует Double Buffering: GPU вычисляет батч N+1, пока CPU пишет батч N. - - Гарантии (исправления по разбору дефектов, см. CHANGELOG.md): - - - F1: записи никогда не выбрасываются. Если очередь полна, ``put_update`` - блокируется на ``put_timeout`` секунд; если за это время место не освободилось, - батч пишется синхронно напрямую в memmap (под локом). Потеря данных исключена. - - F2: писатель под локом ведёт множество строк «в полёте» (в очереди либо - в записи) по каждому имени параметра — только целочисленные индексы. - ``intersects_inflight`` даёт оптимизатору дешёвую проверку пересечения, - ``wait_all_flushed`` — полную очистку очереди. - - F5: индексы сортируются перед записью — локальность IO. - - F6: ``shutdown`` дренирует очередь до конца, ждёт поток с достаточным - таймаутом и при невозможности дописать данные честно бросает ошибку. - Голых ``except: pass`` нет; все ошибки записи фиксируются и всплывают. - """ + """Double-buffered asynchronous writer over ``np.memmap`` files.""" def __init__(self, memmaps: Dict[str, np.memmap], max_queue_size: int = 150, put_timeout: float = 5.0): + if max_queue_size <= 0: + raise ValueError("max_queue_size must be > 0") + if put_timeout < 0: + raise ValueError("put_timeout must be >= 0") self.memmaps = dict(memmaps) - self.queue = queue.Queue(maxsize=max_queue_size) + self.queue: "queue.Queue" = queue.Queue(maxsize=max_queue_size) self._stop_signal = object() - self._lock = threading.Lock() + self._lock = threading.Lock() # protects memmaps and _inflight + self._condition = threading.Condition() # protects the state machine self._put_timeout = float(put_timeout) - # F2: строки в полёте, по имени параметра; защищено self._lock self._inflight: Dict[str, Set[int]] = {} - # F6: честный учёт ошибок вместо молчаливых потерь self._write_errors = [] self._flush_errors = [] self._sync_fallbacks = 0 - self._stopped = False + self._state = WriterState.RUNNING + self._active_puts = 0 + self._stop_sent = False + self._fatal_error: Optional[BaseException] = None self.worker_thread = threading.Thread( target=self._process_queue, name="DSA_AsyncWriter", daemon=True) self.worker_thread.start() - logger.info("💿 AsyncDiskWriter: Фоновый поток инициализирован.") + logger.info("💿 AsyncDiskWriter: фоновый поток инициализирован.") + + # ───────────────────────────── state ───────────────────────────── + + @property + def state(self) -> WriterState: + with self._condition: + return self._state + + @property + def _stopped(self) -> bool: + """Backward-compatible view: True only after a clean shutdown.""" + return self.state is WriterState.STOPPED - # ─────────────────────────── F2: строки в полёте ─────────────────────────── + # ─────────────────────── F2: rows in flight ─────────────────────── def intersects_inflight(self, name: str, indices) -> bool: - """Есть ли среди `indices` строки, запись которых ещё не завершена.""" + """Whether any of ``indices`` is still being written.""" with self._lock: inflight = self._inflight.get(name) if not inflight: @@ -60,40 +86,55 @@ def intersects_inflight(self, name: str, indices) -> bool: return bool(inflight.intersection(indices.tolist())) return any(int(i) in inflight for i in indices) - def wait_all_flushed(self, timeout: float = None): - """Дождаться полной очистки очереди и обнуления множества строк в полёте. + def wait_all_flushed(self, timeout: Optional[float] = None): + """Wait until the queue is empty and no rows are in flight. - `timeout=None` — ждать без ограничений (обычный путь в обучении); - при числовом timeout невыполнение дрена за отведённое время — ошибка. + Raises if the worker died, if a write failed during the session, or + if ``timeout`` expires. A failed write is never hidden from the + training loop. """ - if timeout is None: - self.queue.join() - else: - if not self._wait_queue_drained(timeout): + deadline = None if timeout is None else time.monotonic() + timeout + while True: + if self._write_errors: + raise RuntimeError( + "AsyncDiskWriter: за сессию были ошибки записи; состояние " + f"на диске неполно. Первая ошибка: {self._write_errors[0]}") + if not self.worker_thread.is_alive() and self.queue.unfinished_tasks: + raise RuntimeError( + "AsyncDiskWriter: фоновый поток завершился, но очередь " + f"содержит {self.queue.unfinished_tasks} незавершённых " + "батчей — данные не записаны") + if self.queue.unfinished_tasks == 0: + break + if deadline is not None and time.monotonic() >= deadline: raise RuntimeError( f"AsyncDiskWriter: очередь не опустела за {timeout:.1f} с " f"(осталось батчей: {self.queue.unfinished_tasks}) — " f"данные НЕ дописаны") + time.sleep(0.005) + with self._lock: leftover = {n: sorted(s) for n, s in self._inflight.items() if s} if leftover: raise RuntimeError( - f"AsyncDiskWriter: очередь пуста, но строки остались в полёте: {leftover}") + f"AsyncDiskWriter: очередь пуста, но строки остались в полёте: " + f"{leftover}") def _wait_queue_drained(self, timeout: float) -> bool: - """Ограниченное по времени ожидание доработы очереди (без join-зависаний).""" deadline = time.monotonic() + timeout while True: + if not self.worker_thread.is_alive() and self.queue.unfinished_tasks: + return False if self.queue.unfinished_tasks == 0: return True if time.monotonic() >= deadline: return self.queue.unfinished_tasks == 0 - time.sleep(0.01) + time.sleep(0.005) - # ─────────────────────────── запись ─────────────────────────── + # ─────────────────────────── writing ─────────────────────────── def _apply_write(self, name: str, indices, w_np, m_np, v_np, t_np=None): - """Синхронная запись одного батча в memmap (F5: с сортировкой индексов).""" + """Synchronous write of one batch (F5: sorted indices).""" indices = np.asarray(indices) order = np.argsort(indices) indices = indices[order] @@ -107,30 +148,30 @@ def _apply_write(self, name: str, indices, w_np, m_np, v_np, t_np=None): t_mm = self.memmaps.get(f"{name}_t") if t_mm is None: raise RuntimeError( - f"передан счётчик шага для '{name}', но memmap '{name}_t' не открыт") + f"передан счётчик шага для '{name}', но memmap " + f"'{name}_t' не открыт") t_mm[indices] = t_sorted def _process_queue(self): - """Цикл записи: работает параллельно с основным обучением.""" while True: batch_data = self.queue.get() - if batch_data is self._stop_signal: - # Перед выходом фиксируем все данные на физическом носителе - self._final_flush() - self.queue.task_done() + try: + self._final_flush() + except Exception as exc: # pragma: no cover - defensive + self._flush_errors.append(f"финальный flush: {exc!r}") + finally: + self.queue.task_done() break name, indices, w_np, m_np, v_np, t_np = batch_data try: self._apply_write(name, indices, w_np, m_np, v_np, t_np) - except Exception as e: # F6: ошибка фиксируется и всплывёт на shutdown - msg = f"батч '{name}' ({len(indices)} строк) НЕ записан: {e!r}" + except Exception as exc: # F6: recorded, re-raised on shutdown + msg = f"батч '{name}' ({len(indices)} строк) НЕ записан: {exc!r}" self._write_errors.append(msg) logger.error("❌ IO Error: " + msg) finally: - # F2: строки этого батча больше не в полёте — до task_done, - # чтобы после queue.join() множество гарантированно было пустым with self._lock: inflight = self._inflight.get(name) if inflight is not None: @@ -138,97 +179,155 @@ def _process_queue(self): self.queue.task_done() def _final_flush(self): - """Принудительная запись на диск перед выключением (без голых except).""" - logger.info("💾 AsyncDiskWriter: Финализация данных на диске...") + logger.info("💾 AsyncDiskWriter: финализация данных на диске...") with self._lock: for key, mm in self.memmaps.items(): try: mm.flush() - except Exception as e: - msg = f"flush {key} не выполнен: {e!r}" + except Exception as exc: + msg = f"flush {key} не выполнен: {exc!r}" self._flush_errors.append(msg) logger.error("❌ " + msg) def put_update(self, name: str, indices_pt, w_pt, m_pt, v_pt, t_pt=None): - """ - Подготовка данных для записи. Выполняется в потоке GPU. + """Queue one batch; never drops it (F1).""" + with self._condition: + if self._state is not WriterState.RUNNING: + raise RuntimeError( + f"AsyncDiskWriter в состоянии {self._state.value}: запись " + f"невозможна (shutdown уже начат или завершён)") + self._active_puts += 1 - F1: данные не теряются никогда. При полной очереди ждём место - `put_timeout` секунд; не дождались — пишем синхронно напрямую. - """ - if self._stopped: - raise RuntimeError("AsyncDiskWriter уже остановлен: запись невозможна") - - # Переводим в Numpy на CPU. Это освобождает GPU тензоры немедленно. - idx_np = np.asarray(indices_pt.detach().cpu().numpy()) - w_np = w_pt.detach().cpu().numpy() - m_np = m_pt.detach().cpu().numpy() - v_np = v_pt.detach().cpu().numpy() - t_np = None if t_pt is None else np.asarray( - t_pt.detach().cpu().numpy(), dtype=np.int32) - - idx_set = set(int(i) for i in idx_np.tolist()) - # F2: помечаем строки «в полёте» до постановки в очередь — тогда между - # проверкой пересечения в step() и постановкой не возникает окна гонки - with self._lock: - self._inflight.setdefault(name, set()).update(idx_set) try: - self.queue.put((name, idx_np, w_np, m_np, v_np, t_np), - timeout=self._put_timeout) - except queue.Full: - # F1: диск не успевает — пишем сами, синхронно, данные не теряем. - self._sync_fallbacks += 1 - logger.warning( - f"⚠️ Очередь записи полна дольше {self._put_timeout:.1f} с " - f"('{name}', {len(idx_np)} строк): синхронная запись напрямую.") + idx_np = np.asarray(indices_pt.detach().cpu().numpy()) + w_np = w_pt.detach().cpu().numpy() + m_np = m_pt.detach().cpu().numpy() + v_np = v_pt.detach().cpu().numpy() + t_np = None if t_pt is None else np.asarray( + t_pt.detach().cpu().numpy(), dtype=np.int32) + + idx_set = {int(i) for i in idx_np.tolist()} + with self._lock: + self._inflight.setdefault(name, set()).update(idx_set) try: - self._apply_write(name, idx_np, w_np, m_np, v_np, t_np) - finally: - # строки записаны синхронно — из полёта их снимаем + self.queue.put((name, idx_np, w_np, m_np, v_np, t_np), + timeout=self._put_timeout) + except queue.Full: + # F1: disk is behind — write synchronously, never drop. + self._sync_fallbacks += 1 + logger.warning( + "⚠️ Очередь записи полна дольше %.1f с ('%s', %d строк): " + "синхронная запись напрямую.", + self._put_timeout, name, len(idx_np)) + try: + self._apply_write(name, idx_np, w_np, m_np, v_np, t_np) + finally: + with self._lock: + self._inflight.setdefault(name, set()).difference_update(idx_set) + except BaseException: with self._lock: self._inflight.setdefault(name, set()).difference_update(idx_set) + raise + finally: + with self._condition: + self._active_puts -= 1 + self._condition.notify_all() - def shutdown(self, timeout: float = 300.0): - """ - Честное завершение работы (F6): - 1. Дописывает все оставшиеся батчи из очереди на диск. - 2. Выполняет flush() всех memmap. - 3. Ждёт завершения потока с достаточным таймаутом. - - Если за `timeout` данные дописать не удалось, либо за сессию были - ошибки записи/flush — бросает RuntimeError. «Всё записано» сообщается - только когда это правда. + # ─────────────────────────── shutdown ─────────────────────────── + + def shutdown(self, timeout: float = 300.0, *, raise_on_error: bool = True) -> bool: + """Drain, flush and stop; retryable after a transient failure. + + Returns ``True`` on a clean stop and ``False`` when + ``raise_on_error=False`` and the writer could not finish. A writer + that previously failed and whose thread is still alive is retried, + not reported as stopped (the 0.2.0 bug). """ - if self._stopped: - return - self._stopped = True - logger.info("⏳ AsyncDiskWriter: Дренаж очереди и финальный сброс на диск...") - - # 1. Всё, что уже в очереди, будет записано (дрен тоже ограничен - # таймаутом: застрявший диск — честная ошибка, а не вечное ожидание) - if not self._wait_queue_drained(timeout): - raise RuntimeError( - f"AsyncDiskWriter: за {timeout:.1f} с очередь не опустела " - f"(осталось батчей: {self.queue.unfinished_tasks}) — " - f"данные НЕ дописаны на диск") - # 2. Сигнал остановки: воркер сделает финальный flush и выйдет - # (после дрена очередь пуста, put не блокируется) - self.queue.put(self._stop_signal) - # 3. Ждём завершения фонового потока (финальный msync может быть долгим) - self.worker_thread.join(timeout=timeout) - - if self.worker_thread.is_alive(): - raise RuntimeError( - f"AsyncDiskWriter: фоновый поток не завершился за {timeout:.1f} с — " - f"данные НЕ гарантированно записаны на диск") - - problems = [] - if self._flush_errors: - problems.append(f"ошибки финального flush: {len(self._flush_errors)}") - if self._write_errors: - problems.append(f"невыполненные записи за сессию: {len(self._write_errors)}") - if problems: - first = (self._flush_errors + self._write_errors)[0] - raise RuntimeError( - "AsyncDiskWriter: " + "; ".join(problems) + f"; пример: {first}") - logger.info("✅ AsyncDiskWriter: Все данные успешно сохранены на диск.") + if timeout is None: + raise TypeError("timeout must be a number of seconds") + if timeout < 0: + raise ValueError("timeout must be >= 0") + + with self._condition: + if self._state is WriterState.STOPPED: + return True + if (self._state is WriterState.FAILED + and self._fatal_error is not None + and not self.worker_thread.is_alive()): + if raise_on_error: + raise self._fatal_error + return False + self._state = WriterState.DRAINING + self._condition.notify_all() + + err: Optional[BaseException] = None + deadline = None if timeout is None else time.monotonic() + timeout + + # 1. Wait for put_update calls that started before DRAINING. + while True: + with self._condition: + if self._active_puts == 0: + break + if deadline is not None and time.monotonic() >= deadline: + err = RuntimeError( + f"AsyncDiskWriter: {self._active_puts} вызов(ов) put_update " + f"не завершились за {timeout:.1f} с") + break + time.sleep(0.005) + + # 2. Drain queued batches (bounded: a stuck disk is an error, not a hang). + if err is None and not self._wait_queue_drained(timeout): + if self._write_errors: + err = RuntimeError( + "AsyncDiskWriter: невыполненные записи за сессию: " + f"{len(self._write_errors)}; пример: {self._write_errors[0]}") + else: + err = RuntimeError( + f"AsyncDiskWriter: за {timeout:.1f} с очередь не опустела " + f"(осталось батчей: {self.queue.unfinished_tasks}) — " + f"данные НЕ дописаны") + + # 3. Stop the worker (FIFO: stop is enqueued after all data). + if err is None and self.worker_thread.is_alive() and not self._stop_sent: + self._stop_sent = True + try: + self.queue.put_nowait(self._stop_signal) + except queue.Full: + err = RuntimeError( + "AsyncDiskWriter: очередь заполнилась между дреном и " + "сигналом остановки") + + # 4. Join (msync may be slow). + if err is None and self.worker_thread.is_alive(): + self.worker_thread.join(timeout=timeout) + if err is None and self.worker_thread.is_alive(): + err = RuntimeError( + f"AsyncDiskWriter: фоновый поток не завершился за " + f"{timeout if timeout is not None else '∞'} с — данные НЕ " + f"гарантированно записаны на диск") + + if err is None: + problems = [] + if self._flush_errors: + problems.append(f"ошибки финального flush: {len(self._flush_errors)}") + if self._write_errors: + problems.append(f"невыполненные записи за сессию: {len(self._write_errors)}") + if problems: + first = (self._flush_errors + self._write_errors)[0] + err = RuntimeError( + "AsyncDiskWriter: " + "; ".join(problems) + + f"; пример: {first}") + + if err is not None: + with self._condition: + self._state = WriterState.FAILED + self._fatal_error = err + if raise_on_error: + raise err + return False + + with self._condition: + self._state = WriterState.STOPPED + self._fatal_error = None + logger.info("✅ AsyncDiskWriter: все данные сохранены на диск.") + return True diff --git a/dsa/manifold.py b/dsa/manifold.py new file mode 100644 index 0000000..3657ffd --- /dev/null +++ b/dsa/manifold.py @@ -0,0 +1,164 @@ +"""Poincare ball geometry for curvature -c (c > 0). + +Conventions +----------- +The ball is ``B_c = {x : ||x|| < 1/sqrt(c)}`` with the conformal metric + + g_x = lambda(x)^2 I, lambda(x) = 2 / (1 - c ||x||^2). + +For a Euclidean gradient ``g_E`` the Riemannian gradient is + + grad_R f(x) = lambda(x)^-2 g_E = (1 - c ||x||^2)^2 / 4 * g_E. + +The exponential map and the parallel transport follow the standard +Poincare-ball formulas (the same ones used by geoopt with ``k = -c``). +All functions operate on the last dimension and support batched rows. +""" + +from __future__ import annotations + +import math + +import numpy as np +import torch + +# Relative distance from the boundary kept by the projection. A point at +# radius sqrt((1 - EPS) / c) has conformal factor ~2/EPS, which is already +# numerically extreme; keeping a small margin avoids infinities. +BALL_EPS = 1e-5 + +# tanh saturates to exactly 1.0 in float32 for arguments above ~9, which puts +# the second Moebius operand on the boundary and makes the addition singular. +# Keeping it strictly inside is the numerical equivalent of "finite u". +EXP_TANH_MAX = 1.0 - 1e-7 + + +def sq_norm(x: torch.Tensor) -> torch.Tensor: + """Squared Euclidean norm over the last dimension, shape (..., 1).""" + return (x * x).sum(dim=-1, keepdim=True) + + +def lambda_x(x: torch.Tensor, c: float) -> torch.Tensor: + """Conformal factor ``2 / (1 - c ||x||^2)``.""" + return 2.0 / (1.0 - c * sq_norm(x)).clamp_min(1e-15) + + +def project_to_ball(x: torch.Tensor, c: float, eps: float = BALL_EPS) -> torch.Tensor: + """Project rows into the open ball of radius ``sqrt((1-eps)/c)``. + + Uses a multiplicative scale instead of ``x / ||x||`` so that the zero + row is handled without NaNs. + """ + if c <= 0.0: + return x + max_radius = math.sqrt((1.0 - eps) / c) + norm = x.norm(dim=-1, keepdim=True) + scale = torch.where( + norm > max_radius, + max_radius / norm.clamp_min(1e-30), + torch.ones_like(norm), + ) + return x * scale + + +def project_to_ball_np(w: np.ndarray, c: float, eps: float = BALL_EPS) -> np.ndarray: + """NumPy version of :func:`project_to_ball` (used for memmap chunks). + + Norms are computed in float64 to avoid overflow on corrupted float32 + data; the result is returned in the input dtype. + """ + if c <= 0.0: + return w + max_radius = math.sqrt((1.0 - eps) / c) + norm = np.linalg.norm(w.astype(np.float64), axis=-1) + scale = np.where(norm > max_radius, max_radius / np.maximum(norm, 1e-30), 1.0) + return (w.astype(np.float64) * scale[..., None]).astype(w.dtype, copy=False) + + +def mobius_add(x: torch.Tensor, y: torch.Tensor, c: float) -> torch.Tensor: + """Moebius addition on the ball of curvature -c.""" + if c <= 0.0: + return x + y + x2 = sq_norm(x) + y2 = sq_norm(y) + xy = (x * y).sum(dim=-1, keepdim=True) + num = (1.0 + 2.0 * c * xy + c * y2) * x + (1.0 - c * x2) * y + den = (1.0 + 2.0 * c * xy + c * c * x2 * y2).clamp_min(1e-15) + return num / den + + +def gyration(u: torch.Tensor, v: torch.Tensor, w: torch.Tensor, c: float) -> torch.Tensor: + """Gyration ``gyr[u, v] w`` for curvature -c (geoopt formula).""" + if c <= 0.0: + return w + u2 = sq_norm(u) + v2 = sq_norm(v) + uv = (u * v).sum(dim=-1, keepdim=True) + uw = (u * w).sum(dim=-1, keepdim=True) + vw = (v * w).sum(dim=-1, keepdim=True) + c2 = c * c + a = -c2 * uw * v2 + c * vw + 2.0 * c2 * uv * vw + b = -c2 * vw * u2 - c * uw + den = (1.0 + 2.0 * c * uv + c2 * u2 * v2).clamp_min(1e-15) + return w + 2.0 * (a * u + b * v) / den + + +def expmap(x: torch.Tensor, u: torch.Tensor, c: float) -> torch.Tensor: + """Exponential map ``exp_x(u)`` on the ball of curvature -c. + + ``u`` is a tangent vector in coordinates at ``x``; its Riemannian norm + is ``lambda(x) ||u||``, which is why ``lambda(x)`` appears below. + """ + if c <= 0.0: + return x + u + u_norm = u.norm(dim=-1, keepdim=True).clamp_min(1e-30) + lx = lambda_x(x, c) + sqrt_c = math.sqrt(c) + scale = torch.tanh(sqrt_c * lx * u_norm / 2.0).clamp(max=EXP_TANH_MAX) / ( + sqrt_c * u_norm) + return mobius_add(x, scale * u, c) + + +def parallel_transport( + x: torch.Tensor, y: torch.Tensor, u: torch.Tensor, c: float +) -> torch.Tensor: + """Parallel transport of ``u`` from ``T_x`` to ``T_y`` along the geodesic. + + Uses the gyration form ``gyr[y, -x] u * lambda(x) / lambda(y)``, which is + an isometry between the two tangent spaces (verified numerically). + """ + if c <= 0.0: + return u + return gyration(y, -x, u, c) * lambda_x(x, c) / lambda_x(y, c) + + +def riemannian_gradient(g: torch.Tensor, x: torch.Tensor, c: float) -> torch.Tensor: + """Convert a Euclidean gradient into a Riemannian one: ``lambda(x)^-2 g``.""" + if c <= 0.0: + return g + return g * (1.0 - c * sq_norm(x)).clamp_min(1e-15).pow(2) / 4.0 + + +def distance(x: torch.Tensor, y: torch.Tensor, c: float) -> torch.Tensor: + """Geodesic distance on the ball of curvature -c (for diagnostics/tests). + + Rows outside the ball yield NaN instead of a silently wrong value. + """ + if c <= 0.0: + return (x - y).norm(dim=-1) + x2 = sq_norm(x) + y2 = sq_norm(y) + diff2 = sq_norm(x - y) + den = (1.0 - c * x2) * (1.0 - c * y2) + arg = 1.0 + 2.0 * c * diff2 / den.clamp_min(1e-15) + nan = torch.full_like(arg, float("nan")) + arg = torch.where(den > 0.0, arg.clamp_min(1.0), nan) + return torch.acosh(arg) / math.sqrt(c) + + +def is_in_ball(x: torch.Tensor, c: float, eps: float = BALL_EPS) -> torch.Tensor: + """Boolean mask: rows strictly inside the ball (with margin ``eps``).""" + if c <= 0.0: + return torch.isfinite(x).all(dim=-1) + max_radius = math.sqrt((1.0 - eps) / c) + return torch.isfinite(x).all(dim=-1) & (x.norm(dim=-1) <= max_radius) diff --git a/dsa/optimizer.py b/dsa/optimizer.py index 28d3351..87c8c53 100644 --- a/dsa/optimizer.py +++ b/dsa/optimizer.py @@ -1,61 +1,141 @@ -import os +"""Disk-based sparse Riemannian Adam (DSA). + +The optimizer keeps ``w`` (weights), ``m`` (first moment), ``v`` (second +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``: + +* 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. +""" + +from __future__ import annotations + +import hashlib import json -import math import logging -from datetime import datetime +import math +import os +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, Tuple, Union -import torch import numpy as np +import torch from torch.optim.optimizer import Optimizer -from typing import Dict, Tuple, Optional, Union from .async_writer import AsyncDiskWriter +from .manifold import ( + BALL_EPS, + expmap, + parallel_transport, + project_to_ball, + project_to_ball_np, + riemannian_gradient, +) logger = logging.getLogger("DSA") -# Версия формата дискового состояния (meta-файлы). Несовпадение — явная ошибка. -STATE_FORMAT_VERSION = 1 +# Version of the per-parameter and global metadata formats. +STATE_FORMAT_VERSION = 2 +LEGACY_STATE_FORMAT_VERSIONS = (1,) +ALGORITHM = "riemannian_adam" + +CHECKPOINT_FORMAT = "DSA_checkpoint" +CHECKPOINT_VERSION = 2 +CHECKPOINT_MANIFEST_NAME = "checkpoint.json" +GLOBAL_META_NAME = "dsa_global_meta.json" -# R3: отладочный режим: неизвестный ключ updates — исключение вместо warning. -# Включается переменной окружения DSA_DEBUG=1 либо присвоением -# dsa.optimizer.DEBUG = True. +# Number of rows processed at once by whole-table operations. +CHUNK_ROWS = 65536 + +# R3: debug mode: an unknown key in ``updates`` raises instead of warning. DEBUG = os.environ.get("DSA_DEBUG", "0") not in ("", "0") +WEIGHT_DECAY_MODES = ("l2", "decoupled") + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def _validate_k(k: Any, where: str = "k") -> float: + """Return a finite non-negative curvature; reject NaN/inf/negative early.""" + try: + value = float(k.item() if isinstance(k, torch.Tensor) else k) + except (TypeError, ValueError) as exc: + raise TypeError(f"{where} must be a real number, got {k!r}") from exc + if not math.isfinite(value): + raise ValueError(f"{where} must be finite, got {value!r}") + if value < 0.0: + raise ValueError(f"{where} must be >= 0, got {value!r}") + return value + + +def _validate_mode(mode: str) -> str: + if mode not in WEIGHT_DECAY_MODES: + raise ValueError( + f"weight_decay_mode must be one of {WEIGHT_DECAY_MODES}, got {mode!r}") + return mode + def _project_rows_to_ball(w: torch.Tensor, k: float) -> torch.Tensor: - """Проекция строк в шар Пуанкаре радиуса sqrt((1-1e-5)/k) (F7 и ретракция).""" - max_radius = math.sqrt((1.0 - 1e-5) / k) - norm = torch.norm(w, p=2, dim=-1, keepdim=True) - return torch.where(norm > max_radius, w / norm * max_radius, w) + """Backward-compatible alias for :func:`manifold.project_to_ball`.""" + return project_to_ball(w, k) if k > 0.0 else w class DiskSparseRiemannianAdam(Optimizer): - """ - Disk-based Sparse Riemannian Adam (DSA). - - Специализированный оптимизатор для обучения сверхбольших графовых эмбеддингов. - - Хранит состояния (m, v), веса (w) и пер-строчные счётчики шагов (t) на NVMe - через np.memmap. - - Масштабирует использование памяти от размера батча O(B), а не графа O(N). - - Реализует Riemannian Gradient Scaling и проекцию (Retraction) для шара Пуанкаре. - - Использует фоновую асинхронную запись для устранения простоев GPU. - - Гарантии корректности (см. CHANGELOG.md): записи никогда не выбрасываются - при полной очереди (F1); step() не читает строки, находящиеся в полёте - у писателя, без дрена очереди (F2); повторное открытие каталога состояния - управляется явным `resume` и meta-файлами (F3); - bias-коррекция Adam — по пер-строчному счётчику шага, а не по глобальному t - (F4); индексы сортируются перед IO (F5); чтение весов проецируется в шар (F7); - чекпоинты — save_checkpoint/load_checkpoint, state_dict() честно бросает - NotImplementedError (F8). Дополнительно: не-конечные строки из битого файла - обрабатываются на чтении (R1), переинициализация помечается маркером в meta - (R2), неизвестные ключи updates дают warning со счётчиком (R3). + """Out-of-core sparse Riemannian Adam. + + Parameters + ---------- + params: + ``{name: tensor}`` or a list of tensors. Tensors are used only to + initialise the on-disk state; by default the optimizer does not keep + references to them (``retain_params=False``), which is what makes the + O(batch) memory claim true. Use :meth:`read_rows` to read weights. + k: + Curvature magnitude: ``k > 0`` is the Poincare ball of curvature + ``-k``, ``k = 0`` is Euclidean Adam. A state directory stores its + curvature; on ``resume=True`` the stored value wins over the + constructor argument (a warning is logged if they differ). + resume: + ``None`` (default) refuses to touch an existing state directory and + asks for an explicit choice; ``True`` continues it; ``False`` + reinitialises it from ``params`` (with a warning). A state directory + created by 0.1.0/0.2.0 is migrated on ``resume=True``: weights are + preserved, moments and per-row counters are reset because the old + math is not compatible with tangent-space moments. + weight_decay_mode: + ``"l2"`` (default) adds ``weight_decay * w`` to the Euclidean + gradient before the Riemannian conversion, i.e. it is the + Riemannian gradient of ``f + wd/2 * ||w||^2`` (classic Adam L2 + semantics). ``"decoupled"`` adds ``wd * w`` to the descent + direction in the tangent space, reproducing AdamW exactly when + ``k = 0``. + retain_params: + If ``False`` (default), source tensors are replaced by zero-length + CPU placeholders after initialisation, so the optimizer cannot keep + a full embedding table alive in VRAM. Set to ``True`` only for + compatibility with code that inspects ``param_groups[...]['params']``. """ - def __init__(self, params: Union[list, dict], lr=1e-3, beta1=0.9, beta2=0.999, - eps=1e-8, weight_decay=0, k=1.0, disk_dir="./disk_cache", - max_queue_size=150, resume: bool = False, - put_timeout: float = 5.0): + def __init__(self, params: Union[list, dict], lr: float = 1e-3, + beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, + weight_decay: float = 0.0, k: float = 1.0, + disk_dir: str = "./disk_cache", max_queue_size: int = 150, + resume: Optional[bool] = None, put_timeout: float = 5.0, + weight_decay_mode: str = "l2", + retain_params: bool = False, + init_max_norm_fraction: float = 0.5): if isinstance(params, dict): self.param_names = list(params.keys()) @@ -64,104 +144,293 @@ def __init__(self, params: Union[list, dict], lr=1e-3, beta1=0.9, beta2=0.999, params_list = list(params) self.param_names = [f"param_{i}" for i in range(len(params_list))] - defaults = dict(lr=lr, beta1=beta1, beta2=beta2, eps=eps, - weight_decay=weight_decay, k=k) + if not params_list: + raise ValueError("params must not be empty") + if len(set(self.param_names)) != len(self.param_names): + raise ValueError("parameter names must be unique") + for name in self.param_names: + if not isinstance(name, str) or not name: + raise ValueError("parameter names must be non-empty strings") + if os.sep in name or (os.altsep is not None and os.altsep in name): + raise ValueError( + f"parameter name {name!r} must not contain path separators") + for name, param in zip(self.param_names, params_list): + if not isinstance(param, torch.Tensor): + raise TypeError(f"parameter '{name}' is not a torch.Tensor") + if param.dim() != 2: + raise ValueError( + f"parameter '{name}' must be 2-D (rows, dim), got shape " + f"{tuple(param.shape)}") + if not param.is_floating_point(): + raise TypeError(f"parameter '{name}' must be floating point") - super().__init__(params_list, defaults) + if lr < 0: + raise ValueError(f"lr must be >= 0, got {lr!r}") + if not 0.0 <= beta1 < 1.0: + raise ValueError(f"beta1 must be in [0, 1), got {beta1!r}") + if not 0.0 <= beta2 < 1.0: + raise ValueError(f"beta2 must be in [0, 1), got {beta2!r}") + if eps <= 0: + raise ValueError(f"eps must be > 0, got {eps!r}") + if weight_decay < 0: + raise ValueError(f"weight_decay must be >= 0, got {weight_decay!r}") + if max_queue_size <= 0: + raise ValueError("max_queue_size must be > 0") + if put_timeout < 0: + raise ValueError("put_timeout must be >= 0") + if not 0.0 < init_max_norm_fraction <= 1.0: + raise ValueError( + "init_max_norm_fraction must be in (0, 1], got " + f"{init_max_norm_fraction!r}") + self.init_max_norm_fraction = float(init_max_norm_fraction) - self.disk_dir = os.path.abspath(disk_dir) + requested_k = _validate_k(k) + mode = _validate_mode(weight_decay_mode) + + self.disk_dir = os.path.abspath(os.fspath(disk_dir)) os.makedirs(self.disk_dir, exist_ok=True) - self.state_files = {} - self._paths = {} - # R1/R3: счётчики диагностики (не-конечные строки, неизвестные ключи) + self._paths: Dict[str, Dict[str, str]] = {} + meta_paths: Dict[str, str] = {} + for name, param in zip(self.param_names, params_list): + shape = tuple(int(s) for s in param.shape) + self._paths[name] = { + s: os.path.join(self.disk_dir, f"{name}_{s}.dat") + for s in ("w", "m", "v", "t") + } + meta_paths[name] = os.path.join(self.disk_dir, f"{name}_meta.json") + + global_meta_path = os.path.join(self.disk_dir, GLOBAL_META_NAME) + state_present = os.path.exists(global_meta_path) or any( + os.path.exists(p) for paths in self._paths.values() + for p in paths.values()) or any( + os.path.exists(p) for p in meta_paths.values()) + + if resume is None: + if state_present: + raise RuntimeError( + f"в {self.disk_dir} уже есть состояние DSA. Явно выберите " + f"режим: resume=True — продолжить, resume=False — " + f"переинициализировать из params (старое состояние будет " + f"перезаписано).") + resume = False + if not isinstance(resume, bool): + raise TypeError("resume must be True, False or None") + if resume and not state_present: + raise RuntimeError( + f"resume=True, но в {self.disk_dir} нет состояния DSA: " + f"проверьте disk_dir или создайте состояние с resume=False.") + + stored_k = self._peek_stored_k(meta_paths) if resume else None + effective_k = stored_k if stored_k is not None else requested_k + if stored_k is not None and abs(stored_k - requested_k) > 0.0: + logger.warning( + "DSA: в состоянии сохранена кривизна k=%s; аргумент " + "конструктора k=%s проигнорирован при resume=True. Для смены " + "кривизны вызовите set_curvature() явно.", + stored_k, requested_k) + + defaults = dict(lr=lr, beta1=beta1, beta2=beta2, eps=eps, + weight_decay=weight_decay, k=effective_k, + weight_decay_mode=mode) + super().__init__(params_list, defaults) + + self._param_shapes: Dict[str, Tuple[int, int]] = { + name: tuple(int(s) for s in param.shape) + for name, param in zip(self.param_names, params_list) + } + self._placeholders: List[torch.Tensor] = [] self._nan_rows_replaced = 0 self._unknown_key_warnings = 0 + self._init_rows_projected = 0 + self._resume_rows_projected = 0 + self._global_step = 0 - # ── Инициализация/валидация хранилища на диске (F3) ────────────────── - for name, param in zip(self.param_names, self.param_groups[0]['params']): - shape = tuple(param.shape) - paths = { - 'w': os.path.join(self.disk_dir, f"{name}_w.dat"), - 'm': os.path.join(self.disk_dir, f"{name}_m.dat"), - 'v': os.path.join(self.disk_dir, f"{name}_v.dat"), - 't': os.path.join(self.disk_dir, f"{name}_t.dat"), - } - meta_path = os.path.join(self.disk_dir, f"{name}_meta.json") - self._init_param_storage(name, shape, paths, meta_path, param, resume) - self._paths[name] = paths - - self.state_files[name] = { - 'w': np.memmap(paths['w'], dtype='float32', mode='r+', shape=shape), - 'm': np.memmap(paths['m'], dtype='float32', mode='r+', shape=shape), - 'v': np.memmap(paths['v'], dtype='float32', mode='r+', shape=shape), - # F4: пер-строчные счётчики шага, int32, N×4 байта - 't': np.memmap(paths['t'], dtype='int32', mode='r+', shape=(shape[0],)), - } + self.state_files: Dict[str, Dict[str, np.memmap]] = {} + + for name, param in zip(self.param_names, params_list): + shape = self._param_shapes[name] + self._init_param_storage( + name, shape, self._paths[name], meta_paths[name], param, + resume=resume) + + self.state_files = self._open_memmaps() + + if resume: + global_meta = self._read_global_meta_optional() + if global_meta is not None: + self._global_step = int(global_meta.get("global_step", 0)) + self._validate_ball_on_open() + + # C3: drop references to the (potentially huge) source tensors. + if not retain_params: + self._release_source_params() - # ── Запуск фонового писателя ───────────────────────────────────────── self.writer = AsyncDiskWriter( memmaps={f"{n}_{s}": self.state_files[n][s] - for n in self.param_names for s in ['w', 'm', 'v', 't']}, + for n in self.param_names for s in ("w", "m", "v", "t")}, max_queue_size=max_queue_size, put_timeout=put_timeout, ) + self._write_global_meta() - # ─────────────────────── дисковое состояние (F3) ─────────────────────── + # ───────────────────────── metadata helpers ───────────────────────── - def _write_meta(self, meta_path: str, name: str, shape, paths: dict, - reinit_in_progress: bool = False): - meta = { - 'format_version': STATE_FORMAT_VERSION, - 'param_name': name, - 'shape': list(shape), - 'dtype': 'float32', - 'counter_dtype': 'int32', - 'files': {s: os.path.basename(p) for s, p in paths.items()}, - 'updated': datetime.now().isoformat(timespec='seconds'), - } - if reinit_in_progress: - # R2: маркер незавершённой переинициализации - meta['reinit_in_progress'] = True - tmp = meta_path + '.tmp' - with open(tmp, 'w', encoding='utf-8') as f: - json.dump(meta, f, ensure_ascii=False, indent=2) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, meta_path) + def _current_k(self) -> float: + return float(self.param_groups[0]["k"]) - def _load_meta(self, meta_path: str, name: str) -> dict: + def _meta_path(self, name: str) -> str: + return os.path.join(self.disk_dir, f"{name}_meta.json") + + def _legacy_tag(self, name: str): + meta = self._read_meta_optional(self._meta_path(name)) + return None if meta is None else meta.get("legacy_migrated_from") + + def _read_meta_optional(self, meta_path: str) -> Optional[dict]: if not os.path.exists(meta_path): - raise RuntimeError( - f"для параметра '{name}' в {self.disk_dir} есть файлы состояния, но нет " - f"meta-файла {os.path.basename(meta_path)}: происхождение состояния " - f"неизвестно. Удалите каталог или восстановите meta.") + return None try: - with open(meta_path, 'r', encoding='utf-8') as f: + with open(meta_path, "r", encoding="utf-8") as f: + return json.load(f) + except (OSError, ValueError) as exc: + raise RuntimeError(f"meta-файл {meta_path} не читается: {exc!r}") from exc + + def _peek_stored_k(self, meta_paths: Dict[str, str]) -> Optional[float]: + global_meta = self._read_global_meta_optional() + if global_meta is not None and "k" in global_meta: + return _validate_k(global_meta["k"], where="сохранённый k") + for name in self.param_names: + meta = self._read_meta_optional(meta_paths[name]) + if meta is not None and "k" in meta: + return _validate_k(meta["k"], where=f"сохранённый k ({name})") + return None + + def _read_global_meta_optional(self) -> Optional[dict]: + path = os.path.join(self.disk_dir, GLOBAL_META_NAME) + if not os.path.exists(path): + return None + try: + with open(path, "r", encoding="utf-8") as f: meta = json.load(f) - except (OSError, ValueError) as e: - raise RuntimeError(f"meta-файл {meta_path} не читается: {e!r}") from e - if int(meta.get('format_version', -1)) != STATE_FORMAT_VERSION: + 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: raise RuntimeError( - f"meta-файл {meta_path}: версия формата {meta.get('format_version')!r} " + f"глобальный meta-файл {path}: версия формата {version!r} " f"не поддерживается (ожидается {STATE_FORMAT_VERSION})") + if "global_step" in meta: + try: + step = int(meta["global_step"]) + except (TypeError, ValueError) as exc: + raise RuntimeError( + f"глобальный meta-файл {path}: global_step=" + f"{meta['global_step']!r} не является целым") from exc + if step < 0: + raise RuntimeError( + f"глобальный meta-файл {path}: global_step={step} < 0") return meta - def _verify_state_file_sizes(self, name, shape, paths): - """Сверка фактических размеров файлов состояния с формой из meta (F3). + def _write_global_meta(self) -> str: + path = os.path.join(self.disk_dir, GLOBAL_META_NAME) + meta = { + "format_version": STATE_FORMAT_VERSION, + "algorithm": ALGORITHM, + "global_step": int(self._global_step), + "k": float(self._current_k()), + "weight_decay_mode": self.param_groups[0]["weight_decay_mode"], + "param_names": list(self.param_names), + "updated": _utc_now(), + } + self._atomic_json_write(path, meta) + return path - np.memmap(dtype float32, mode='r+', shape) молча продлевает короткий - файл нулями до запрошенного размера (и молча отображает только начало - слишком длинного), поэтому обрезанный при сбое диска/копирования файл - без этой проверки открылся бы штатно: потерянные строки читались бы - как нули и тихо смешивались со старыми. Та же проверка выполняется - в load_checkpoint (_verify_checkpoint_file); здесь — штатный путь - открытия с resume=True. - """ - row_bytes = int(np.prod(shape)) * 4 # float32: w, m, v - expected = {'w': row_bytes, 'm': row_bytes, 'v': row_bytes, - 't': int(shape[0]) * 4} # int32: пер-строчные счётчики + @staticmethod + def _atomic_json_write(path: str, payload: dict) -> None: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False, indent=2) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + def _write_meta(self, meta_path: str, name: str, shape, + paths: Dict[str, str], *, + reinit_in_progress: bool = False, + legacy_from: Optional[Union[int, str]] = None, + ball_validated: bool = True) -> None: + meta = { + "format_version": STATE_FORMAT_VERSION, + "algorithm": ALGORITHM, + "param_name": name, + "shape": list(shape), + "dtype": "float32", + "counter_dtype": "int32", + "k": float(self._current_k()), + "weight_decay_mode": self.param_groups[0]["weight_decay_mode"], + "files": {s: os.path.basename(p) for s, p in paths.items()}, + "ball_validated": bool(ball_validated), + "updated": _utc_now(), + } + if reinit_in_progress: + meta["reinit_in_progress"] = True + if legacy_from is not None: + meta["legacy_migrated_from"] = legacy_from + self._atomic_json_write(meta_path, meta) + + def _check_meta_compat(self, meta: dict, name: str, shape, meta_path: str, + paths: Dict[str, str]) -> int: + version = int(meta.get("format_version", -1)) + if version != STATE_FORMAT_VERSION and version not in LEGACY_STATE_FORMAT_VERSIONS: + raise RuntimeError( + f"meta-файл {meta_path}: версия формата {version!r} не " + f"поддерживается (поддерживаются " + f"{list(LEGACY_STATE_FORMAT_VERSIONS) + [STATE_FORMAT_VERSION]})") + disk_shape = tuple(meta.get("shape", [])) + if disk_shape != tuple(shape): + raise RuntimeError( + f"несовпадение формы состояния DSA для '{name}': на диске " + f"{list(disk_shape)}, в параметрах {list(shape)}. Удалите " + f"{self.disk_dir} или укажите новый disk_dir.") + if meta.get("dtype") != "float32" or meta.get("counter_dtype") != "int32": + raise RuntimeError( + f"meta-файл {meta_path}: неожиданные типы данных " + f"(w/m/v={meta.get('dtype')!r}, t={meta.get('counter_dtype')!r}), " + f"ожидались float32/int32.") + algorithm = meta.get("algorithm") + if algorithm is not None and algorithm != ALGORITHM: + raise RuntimeError( + f"meta-файл {meta_path}: алгоритм {algorithm!r} не совпадает " + f"с {ALGORITHM!r}; состояние создано другой версией DSA.") + if "k" in meta: + disk_k = _validate_k(meta["k"], where=f"сохранённый k ({name})") + if abs(disk_k - self._current_k()) > 0.0: + raise RuntimeError( + f"meta-файл {meta_path}: кривизна k={disk_k} не совпадает " + f"с k={self._current_k()}; смешивать геометрии нельзя.") + files = meta.get("files") or {} + for suffix, path in paths.items(): + expected = os.path.basename(path) + if suffix in files and files[suffix] != expected: + raise RuntimeError( + f"meta-файл {meta_path}: файл '{suffix}' называется " + f"{files[suffix]!r}, ожидался {expected!r}.") + return version + + # ───────────────────────── state initialisation ───────────────────────── + + def _verify_state_file_sizes(self, name, shape, paths, *, include_t: bool = True): + """Check real file sizes before memmap opens them (r+ zero-extends).""" + row_bytes = int(np.prod(shape)) * 4 + expected = {"w": row_bytes, "m": row_bytes, "v": row_bytes} + if include_t: + expected["t"] = int(shape[0]) * 4 for suffix, exp in expected.items(): + if not os.path.exists(paths[suffix]): + raise RuntimeError( + f"частичное состояние DSA для '{name}': отсутствует файл " + f"{paths[suffix]}") actual = os.path.getsize(paths[suffix]) if actual != exp: raise RuntimeError( @@ -172,123 +441,326 @@ def _verify_state_file_sizes(self, name, shape, paths): f"резервной копии/чекпоинта или переоткройтесь с " f"resume=False (полная переинициализация).") + def _project_param_to_ball(self, name: str, arr: np.ndarray) -> np.ndarray: + k = self._current_k() + if k <= 0.0: + return arr + if arr.shape[0] == 0: + return arr + radius = math.sqrt((1.0 - BALL_EPS) / k) + max_norm = float(np.linalg.norm(arr.astype(np.float64), axis=1).max()) + if max_norm <= 0.0: + return arr + limit = self.init_max_norm_fraction * radius + if max_norm > limit and self.init_max_norm_fraction < 1.0: + scale = limit / max_norm + self._init_rows_projected += arr.shape[0] + logger.warning( + "DSA '%s': начальные веса выходят за %.0f%% радиуса шара " + "(max |w| = %.6g, радиус %.6g); весь тензор умножен на %.6g, " + "чтобы обучение не началось на границе (λ² ≈ 0). Передайте " + "init_max_norm_fraction=1.0, если нужна только проекция.", + name, 100.0 * self.init_max_norm_fraction, max_norm, radius, + scale) + arr = (arr.astype(np.float64) * scale).astype(arr.dtype, copy=False) + # Even with fraction=1.0 a row may be outside; project individually. + norm = np.linalg.norm(arr.astype(np.float64), axis=1) + bad = norm > radius + if bad.any(): + self._init_rows_projected += int(np.count_nonzero(bad)) + logger.warning( + "DSA '%s': %d строк начальных весов лежат вне шара " + "(радиус %.6g) и спроецированы.", name, int(bad.sum()), radius) + arr = arr.copy() + arr[bad] = project_to_ball_np(arr[bad], k) + return arr + def _write_fresh_state(self, name, shape, paths, param): - """Полная (пере)инициализация: w из params, m=v=0, счётчики 0, свежий meta.""" - for suffix in ('w', 'm', 'v'): - mm = np.memmap(paths[suffix], dtype='float32', mode='w+', shape=shape) - if suffix == 'w': - mm[:] = param.detach().cpu().numpy().astype('float32') + arr = param.detach().to(device="cpu", dtype=torch.float32).contiguous().numpy() + if not np.isfinite(arr).all(): + bad_rows = np.nonzero(~np.isfinite(arr).all(axis=1))[0][:10].tolist() + raise ValueError( + f"начальные веса '{name}' содержат NaN/inf (строки {bad_rows}); " + f"исправьте инициализацию.") + arr = self._project_param_to_ball(name, arr) + + for suffix in ("w", "m", "v"): + mm = np.memmap(paths[suffix], dtype="float32", mode="w+", shape=shape) + if suffix == "w": + mm[:] = arr else: mm[:] = 0.0 mm.flush() del mm - t = np.memmap(paths['t'], dtype='int32', mode='w+', shape=(shape[0],)) + t = np.memmap(paths["t"], dtype="int32", mode="w+", shape=(shape[0],)) + t[:] = 0 + t.flush() + del t + + def _reset_moments(self, name, shape, paths): + for suffix in ("m", "v"): + mm = np.memmap(paths[suffix], dtype="float32", mode="r+", shape=shape) + mm[:] = 0.0 + mm.flush() + del mm + t = np.memmap(paths["t"], dtype="int32", mode="r+", shape=(shape[0],)) t[:] = 0 t.flush() del t - def _init_param_storage(self, name, shape, paths, meta_path, param, resume: bool): + def _init_param_storage(self, name, shape, paths, meta_path, param, + *, resume: bool) -> None: exists = {s: os.path.exists(p) for s, p in paths.items()} + meta = self._read_meta_optional(meta_path) + any_data = any(exists.values()) - if not any(exists.values()): - # Новый каталог — честная инициализация с нуля + if not resume: + # resume=False is the documented reinitialisation path: it must + # work even for a partial/legacy/truncated directory (C1), but a + # trustworthy shape mismatch is still an explicit error. Curvature + # and algorithm may differ: the whole state is being replaced. + if any_data or meta is not None: + logger.warning( + "DSA '%s': в %s найдено состояние; resume=False → полная " + "переинициализация: w берётся из params, m=v=0, счётчики " + "шагов сброшены в 0.", name, self.disk_dir) + if meta is not None: + disk_shape = tuple(meta.get("shape", [])) + if disk_shape != tuple(shape): + raise RuntimeError( + f"несовпадение формы состояния DSA для '{name}': на " + f"диске {list(disk_shape)}, в параметрах {list(shape)}. " + f"Удалите {self.disk_dir} или укажите новый disk_dir.") self._write_fresh_state(name, shape, paths, param) self._write_meta(meta_path, name, shape, paths) return - # На диске что-то есть — валидируем строго (F3: никакого тихого переиспользования) - if not all(exists.values()): - missing = [s for s, e in exists.items() if not e] + # resume=True + if not any_data and meta is None: raise RuntimeError( - f"частичное состояние DSA для '{name}' в {self.disk_dir}: отсутствуют " - f"{missing}. Состояние ненадёжно — удалите каталог целиком или " - f"восстановите все файлы ({', '.join(paths.values())}).") - - meta = self._load_meta(meta_path, name) - if meta.get('reinit_in_progress'): - # R2: процесс умер посреди переинициализации — часть файлов уже - # перезаписана, часть осталась от старого состояния. Принимать - # такое молча нельзя: resume=True — явная ошибка, resume=False — - # штатное восстановление полной переинициализацией. - if resume: - raise RuntimeError( - f"meta-файл {meta_path}: состояние помечено " - f"'reinit_in_progress' — предыдущая переинициализация не " - f"завершилась (сбой процесса между записью файлов и meta), " - f"файлы могут быть смешанными. resume=True невозможен: " - f"удалите {self.disk_dir} целиком или переоткройте с " - f"resume=False — полная переинициализация пересоздаст " - f"файлы и снимет маркер.") - logger.warning( - f"DSA '{name}': предыдущая переинициализация не завершилась " - f"(маркер 'reinit_in_progress'); resume=False → полная " - f"переинициализация, маркер будет снят.") - disk_shape = tuple(meta.get('shape', [])) - if disk_shape != tuple(shape): + f"resume=True, но состояния для '{name}' в {self.disk_dir} нет") + if meta is None: + # 0.1.0: only w/m/v, no meta, no per-row counters. + if exists["w"] and exists["m"] and exists["v"] and not exists["t"]: + self._verify_state_file_sizes(name, shape, paths, include_t=False) + logger.warning( + "DSA '%s': найдено состояние 0.1.0 (без meta и без " + "счётчиков t). Веса сохранены, моменты m/v и счётчики " + "сброшены: старая математика несовместима с моментами в " + "касательном пространстве.", name) + # Create the missing counter file first, then zero m/v/t. + t = np.memmap(paths["t"], dtype="int32", mode="w+", + shape=(shape[0],)) + t[:] = 0 + t.flush() + del t + self._reset_moments(name, shape, paths) + self._write_meta(meta_path, name, shape, paths, + legacy_from="0.1.0", ball_validated=False) + return + missing = [s for s, e in exists.items() if not e] raise RuntimeError( - f"несовпадение формы состояния DSA для '{name}': на диске " - f"{list(disk_shape)}, в параметрах {list(shape)}. Удалите " - f"{self.disk_dir} или укажите новый disk_dir.") - if meta.get('dtype') != 'float32' or meta.get('counter_dtype') != 'int32': + f"для параметра '{name}' в {self.disk_dir} есть файлы " + f"состояния, но нет meta-файла {os.path.basename(meta_path)}: " + f"происхождение состояния неизвестно (отсутствуют {missing}). " + f"Удалите каталог или восстановите meta.") + + version = self._check_meta_compat(meta, name, shape, meta_path, paths) + if meta.get("reinit_in_progress"): raise RuntimeError( - f"meta-файл {meta_path}: неожиданные типы данных " - f"(w/m/v={meta.get('dtype')!r}, t={meta.get('counter_dtype')!r}), " - f"ожидались float32/int32.") + f"meta-файл {meta_path}: состояние помечено " + f"'reinit_in_progress' — предыдущая переинициализация не " + f"завершилась (сбой процесса между записью файлов и meta), " + f"файлы могут быть смешанными. resume=True невозможен: " + f"удалите {self.disk_dir} целиком или переоткройте с " + f"resume=False — полная переинициализация пересоздаст файлы " + f"и снимет маркер.") - if resume: - # Явное продолжение: существующее состояние используется как есть, - # w из params НЕ перезаписывает диск, счётчики шагов сохраняются. - # До открытия memmap сверяем фактические размеры файлов: np.memmap - # в режиме r+ молча продлевает обрезанный файл нулями (сбой диска - # или копирования), и потерянные строки тихо смешались бы со - # старыми как нулевые (начало шара) — без ошибки и без лога. - self._verify_state_file_sizes(name, shape, paths) - logger.info( - f"DSA '{name}': продолжаем существующее состояние {list(shape)} " - f"из {self.disk_dir} (resume=True)") - else: + self._verify_state_file_sizes(name, shape, paths) + if version == 1: logger.warning( - f"DSA '{name}': в {self.disk_dir} найдено состояние {list(shape)}; " - f"resume=False → полная переинициализация: w берётся из params, " - f"m=v=0, счётчики шагов сброшены в 0.") - # R2: маркер ставится ДО перезаписи файлов и снимается финальной - # записью meta — смерть процесса посреди переинициализации не - # оставит смешанное состояние, выглядящее как валидное + "DSA '%s': состояние 0.2.0 (format_version=1). Веса сохранены, " + "моменты m/v и счётчики t сброшены: 0.3.0 хранит моменты в " + "касательном пространстве и переносит их параллельно, старые " + "значения к новой математике неприменимы.", name) + self._reset_moments(name, shape, paths) 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) + legacy_from=1, ball_validated=False) + else: + logger.info( + "DSA '%s': продолжаем состояние %s из %s (resume=True)", + name, list(shape), self.disk_dir) + + def _open_memmaps(self) -> Dict[str, Dict[str, np.memmap]]: + files: Dict[str, Dict[str, np.memmap]] = {} + for name in self.param_names: + shape = self._param_shapes[name] + files[name] = { + "w": np.memmap(self._paths[name]["w"], dtype="float32", + mode="r+", shape=shape), + "m": np.memmap(self._paths[name]["m"], dtype="float32", + mode="r+", shape=shape), + "v": np.memmap(self._paths[name]["v"], dtype="float32", + mode="r+", shape=shape), + "t": np.memmap(self._paths[name]["t"], dtype="int32", + mode="r+", shape=(shape[0],)), + } + return files + + @staticmethod + def _close_memmaps(files: Dict[str, Dict[str, np.memmap]]) -> None: + for per_param in files.values(): + for mm in per_param.values(): + try: + mm.flush() + except Exception: # pragma: no cover - best effort on broken mmap + pass + try: + mm._mmap.close() + except Exception: # pragma: no cover + pass + + def _release_source_params(self) -> None: + """Replace source tensors with zero-length CPU placeholders (C3).""" + placeholders = [torch.empty(0, dtype=torch.float32) for _ in self.param_names] + self.param_groups[0]["params"] = placeholders + self._placeholders = placeholders + # No optimizer state is keyed by the source tensors (DSA keeps it on + # disk), but clear the mapping defensively. + self.state.clear() + + def _validate_ball_on_open(self) -> None: + """On resume, project finite rows that fall outside the stored ball.""" + k = self._current_k() + if k <= 0.0: + return + radius = math.sqrt((1.0 - BALL_EPS) / k) + projected = 0 + nonfinite = 0 + for name in self.param_names: + mm = self.state_files[name]["w"] + n = mm.shape[0] + for start in range(0, n, CHUNK_ROWS): + stop = min(start + CHUNK_ROWS, n) + w = np.array(mm[start:stop], dtype=np.float32, copy=True) + finite = np.isfinite(w).all(axis=1) + nonfinite += int(np.count_nonzero(~finite)) + if finite.any(): + norm = np.linalg.norm(w[finite].astype(np.float64), axis=1) + bad = np.nonzero(norm > radius)[0] + if bad.size: + rows = np.nonzero(finite)[0][bad] + w[rows] = project_to_ball_np(w[rows], k) + mm[start:stop] = w + projected += int(rows.size) + mm.flush() + self._write_meta(self._meta_path(name), name, + self._param_shapes[name], self._paths[name], + legacy_from=self._legacy_tag(name), + ball_validated=True) + self._resume_rows_projected = projected + if projected: + logger.warning( + "DSA: при возобновлении %d строк весов лежали вне шара " + "k=%s и были спроецированы.", projected, k) + if nonfinite: + logger.warning( + "DSA: при возобновлении найдено %d не-конечных строк; они " + "будут полностью сброшены при первом обновлении.", nonfinite) - # ─────────────────────────────── шаг ─────────────────────────────── + # ───────────────────────── indices / gradients ───────────────────────── + + def _validate_indices(self, indices: Any, n: int, name: str) -> torch.Tensor: + if isinstance(indices, torch.Tensor): + if indices.dim() != 1: + raise ValueError( + f"indices для '{name}' должны быть 1-D, получено " + f"{tuple(indices.shape)}") + if indices.dtype not in (torch.int64, torch.int32, torch.int16): + raise TypeError( + f"indices для '{name}' должны быть целочисленными, " + f"получен dtype {indices.dtype}") + idx = indices.detach().to(device="cpu", dtype=torch.int64) + else: + arr = np.asarray(indices) + if arr.ndim != 1: + raise ValueError( + f"indices для '{name}' должны быть 1-D, получено " + f"{arr.shape}") + if not np.issubdtype(arr.dtype, np.integer): + raise TypeError( + f"indices для '{name}' должны быть целочисленными, " + f"получен dtype {arr.dtype}") + idx = torch.from_numpy(arr.astype(np.int64, copy=False)) + if idx.numel(): + if int(idx.min()) < 0: + raise IndexError( + f"indices для '{name}' содержат отрицательный индекс " + f"{int(idx.min())}: отрицательные индексы не поддерживаются") + if int(idx.max()) >= n: + raise IndexError( + f"indices для '{name}' выходят за границы: максимум " + f"{int(idx.max())}, строк {n}") + return idx + + def _validate_grads(self, grads: torch.Tensor, count: int, dim: int, + name: str) -> torch.Tensor: + if not isinstance(grads, torch.Tensor): + grads = torch.as_tensor(grads) + if grads.is_sparse: + raise TypeError( + f"градиент '{name}' разреженный (sparse); DSA принимает " + f"плотные градиенты батча") + if tuple(grads.shape) != (count, dim): + raise ValueError( + f"градиент '{name}' имеет форму {tuple(grads.shape)}, " + f"ожидалась ({count}, {dim})") + if not grads.is_floating_point(): + raise TypeError( + f"градиент '{name}' должен быть вещественным, получен " + f"dtype {grads.dtype}") + g = grads.detach().to(device="cpu", dtype=torch.float64) + if not torch.isfinite(g).all(): + bad = torch.nonzero(~torch.isfinite(g).all(dim=-1)).flatten()[:10] + raise ValueError( + f"градиент '{name}' содержит NaN/inf (позиции {bad.tolist()}); " + f"обучение остановлено, чтобы не испортить моменты.") + return g + + # ─────────────────────────────── step ─────────────────────────────── @torch.no_grad() def step(self, updates: Dict[str, Tuple[torch.Tensor, torch.Tensor]], current_k: Optional[Union[float, torch.Tensor]] = None): + """Apply one optimizer step. + + ``updates`` maps a parameter name to ``(indices, grads)``. Indices + must be 1-D integer tensors/arrays in ``[0, N)``; gradients must have + shape ``(len(indices), dim)`` and be finite. Duplicate indices are + summed. ``current_k`` changes the curvature for this and subsequent + steps (all rows are reprojected, because the ball itself changes). """ - updates: словарь {'param_name': (indices_tensor, grads_tensor)} - """ - # Глобальный счетчик шагов хранится в первом параметре для совместимости - # (в bias-коррекции он больше не участвует — она пер-строчная, F4) - param_base = self.param_groups[0]['params'][0] - state = self.state[param_base] - if len(state) == 0: - state['step'] = 0 - state['step'] += 1 + if current_k is not None: + k_new = _validate_k(current_k, where="current_k") + if k_new != self._current_k(): + self.set_curvature(k_new) - group = self.param_groups[0] - lr = group['lr'] - beta1, beta2 = group['beta1'], group['beta2'] - eps = group['eps'] - weight_decay = group['weight_decay'] + if not isinstance(updates, dict): + raise TypeError("updates must be a dict {name: (indices, grads)}") + + self._global_step += 1 - k = current_k.item() if isinstance(current_k, torch.Tensor) else ( - current_k if current_k is not None else group['k']) + group = self.param_groups[0] + lr = float(group["lr"]) + beta1 = float(group["beta1"]) + beta2 = float(group["beta2"]) + eps = float(group["eps"]) + weight_decay = float(group["weight_decay"]) + mode = group["weight_decay_mode"] + k = self._current_k() - for name, (indices, grads) in updates.items(): + for name, pair in updates.items(): if name not in self.state_files: - # R3: молча отброшенное обновление — потерянные данные без - # сообщения. Warning со счётчиком; в отладочном режиме — - # исключение (счётчик не растёт: warning не выдавался). msg = (f"step(updates): неизвестный ключ '{name}' — обновление " f"пропущено (известные ключи: {sorted(self.state_files)}); " f"вероятна опечатка в имени параметра") @@ -297,227 +769,543 @@ def step(self, updates: Dict[str, Tuple[torch.Tensor, torch.Tensor]], self._unknown_key_warnings += 1 logger.warning(msg) continue - if indices.numel() == 0: + if not isinstance(pair, (tuple, list)) or len(pair) != 2: + raise TypeError( + f"updates['{name}'] должен быть (indices, grads), " + f"получено {type(pair).__name__}") + indices, grads = pair + n, dim = self._param_shapes[name] + idx = self._validate_indices(indices, n, name) + if idx.numel() == 0: continue + g_euc = self._validate_grads(grads, int(idx.numel()), dim, name).double() - # 1. Агрегация градиентов для дублирующихся индексов в батче - if indices.numel() != torch.unique(indices).numel(): - unique_idx, inv = torch.unique(indices, return_inverse=True) - agg_grads = torch.zeros((len(unique_idx), grads.shape[1]), - device=grads.device, dtype=grads.dtype) - agg_grads.index_add_(0, inv, grads) - active_indices, active_grads = unique_idx, agg_grads - else: - active_indices, active_grads = indices, grads + # 1. Sum duplicate indices (hub nodes) within the batch. + unique, inverse = torch.unique(idx, return_inverse=True) + if unique.numel() != idx.numel(): + agg = torch.zeros((unique.numel(), dim), dtype=g_euc.dtype) + agg.index_add_(0, inverse, g_euc) + idx, g_euc = unique, agg - # F5: сортировка индексов — локальность и чтения, и записи - if active_indices.numel() > 1: - order = torch.argsort(active_indices) - active_indices = active_indices[order] - active_grads = active_grads[order] + # 2. Sequential IO: sorted, unique indices. + if idx.numel() > 1 and not bool(torch.all(idx[1:] > idx[:-1])): + order = torch.argsort(idx) + idx = idx[order] + g_euc = g_euc[order] - idx_np = active_indices.cpu().numpy() + idx_np = idx.numpy() states = self.state_files[name] - # F2: не читать строки, которые писатель ещё не дописал. - # Пересечение пусто → читаем сразу (асинхронность сохраняется); - # непусто → ждём полной очистки очереди и только потом читаем. + # 3. Never read rows the writer has not finished writing (F2). if self.writer.intersects_inflight(name, idx_np): self.writer.wait_all_flushed() - # 2. Чтение из mmap в оперативную память (только нужные строки) - w = torch.from_numpy(states['w'][idx_np].copy()) - m = torch.from_numpy(states['m'][idx_np].copy()) - v = torch.from_numpy(states['v'][idx_np].copy()) - # F4: пер-строчный счётчик шага; +1 = этот шаг - t_row = torch.from_numpy( - states['t'][idx_np].copy().astype('int64')) + 1 - g_euc = active_grads.cpu().float() - - # R1: не-конечные строки из битого файла не проходят молча: - # torch.where(norm > max) при NaN даёт False, и NaN расползся бы - # в m/v и далее в обучение. Строка обнуляется (как у свежей строки), - # проекция в шар ниже удерживает геометрию. - bad = torch.nonzero( - ~(torch.isfinite(w).all(dim=-1) & torch.isfinite(m).all(dim=-1) - & torch.isfinite(v).all(dim=-1))).flatten() - if bad.numel() > 0: - self._nan_rows_replaced += int(bad.numel()) + # Batch math is done in float64. The chart is numerically + # ill-conditioned near the boundary; float32 intermediate values + # lose the cancellation in the gyration formula and can overflow + # the moments even though the exact transport is an isometry. + # Only O(batch) rows are involved, so the cost is negligible + # compared with disk I/O; storage remains float32. + w = torch.from_numpy(states["w"][idx_np].copy()).double() + m = torch.from_numpy(states["m"][idx_np].copy()).double() + v = torch.from_numpy(states["v"][idx_np].copy()).double() + t_prev = torch.from_numpy(states["t"][idx_np].copy().astype("int64")) + + # 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). + finite = (torch.isfinite(w).all(dim=-1) + & torch.isfinite(m).all(dim=-1) + & torch.isfinite(v).all(dim=-1) + & (v >= 0).all(dim=-1)) + bad = ~finite + if bool(bad.any()): + count = int(bad.sum()) + self._nan_rows_replaced += count logger.warning( - f"DSA '{name}': {int(bad.numel())} строк с NaN/inf в файлах " - f"состояния (строки диска {active_indices[bad][:10].tolist()}); " - f"обнулены и спроецированы в шар") - w = torch.nan_to_num(w, nan=0.0, posinf=0.0, neginf=0.0) - m = torch.nan_to_num(m, nan=0.0, posinf=0.0, neginf=0.0) - v = torch.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0) - - # F7: защита чтения — битый файл не должен расходить геометрию: - # прочитанные веса сразу проецируются в шар - if k > 1e-9: - w = _project_rows_to_ball(w, k) - - # 3. Riemannian Gradient Scaling (Poincaré Ball) - if k > 1e-9: - w_sq_norm = torch.sum(w * w, dim=-1, keepdim=True) - max_radius_sq = (1.0 - 1e-5) / k - w_sq_norm = torch.clamp(w_sq_norm, max=max_radius_sq) - rescale_factor = ((1.0 - k * w_sq_norm) ** 2) / 4.0 - g_riem = g_euc * rescale_factor - else: - g_riem = g_euc + "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 + t_prev[bad] = 0 + + # A zero/negative counter means the row has never had an Adam + # step; any moments on disk are stale (e.g. a lost counter) and + # must not leak into the fresh bias correction. + fresh = t_prev <= 0 + if bool(fresh.any()): + m[fresh] = 0.0 + v[fresh] = 0.0 + + # 5. Defensive projection of read weights (corrupted file, F7/H1). + if k > 0.0: + w = project_to_ball(w, k) - # F4: bias-коррекция по счётчику СТРОКИ, а не по глобальному шагу. - # Первый шаг свежей строки даёт lr*sign(g): m_hat=g, v_hat=g^2. - bc1 = (1.0 - torch.pow(torch.full_like(t_row, beta1, dtype=torch.float64), - t_row.double())).to(torch.float32).unsqueeze(-1) - bc2 = (1.0 - torch.pow(torch.full_like(t_row, beta2, dtype=torch.float64), - t_row.double())).to(torch.float32).unsqueeze(-1) + # 6. Riemannian gradient. "l2" is the Riemannian gradient of + # f + wd/2 * ||w||^2; in the Euclidean case it is classic Adam L2. + if weight_decay > 0.0 and mode == "l2": + g_euc = g_euc + weight_decay * w + g_riem = riemannian_gradient(g_euc, w, k) - # 4. Adam Step - if weight_decay > 0: - w.mul_(1.0 - lr * weight_decay) + # 7. Adam moments in the tangent space at w. + 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) + direction = (m / bc1) / (torch.sqrt(v / bc2) + eps) + + # 8. Decoupled decay is a tangent-space term, so in the Euclidean + # case it is exactly AdamW and in the hyperbolic case it is part + # of the same geodesic step. + 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). + 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() + else: + w_new = w - lr * direction + m_new, v_new = m, v - m_hat = m / bc1 - v_hat = v / bc2 - step_size = lr * m_hat / (torch.sqrt(v_hat) + eps) - w.sub_(step_size) + if not torch.isfinite(w_new).all(): + raise RuntimeError( + f"DSA '{name}': шаг дал не-конечные веса (k={k}, lr={lr}); " + f"уменьшите lr или проверьте градиенты.") - # 5. Retraction (удержание на многообразии) - if k > 1e-9: - w = _project_rows_to_ball(w, k) + # 10. Queue the consistent row update (storage is float32). + w_out = w_new.float() + m_out = m_new.float() + v_out = v_new.float() + if not (torch.isfinite(w_out).all() and torch.isfinite(m_out).all() + and torch.isfinite(v_out).all()): + raise RuntimeError( + f"DSA '{name}': состояние момента не представимо во " + f"float32 (k={k}); точка слишком близка к границе шара. " + f"Уменьшите lr или увеличьте зазор BALL_EPS.") + self.writer.put_update(name, idx, w_out, m_out, v_out, + torch.clamp(t_new, max=2**31 - 1).to(torch.int32)) - # 6. Передача на асинхронную фоновую запись (счётчик — вместе - # с w/m/v, одной очередью, чтобы состояние осталось согласованным) - self.writer.put_update(name, active_indices, w, m, v, - t_row.to(torch.int32)) + # ───────────────────────── curvature / rows ───────────────────────── - # ─────────────────────── чекпоинты (F8) ─────────────────────── + def set_curvature(self, k: float, *, reset_moments: bool = False) -> None: + """Change curvature and restore the ball invariant for all rows. - def save_checkpoint(self, path: str) -> str: + Changing curvature changes the manifold, so the whole table is + scanned and out-of-ball rows are projected. Moments live in tangent + spaces of the old geometry; by default they are kept (a documented + approximation for annealing). Pass ``reset_moments=True`` for a + mathematically clean restart of the adaptive state. """ - Явный чекпоинт: дождаться записи всей очереди, flush всех файлов - состояния (w/m/v/t) и атомарно записать манифест путей в `path`. - Возвращает абсолютный путь манифеста. + k = _validate_k(k) + self.writer.wait_all_flushed() + old_k = self._current_k() + self.param_groups[0]["k"] = k + self.defaults["k"] = k + + if reset_moments: + with self.writer._lock: + for name in self.param_names: + self._reset_moments(name, self._param_shapes[name], + self._paths[name]) + logger.warning( + "DSA: set_curvature(%s -> %s, reset_moments=True): моменты и " + "счётчики сброшены.", old_k, k) + else: + logger.info("DSA: кривизна изменена %s -> %s", old_k, k) + + self.reproject_all() + self._write_global_meta() + + def reproject_all(self, name: Optional[str] = None) -> int: + """Project every finite row into the current ball (O(N), honest). + + Non-finite rows are fully reset (``w=m=v=0, t=0``). Returns the + number of projected rows. """ - path = os.path.abspath(path) self.writer.wait_all_flushed() - param_base = self.param_groups[0]['params'][0] + k = self._current_k() + names = [name] if name is not None else list(self.param_names) + if name is not None and name not in self.state_files: + raise KeyError(f"неизвестный параметр '{name}'") + projected = 0 with self.writer._lock: - manifest = { - 'format': 'DSA_checkpoint', - 'version': 1, - 'saved_at': datetime.now().isoformat(timespec='seconds'), - 'global_step': int(self.state[param_base].get('step', 0)), - 'params': {}, - } - for name in self.param_names: - files = {} - for suffix, mm in self.state_files[name].items(): + for pname in names: + shape = self._param_shapes[pname] + states = self.state_files[pname] + for start in range(0, shape[0], CHUNK_ROWS): + stop = min(start + CHUNK_ROWS, shape[0]) + w = np.array(states["w"][start:stop], copy=True) + m = np.array(states["m"][start:stop], copy=True) + v = np.array(states["v"][start:stop], copy=True) + t = np.array(states["t"][start:stop], copy=True) + finite = np.isfinite(w).all(axis=1) & \ + np.isfinite(m).all(axis=1) & \ + np.isfinite(v).all(axis=1) & (v >= 0).all(axis=1) + if (~finite).any(): + w[~finite] = 0.0 + m[~finite] = 0.0 + v[~finite] = 0.0 + t[~finite] = 0 + if k > 0.0 and finite.any(): + norm = np.linalg.norm(w[finite].astype(np.float64), axis=1) + radius = math.sqrt((1.0 - BALL_EPS) / k) + bad = norm > radius + if bad.any(): + rows = np.nonzero(finite)[0][bad] + w[rows] = project_to_ball_np(w[rows], k) + projected += int(rows.size) + states["w"][start:stop] = w + states["m"][start:stop] = m + states["v"][start:stop] = v + states["t"][start:stop] = t + for mm in states.values(): mm.flush() - files[suffix] = self._paths[name][suffix] - manifest['params'][name] = { - 'shape': list(self.state_files[name]['w'].shape), - 'files': files, - } - os.makedirs(os.path.dirname(path) or '.', exist_ok=True) - tmp = path + '.tmp' - with open(tmp, 'w', encoding='utf-8') as f: - json.dump(manifest, f, ensure_ascii=False, indent=2) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) - return path + self._write_meta(self._meta_path(pname), pname, shape, + self._paths[pname], + legacy_from=self._legacy_tag(pname), + ball_validated=True) + if projected: + logger.warning("DSA: %d строк спроецированы в шар k=%s.", + projected, k) + return projected - def load_checkpoint(self, path: str) -> dict: + def read_rows(self, indices, name: Optional[str] = None) -> torch.Tensor: + """Read rows of ``w`` for inference (projected into the ball). + + This is the public alternative to reading ``state_files`` directly. """ - Восстановление из манифеста save_checkpoint. Файлы чекпоинта должны - лежать по тем же путям, на которых открыт оптимизатор (обычный сценарий: - создать DiskSparseRiemannianAdam(..., resume=True) на том же disk_dir и - вызвать load_checkpoint(манифест) для проверки целостности и переоткрытия - memmap). Чекпоинт из другого каталога не подхватывается молча — это - явная ошибка: скопируйте файлы в disk_dir и откройтесь с resume=True. - Проверяются существование файлов, их размеры и формы. + if name is None: + name = self.param_names[0] + if name not in self.state_files: + raise KeyError(f"неизвестный параметр '{name}'") + n, _ = self._param_shapes[name] + idx = self._validate_indices(indices, n, name) + if idx.numel() == 0: + return torch.empty((0, self._param_shapes[name][1]), + dtype=torch.float32) + idx_np = idx.numpy() + if self.writer.intersects_inflight(name, idx_np): + self.writer.wait_all_flushed() + w = torch.from_numpy(self.state_files[name]["w"][idx_np].copy()) + if not torch.isfinite(w).all(): + raise RuntimeError( + f"DSA '{name}': прочитанные строки содержат NaN/inf; файл " + f"состояния повреждён (лечится обновлением этих строк).") + k = self._current_k() + if k > 0.0: + w = project_to_ball(w, k) + return w + + # ───────────────────────── checkpoints (F8/C2) ───────────────────────── + + @staticmethod + def _copy_and_hash(src: str, dst: str) -> Tuple[int, str]: + digest = hashlib.sha256() + size = 0 + with open(src, "rb") as fin, open(dst, "wb") as fout: + while True: + chunk = fin.read(1 << 20) + if not chunk: + break + digest.update(chunk) + fout.write(chunk) + size += len(chunk) + fout.flush() + os.fsync(fout.fileno()) + return size, digest.hexdigest() + + def export_checkpoint(self, directory: str) -> str: + """Write a self-contained checkpoint directory (C2). + + Data files are copied into ``directory`` under unique names; the + manifest (``checkpoint.json``) stores relative file names, sizes and + SHA-256 checksums. The directory can be moved to another machine and + imported with :meth:`import_checkpoint`. """ - path = os.path.abspath(path) + directory = os.path.abspath(os.fspath(directory)) + os.makedirs(directory, exist_ok=True) + self.writer.wait_all_flushed() + manifest_path = os.path.join(directory, CHECKPOINT_MANIFEST_NAME) + + old_files = set() + if os.path.exists(manifest_path): + try: + with open(manifest_path, "r", encoding="utf-8") as f: + old = json.load(f) + for info in old.get("params", {}).values(): + for rec in info.get("files", {}).values(): + if isinstance(rec, dict) and "name" in rec: + old_files.add(rec["name"]) + except (OSError, ValueError): + old_files = set() + + stamp = (datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S") + + f"-{os.getpid()}-{os.urandom(3).hex()}") + manifest = { + "format": CHECKPOINT_FORMAT, + "version": CHECKPOINT_VERSION, + "algorithm": ALGORITHM, + "saved_at": _utc_now(), + "global_step": int(self._global_step), + "k": float(self._current_k()), + "weight_decay_mode": self.param_groups[0]["weight_decay_mode"], + "param_names": list(self.param_names), + "params": {}, + } + + new_files = set() + with self.writer._lock: + for name in self.param_names: + shape = self._param_shapes[name] + info = {"shape": list(shape), "files": {}} + for suffix in ("w", "m", "v", "t"): + mm = self.state_files[name][suffix] + mm.flush() + filename = f"{name}_{suffix}_{stamp}.dat" + dst = os.path.join(directory, filename) + size, sha = self._copy_and_hash(self._paths[name][suffix], dst) + info["files"][suffix] = { + "name": filename, "size": int(size), "sha256": sha} + new_files.add(filename) + manifest["params"][name] = info + + self._atomic_json_write(manifest_path, manifest) + + # Remove files of the previous export only after the new manifest is + # durable: an interrupted export leaves the old checkpoint valid. + for filename in sorted(old_files - new_files): + try: + os.remove(os.path.join(directory, filename)) + except OSError: + pass + return manifest_path + + def _resolve_checkpoint(self, path: str) -> Tuple[str, str]: + path = os.path.abspath(os.fspath(path)) + if os.path.isdir(path): + return os.path.join(path, CHECKPOINT_MANIFEST_NAME), path + if os.path.isfile(path): + if os.path.basename(path) == CHECKPOINT_MANIFEST_NAME: + return path, os.path.dirname(path) + raise RuntimeError( + f"{path} — файл, а не каталог чекпоинта. Начиная с 0.3.0 " + f"чекпоинт — это каталог с checkpoint.json и файлами данных; " + f"используйте export_checkpoint(dir)/import_checkpoint(dir). " + f"Манифесты 0.2.0 с абсолютными путями не поддерживаются.") + raise RuntimeError(f"чекпоинт {path} не найден") + + def import_checkpoint(self, path: str, *, verify_hashes: bool = True) -> dict: + """Import a checkpoint directory exported by :meth:`export_checkpoint`.""" + manifest_path, directory = self._resolve_checkpoint(path) try: - with open(path, 'r', encoding='utf-8') as f: + with open(manifest_path, "r", encoding="utf-8") as f: manifest = json.load(f) - except (OSError, ValueError) as e: - raise RuntimeError(f"чекпоинт {path} не читается: {e!r}") from e - if manifest.get('format') != 'DSA_checkpoint': - raise RuntimeError(f"{path} не является чекпоинтом DSA " - f"(format={manifest.get('format')!r})") + except (OSError, ValueError) as exc: + raise RuntimeError(f"чекпоинт {manifest_path} не читается: {exc!r}") from exc + if manifest.get("format") != CHECKPOINT_FORMAT: + raise RuntimeError( + f"{manifest_path} не является чекпоинтом DSA " + f"(format={manifest.get('format')!r})") + if int(manifest.get("version", -1)) != CHECKPOINT_VERSION: + raise RuntimeError( + f"чекпоинт {manifest_path}: версия {manifest.get('version')!r} " + f"не поддерживается (ожидается {CHECKPOINT_VERSION}). " + f"Чекпоинты 0.2.0 содержали абсолютные пути и не переносимы; " + f"создайте новый чекпоинт.") + if manifest.get("algorithm") != ALGORITHM: + raise RuntimeError( + f"чекпоинт {manifest_path}: алгоритм " + f"{manifest.get('algorithm')!r}, ожидался {ALGORITHM!r}") + manifest_params = manifest.get("params") + if not isinstance(manifest_params, dict): + raise RuntimeError(f"чекпоинт {manifest_path}: поле 'params' повреждено") + saved_names = list(manifest_params.keys()) + if set(saved_names) != set(self.param_names): + raise RuntimeError( + f"чекпоинт содержит параметры {sorted(saved_names)}, " + f"оптимизатор — {sorted(self.param_names)}") + k = _validate_k(manifest.get("k"), where="k чекпоинта") self.writer.wait_all_flushed() - new_files = {} - with self.writer._lock: - for name, param in zip(self.param_names, self.param_groups[0]['params']): - info = manifest['params'].get(name) - if info is None: - raise RuntimeError(f"чекпоинт {path} не содержит параметр '{name}'") - shape = tuple(param.shape) - if tuple(info['shape']) != shape: + staged: Dict[Tuple[str, str], str] = {} + try: + for name in self.param_names: + info = manifest_params[name] + if not isinstance(info, dict): + raise RuntimeError( + f"чекпоинт {manifest_path}: описание '{name}' повреждено") + shape = self._param_shapes[name] + saved_shape = info.get("shape") + if not isinstance(saved_shape, (list, tuple)): + raise RuntimeError( + f"чекпоинт {manifest_path}: форма '{name}' повреждена") + if tuple(saved_shape) != shape: raise RuntimeError( f"несовпадение формы для '{name}': чекпоинт " - f"{list(info['shape'])}, оптимизатор {list(shape)}") - for suffix in ('w', 'm', 'v'): - expected = int(np.prod(shape)) * 4 # float32 - self._verify_checkpoint_file(name, suffix, info, expected) - self._verify_checkpoint_file(name, 't', info, shape[0] * 4) # int32 - - new_files[name] = { - s: np.memmap(self._paths[name][s], - dtype='float32' if s != 't' else 'int32', - mode='r+', - shape=shape if s != 't' else (shape[0],)) - for s in ('w', 'm', 'v', 't') - } - self.state_files = new_files - self.writer.memmaps = {f"{n}_{s}": mm - for n, mm2 in new_files.items() - for s, mm in mm2.items()} + f"{info.get('shape')}, оптимизатор {list(shape)}") + info_files = info.get("files") + if not isinstance(info_files, dict): + raise RuntimeError( + f"чекпоинт {manifest_path}: файлы '{name}' повреждены") + for suffix in ("w", "m", "v", "t"): + rec = info_files.get(suffix) + if not isinstance(rec, dict) or "name" not in rec: + raise RuntimeError( + f"чекпоинт не содержит файл '{suffix}' для '{name}'") + filename = rec["name"] + if not isinstance(filename, str): + raise RuntimeError( + f"недопустимое имя файла в чекпоинте: {filename!r}") + if os.path.basename(filename) != filename: + raise RuntimeError( + f"недопустимое имя файла в чекпоинте: {filename!r}") + src = os.path.join(directory, filename) + if not os.path.exists(src): + raise RuntimeError(f"файл чекпоинта отсутствует: {src}") + actual_size = os.path.getsize(src) + if "size" in rec and int(rec["size"]) != actual_size: + raise RuntimeError( + f"файл {src}: размер {actual_size} байт, в манифесте " + f"{rec['size']} — чекпоинт повреждён") + expected_size = (shape[0] * shape[1] * 4 if suffix != "t" + else shape[0] * 4) + if actual_size != expected_size: + raise RuntimeError( + f"файл {src}: размер {actual_size} байт, по форме " + f"{list(shape)} ожидается {expected_size}") + tmp = os.path.join( + self.disk_dir, f".import-{os.getpid()}-{name}-{suffix}") + size, sha = self._copy_and_hash(src, tmp) + if "size" in rec and size != int(rec["size"]): + raise RuntimeError(f"файл {src}: размер изменился при копировании") + if verify_hashes and "sha256" in rec and sha != rec["sha256"]: + raise RuntimeError( + f"файл {src}: SHA-256 не совпадает с манифестом — " + f"чекпоинт повреждён") + staged[(name, suffix)] = tmp + + # Mark the state transition. If the process dies between the + # individual os.replace calls, the next resume=True must refuse + # the mixed state instead of accepting it (same rule as R2). + for name in self.param_names: + self._write_meta(self._meta_path(name), name, + self._param_shapes[name], self._paths[name], + reinit_in_progress=True, ball_validated=False) + + # Replace the live state only after every file has been staged and + # verified. Close memmaps first: on Windows an open file cannot + # be replaced. + try: + with self.writer._lock: + self._close_memmaps(self.state_files) + for (name, suffix), tmp in staged.items(): + os.replace(tmp, self._paths[name][suffix]) + self.state_files = self._open_memmaps() + self.writer.memmaps = { + f"{n}_{s}": mm + for n, per_param in self.state_files.items() + for s, mm in per_param.items() + } + except Exception: + # Reopen whatever is on disk so the object is not left with + # closed memmaps; the reinit marker still blocks resume=True. + try: + self.state_files = self._open_memmaps() + self.writer.memmaps = { + f"{n}_{s}": mm + for n, per_param in self.state_files.items() + for s, mm in per_param.items() + } + except Exception: + pass + raise + except Exception: + for tmp in staged.values(): + try: + os.remove(tmp) + except OSError: + pass + raise + + self.param_groups[0]["k"] = k + self.defaults["k"] = k + self._global_step = int(manifest.get("global_step", 0)) + for name in self.param_names: + self._write_meta(self._meta_path(name), name, + self._param_shapes[name], self._paths[name], + ball_validated=False) + self._write_global_meta() + self._validate_ball_on_open() + logger.info( + "DSA: импортирован чекпоинт %s (global_step=%d, k=%s)", + manifest_path, self._global_step, k) return manifest - def _verify_checkpoint_file(self, name, suffix, info, expected_bytes: int): - p = info['files'].get(suffix) - mine = os.path.abspath(self._paths[name][suffix]) - if p is None: - raise RuntimeError(f"чекпоинт не содержит файл '{suffix}' для '{name}'") - if os.path.abspath(p) != mine: - raise RuntimeError( - f"чекпоинт ссылается на {p}, а оптимизатор открыт на {mine}: " - f"скопируйте файлы чекпоинта в disk_dir и создайте оптимизатор " - f"с resume=True") - if not os.path.exists(p): - raise RuntimeError(f"файл чекпоинта отсутствует: {p}") - actual = os.path.getsize(p) - if actual != expected_bytes: - raise RuntimeError( - f"файл {p}: размер {actual} байт, ожидалось {expected_bytes} " - f"— состояние повреждено") + def save_checkpoint(self, path: str) -> str: + """Deprecated alias: ``path`` is a checkpoint directory (C2).""" + path = os.path.abspath(os.fspath(path)) + if os.path.isfile(path) or (path.endswith(".json") and not os.path.isdir(path)): + raise ValueError( + "save_checkpoint(path) в 0.3.0 принимает КАТАЛОГ, а не файл " + "манифеста: используйте save_checkpoint('checkpoints/run1') " + "или export_checkpoint('checkpoints/run1').") + return self.export_checkpoint(path) + + def load_checkpoint(self, path: str, **kwargs) -> dict: + """Deprecated alias: ``path`` is a checkpoint directory (C2).""" + return self.import_checkpoint(path, **kwargs) + + # ───────────────────────── compatibility ───────────────────────── - # F8: молчаливо неверного state_dict не отдаём — только честная ошибка def state_dict(self): raise NotImplementedError( "DSA хранит состояние на диске (файлы _w/_m/_v/_t.dat в " f"{self.disk_dir}); сериализовать его через state_dict() нельзя. " - "Используйте save_checkpoint()/load_checkpoint() либо shutdown() " - "и повторное открытие с resume=True.") + "Используйте export_checkpoint()/import_checkpoint() либо " + "shutdown() и повторное открытие с resume=True.") def load_state_dict(self, state_dict): raise NotImplementedError( "DSA восстанавливается с диска: создайте оптимизатор с resume=True " - "на том же disk_dir (и при необходимости load_checkpoint()). " + "на том же disk_dir (и при необходимости import_checkpoint()). " "load_state_dict() не поддерживается.") def shutdown(self): - """Вызывать в конце обучения для финализации записи на диск.""" - if hasattr(self, 'writer'): + """Drain the writer, flush the state and persist the global step. + + Retryable: if a previous call failed (e.g. a stuck disk), calling it + again resumes the drain. Do not use the optimizer for further steps + after a successful shutdown. + """ + error = None + try: self.writer.shutdown() + except BaseException as exc: # preserve the writer error + error = exc + try: + self._write_global_meta() + except BaseException as exc: + if error is None: + error = exc + else: + logger.error("DSA: не удалось обновить глобальный meta: %r", exc) + if error is not None: + raise error def zero_grad(self, set_to_none: bool = False): - """Заглушка для совместимости, так как мы работаем с явными градиентами.""" + """No-op: gradients are passed explicitly to :meth:`step`.""" pass diff --git a/pyproject.toml b/pyproject.toml index 488b30d..92b58cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta" [project] name = "dsa" -version = "0.2.0" -description = "Disk Sparse Adam (DSA): An Out-of-Core Sparse Adam optimizer for PyTorch" +version = "0.3.0" +description = "Disk Sparse Adam (DSA): an out-of-core sparse Riemannian Adam optimizer for PyTorch" readme = "README.md" license = "MIT" license-files = ["LICENSE"] @@ -14,8 +14,14 @@ authors = [ { name = "Maksim Khotinsky" } ] dependencies = [ - "torch>=1.12.0", - "numpy>=1.20.0" + "torch>=2.0.0", + "numpy>=1.21.0" +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering :: Artificial Intelligence", ] [project.urls] diff --git a/setup.py b/setup.py deleted file mode 100644 index 9e8719a..0000000 --- a/setup.py +++ /dev/null @@ -1,24 +0,0 @@ -from setuptools import setup, find_packages - -setup( - name="disk-sparse-adam", - version="0.2.0", - description="Drop-in Out-of-Core Sparse Adam optimizer for Large-Scale PyTorch Graphs.", - long_description=open("README.md", encoding="utf-8").read(), - long_description_content_type="text/markdown", - author="Maksim Khotinskiy", - url="https://github.com/Assistentus/DSA", - packages=find_packages(), - install_requires=[ - "torch>=2.0.0", - "numpy>=1.21.0" - ], - classifiers=[ - "Development Status :: 4 - Beta", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: MIT License", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Programming Language :: Python :: 3.12", - ], -) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/reference.py b/tests/reference.py new file mode 100644 index 0000000..1d85a44 --- /dev/null +++ b/tests/reference.py @@ -0,0 +1,132 @@ +"""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. +""" + +from __future__ import annotations + +import numpy as np + + +def _sq(x): + return np.sum(x * x, axis=-1, keepdims=True) + + +def _lambda(x, k): + return 2.0 / (1.0 - k * _sq(x)) + + +def _mobius_add(x, y, k): + x2, y2 = _sq(x), _sq(y) + xy = np.sum(x * y, axis=-1, keepdims=True) + num = (1.0 + 2.0 * k * xy + k * y2) * x + (1.0 - k * x2) * y + den = 1.0 + 2.0 * k * xy + k * k * x2 * y2 + return num / den + + +def _gyration(u, v, w, k): + u2, v2 = _sq(u), _sq(v) + uv = np.sum(u * v, axis=-1, keepdims=True) + uw = np.sum(u * w, axis=-1, keepdims=True) + vw = np.sum(v * w, axis=-1, keepdims=True) + k2 = k * k + a = -k2 * uw * v2 + k * vw + 2.0 * k2 * uv * vw + b = -k2 * vw * u2 - k * uw + den = 1.0 + 2.0 * k * uv + k2 * u2 * v2 + return w + 2.0 * (a * u + b * v) / den + + +def expmap(x, u, k): + if k <= 0.0: + return x + u + un = np.linalg.norm(u, axis=-1, keepdims=True) + un = np.maximum(un, 1e-300) + scale = np.minimum(np.tanh(np.sqrt(k) * _lambda(x, k) * un / 2.0), + 1.0 - 1e-7) / (np.sqrt(k) * un) + return _mobius_add(x, scale * u, k) + + +def logmap(x, y, k): + if k <= 0.0: + return y - x + v = _mobius_add(-x, y, k) + vn = np.linalg.norm(v, axis=-1, keepdims=True) + vn = np.maximum(vn, 1e-300) + return (2.0 / (np.sqrt(k) * _lambda(x, k))) * np.arctanh(np.sqrt(k) * vn) * v / vn + + +def parallel_transport(x, y, u, k): + if k <= 0.0: + return u + return _gyration(y, -x, u, k) * _lambda(x, k) / _lambda(y, k) + + +def distance(x, y, k): + if k <= 0.0: + return np.linalg.norm(x - y, axis=-1) + x2, y2 = _sq(x), _sq(y) + diff2 = _sq(x - y) + den = (1.0 - k * x2) * (1.0 - k * y2) + arg = 1.0 + 2.0 * k * diff2 / den + return np.arccosh(np.maximum(arg, 1.0)) / np.sqrt(k) + + +class RiemannianAdamReference: + """In-memory float64 tangent-space Adam with parallel transport.""" + + def __init__(self, w0, lr, beta1=0.9, beta2=0.999, eps=1e-8, + k=0.0, weight_decay=0.0, weight_decay_mode="l2"): + self.w = np.array(w0, dtype=np.float64) + self.n, self.d = self.w.shape + self.m = np.zeros_like(self.w) + self.v = np.zeros_like(self.w) + self.t = np.zeros(self.n, dtype=np.int64) + self.lr = float(lr) + self.b1, self.b2, self.eps = float(beta1), float(beta2), float(eps) + self.k = float(k) + self.wd = float(weight_decay) + self.wd_mode = weight_decay_mode + + def step(self, indices, grads): + idx = np.asarray(indices, dtype=np.int64) + g = np.asarray(grads, dtype=np.float64) + unique, inverse = np.unique(idx, return_inverse=True) + if len(unique) != len(idx): + agg = np.zeros((len(unique), g.shape[1])) + np.add.at(agg, inverse, g) + idx, g = unique, agg + order = np.argsort(idx) + idx, g = idx[order], g[order] + + x = self.w[idx] + if self.wd > 0.0 and self.wd_mode == "l2": + g = g + self.wd * x + if self.k > 0.0: + g = g * (1.0 - self.k * _sq(x)) ** 2 / 4.0 + + 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 + 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": + direction = direction + self.wd * x + + 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 + else: + x_new = x - self.lr * direction + + self.w[idx] = x_new + self.m[idx] = m + self.v[idx] = v diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 3fe3ca8..5678e3c 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1,129 +1,74 @@ -"""Тесты исправлений дефектов F1–F8 из разбора DSA (см. ../CHANGELOG.md). - -Каждый тест привязан к дефекту: - F1 — очередь не теряет записи при медленном диске; синхронный fallback. - F2 — два последовательных шага с теми же строками не читают устаревшее - состояние; при непересекающихся строках дрена нет. - F3 — resume/reinit/shape-mismatch/частичное состояние/отсутствие meta. - F4 — пер-строчная bias-коррекция; первый шаг свежей строки = lr*sign(g). - F5 — индексы сортируются перед чтением и записью. - F6 — честный shutdown: ошибка вместо «всё записано». - F7 — прочитанные из файла w проецируются в шар. - F8 — save_checkpoint/load_checkpoint; state_dict честно бросает исключение. - R1 — не-конечные строки (NaN/inf) из битого файла обрабатываются на чтении. - R2 — маркер незавершённой переинициализации в meta. - R3 — неизвестный ключ updates: warning со счётчиком, исключение в debug. -Плюс математические проверки: ретракция, дедупликация, k=0 (евклидов путь), -формула масштабирования (1-k*||w||^2)^2/4 == egrad/lambda^2. +"""Regression and acceptance tests for DSA 0.3.0. + +The suite is organised by the audit finding it covers: + + C1 legacy 0.1.0 state can be reinitialised and migrated + C2 portable checkpoint directory (no absolute paths) + C3 source tensors are not retained (O(batch) memory contract) + C4 explicit weight-decay semantics (L2 vs decoupled) + C5 one packaging metadata source + A1 shutdown is retryable and never reports a false success + A2 global_step survives resume and checkpoint + A3 negative/out-of-range indices are rejected + A4 a non-finite row is fully reset (w, m, v, t) + A9 gradient shape/finiteness is validated + H1 initial/resumed weights are inside the ball + H2 full Riemannian Adam (tangent moments, exp map, parallel transport) + H3 curvature changes reproject the whole table + H4 k is validated early (NaN/inf/negative) + F1-F8, R1-R3 the 0.2.0 correctness guarantees """ + +import gc import json +import logging +import math import os import shutil import tempfile import threading import time +import weakref import numpy as np import pytest import torch +from dsa import manifold from dsa.optimizer import DiskSparseRiemannianAdam +from tests.reference import RiemannianAdamReference, expmap BETA1, BETA2, EPS = 0.9, 0.999, 1e-8 -# ───────────────────────── эталонная модель ───────────────────────── - -class RefAdam: - """Та же математика, что в DiskSparseRiemannianAdam.step, но во float64 - и в оперативной памяти. Используется как эталон для аналитических проверок.""" - - def __init__(self, w0, lr, beta1=BETA1, beta2=BETA2, eps=EPS, k=0.0, wd=0.0): - self.w = np.array(w0, dtype=np.float64) - self.n, self.d = self.w.shape - self.m = np.zeros_like(self.w) - self.v = np.zeros_like(self.w) - self.t = np.zeros(self.n, dtype=np.int64) - self.lr, self.b1, self.b2, self.eps = lr, beta1, beta2, eps - self.k, self.wd = k, wd - - @classmethod - def from_state(cls, w, m, v, t, **kw): - obj = cls(w, lr=kw.get('lr', 1e-3)) - obj.w = np.array(w, dtype=np.float64) - obj.m = np.array(m, dtype=np.float64) - obj.v = np.array(v, dtype=np.float64) - obj.t = np.array(t, dtype=np.int64) - obj.n, obj.d = obj.w.shape - obj.b1 = kw.get('beta1', obj.b1) - obj.b2 = kw.get('beta2', obj.b2) - obj.eps = kw.get('eps', obj.eps) - obj.k = kw.get('k', obj.k) - obj.wd = kw.get('wd', obj.wd) - return obj - - def _project(self, w): - r = np.sqrt((1.0 - 1e-5) / self.k) - n = np.linalg.norm(w, axis=-1, keepdims=True) - return np.where(n > r, w / n * r, w) - - def step(self, idx, g): - idx = np.asarray(idx, dtype=np.int64) - g = np.asarray(g, dtype=np.float64) - # дедупликация как в оптимизаторе (сумма градиентов дубликатов) - uniq, inv = np.unique(idx, return_inverse=True) - agg = np.zeros((len(uniq), g.shape[1])) - for j, i in enumerate(inv): - agg[i] += g[j] - idx, g = uniq, agg - order = np.argsort(idx) - idx, g = idx[order], g[order] - - self.t[idx] += 1 - w = self.w[idx] - if self.k > 1e-9: - w = self._project(w) - sq = np.sum(w * w, axis=-1, keepdims=True) - sq = np.minimum(sq, (1.0 - 1e-5) / self.k) - g = g * ((1.0 - self.k * sq) ** 2 / 4.0) - if self.wd > 0: - w = w * (1.0 - self.lr * self.wd) - m, v = self.m[idx], self.v[idx] - t = self.t[idx].astype(np.float64) - m = self.b1 * m + (1 - self.b1) * g - v = self.b2 * v + (1 - self.b2) * g * g - m_hat = m / (1 - self.b1 ** t)[:, None] - v_hat = v / (1 - self.b2 ** t)[:, None] - w = w - self.lr * m_hat / (np.sqrt(v_hat) + self.eps) - if self.k > 1e-9: - w = self._project(w) - self.w[idx], self.m[idx], self.v[idx] = w, m, v - - -# ───────────────────────── помощники ───────────────────────── - -def make_opt(disk_dir, w0, lr=0.05, k=0.0, **kw): +# ───────────────────────── helpers ───────────────────────── + +def make_opt(disk_dir, w0, lr=0.05, k=0.0, resume=False, **kw): params = torch.from_numpy(np.array(w0, dtype=np.float32)) - kw.setdefault('lr', lr) - kw.setdefault('k', k) - return DiskSparseRiemannianAdam(params={'ent': params}, disk_dir=disk_dir, **kw) + kw.setdefault("lr", lr) + kw.setdefault("k", k) + # Numeric tests need exactly the weights they passed; H1 tests exercise + # the default initial scaling explicitly. + kw.setdefault("init_max_norm_fraction", 1.0) + return DiskSparseRiemannianAdam(params={"ent": params}, + disk_dir=disk_dir, resume=resume, **kw) def do_step(opt, idx, g): idx_t = torch.tensor(np.asarray(idx, dtype=np.int64)) g_t = torch.tensor(np.asarray(g, dtype=np.float32)) - opt.step(updates={'ent': (idx_t, g_t)}) + opt.step(updates={"ent": (idx_t, g_t)}) -def read_state(opt, name='ent'): +def read_state(opt, name="ent"): opt.writer.wait_all_flushed() sf = opt.state_files[name] - return (np.array(sf['w']), np.array(sf['m']), np.array(sf['v']), - np.array(sf['t'])) + return (np.array(sf["w"]), np.array(sf["m"]), np.array(sf["v"]), + np.array(sf["t"])) -def assert_matches_reference(opt, ref, name='ent', rows=None, atol=1e-5): - w, m, v, t = read_state(opt, name) +def assert_matches_reference(opt, ref, rows=None, atol=1e-5): + w, m, v, t = read_state(opt) rows = np.arange(ref.n) if rows is None else np.asarray(rows) assert np.array_equal(t[rows], ref.t[rows]), \ f"пер-строчные счётчики разошлись: {t[rows]} vs {ref.t[rows]}" @@ -133,7 +78,6 @@ def assert_matches_reference(opt, ref, name='ent', rows=None, atol=1e-5): def slow_down_writer(opt, delay): - """Подставной медленный писатель: задержка перед каждой записью батча.""" orig = opt.writer._apply_write def slow(name, indices, w_np, m_np, v_np, t_np=None): @@ -144,7 +88,6 @@ def slow(name, indices, w_np, m_np, v_np, t_np=None): def gate_writer(opt): - """Полная остановка писателя; возвращает функцию открытия.""" orig = opt.writer._apply_write gate = threading.Event() @@ -156,19 +99,45 @@ def gated(name, indices, w_np, m_np, v_np, t_np=None): return gate.set -# ───────────────────────── F1: потери записей ───────────────────────── +def _write_legacy_010(directory, w, m, v): + for suffix, arr in (("w", w), ("m", m), ("v", v)): + mm = np.memmap(os.path.join(directory, f"ent_{suffix}.dat"), + dtype="float32", mode="w+", shape=arr.shape) + mm[:] = arr.astype(np.float32) + mm.flush() + del mm + + +def _write_legacy_020(directory, w, m, v, t): + _write_legacy_010(directory, w, m, v) + mm = np.memmap(os.path.join(directory, "ent_t.dat"), dtype="int32", + mode="w+", shape=(len(t),)) + mm[:] = t + mm.flush() + del mm + meta = { + "format_version": 1, + "param_name": "ent", + "shape": list(w.shape), + "dtype": "float32", + "counter_dtype": "int32", + "files": {s: f"ent_{s}.dat" for s in ("w", "m", "v", "t")}, + } + with open(os.path.join(directory, "ent_meta.json"), "w") as f: + json.dump(meta, f) + + +# ───────────────────────── F1: no dropped batches ───────────────────────── def test_f1_slow_disk_loses_nothing(tmp_path): - """Медленный диск (подставной писатель): ни один батч не теряется, - включая шаги с перекрытием строк (дрены по F2).""" - n, d = 400, 4 + n, d = 300, 4 rng = np.random.default_rng(0) w0 = rng.normal(0, 0.05, (n, d)) opt = make_opt(str(tmp_path), w0, lr=0.05, k=0.0, max_queue_size=2) - slow_down_writer(opt, 0.02) + slow_down_writer(opt, 0.01) - ref = RefAdam(w0, lr=0.05, k=0.0) - for s in range(15): + ref = RiemannianAdamReference(w0, lr=0.05, k=0.0) + for s in range(12): rows = (s * 25 + np.arange(40)) % n g = rng.normal(0, 1.0, (len(rows), d)) do_step(opt, rows, g) @@ -179,8 +148,6 @@ def test_f1_slow_disk_loses_nothing(tmp_path): def test_f1_sync_fallback_on_stalled_writer(tmp_path): - """Очередь полна и писатель стоит: put_update уходит в синхронную запись, - данные не теряются (счётчик fallbacks растёт, состояние совпадает).""" n, d = 200, 4 rng = np.random.default_rng(1) w0 = rng.normal(0, 0.05, (n, d)) @@ -188,31 +155,29 @@ def test_f1_sync_fallback_on_stalled_writer(tmp_path): max_queue_size=1, put_timeout=0.15) open_gate = gate_writer(opt) - ref = RefAdam(w0, lr=0.05, k=0.0) - threading.Timer(0.5, open_gate).start() # открыть шлюз через 0.5 с + ref = RiemannianAdamReference(w0, lr=0.05, k=0.0) + threading.Timer(0.5, open_gate).start() - for s in range(5): # 5 батчей по 20 непересекающихся строк + for s in range(5): rows = np.arange(s * 20, s * 20 + 20) g = rng.normal(0, 1.0, (len(rows), d)) do_step(opt, rows, g) ref.step(rows, g) opt.shutdown() - assert opt.writer._sync_fallbacks >= 1, "синхронный fallback не сработал" + assert opt.writer._sync_fallbacks >= 1 assert_matches_reference(opt, ref, rows=np.arange(100)) -# ───────────────────────── F2: гонка step/писатель ───────────────────────── +# ───────────────────────── F2: step/writer race ───────────────────────── def test_f2_repeated_rows_never_stale(tmp_path): - """Два+ последовательных шага по тем же строкам обязаны видеть результаты - предыдущего шага (дрен), а не устаревшее состояние с диска.""" n, d = 16, 3 w0 = np.zeros((n, d)) opt = make_opt(str(tmp_path), w0, lr=0.05, k=0.0) - slow_down_writer(opt, 0.05) + slow_down_writer(opt, 0.03) - ref = RefAdam(w0, lr=0.05, k=0.0) + ref = RiemannianAdamReference(w0, lr=0.05, k=0.0) rows = [3, 7] for g_val in (1.0, -2.0, 0.5): g = np.full((len(rows), d), g_val) @@ -224,49 +189,47 @@ def test_f2_repeated_rows_never_stale(tmp_path): def test_f2_disjoint_rows_do_not_drain(tmp_path): - """Пересечение с строками в полёте пусто → step() не ждёт дрен очереди - (асинхронность сохраняется в общем случае); непусто → ждёт.""" n, d = 32, 3 w0 = np.zeros((n, d)) opt = make_opt(str(tmp_path), w0, lr=0.05, k=0.0) - open_gate = gate_writer(opt) # первый батч гарантированно висит в полёте + open_gate = gate_writer(opt) do_step(opt, [0, 1], np.ones((2, d))) - assert opt.writer.intersects_inflight('ent', np.array([0, 1])) - assert not opt.writer.intersects_inflight('ent', np.array([5, 6])) + assert opt.writer.intersects_inflight("ent", np.array([0, 1])) + assert not opt.writer.intersects_inflight("ent", np.array([5, 6])) - drains = {'n': 0} + drains = {"n": 0} def record_only_wait(*a, **kw): - # заглушка: фиксируем сам факт запроса дрена; реальный join при - # закрытом шлюзе завис бы, а проверяется здесь только логика step() - drains['n'] += 1 + drains["n"] += 1 opt.writer.wait_all_flushed = record_only_wait - do_step(opt, [5, 6], np.ones((2, d))) # непересекающиеся → без дрена - assert drains['n'] == 0, "лишний дрен на непересекающихся строках" - do_step(opt, [0, 1], np.ones((2, d))) # пересечение → дрен запрошен - assert drains['n'] == 1 + do_step(opt, [5, 6], np.ones((2, d))) + assert drains["n"] == 0 + do_step(opt, [0, 1], np.ones((2, d))) + assert drains["n"] == 1 open_gate() opt.shutdown() -# ───────────────────────── F3: состояние на диске ───────────────────────── +# ───────────────────────── F3: state lifecycle ───────────────────────── def test_f3_meta_written_on_fresh_init(tmp_path): n, d = 10, 2 opt = make_opt(str(tmp_path), np.zeros((n, d))) - meta_path = os.path.join(str(tmp_path), 'ent_meta.json') - assert os.path.exists(meta_path) + meta_path = os.path.join(str(tmp_path), "ent_meta.json") with open(meta_path) as f: meta = json.load(f) - assert meta['shape'] == [n, d] - assert meta['dtype'] == 'float32' and meta['counter_dtype'] == 'int32' + assert meta["format_version"] == 2 + assert meta["algorithm"] == "riemannian_adam" + assert meta["shape"] == [n, d] + assert meta["dtype"] == "float32" and meta["counter_dtype"] == "int32" + assert meta["k"] == 0.0 and meta["ball_validated"] is True opt.shutdown() -def test_f3_resume_true_preserves_state(tmp_path): +def test_f3_resume_true_preserves_state_and_global_step(tmp_path): n, d = 12, 3 rng = np.random.default_rng(2) w0 = rng.normal(0, 0.1, (n, d)) @@ -277,60 +240,69 @@ def test_f3_resume_true_preserves_state(tmp_path): opt.shutdown() w1, m1, v1, t1 = read_state(opt) - assert t1[rows].tolist() == [2, 2, 2] and t1.sum() == 6 + assert t1[rows].tolist() == [2, 2, 2] - # продолжаем: w из params НЕ должен перезаписать диск other = rng.normal(0, 5.0, (n, d)) opt2 = make_opt(str(tmp_path), other, lr=0.05, k=0.0, resume=True) + assert opt2._global_step == 2 w2, m2, v2, t2 = read_state(opt2) assert np.array_equal(w1, w2) and np.array_equal(m1, m2) assert np.array_equal(v1, v2) and np.array_equal(t1, t2) - # следующий шаг продолжает с того же места (счётчик строки = 3) - ref = RefAdam.from_state(w1, m1, v1, t1, lr=0.05, k=0.0) + ref = RiemannianAdamReference.__new__(RiemannianAdamReference) + ref.w, ref.m, ref.v, ref.t = (np.array(w1, dtype=np.float64), + np.array(m1, dtype=np.float64), + np.array(v1, dtype=np.float64), + np.array(t1, dtype=np.int64)) + ref.n, ref.d = n, d + ref.lr, ref.b1, ref.b2, ref.eps = 0.05, BETA1, BETA2, EPS + ref.k, ref.wd, ref.wd_mode = 0.0, 0.0, "l2" do_step(opt2, rows, np.full((len(rows), d), 2.0)) ref.step(rows, np.full((len(rows), d), 2.0)) opt2.shutdown() assert_matches_reference(opt2, ref, rows=rows) -def test_f3_default_reinits_and_warns(tmp_path, caplog): - import logging +def test_f3_existing_state_requires_explicit_resume(tmp_path): + n, d = 8, 3 + opt = make_opt(str(tmp_path), np.zeros((n, d))) + opt.shutdown() + with pytest.raises(RuntimeError, match="resume=True"): + make_opt(str(tmp_path), np.zeros((n, d)), resume=None) + + +def test_f3_resume_false_reinitialises_and_warns(tmp_path, caplog): n, d = 12, 3 w0 = np.full((n, d), 0.25) opt = make_opt(str(tmp_path), w0, lr=0.05, k=0.0) do_step(opt, [0, 1], np.ones((2, d))) opt.shutdown() - fresh = rng_vals(3, n, d) + fresh = np.random.default_rng(3).normal(0, 0.3, (n, d)) with caplog.at_level(logging.WARNING, logger="DSA"): - opt2 = make_opt(str(tmp_path), fresh, lr=0.05, k=0.0) # resume=False по умолчанию + opt2 = make_opt(str(tmp_path), fresh, lr=0.05, k=0.0, resume=False) assert any("переинициализация" in r.message for r in caplog.records) w, m, v, t = read_state(opt2) - assert np.allclose(w, fresh, atol=1e-7), "w должен быть взят заново из params" + assert np.allclose(w, fresh, atol=1e-7) assert np.all(m == 0) and np.all(v == 0) and np.all(t == 0) + assert opt2._global_step == 0 opt2.shutdown() -def rng_vals(seed, n, d): - return np.random.default_rng(seed).normal(0, 0.3, (n, d)) - - def test_f3_shape_mismatch_raises(tmp_path): opt = make_opt(str(tmp_path), np.zeros((50, 4))) opt.shutdown() with pytest.raises(RuntimeError, match="несовпадение формы"): make_opt(str(tmp_path), np.zeros((60, 4)), resume=True) - # то же и при resume=False: тихо выбрасывать чужое состояние нельзя with pytest.raises(RuntimeError, match="несовпадение формы"): - make_opt(str(tmp_path), np.zeros((50, 8))) + make_opt(str(tmp_path), np.zeros((50, 8)), resume=False) def test_f3_partial_state_raises(tmp_path): opt = make_opt(str(tmp_path), np.zeros((10, 2))) opt.shutdown() - os.remove(os.path.join(str(tmp_path), 'ent_v.dat')) + os.remove(os.path.join(str(tmp_path), "ent_v.dat")) with pytest.raises(RuntimeError, match="частичное состояние"): make_opt(str(tmp_path), np.zeros((10, 2)), resume=True) @@ -338,177 +310,625 @@ def test_f3_partial_state_raises(tmp_path): def test_f3_missing_meta_raises(tmp_path): opt = make_opt(str(tmp_path), np.zeros((10, 2))) opt.shutdown() - os.remove(os.path.join(str(tmp_path), 'ent_meta.json')) + os.remove(os.path.join(str(tmp_path), "ent_meta.json")) with pytest.raises(RuntimeError, match="meta"): make_opt(str(tmp_path), np.zeros((10, 2)), resume=True) @pytest.mark.parametrize("suffix", ["w", "m", "v", "t"]) def test_f3_truncated_data_file_raises_on_resume(tmp_path, suffix): - # Контрпример ревью: валидный meta + обрезанный файл данных (сбой диска - # или копирования). np.memmap в r+ молча продлевает файл нулями — - # открытие обязано быть явной ошибкой, файл не должен «вырастать обратно». n, d = 10, 2 - opt = make_opt(str(tmp_path), rng_vals(5, n, d), lr=0.05, k=0.0) + opt = make_opt(str(tmp_path), np.random.default_rng(5).normal(0, 0.1, (n, d)), + lr=0.05, k=0.0) do_step(opt, [0, 5], np.ones((2, d))) opt.shutdown() - p = os.path.join(str(tmp_path), f'ent_{suffix}.dat') + p = os.path.join(str(tmp_path), f"ent_{suffix}.dat") sz = os.path.getsize(p) - with open(p, 'r+b') as f: + with open(p, "r+b") as f: f.truncate(sz // 2) with pytest.raises(RuntimeError, match="повреждён: размер"): make_opt(str(tmp_path), np.zeros((n, d)), resume=True) - # memmap так и не открылся — файл НЕ продлился нулями молча assert os.path.getsize(p) == sz // 2 def test_f3_oversized_data_file_raises_on_resume(tmp_path): - # Слишком длинный файл (обрыв копирования поверх старого состояния) - # np.memmap отобразил бы молча только начало — тоже явная ошибка. n, d = 10, 2 opt = make_opt(str(tmp_path), np.zeros((n, d))) opt.shutdown() - p = os.path.join(str(tmp_path), 'ent_w.dat') - with open(p, 'ab') as f: - f.write(b'\x00' * 16) + p = os.path.join(str(tmp_path), "ent_w.dat") + with open(p, "ab") as f: + f.write(b"\x00" * 16) with pytest.raises(RuntimeError, match="повреждён: размер"): make_opt(str(tmp_path), np.zeros((n, d)), resume=True) def test_f3_truncated_file_resume_false_reinits(tmp_path): - # resume=False — документированный путь восстановления: полная - # переинициализация пересоздаёт файлы правильного размера. n, d = 10, 2 opt = make_opt(str(tmp_path), np.zeros((n, d))) opt.shutdown() - p = os.path.join(str(tmp_path), 'ent_v.dat') - with open(p, 'r+b') as f: + p = os.path.join(str(tmp_path), "ent_v.dat") + with open(p, "r+b") as f: f.truncate(os.path.getsize(p) // 3) - fresh = rng_vals(7, n, d) - opt2 = make_opt(str(tmp_path), fresh, lr=0.05, k=0.0) # resume=False + fresh = np.random.default_rng(7).normal(0, 0.3, (n, d)) + opt2 = make_opt(str(tmp_path), fresh, lr=0.05, k=0.0, resume=False) w, m, v, t = read_state(opt2) - assert os.path.getsize(p) == n * d * 4, "файл должен быть пересоздан целиком" + assert os.path.getsize(p) == n * d * 4 assert np.allclose(w, fresh, atol=1e-7) assert np.all(m == 0) and np.all(v == 0) and np.all(t == 0) opt2.shutdown() -# ───────────────────────── F4: пер-строчная bias-коррекция ───────────────────────── +# ───────────────────────── C1: legacy 0.1.0 / 0.2.0 ───────────────────────── + +def test_c1_legacy_010_resume_true_migrates(tmp_path): + n, d = 6, 3 + w = np.full((n, d), 0.05) + _write_legacy_010(str(tmp_path), w, np.full((n, d), 0.3), + np.full((n, d), 0.4)) + + 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.1.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: + meta = json.load(f) + assert meta["format_version"] == 2 + assert meta["legacy_migrated_from"] == "0.1.0" + opt.shutdown() + + +def test_c1_legacy_010_resume_false_reinitialises(tmp_path): + n, d = 6, 3 + _write_legacy_010(str(tmp_path), np.full((n, d), 0.05), + np.full((n, d), 0.3), np.full((n, d), 0.4)) + fresh = np.full((n, d), 0.02) + opt = make_opt(str(tmp_path), fresh, k=0.0, resume=False) + st = read_state(opt) + assert np.allclose(st[0], fresh, atol=1e-7) + assert np.all(st[1] == 0) and np.all(st[2] == 0) and np.all(st[3] == 0) + opt.shutdown() + + +def test_c1_legacy_010_resume_none_raises(tmp_path): + n, d = 6, 3 + _write_legacy_010(str(tmp_path), np.zeros((n, d)), np.zeros((n, d)), + np.zeros((n, d))) + with pytest.raises(RuntimeError, match="resume=True"): + make_opt(str(tmp_path), np.zeros((n, d)), resume=None) + + +def test_c1_legacy_020_resets_incompatible_moments(tmp_path): + n, d = 4, 2 + w = np.full((n, d), 0.04) + _write_legacy_020(str(tmp_path), w, np.full((n, d), 0.5), + np.full((n, d), 0.6), np.full(n, 7)) + 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) + 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: + assert json.load(f)["legacy_migrated_from"] == 1 + opt.shutdown() + + +# ───────────────────────── F4 / C4: Adam and weight decay ───────────────────────── def test_f4_first_step_of_fresh_row_is_lr_sign(tmp_path): - """Аналитически: первый шаг свежей строки обязан давать w -= lr*sign(g), - точнее w -= lr*g/(|g|+eps) — m_hat=g, v_hat=g^2 при t=1.""" n, d = 10, 3 lr = 0.1 - w0 = np.zeros((n, d)) - opt = make_opt(str(tmp_path), w0, lr=lr, k=0.0) - + opt = make_opt(str(tmp_path), np.zeros((n, d)), lr=lr, k=0.0) g = np.array([[1.0, -2.0, 0.5]]) do_step(opt, [4], g) opt.shutdown() w, m, v, t = read_state(opt) expected_delta = lr * g[0] / (np.abs(g[0]) + EPS) - assert np.allclose(w[4], w0[4] - expected_delta, atol=1e-7) + assert np.allclose(w[4], -expected_delta, atol=1e-7) assert np.allclose(m[4], (1 - BETA1) * g[0], atol=1e-7) assert np.allclose(v[4], (1 - BETA2) * g[0] ** 2, atol=1e-7) assert t[4] == 1 - assert np.all(np.sign(w0[4] - w[4]) == np.sign(g[0])), \ - "направление первого шага должно быть -sign(g)" def test_f4_row_counter_independent_of_global_steps(tmp_path): - """Строка, впервые обновлённая на глобальном шаге 6, получает bias-коррекцию - по своему счётчику (t=1), а не по глобальному: движение = lr*g/(|g|+eps). - Старый код с глобальным t давал бы ~lr*0.213*g — тест ловит регрессию.""" n, d = 8, 3 lr = 0.1 - w0 = np.zeros((n, d)) - opt = make_opt(str(tmp_path), w0, lr=lr, k=0.0) - - for _ in range(5): # пять шагов только по строке 0 + opt = make_opt(str(tmp_path), np.zeros((n, d)), lr=lr, k=0.0) + for _ in range(5): do_step(opt, [0], np.ones((1, d))) - do_step(opt, [1], np.ones((1, d))) # строка 1 — впервые, глобальный шаг 6 + do_step(opt, [1], np.ones((1, d))) opt.shutdown() w, _, _, t = read_state(opt) - expected = lr * 1.0 / (1.0 + EPS) + expected = lr / (1.0 + EPS) assert t[0] == 5 and t[1] == 1 and t[2:].sum() == 0 - assert np.allclose(w[1], -expected, atol=1e-7), \ - f"первый шаг свежей строки: {w[1]}, ожидалось {-expected}" + assert np.allclose(w[1], -expected, atol=1e-7) -def test_f4_full_adam_reference_euclidean(tmp_path): - """k=0 — чистый евклидов путь: несколько шагов с дублями, weight_decay, - случайными градиентами — полное совпадение с эталоном.""" - n, d = 100, 8 +@pytest.mark.parametrize("mode", ["l2", "decoupled"]) +def test_c4_euclidean_weight_decay_matches_reference(tmp_path, mode): + n, d = 40, 5 rng = np.random.default_rng(7) w0 = rng.normal(0, 0.1, (n, d)) - lr, wd = 0.05, 0.01 - opt = make_opt(str(tmp_path), w0, lr=lr, k=0.0, weight_decay=wd) - - ref = RefAdam(w0, lr=lr, k=0.0, wd=wd) - torch.manual_seed(7) - for _ in range(6): - rows = torch.randint(0, n, (37,)).numpy() - g = torch.randn(len(rows), d).numpy() + lr, wd = 0.05, 0.1 + opt = make_opt(str(tmp_path), w0, lr=lr, k=0.0, weight_decay=wd, + weight_decay_mode=mode) + ref = RiemannianAdamReference(w0, lr=lr, k=0.0, weight_decay=wd, + weight_decay_mode=mode) + for _ in range(5): + rows = rng.integers(0, n, 20) + g = rng.normal(0, 1.0, (len(rows), d)) do_step(opt, rows, g) ref.step(rows, g) - opt.shutdown() assert_matches_reference(opt, ref) -# ───────────────────────── ретракция и F7 ───────────────────────── +def test_c4_l2_is_adam_l2_not_adamw(tmp_path): + """The default must be classic Adam L2 (decay inside the gradient), not AdamW.""" + n, d = 6, 2 + w0 = np.full((n, d), 0.5) + lr, wd = 0.1, 0.1 + g = np.ones((1, d)) + + opt = make_opt(str(tmp_path), w0, lr=lr, k=0.0, weight_decay=wd, + weight_decay_mode="l2") + do_step(opt, [0], g) + opt.shutdown() + w, *_ = read_state(opt) + + ref = RiemannianAdamReference(w0, lr=lr, k=0.0, weight_decay=wd, + weight_decay_mode="l2") + ref.step([0], g) + assert np.allclose(w[0], ref.w[0], atol=1e-6) + + # AdamW with the same coefficients differs substantially. + x = w0[0].copy() + m = np.zeros_like(x) + v = np.zeros_like(x) + m = BETA1 * m + (1 - BETA1) * g[0] + v = BETA2 * v + (1 - BETA2) * g[0] ** 2 + direction = (m / (1 - BETA1)) / (np.sqrt(v / (1 - BETA2)) + EPS) + adamw = x * (1 - lr * wd) - lr * direction + assert np.max(np.abs(w[0] - adamw)) > 1e-3, \ + "L2-режим не должен совпадать с AdamW" + + +def test_c4_hyperbolic_l2_is_riemannian_gradient(tmp_path): + """L2 decay is part of the gradient, not a coordinate shrink. + + With ``g_E = -wd * w`` the L2 gradient is exactly zero, so Adam has no + direction and the point must not move. The decoupled variant has no such + cancellation: it applies ``wd * w`` as a separate tangent term. + """ + n, d = 2, 2 + k, lr, wd = 1.0, 0.1, 0.5 + w0 = np.array([[0.05, 0.0], [0.9, 0.0]]) + g = -wd * w0 + + opt_l2 = make_opt(str(tmp_path / "l2"), w0, lr=lr, k=k, + weight_decay=wd, weight_decay_mode="l2") + do_step(opt_l2, [0, 1], g) + opt_l2.shutdown() + w_l2, *_ = read_state(opt_l2) + assert np.allclose(w_l2, w0, atol=1e-6), \ + "L2-штраф обязан сократиться с g = -wd*w" + + opt_dec = make_opt(str(tmp_path / "dec"), w0, lr=lr, k=k, + weight_decay=wd, weight_decay_mode="decoupled") + do_step(opt_dec, [0, 1], g) + opt_dec.shutdown() + w_dec, *_ = read_state(opt_dec) + assert not np.allclose(w_dec, w0, atol=1e-4), \ + "decoupled-штраф не участвует в градиенте и не должен сокращаться" + + +def test_c4_invalid_weight_decay_mode(tmp_path): + with pytest.raises(ValueError, match="weight_decay_mode"): + make_opt(str(tmp_path), np.zeros((3, 2)), weight_decay_mode="adamw") + + +# ───────────────────────── F7 / H1: ball invariants ───────────────────────── + +def test_h1_init_is_scaled_into_ball(tmp_path): + n, d = 4, 8 + w0 = np.full((n, d), 0.5) # norm 1.414 > radius 1 + opt = DiskSparseRiemannianAdam( + {"ent": torch.tensor(w0)}, lr=0.01, k=1.0, disk_dir=str(tmp_path), + resume=False) # default init_max_norm_fraction=0.5 + w = np.array(opt.state_files["ent"]["w"]) + radius = math.sqrt((1 - manifold.BALL_EPS) / 1.0) + norms = np.linalg.norm(w, axis=1) + assert np.all(norms <= 0.5 * radius + 1e-6) + assert np.isfinite(manifold.distance( + torch.tensor(w[:1]), torch.tensor(w[1:2]), 1.0).numpy()).all() + opt.shutdown() + + +def test_h1_init_fraction_one_only_projects(tmp_path): + n, d = 4, 8 + w0 = np.full((n, d), 0.5) + opt = DiskSparseRiemannianAdam( + {"ent": torch.tensor(w0)}, lr=0.01, k=1.0, disk_dir=str(tmp_path), + resume=False, init_max_norm_fraction=1.0) + w = np.array(opt.state_files["ent"]["w"]) + radius = math.sqrt((1 - manifold.BALL_EPS) / 1.0) + assert np.all(np.linalg.norm(w, axis=1) <= radius + 1e-6) + opt.shutdown() + + +def test_h1_read_rows_projects(tmp_path): + n, d, k = 8, 4, 1.0 + opt = make_opt(str(tmp_path), np.full((n, d), 0.1), k=k) + opt.state_files["ent"]["w"][2] = np.full(d, 2.5, dtype=np.float32) + opt.state_files["ent"]["w"].flush() + + rows = opt.read_rows([2]).numpy() + radius = math.sqrt((1 - manifold.BALL_EPS) / k) + assert np.linalg.norm(rows[0]) <= radius + 1e-6 + opt.shutdown() + + +def test_h1_resume_projects_corrupted_rows(tmp_path): + n, d, k = 6, 3, 1.0 + opt = make_opt(str(tmp_path), np.full((n, d), 0.1), k=k) + opt.shutdown() + with open(os.path.join(str(tmp_path), "ent_w.dat"), "r+b") as f: + f.seek(0) + f.write(np.full(d, 3.0, dtype=np.float32).tobytes()) + f.write(np.full(d, -3.0, dtype=np.float32).tobytes()) + + opt2 = make_opt(str(tmp_path), np.zeros((n, d)), k=k, resume=True) + w = np.array(opt2.state_files["ent"]["w"]) + radius = math.sqrt((1 - manifold.BALL_EPS) / k) + assert np.all(np.linalg.norm(w, axis=1) <= radius + 1e-6) + assert opt2._resume_rows_projected == 2 + opt2.shutdown() + def test_retraction_keeps_norm_in_ball(tmp_path): - """Ретракция держит норму каждой строки ≤ sqrt((1-1e-5)/k) после каждого шага.""" - n, d = 6, 4 - k = 1.0 - max_radius = np.sqrt((1.0 - 1e-5) / k) - w0 = np.full((n, d), 0.45) # норма 0.9 — внутри шара + """An extreme outward step keeps every row inside the ball and finite. + + The float64 reference cannot be used here: without projection it reaches + the boundary exactly (lambda = inf), which is not a point of the manifold. + Exact numerical agreement is checked by the H2 tests away from the + boundary; this test checks the invariant that matters at the boundary. + """ + n, d, k = 6, 4, 1.0 + max_radius = math.sqrt((1.0 - manifold.BALL_EPS) / k) + w0 = np.full((n, d), 0.45) opt = make_opt(str(tmp_path), w0, lr=10.0, k=k) - - ref = RefAdam(w0, lr=10.0, k=k) - for _ in range(8): + for step in range(8): g = np.full((n, d), 100.0) do_step(opt, np.arange(n), g) - ref.step(np.arange(n), g) - w, *_ = read_state(opt) # read_state ждёт дрен — тот же приём, что step - norms = np.linalg.norm(w, axis=1) - assert np.all(norms <= max_radius + 1e-6), \ - f"норма вышла за шар: {norms} > {max_radius}" + w, m, v, t = read_state(opt) + assert np.all(np.linalg.norm(w, axis=1) <= max_radius + 1e-6) + assert np.isfinite(w).all() and np.isfinite(m).all() and np.isfinite(v).all() + assert np.all(t == step + 1), "строка была сброшена вместо шага" + opt.shutdown() + + +def test_f7_corrupted_row_is_projected_on_step(tmp_path): + n, d, k = 8, 4, 1.0 + max_radius = math.sqrt((1.0 - manifold.BALL_EPS) / k) + opt = make_opt(str(tmp_path), np.full((n, d), 0.1), lr=0.01, k=k) + opt.state_files["ent"]["w"][2] = np.full(d, 2.5, dtype=np.float32) + opt.state_files["ent"]["w"].flush() + + do_step(opt, [2], np.full((1, d), 0.1)) + opt.shutdown() + w, m, v, t = read_state(opt) + assert np.linalg.norm(w[2]) <= max_radius + 1e-6 + assert np.isfinite(w[2]).all() and np.isfinite(m[2]).all() + assert t[2] == 1 + + +# ───────────────────────── H2: full Riemannian Adam ───────────────────────── + +@pytest.mark.parametrize("k", [0.0, 0.7]) +@pytest.mark.parametrize("mode", ["l2", "decoupled"]) +def test_h2_matches_independent_reference(tmp_path, k, mode): + n, d = 30, 5 + rng = np.random.default_rng(9) + w0 = rng.normal(0, 0.1, (n, d)) + if k > 0: + w0 *= 0.8 * math.sqrt(1 / k) / np.linalg.norm(w0, axis=1).max() + lr, wd = 0.03, 0.02 + opt = make_opt(str(tmp_path), w0, lr=lr, k=k, weight_decay=wd, + weight_decay_mode=mode) + ref = RiemannianAdamReference(w0, lr=lr, k=k, weight_decay=wd, + weight_decay_mode=mode) + for _ in range(6): + rows = rng.integers(0, n, 12) + g = rng.normal(0, 0.5, (len(rows), d)) + do_step(opt, rows, g) + ref.step(rows, g) + opt.shutdown() + assert_matches_reference(opt, ref, atol=1e-5) + + +def test_h2_moments_are_transported_isometrically(tmp_path): + n, d, k = 1, 3, 0.7 + lr = 0.05 + w0 = np.array([[0.2, -0.1, 0.15]]) + g = np.array([[0.3, -0.4, 0.2]]) + opt = make_opt(str(tmp_path), w0, lr=lr, k=k) + do_step(opt, [0], g) + opt.shutdown() + w, m, v, _ = read_state(opt) + g_riem = g[0] * (1 - k * np.dot(w0[0], w0[0])) ** 2 / 4 + m_before = (1 - BETA1) * g_riem + lam0 = 2.0 / (1 - k * np.dot(w0[0], w0[0])) + 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) + + +def test_h2_first_step_from_origin_closed_form(tmp_path): + n, d, k = 1, 4, 0.7 + lr = 0.1 + g = np.array([[0.5, -1.0, 2.0, -0.25]]) + opt = make_opt(str(tmp_path), np.zeros((n, d)), lr=lr, k=k) + do_step(opt, [0], g) opt.shutdown() - assert_matches_reference(opt, ref, atol=1e-4) + w, *_ = read_state(opt) + direction = g[0] / (np.abs(g[0]) + EPS) # first Adam step at t=1 + # exp_0(-lr * direction) = -tanh(sqrt(k)*lr*||direction||) / + # (sqrt(k)*||direction||) * direction + norm = np.linalg.norm(direction) + expected = (-math.tanh(math.sqrt(k) * lr * norm) + / (math.sqrt(k) * norm) * direction) + assert np.allclose(w[0], expected, atol=1e-6) -def test_f7_corrupted_row_is_projected_on_read(tmp_path): - """Битая строка (норма вне шара) в файле w не расходить геометрию не должна: - после шага строка снова внутри шара, без NaN.""" - n, d = 8, 4 - k = 1.0 - max_radius = np.sqrt((1.0 - 1e-5) / k) + +def test_h2_zero_gradient_is_identity(tmp_path): + n, d, k = 3, 4, 1.0 w0 = np.full((n, d), 0.1) - opt = make_opt(str(tmp_path), w0, lr=0.01, k=k) + opt = make_opt(str(tmp_path), w0, lr=0.1, k=k) + do_step(opt, np.arange(n), np.zeros((n, d))) + opt.shutdown() + w, m, v, t = read_state(opt) + assert np.allclose(w, w0, atol=1e-6) + assert np.all(m == 0) and np.all(v == 0) and np.all(t == 1) - # портим строку 2 напрямую в memmap (имитация битого файла) - opt.state_files['ent']['w'][2] = np.full(d, 2.5, dtype=np.float32) # норма 5 - opt.state_files['ent']['w'].flush() - do_step(opt, [2], np.full((1, d), 0.1)) +# ───────────────────────── H3: curvature changes ───────────────────────── + +def test_h3_set_curvature_reprojects_whole_table(tmp_path): + n, d = 8, 3 + k0, k1 = 0.25, 4.0 + radius0 = math.sqrt((1 - manifold.BALL_EPS) / k0) + w0 = np.full((n, d), 0.95 * radius0 / math.sqrt(d)) + opt = make_opt(str(tmp_path), w0, k=k0) + opt.set_curvature(k1) + w = np.array(opt.state_files["ent"]["w"]) + radius1 = math.sqrt((1 - manifold.BALL_EPS) / k1) + assert np.all(np.linalg.norm(w, axis=1) <= radius1 + 1e-6) + assert opt.param_groups[0]["k"] == k1 + with open(os.path.join(str(tmp_path), "dsa_global_meta.json")) as f: + assert json.load(f)["k"] == k1 + opt.shutdown() + + +def test_h3_current_k_reprojects_whole_table(tmp_path): + n, d = 6, 3 + k0, k1 = 0.25, 4.0 + radius0 = math.sqrt((1 - manifold.BALL_EPS) / k0) + w0 = np.full((n, d), 0.9 * radius0 / math.sqrt(d)) + opt = make_opt(str(tmp_path), w0, k=k0) + do_step(opt, [0], np.ones((1, d))) + opt.step(updates={"ent": (torch.tensor([1]), torch.ones(1, d))}, + current_k=k1) + w = np.array(opt.state_files["ent"]["w"]) + radius1 = math.sqrt((1 - manifold.BALL_EPS) / k1) + assert np.all(np.linalg.norm(w, axis=1) <= radius1 + 1e-6) + opt.shutdown() + + +def test_h3_set_curvature_can_reset_moments(tmp_path): + n, d = 4, 3 + opt = make_opt(str(tmp_path), np.full((n, d), 0.05), k=0.5) + do_step(opt, [0, 1], np.ones((2, d))) + opt.set_curvature(0.5, reset_moments=True) + m, v, t = read_state(opt)[1:] + assert np.all(m == 0) and np.all(v == 0) and np.all(t == 0) + opt.shutdown() + + +# ───────────────────────── H4: curvature validation ───────────────────────── + +@pytest.mark.parametrize("bad", [-1.0, float("nan"), float("inf"), -float("inf")]) +def test_h4_invalid_k_rejected_at_construction(tmp_path, bad): + with pytest.raises((ValueError, TypeError)): + make_opt(str(tmp_path), np.zeros((3, 2)), k=bad) + + +@pytest.mark.parametrize("bad", [-1.0, float("nan"), float("inf")]) +def test_h4_invalid_k_rejected_at_runtime(tmp_path, bad): + opt = make_opt(str(tmp_path), np.zeros((3, 2)), k=0.0) + with pytest.raises((ValueError, TypeError)): + opt.set_curvature(bad) + with pytest.raises((ValueError, TypeError)): + opt.step(updates={"ent": (torch.tensor([0]), torch.ones(1, 2))}, + current_k=bad) + opt.shutdown() + + +# ───────────────────────── A1: retryable shutdown ───────────────────────── + +def test_a1_shutdown_retry_after_timeout(tmp_path): + n, d = 20, 3 + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) + open_gate = gate_writer(opt) + do_step(opt, [0, 1], np.ones((2, d))) + do_step(opt, [2, 3], np.ones((2, d))) + + with pytest.raises(RuntimeError, match="НЕ дописаны"): + opt.writer.shutdown(timeout=0.2) + assert opt.writer.state.name == "FAILED" + assert opt.writer.worker_thread.is_alive() + assert opt.writer.queue.unfinished_tasks == 2 + + open_gate() + assert opt.writer.shutdown(timeout=10.0) is True + assert opt.writer.state.name == "STOPPED" + assert not opt.writer.worker_thread.is_alive() + assert opt.writer.queue.unfinished_tasks == 0 + + # Idempotent after a clean stop. + assert opt.writer.shutdown(timeout=1.0) is True + + +def test_a1_shutdown_reports_persistent_write_error(tmp_path): + n, d = 8, 2 + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) + + def broken(*a, **kw): + raise OSError("simulated disk failure") + + opt.writer._apply_write = broken + do_step(opt, [0, 1], np.ones((2, d))) + with pytest.raises(RuntimeError, match="НЕ записан"): + opt.shutdown() + assert opt.writer.state.name == "FAILED" + assert not opt.writer.worker_thread.is_alive(), \ + "после сбора ошибок воркер должен быть остановлен" + # A second call must not pretend that the failed write succeeded. + with pytest.raises(RuntimeError, match="НЕ записан"): + opt.shutdown() + + +def test_a1_put_after_shutdown_raises(tmp_path): + n, d = 6, 2 + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) + opt.shutdown() + with pytest.raises(RuntimeError, match="запись невозможна"): + opt.writer.put_update("ent", torch.tensor([0]), + torch.zeros(1, d), torch.zeros(1, d), + torch.zeros(1, d), torch.tensor([1])) + + +# ───────────────────────── A2: global step ───────────────────────── + +def test_a2_global_step_survives_resume_and_checkpoint(tmp_path): + n, d = 6, 2 + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) + for _ in range(3): + do_step(opt, [0], np.ones((1, d))) + assert opt._global_step == 3 + opt.shutdown() + + opt2 = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0, resume=True) + assert opt2._global_step == 3 + + ckpt_dir = str(tmp_path / "ckpt") + opt2.export_checkpoint(ckpt_dir) + opt2.shutdown() + + opt3 = make_opt(str(tmp_path / "fresh"), np.zeros((n, d)), k=0.0) + manifest = opt3.import_checkpoint(ckpt_dir) + assert manifest["global_step"] == 3 + assert opt3._global_step == 3 + opt3.shutdown() + + +# ───────────────────────── A3: index validation ───────────────────────── + +def test_a3_negative_and_out_of_range_indices_rejected(tmp_path): + n, d = 6, 3 + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) + with pytest.raises(IndexError, match="отрицательный"): + do_step(opt, [-1], np.ones((1, d))) + with pytest.raises(IndexError, match="за границы"): + do_step(opt, [n], np.ones((1, d))) + with pytest.raises(TypeError, match="целочисленными"): + opt.step(updates={"ent": (torch.tensor([0.0]), torch.ones(1, d))}) + with pytest.raises(ValueError, match="1-D"): + opt.step(updates={"ent": (torch.tensor([[0]]), torch.ones(1, d))}) + _, _, _, t = read_state(opt) + assert np.all(t == 0), "невалидный индекс не должен менять состояние" + opt.shutdown() + + +# ───────────────────────── A4: NaN row reset ───────────────────────── + +def test_a4_nan_row_is_fully_reset(tmp_path): + n, d = 6, 3 + lr = 0.1 + opt = make_opt(str(tmp_path), np.zeros((n, d)), lr=lr, k=0.0) + opt.writer.wait_all_flushed() + + sf = opt.state_files["ent"] + sf["w"][2] = np.nan + sf["m"][2] = 7.0 + sf["v"][2] = 9.0 + sf["t"][2] = 50 + for mm in sf.values(): + mm.flush() + + g = np.array([[1.0, -2.0, 0.5]]) + do_step(opt, [2], g) opt.shutdown() w, m, v, t = read_state(opt) - assert np.linalg.norm(w[2]) <= max_radius + 1e-6 - assert not np.isnan(w[2]).any() and not np.isnan(m[2]).any() + expected = -lr * g[0] / (np.abs(g[0]) + EPS) + assert np.allclose(w[2], expected, atol=1e-7) + assert np.allclose(m[2], (1 - BETA1) * g[0], atol=1e-7) + assert np.allclose(v[2], (1 - BETA2) * g[0] ** 2, atol=1e-7) + assert t[2] == 1, f"счётчик строки должен перезапуститься, получено {t[2]}" + assert opt._nan_rows_replaced == 1 + + +def test_a4_zero_counter_discards_stale_moments(tmp_path): + """A row with t=0 is fresh; leftover m/v from a lost counter must not leak.""" + n, d = 4, 2 + opt = make_opt(str(tmp_path), np.zeros((n, d)), lr=0.1, k=0.0) + opt.writer.wait_all_flushed() + sf = opt.state_files["ent"] + sf["m"][1] = 5.0 + sf["v"][1] = 7.0 + sf["t"][1] = 0 + for mm in sf.values(): + mm.flush() + + g = np.array([[1.0, -2.0]]) + do_step(opt, [1], g) + opt.shutdown() + _, m, v, t = read_state(opt) + assert np.allclose(m[1], (1 - BETA1) * g[0], atol=1e-7) + assert np.allclose(v[1], (1 - BETA2) * g[0] ** 2, atol=1e-7) + assert t[1] == 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) + g = np.ones((1, d)) + g[0, 1] = np.nan + with pytest.raises(ValueError, match="NaN/inf"): + do_step(opt, [0], g) + _, _, _, t = read_state(opt) + assert np.all(t == 0) + opt.shutdown() + + +def test_a9_gradient_shape_validation(tmp_path): + n, d = 4, 2 + opt = make_opt(str(tmp_path), np.zeros((n, d)), k=0.0) + with pytest.raises(ValueError, match="форму"): + do_step(opt, [0], np.ones((2, d))) + with pytest.raises(ValueError, match="форму"): + do_step(opt, [0], np.ones((1, d + 1))) + opt.shutdown() -# ───────────────────────── дедупликация и формула масштабирования ───────────────────────── +# ───────────────────────── dedup / scaling ───────────────────────── def test_dedup_equals_manual_aggregation(tmp_path): - """Батч с дублями [3,3,5] и градиентами g1,g2,g3 == батч [3,5] с g1+g2, g3. - Проверено для k=0 и k=1 (масштабирование не ломает эквивалентность).""" for k in (0.0, 1.0): dirA = tempfile.mkdtemp() dirB = tempfile.mkdtemp() @@ -516,12 +936,10 @@ def test_dedup_equals_manual_aggregation(tmp_path): g1, g2, g3 = (np.array([[0.3, -1.2, 0.7]]), np.array([[-0.4, 0.2, 1.1]]), np.array([[1.5, -0.6, 0.2]])) - w0 = np.array([[0.05, -0.02, 0.03] * 8]).reshape(8, 3) * np.arange(1, 9)[:, None] - + w0 = np.array([[0.05, -0.02, 0.03] * 8]).reshape(8, 3) optA = make_opt(dirA, w0, lr=0.05, k=k) do_step(optA, [3, 3, 5], np.concatenate([g1, g2, g3])) optA.shutdown() - optB = make_opt(dirB, w0, lr=0.05, k=k) do_step(optB, [3, 5], np.concatenate([g1 + g2, g3])) optB.shutdown() @@ -537,35 +955,30 @@ def test_dedup_equals_manual_aggregation(tmp_path): shutil.rmtree(dirB, ignore_errors=True) -def test_riemannian_scaling_formula(tmp_path): - """Масштабирование градиента: g_riem = g*(1-k||w||^2)^2/4 == egrad/lambda^2, - где lambda = 2/(1-k||w||^2). Проверяем по m после первого шага и по сдвигу w.""" - n, d = 4, 3 - k = 0.7 +def test_riemannian_gradient_scaling_formula(tmp_path): + n, d, k = 4, 3, 0.7 lr = 0.05 w0 = np.zeros((n, d)) - w0[0] = np.array([0.3, -0.4, 0.2657]) # норма ≈0.57, внутри шара + w0[0] = np.array([0.3, -0.4, 0.2657]) opt = make_opt(str(tmp_path), w0, lr=lr, k=k) - g = np.array([[1.0, -0.5, 2.0]]) do_step(opt, [0], g) opt.shutdown() + w, m, _, _ = read_state(opt) - w, m, v, t = read_state(opt) - sq = float(np.sum(w0[0] ** 2)) # 0.36, внутри шара — проекция не меняет - rescale = (1.0 - k * sq) ** 2 / 4.0 - lam_sq = 4.0 / (1.0 - k * sq) ** 2 # λ² для шара Пуанкаре с кривизной k - g_riem_formula = g[0] * rescale - g_riem_via_lambda = g[0] / lam_sq - assert np.allclose(g_riem_formula, g_riem_via_lambda, rtol=0, atol=0), \ - "формулы (1-k||w||^2)^2/4 и egrad/λ^2 обязаны совпадать" + sq = float(np.sum(w0[0] ** 2)) + g_riem = g[0] * (1.0 - k * sq) ** 2 / 4.0 + lam0 = 2.0 / (1 - k * sq) + lam1 = 2.0 / (1 - k * np.sum(w[0] ** 2)) + 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] + assert np.allclose(w[0], expected, atol=1e-6) - assert np.allclose(m[0], (1 - BETA1) * g_riem_formula, atol=1e-7) - expected_w = w0[0] - lr * g_riem_formula / (np.abs(g_riem_formula) + EPS) - assert np.allclose(w[0], expected_w, atol=1e-6) - -# ───────────────────────── F5: сортировка индексов ───────────────────────── +# ───────────────────────── F5: sorted indices ───────────────────────── def test_f5_indices_sorted_before_write(tmp_path): n, d = 16, 3 @@ -578,8 +991,8 @@ def capturing(name, indices, w_np, m_np, v_np, t_np=None): return orig(name, indices, w_np, m_np, v_np, t_np) opt.writer._apply_write = capturing - do_step(opt, [9, 2, 7, 4, 0], np.ones((5, d))) # без дублей, несортировано - do_step(opt, [5, 1, 5, 3], np.ones((4, d))) # с дублями + do_step(opt, [9, 2, 7, 4, 0], np.ones((5, d))) + do_step(opt, [5, 1, 5, 3], np.ones((4, d))) opt.shutdown() assert len(captured) == 2 @@ -588,158 +1001,213 @@ def capturing(name, indices, w_np, m_np, v_np, t_np=None): def test_f5_apply_write_maps_sorted_correctly(tmp_path): - """Прямой вызов _apply_write с несортированными индексами пишет строки - в правильные места (внутренняя сортировка не путает данные).""" n, d = 10, 3 opt = make_opt(str(tmp_path), np.zeros((n, d))) opt.writer.wait_all_flushed() idx = np.array([7, 2]) w_rows = np.array([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]], dtype=np.float32) - m_rows = np.zeros_like(w_rows) - v_rows = np.zeros_like(w_rows) + zeros = np.zeros_like(w_rows) t_rows = np.array([3, 4], dtype=np.int32) - opt.writer._apply_write('ent', idx, w_rows, m_rows, v_rows, t_rows) - - w, m, v, t = read_state(opt) + opt.writer._apply_write("ent", idx, w_rows, zeros, zeros, t_rows) + w, _, _, t = read_state(opt) assert np.allclose(w[7], 1.0) and np.allclose(w[2], 2.0) assert t[7] == 3 and t[2] == 4 opt.shutdown() -# ───────────────────────── F6: честный shutdown ───────────────────────── +# ───────────────────────── F6: honest shutdown ───────────────────────── def test_f6_shutdown_raises_when_cannot_finish(tmp_path): - """Писатель стоит, очередь не опустела за таймаут → shutdown честно бросает - ошибку вместо «всё записано»; после открытия шлюза данные дописываются.""" n, d = 40, 3 w0 = np.zeros((n, d)) opt = make_opt(str(tmp_path), w0, lr=0.05, k=0.0) open_gate = gate_writer(opt) - - ref = RefAdam(w0, lr=0.05, k=0.0) - do_step(opt, np.arange(0, 10), np.ones((10, d))) # висит в полёте - ref.step(np.arange(0, 10), np.ones((10, d))) - do_step(opt, np.arange(10, 20), np.ones((10, d))) # стоит в очереди - ref.step(np.arange(10, 20), np.ones((10, d))) + do_step(opt, np.arange(0, 10), np.ones((10, d))) + do_step(opt, np.arange(10, 20), np.ones((10, d))) with pytest.raises(RuntimeError, match="НЕ дописаны"): opt.writer.shutdown(timeout=0.3) - - # прибираемся вручную: открываем шлюз, ждём дрен, останавливаем поток open_gate() - opt.writer.queue.join() - opt.writer.queue.put(opt.writer._stop_signal) - opt.writer.worker_thread.join(timeout=10) + assert opt.writer.shutdown(timeout=10.0) is True - assert_matches_reference(opt, ref, rows=np.arange(20)) +# ───────────────────────── F8 / C2: checkpoints ───────────────────────── -# ───────────────────────── F8: чекпоинты ───────────────────────── - -def test_f8_checkpoint_save_load_and_state_dict(tmp_path): +def test_c2_checkpoint_is_portable_and_self_contained(tmp_path): n, d = 24, 3 rng = np.random.default_rng(11) w0 = rng.normal(0, 0.1, (n, d)) - opt = make_opt(str(tmp_path), w0, lr=0.05, k=0.0) + opt = make_opt(str(tmp_path / "state"), w0, lr=0.05, k=0.0) rows = [2, 5, 11] for g_val in (1.0, -0.7): do_step(opt, rows, np.full((len(rows), d), g_val)) - - ckpt = opt.save_checkpoint(str(tmp_path / 'ckpt.json')) - assert os.path.exists(ckpt) - with open(ckpt) as f: + ckpt_dir = str(tmp_path / "checkpoint") + manifest_path = opt.export_checkpoint(ckpt_dir) + with open(manifest_path) as f: manifest = json.load(f) - assert manifest['format'] == 'DSA_checkpoint' - assert manifest['params']['ent']['shape'] == [n, d] + assert manifest["format"] == "DSA_checkpoint" + assert manifest["version"] == 2 + assert manifest["global_step"] == 2 + for info in manifest["params"].values(): + for rec in info["files"].values(): + assert os.path.basename(rec["name"]) == rec["name"] + assert not os.path.isabs(rec["name"]) opt.shutdown() - w1, m1, v1, t1 = read_state(opt) - # state_dict/load_state_dict — честные NotImplementedError, не мусор - with pytest.raises(NotImplementedError): - opt.state_dict() - with pytest.raises(NotImplementedError): - opt.load_state_dict({}) + # Move the checkpoint: absolute paths would break here. + moved = str(tmp_path / "moved") + shutil.copytree(ckpt_dir, moved) - # восстановление: resume=True + load_checkpoint - opt2 = make_opt(str(tmp_path), rng.normal(0, 9.0, (n, d)), - lr=0.05, k=0.0, resume=True) - opt2.load_checkpoint(ckpt) + opt2 = make_opt(str(tmp_path / "restored"), np.zeros((n, d)), lr=0.05, k=0.0) + opt2.import_checkpoint(moved) w2, m2, v2, t2 = read_state(opt2) assert np.array_equal(w1, w2) and np.array_equal(m1, m2) assert np.array_equal(v1, v2) and np.array_equal(t1, t2) - - # продолжение обучения с чекпоинта сходится с эталоном - ref = RefAdam.from_state(w1, m1, v1, t1, lr=0.05, k=0.0) + assert opt2._global_step == 2 + + ref = RiemannianAdamReference.__new__(RiemannianAdamReference) + ref.w, ref.m, ref.v, ref.t = (np.array(w1, dtype=np.float64), + np.array(m1, dtype=np.float64), + np.array(v1, dtype=np.float64), + np.array(t1, dtype=np.int64)) + ref.n, ref.d = n, d + ref.lr, ref.b1, ref.b2, ref.eps = 0.05, BETA1, BETA2, EPS + ref.k, ref.wd, ref.wd_mode = 0.0, 0.0, "l2" do_step(opt2, rows, np.full((len(rows), d), 3.0)) ref.step(rows, np.full((len(rows), d), 3.0)) opt2.shutdown() assert_matches_reference(opt2, ref, rows=rows) -def test_f8_load_checkpoint_foreign_paths_rejected(tmp_path): - n, d = 10, 2 +def test_c2_corrupted_checkpoint_rejected(tmp_path): + n, d = 6, 2 + opt = make_opt(str(tmp_path / "state"), np.zeros((n, d))) + ckpt_dir = str(tmp_path / "ckpt") + opt.export_checkpoint(ckpt_dir) + opt.shutdown() + + manifest = json.load(open(os.path.join(ckpt_dir, "checkpoint.json"))) + victim = os.path.join(ckpt_dir, + manifest["params"]["ent"]["files"]["w"]["name"]) + with open(victim, "r+b") as f: + f.seek(0) + f.write(b"\x01\x02\x03\x04") + + opt2 = make_opt(str(tmp_path / "restored"), np.zeros((n, d))) + before = np.array(opt2.state_files["ent"]["w"]) + with pytest.raises(RuntimeError, match="SHA-256"): + opt2.import_checkpoint(ckpt_dir) + after = np.array(opt2.state_files["ent"]["w"]) + assert np.array_equal(before, after), "неудачный импорт изменил живое состояние" + opt2.shutdown() + + +def test_c2_old_manifest_and_json_paths_rejected(tmp_path): + n, d = 6, 2 opt = make_opt(str(tmp_path), np.zeros((n, d))) + old = tmp_path / "old.json" + old.write_text(json.dumps({"format": "DSA_checkpoint", "version": 1, + "params": {}})) + with pytest.raises(RuntimeError, match="каталог"): + opt.import_checkpoint(str(old)) + with pytest.raises(ValueError, match="КАТАЛОГ"): + opt.save_checkpoint(str(old)) + opt.shutdown() + + +def test_c2_checkpoint_immediately_after_init_is_valid(tmp_path): + n, d, k = 5, 4, 1.0 + opt = DiskSparseRiemannianAdam( + {"ent": torch.full((n, d), 0.5)}, lr=0.01, k=k, + disk_dir=str(tmp_path / "state"), resume=False) + ckpt_dir = str(tmp_path / "ckpt") + opt.export_checkpoint(ckpt_dir) + opt.shutdown() + + opt2 = make_opt(str(tmp_path / "restored"), np.zeros((n, d)), k=k) + opt2.import_checkpoint(ckpt_dir) + w = np.array(opt2.state_files["ent"]["w"]) + radius = math.sqrt((1 - manifold.BALL_EPS) / k) + assert np.all(np.linalg.norm(w, axis=1) <= radius + 1e-6) + opt2.shutdown() + + +def test_state_dict_raises(tmp_path): + opt = make_opt(str(tmp_path), np.zeros((4, 2))) + with pytest.raises(NotImplementedError): + opt.state_dict() + with pytest.raises(NotImplementedError): + opt.load_state_dict({}) opt.shutdown() - foreign = { - 'format': 'DSA_checkpoint', - 'params': {'ent': {'shape': [n, d], - 'files': {'w': '/nonexistent/x_w.dat', - 'm': '/nonexistent/x_m.dat', - 'v': '/nonexistent/x_v.dat', - 't': '/nonexistent/x_t.dat'}}}, - } - bad = tmp_path / 'foreign.json' - bad.write_text(json.dumps(foreign)) - with pytest.raises(RuntimeError, match="скопируйте файлы"): - opt.load_checkpoint(str(bad)) -# ───────────────────────── совместимость публичного API ───────────────────────── +# ───────────────────────── C3: source tensors released ───────────────────────── + +def test_c3_source_tensor_is_not_retained(tmp_path): + source = torch.zeros(8, 4) + ref = weakref.ref(source) + opt = DiskSparseRiemannianAdam({"ent": source}, lr=0.01, k=0.0, + disk_dir=str(tmp_path), resume=False) + del source + gc.collect() + assert ref() is None, "оптимизатор удержал исходный тензор (O(N) память)" + assert opt.param_groups[0]["params"][0].numel() == 0 + opt.shutdown() + + +def test_c3_retain_params_opt_in(tmp_path): + source = torch.zeros(8, 4) + ref = weakref.ref(source) + opt = DiskSparseRiemannianAdam({"ent": source}, lr=0.01, k=0.0, + disk_dir=str(tmp_path), resume=False, + retain_params=True) + del source + gc.collect() + assert ref() is not None + assert opt.param_groups[0]["params"][0].shape == (8, 4) + opt.shutdown() + del opt + gc.collect() + assert ref() is None + + +# ───────────────────────── public API ───────────────────────── def test_public_api_compat(tmp_path): - """Поверхность, которую используют внешние интеграции: - конструктор (params dict, lr, k, disk_dir), state_files, writer._lock, - writer.queue.join(), param_groups, step(updates=...), shutdown().""" n, d = 10, 2 table = torch.randn(n, d) opt = DiskSparseRiemannianAdam( - params={'ent': table}, lr=0.01, k=0.5, disk_dir=str(tmp_path)) - assert hasattr(opt, 'state_files') and 'ent' in opt.state_files - assert {'w', 'm', 'v'} <= set(opt.state_files['ent']) - assert hasattr(opt.writer, '_lock') and hasattr(opt.writer, 'queue') - opt.writer.queue.join() # wait_idle() обёртки + params={"ent": table}, lr=0.01, k=0.5, disk_dir=str(tmp_path)) + assert hasattr(opt, "state_files") and "ent" in opt.state_files + assert {"w", "m", "v", "t"} <= set(opt.state_files["ent"]) + assert hasattr(opt.writer, "_lock") and hasattr(opt.writer, "queue") + opt.writer.queue.join() - opt.param_groups[0]['lr'] *= 0.5 - assert opt.param_groups[0]['lr'] == 0.005 + opt.param_groups[0]["lr"] *= 0.5 + assert opt.param_groups[0]["lr"] == 0.005 idx = torch.tensor([1, 2]) - opt.step(updates={'ent': (idx, torch.ones(2, d))}) + opt.step(updates={"ent": (idx, torch.ones(2, d))}) opt.zero_grad() opt.shutdown() - # список параметров тоже поддерживается (как в исходном API) opt2 = DiskSparseRiemannianAdam([torch.randn(6, 3)], lr=0.01, k=0.0, - disk_dir=str(tmp_path / 'b')) - opt2.step(updates={'param_0': (torch.tensor([0]), torch.ones(1, 3))}) + disk_dir=str(tmp_path / "b")) + opt2.step(updates={"param_0": (torch.tensor([0]), torch.ones(1, 3))}) opt2.shutdown() -# ───────────────────────── R1: NaN/inf строки в файле состояния ───────────────────────── +# ───────────────────────── R1: non-finite rows ───────────────────────── -def test_r1_nan_row_is_sanitized_on_step(tmp_path, caplog): - """Битая строка w (NaN/inf) не проходит проекцию молча (torch.where при - NaN даёт False) и расползлась бы в m/v и обучение. После одного шага - в w/m/v нет NaN/inf, строка внутри шара; в логе warning, счётчик растёт.""" - import logging +def test_r1_nan_row_is_reset_on_step(tmp_path, caplog): n, d, k = 8, 4, 1.0 - max_radius = np.sqrt((1.0 - 1e-5) / k) + max_radius = math.sqrt((1.0 - manifold.BALL_EPS) / k) w0 = np.full((n, d), 0.1) opt = make_opt(str(tmp_path), w0, lr=0.01, k=k) opt.shutdown() - # портим файл w напрямую (имитация битого файла): строка 2 = NaN, 3 = inf - with open(os.path.join(str(tmp_path), 'ent_w.dat'), 'r+b') as f: + with open(os.path.join(str(tmp_path), "ent_w.dat"), "r+b") as f: f.seek(2 * d * 4) f.write(np.full(d, np.nan, dtype=np.float32).tobytes()) f.seek(3 * d * 4) @@ -752,59 +1220,47 @@ def test_r1_nan_row_is_sanitized_on_step(tmp_path, caplog): w, m, v, t = read_state(opt2) for state in (w, m, v): - assert np.isfinite(state[[2, 3]]).all(), "NaN/inf выжил после шага" + assert np.isfinite(state[[2, 3]]).all() assert np.linalg.norm(w[2]) <= max_radius + 1e-6 assert np.linalg.norm(w[3]) <= max_radius + 1e-6 assert t[2] == 1 and t[3] == 1 - assert opt2._nan_rows_replaced == 2, "счётчик не-конечных строк не вырос" - assert any("NaN" in r.message for r in caplog.records), "нет warning в логе" + assert opt2._nan_rows_replaced == 2 + assert any("NaN" in r.message for r in caplog.records) -# ───────────────────────── R2: маркер переинициализации ───────────────────────── +# ───────────────────────── R2: reinit marker ───────────────────────── def test_r2_reinit_marker_blocks_resume(tmp_path): - """Смешанное состояние после смерти процесса посреди переинициализации - (маркер 'reinit_in_progress' в meta + частично перезаписанные файлы, - воспроизведённые руками): resume=True обязан дать явную ошибку, а не - принять его молча.""" n, d = 8, 3 opt = make_opt(str(tmp_path), np.zeros((n, d))) do_step(opt, [1, 2], np.ones((2, d))) opt.shutdown() - meta_path = os.path.join(str(tmp_path), 'ent_meta.json') - with open(meta_path) as f: - meta = json.load(f) - meta['reinit_in_progress'] = True - with open(meta_path, 'w', encoding='utf-8') as f: + meta_path = os.path.join(str(tmp_path), "ent_meta.json") + meta = json.load(open(meta_path)) + meta["reinit_in_progress"] = True + with open(meta_path, "w", encoding="utf-8") as f: json.dump(meta, f) - with open(os.path.join(str(tmp_path), 'ent_w.dat'), 'r+b') as f: - f.seek(0) - f.write(np.full((n, d), 0.5, dtype=np.float32).tobytes()) with pytest.raises(RuntimeError, match="reinit_in_progress"): make_opt(str(tmp_path), np.zeros((n, d)), resume=True) def test_r2_reinit_marker_cleared_by_full_reinit(tmp_path): - """resume=False — документированный путь восстановления: полная - переинициализация пересоздаёт файлы и снимает маркер; после неё - resume=True снова работает.""" n, d = 8, 3 opt = make_opt(str(tmp_path), np.zeros((n, d))) opt.shutdown() - meta_path = os.path.join(str(tmp_path), 'ent_meta.json') - with open(meta_path) as f: - meta = json.load(f) - meta['reinit_in_progress'] = True - with open(meta_path, 'w', encoding='utf-8') as f: + meta_path = os.path.join(str(tmp_path), "ent_meta.json") + meta = json.load(open(meta_path)) + meta["reinit_in_progress"] = True + with open(meta_path, "w", encoding="utf-8") as f: json.dump(meta, f) - fresh = rng_vals(13, n, d) - opt2 = make_opt(str(tmp_path), fresh, lr=0.05, k=0.0) # resume=False + fresh = np.random.default_rng(13).normal(0, 0.3, (n, d)) + opt2 = make_opt(str(tmp_path), fresh, lr=0.05, k=0.0, resume=False) opt2.shutdown() with open(meta_path) as f: - assert 'reinit_in_progress' not in json.load(f), "маркер не снят" + assert "reinit_in_progress" not in json.load(f) w, m, v, t = read_state(opt2) assert np.allclose(w, fresh, atol=1e-7) assert np.all(m == 0) and np.all(v == 0) and np.all(t == 0) @@ -813,34 +1269,29 @@ def test_r2_reinit_marker_cleared_by_full_reinit(tmp_path): opt3.shutdown() -# ───────────────────────── R3: неизвестный ключ updates ───────────────────────── +# ───────────────────────── R3: unknown update key ───────────────────────── def test_r3_unknown_updates_key_warns_and_counts(tmp_path, caplog): - """Ключ-опечатка в updates не отбрасывается молча: warning в логе, - счётчик растёт, чужие строки не изменяются.""" - import logging n, d = 6, 3 opt = make_opt(str(tmp_path), np.full((n, d), 0.1)) before = opt._unknown_key_warnings with caplog.at_level(logging.WARNING, logger="DSA"): - do_step(opt, [1], np.ones((1, d))) # валидный ключ - opt.step(updates={'entt': (torch.tensor([2]), torch.ones(1, d))}) + do_step(opt, [1], np.ones((1, d))) + opt.step(updates={"entt": (torch.tensor([2]), torch.ones(1, d))}) opt.shutdown() - assert opt._unknown_key_warnings == before + 1, "счётчик не вырос" - assert any("entt" in r.message for r in caplog.records), "нет warning в логе" + assert opt._unknown_key_warnings == before + 1 + assert any("entt" in r.message for r in caplog.records) w, *_ = read_state(opt) - assert np.allclose(w[2], 0.1, atol=1e-7), "строка 2 не должна была измениться" + assert np.allclose(w[2], 0.1, atol=1e-7) def test_r3_unknown_key_raises_in_debug_mode(tmp_path, monkeypatch): - """В отладочном режиме (dsa.optimizer.DEBUG, env DSA_DEBUG=1) неизвестный - ключ updates — исключение KeyError, а не warning.""" import dsa.optimizer as dsa_opt n, d = 6, 3 opt = make_opt(str(tmp_path), np.zeros((n, d))) - monkeypatch.setattr(dsa_opt, 'DEBUG', True) + monkeypatch.setattr(dsa_opt, "DEBUG", True) with pytest.raises(KeyError, match="entt"): - opt.step(updates={'entt': (torch.tensor([2]), torch.ones(1, d))}) - assert opt._unknown_key_warnings == 0 # до warning дело не дошло + opt.step(updates={"entt": (torch.tensor([2]), torch.ones(1, d))}) + assert opt._unknown_key_warnings == 0 opt.shutdown() diff --git a/tests/test_geoopt_golden.py b/tests/test_geoopt_golden.py new file mode 100644 index 0000000..09f1273 --- /dev/null +++ b/tests/test_geoopt_golden.py @@ -0,0 +1,52 @@ +"""Golden values produced by geoopt's stereographic Poincare math. + +The reference numbers below were generated with the actual +``geoopt/manifolds/stereographic/math.py`` (master) using its convention +``k = -c`` (geoopt's signed curvature) against the DSA convention ``c > 0``. +They are hard-coded so the repository does not need geoopt installed and so a +future change in ``dsa.manifold`` is checked against an external oracle, not +only against another formula written in this repository. +""" + +import torch + +from dsa import manifold + + +C = 0.7 +X = torch.tensor([[0.2, -0.1]], dtype=torch.float64) +Y = torch.tensor([[-0.3, 0.4]], dtype=torch.float64) +U = torch.tensor([[0.5, -0.2]], dtype=torch.float64) + + +def test_lambda_matches_geoopt(): + assert torch.allclose(manifold.lambda_x(X, C), + torch.tensor([[2.072538860103627]], + dtype=torch.float64), + atol=1e-14) + + +def test_mobius_add_matches_geoopt(): + expected = torch.tensor([[-0.09525184009236541, 0.3261653918314331]], + dtype=torch.float64) + assert torch.allclose(manifold.mobius_add(X, Y, C), expected, atol=1e-14) + + +def test_expmap_matches_geoopt(): + expected = torch.tensor([[0.6304294392917015, -0.27530547641281544]], + dtype=torch.float64) + assert torch.allclose(manifold.expmap(X, U, C), expected, atol=1e-12) + + +def test_parallel_transport_matches_geoopt(): + expected = torch.tensor([[0.41537343956347617, -0.19855381136144154]], + dtype=torch.float64) + assert torch.allclose( + manifold.parallel_transport(X, Y, U, C), expected, atol=1e-14) + + +def test_distance_matches_geoopt(): + assert torch.allclose(manifold.distance(X, Y, C), + torch.tensor([[1.487174390441505]], + dtype=torch.float64), + atol=1e-14) diff --git a/tests/test_manifold.py b/tests/test_manifold.py new file mode 100644 index 0000000..e63adf4 --- /dev/null +++ b/tests/test_manifold.py @@ -0,0 +1,167 @@ +"""Geometric properties of the Poincare-ball implementation. + +These tests do not compare two copies of the same formula: they check the +defining properties of the exponential map, parallel transport and the +Riemannian gradient. The same properties are checked for the independent +float64 reference used in ``test_fixes.py``. +""" + +import math + +import numpy as np +import pytest +import torch + +from dsa import manifold +from tests import reference + + +K = 0.7 +D = 6 + + +def _inside(rows, radius_fraction=0.8): + radius = math.sqrt(1.0 / K) + x = np.asarray(rows, dtype=np.float64) + scale = min(1.0, radius_fraction * radius / np.linalg.norm(x, axis=-1).max()) + return x * scale + + +@pytest.fixture(scope="module") +def points(): + rng = np.random.default_rng(0) + x = _inside(rng.normal(0, 0.3, (5, D))) + y = _inside(rng.normal(0, 0.3, (5, D))) + u = rng.normal(0, 0.2, (5, D)) + v = rng.normal(0, 0.5, (5, D)) + return x, y, u, v + + +def _torch(a): + return torch.tensor(a, dtype=torch.float64) + + +def test_expmap_zero_is_identity(points): + x, _, _, _ = points + out = manifold.expmap(_torch(x), torch.zeros_like(_torch(x)), K).numpy() + assert np.allclose(out, x, atol=1e-14) + + +def test_expmap_stays_inside_ball(points): + x, _, u, _ = points + out = manifold.expmap(_torch(x), _torch(u), K).numpy() + radius = math.sqrt((1.0 - manifold.BALL_EPS) / K) + assert np.all(np.linalg.norm(out, axis=-1) <= radius + 1e-12) + + +def test_expmap_distance_equals_tangent_norm(points): + """dist(x, exp_x(u)) == lambda(x) * ||u|| for small u.""" + x, _, u, _ = points + out = manifold.expmap(_torch(x), _torch(u), K) + lhs = manifold.distance(_torch(x), out, K).numpy().ravel() + rhs = (manifold.lambda_x(_torch(x), K) + * torch.linalg.norm(_torch(u), dim=-1, keepdim=True)).numpy().ravel() + assert np.allclose(lhs, rhs, rtol=1e-12, atol=1e-12) + + +def test_log_exp_roundtrip(points): + x, _, u, _ = points + y = manifold.expmap(_torch(x), _torch(u), K) + back = reference.logmap(x, y.numpy(), K) + assert np.allclose(back, u, atol=1e-12) + + +def test_parallel_transport_identity(points): + x, _, _, v = points + out = manifold.parallel_transport(_torch(x), _torch(x), _torch(v), K).numpy() + assert np.allclose(out, v, atol=1e-14) + + +def test_parallel_transport_is_isometry(points): + x, y, _, v = points + out = manifold.parallel_transport(_torch(x), _torch(y), _torch(v), K) + lhs = (manifold.lambda_x(_torch(y), K) + * torch.linalg.norm(out, dim=-1, keepdim=True)).numpy().ravel() + rhs = (manifold.lambda_x(_torch(x), K) + * torch.linalg.norm(_torch(v), dim=-1, keepdim=True)).numpy().ravel() + assert np.allclose(lhs, rhs, rtol=1e-12, atol=1e-12) + + +def test_parallel_transport_roundtrip(points): + x, y, _, v = points + forward = manifold.parallel_transport(_torch(x), _torch(y), _torch(v), K) + back = manifold.parallel_transport(_torch(y), _torch(x), forward, K).numpy() + assert np.allclose(back, v, atol=1e-12) + + +def test_parallel_transport_linear(points): + x, y, _, v = points + a = 0.37 + lhs = manifold.parallel_transport(_torch(x), _torch(y), _torch(a * v), K).numpy() + rhs = a * manifold.parallel_transport(_torch(x), _torch(y), _torch(v), K).numpy() + assert np.allclose(lhs, rhs, atol=1e-13) + + +def test_riemannian_gradient_is_inverse_metric(points): + x, _, _, _ = points + g = _torch(np.ones_like(x)) + rg = manifold.riemannian_gradient(g, _torch(x), K) + lam = manifold.lambda_x(_torch(x), K) + assert np.allclose(rg.numpy(), (g / lam ** 2).numpy(), atol=1e-14) + + +def test_projection_and_distance_validity(points): + x, y, _, _ = points + radius = math.sqrt((1.0 - manifold.BALL_EPS) / K) + outside = torch.full((2, D), 10.0) + projected = manifold.project_to_ball(outside, K) + assert torch.all(torch.linalg.norm(projected, dim=-1) <= radius + 1e-12) + assert np.isfinite(manifold.distance( + projected[:1], projected[1:], K).numpy()).all() + # an out-of-ball point yields NaN, not a silently wrong distance + assert np.isnan(manifold.distance( + outside[:1], _torch(y[:1]), K).numpy()).all() + + +def test_reference_satisfies_same_properties(points): + x, y, u, v = points + e = reference.expmap(x, u, K) + assert np.allclose(reference.expmap(x, np.zeros_like(x), K), x, atol=1e-14) + lhs = reference.distance(x, e, K).ravel() + rhs = reference._lambda(x, K).ravel() * np.linalg.norm(u, axis=-1) + assert np.allclose(lhs, rhs, rtol=1e-12, atol=1e-12) + assert np.allclose(reference.logmap(x, e, K), u, atol=1e-12) + + pt = reference.parallel_transport(x, y, v, K) + assert np.allclose( + reference._lambda(y, K).ravel() * np.linalg.norm(pt, axis=-1), + reference._lambda(x, K).ravel() * np.linalg.norm(v, axis=-1), + rtol=1e-12, atol=1e-12) + assert np.allclose(reference.parallel_transport(y, x, pt, K), v, atol=1e-12) + + +def test_euclidean_branches_are_identity(points): + """k=0 must reduce every geometry helper to its Euclidean counterpart.""" + x, y, u, v = points + tx, ty, tu, tv = _torch(x), _torch(y), _torch(u), _torch(v) + assert torch.allclose(manifold.project_to_ball(tx, 0.0), tx) + assert np.allclose(manifold.project_to_ball_np(x, 0.0), x) + assert torch.allclose(manifold.mobius_add(tx, ty, 0.0), tx + ty) + assert torch.allclose(manifold.gyration(tx, ty, tv, 0.0), tv) + assert torch.allclose(manifold.expmap(tx, tu, 0.0), tx + tu) + assert torch.allclose(manifold.parallel_transport(tx, ty, tv, 0.0), tv) + assert torch.allclose(manifold.riemannian_gradient(tv, tx, 0.0), tv) + assert np.allclose(manifold.distance(tx, ty, 0.0).numpy(), + np.linalg.norm(x - y, axis=-1)) + + +def test_is_in_ball_mask(): + inside = torch.tensor([[0.0, 0.0], [0.5, 0.0]], dtype=torch.float64) + outside = torch.tensor([[2.0, 0.0]], dtype=torch.float64) + nan = torch.tensor([[float("nan"), 0.0]], dtype=torch.float64) + assert manifold.is_in_ball(inside, 1.0).tolist() == [True, True] + assert manifold.is_in_ball(outside, 1.0).tolist() == [False] + assert manifold.is_in_ball(nan, 1.0).tolist() == [False] + assert manifold.is_in_ball(outside, 0.0).tolist() == [True] + assert manifold.is_in_ball(nan, 0.0).tolist() == [False] + diff --git a/tests/test_packaging.py b/tests/test_packaging.py new file mode 100644 index 0000000..d6f5c02 --- /dev/null +++ b/tests/test_packaging.py @@ -0,0 +1,49 @@ +"""C5: one metadata source and consistent version/requirements.""" + +import os +import re + +import dsa + + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def _pyproject(): + path = os.path.join(ROOT, "pyproject.toml") + text = open(path, encoding="utf-8").read() + try: + import tomllib + return tomllib.loads(text), text + except ModuleNotFoundError: # Python < 3.11 + return None, text + + +def test_setup_py_removed(): + assert not os.path.exists(os.path.join(ROOT, "setup.py")), ( + "setup.py — второй источник метаданных; оставьте только pyproject.toml") + + +def test_version_matches_package(): + data, _ = _pyproject() + if data is not None: + 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) + + +def test_torch_floor_matches_readme_badge(): + data, text = _pyproject() + if data is not None: + deps = data["project"]["dependencies"] + torch_req = next(d for d in deps if d.startswith("torch>=")) + assert torch_req == "torch>=2.0.0" + readme = open(os.path.join(ROOT, "README.md"), encoding="utf-8").read() + assert "pytorch-2.0+" in readme + + +def test_python_requires_is_declared(): + data, _ = _pyproject() + if data is not None: + assert data["project"]["requires-python"] diff --git a/tests/test_validation.py b/tests/test_validation.py new file mode 100644 index 0000000..d240868 --- /dev/null +++ b/tests/test_validation.py @@ -0,0 +1,182 @@ +"""Validation and corruption paths that the behavioural suite does not hit.""" + +import json +import os + +import numpy as np +import pytest +import torch + +from dsa.optimizer import DiskSparseRiemannianAdam + + +def _opt(tmp_path, **kwargs): + kwargs.setdefault("disk_dir", str(tmp_path)) + return DiskSparseRiemannianAdam({"e": torch.zeros(3, 2)}, **kwargs) + + +@pytest.mark.parametrize("kwargs,exc", [ + ({"lr": -1}, ValueError), + ({"beta1": 1.0}, ValueError), + ({"beta2": -0.1}, ValueError), + ({"eps": 0.0}, ValueError), + ({"weight_decay": -1.0}, ValueError), + ({"max_queue_size": 0}, ValueError), + ({"put_timeout": -1.0}, ValueError), + ({"init_max_norm_fraction": 0.0}, ValueError), + ({"init_max_norm_fraction": 1.1}, ValueError), + ({"weight_decay_mode": "adamw"}, ValueError), + ({"k": -1.0}, ValueError), + ({"k": float("nan")}, ValueError), + ({"k": float("inf")}, ValueError), + ({"resume": "yes"}, TypeError), +]) +def test_constructor_rejects_invalid_arguments(tmp_path, kwargs, exc): + with pytest.raises(exc): + _opt(tmp_path, **kwargs) + + +def test_constructor_rejects_bad_params(tmp_path): + with pytest.raises(ValueError, match="empty"): + DiskSparseRiemannianAdam({}, disk_dir=str(tmp_path)) + with pytest.raises(ValueError, match="2-D"): + DiskSparseRiemannianAdam({"e": torch.zeros(3)}, disk_dir=str(tmp_path)) + with pytest.raises(TypeError, match="floating"): + DiskSparseRiemannianAdam({"e": torch.zeros(3, 2, dtype=torch.int64)}, + disk_dir=str(tmp_path)) + with pytest.raises(TypeError, match="Tensor"): + DiskSparseRiemannianAdam({"e": np.zeros((3, 2))}, disk_dir=str(tmp_path)) + with pytest.raises(ValueError, match="path separators"): + DiskSparseRiemannianAdam({"a/b": torch.zeros(3, 2)}, + disk_dir=str(tmp_path)) + + +def test_nonfinite_initial_weights_rejected(tmp_path): + with pytest.raises(ValueError, match="NaN/inf"): + DiskSparseRiemannianAdam({"e": torch.full((3, 2), float("nan"))}, + disk_dir=str(tmp_path)) + + +def test_resume_true_on_empty_directory_raises(tmp_path): + with pytest.raises(RuntimeError, match="нет состояния"): + _opt(tmp_path, resume=True) + + +def _valid_state(tmp_path): + opt = _opt(tmp_path) + opt.shutdown() + return os.path.join(str(tmp_path), "e_meta.json") + + +@pytest.mark.parametrize("mutator,match", [ + (lambda m: m.__setitem__("format_version", 99), "версия формата"), + (lambda m: m.__setitem__("algorithm", "other"), "алгоритм"), + (lambda m: m.__setitem__("k", 0.5), "кривизна"), + (lambda m: m.__setitem__("dtype", "float64"), "типы данных"), +]) +def test_corrupt_meta_is_rejected(tmp_path, mutator, match): + meta_path = _valid_state(tmp_path) + meta = json.load(open(meta_path)) + mutator(meta) + with open(meta_path, "w") as f: + json.dump(meta, f) + with pytest.raises(RuntimeError, match=match): + _opt(tmp_path, resume=True) + + +def test_unreadable_meta_is_rejected(tmp_path): + meta_path = _valid_state(tmp_path) + with open(meta_path, "w") as f: + f.write("{not json") + with pytest.raises(RuntimeError, match="не читается"): + _opt(tmp_path, resume=True) + + +@pytest.mark.parametrize("mutator,match", [ + (lambda m: m.__setitem__("format_version", 99), "глобальный meta"), + (lambda m: m.__setitem__("global_step", "many"), "global_step"), +]) +def test_corrupt_global_meta_is_rejected(tmp_path, mutator, match): + _valid_state(tmp_path) + path = os.path.join(str(tmp_path), "dsa_global_meta.json") + meta = json.load(open(path)) + mutator(meta) + with open(path, "w") as f: + json.dump(meta, f) + with pytest.raises((RuntimeError, ValueError), match=match): + _opt(tmp_path, resume=True) + + +def test_sparse_gradient_is_rejected(tmp_path): + opt = _opt(tmp_path) + idx = torch.tensor([0]) + sparse = torch.sparse_coo_tensor( + torch.tensor([[0], [0]]), torch.tensor([1.0]), size=(1, 2)) + with pytest.raises(TypeError, match="разреженный"): + opt.step(updates={"e": (idx, sparse)}) + opt.shutdown() + + +def test_updates_contract_and_numpy_inputs(tmp_path): + opt = _opt(tmp_path) + with pytest.raises(TypeError, match="dict"): + opt.step(updates=[]) + with pytest.raises(TypeError, match="indices, grads"): + opt.step(updates={"e": 5}) + with pytest.raises(ValueError, match="1-D"): + opt.step(updates={"e": (np.array([[0]]), + np.ones((1, 2), dtype=np.float32))}) + with pytest.raises(TypeError, match="целочисленными"): + opt.step(updates={"e": (np.array([0.0]), + np.ones((1, 2), dtype=np.float32))}) + # numpy indices and gradients are accepted + opt.step(updates={"e": (np.array([0]), np.ones((1, 2), dtype=np.float32))}) + opt.step(updates={"e": (torch.tensor([1]), np.ones((1, 2), dtype=np.float32))}) + w = np.array(opt.state_files["e"]["w"]) + assert np.isfinite(w).all() + opt.shutdown() + + +def test_zero_initial_weights_hyperbolic(tmp_path): + opt = DiskSparseRiemannianAdam( + {"e": torch.zeros(3, 2)}, k=1.0, disk_dir=str(tmp_path)) + assert np.all(np.array(opt.state_files["e"]["w"]) == 0) + opt.shutdown() + + +def test_writer_timeout_and_shutdown_validation(tmp_path): + import threading + opt = _opt(tmp_path) + gate = threading.Event() + orig = opt.writer._apply_write + + def gated(*args, **kwargs): + gate.wait() + return orig(*args, **kwargs) + + opt.writer._apply_write = gated + opt.step(updates={"e": (torch.tensor([0]), torch.ones(1, 2))}) + with pytest.raises(RuntimeError, match="очередь не опустела"): + opt.writer.wait_all_flushed(timeout=0.2) + with pytest.raises(TypeError, match="timeout"): + opt.writer.shutdown(timeout=None) + with pytest.raises(ValueError, match="timeout"): + opt.writer.shutdown(timeout=-1) + gate.set() + opt.shutdown() + + +def test_malformed_checkpoint_manifest_is_rejected(tmp_path): + opt = _opt(tmp_path) + ckpt = tmp_path / "ckpt" + opt.export_checkpoint(str(ckpt)) + opt.shutdown() + + manifest_path = ckpt / "checkpoint.json" + manifest = json.loads(manifest_path.read_text()) + manifest["params"] = [] + manifest_path.write_text(json.dumps(manifest)) + opt2 = _opt(tmp_path / "new") + with pytest.raises(RuntimeError, match="params"): + opt2.import_checkpoint(str(ckpt)) + opt2.shutdown()