Match GPU NMF denominator handling to sklearn

README.md

@@ -22,6 +22,25 @@ pip install harmonypy
 pip install scikit-misc
 ```
 
+# GPU NMF parity tests
+
+The regular GPU-NMF test suite compares sklearn Frobenius MU with the PyTorch
+kernel in fp64 and fp32 using identical inputs, initialization seeds, and MU
+iteration counts. CUDA cases run when a CUDA device is available:
+
+```bash
+pytest -q tests/test_nmf_gpu.py -k sklearn_mu_matches
+```
+
+The 100,000-cell by 20,000-gene cases are opt-in because the dense input alone
+uses 7.45 GiB in fp32 or 14.90 GiB in fp64. They run one rank-2 MU iteration
+and check host/device memory before allocating:
+
+```bash
+CNMF_RUN_LARGE_GPU_PARITY=1 \
+pytest -q tests/test_nmf_gpu.py -k 100k_by_20k
+```
+
 # Running cNMF
 
 cNMF can be run from the command line without any parallelization using the example commands below:

src/cnmf/nmf_gpu.py

@@ -35,12 +35,15 @@
 }
 
 
+_SKLEARN_EPSILON = float(np.finfo(np.float32).eps)
+
+
 DEFAULT_GPU = {
     "device": "auto",
     "dtype": "auto",
     "allow_tf32": False,
     "compile": False,
-    "eps": 1e-9,
+    "eps": _SKLEARN_EPSILON,
     "check_every": 10,
     "compile_block": 1,
     "batch": 1,             # factorize replicates per launch
@@ -60,7 +63,7 @@ def parse_gpu_args(parser):
     group.add_argument("--gpu-dtype", type=str.lower, choices=["auto", "fp32", "fp64", "bf16"], help="[factorize,consensus,gpu] Storage and matmul dtype for GPU NMF (default auto)")
     group.add_argument("--gpu-allow-tf32", action="store_const", const=True, help="[factorize,consensus,gpu] Allow TF32 for CUDA fp32 matrix multiplication")
     group.add_argument("--gpu-compile", action="store_const", const=True, help="[factorize,consensus,gpu] Enable torch.compile for the GPU NMF update step")
-    group.add_argument("--gpu-eps", type=float, help="[factorize,consensus,gpu] Multiplicative-update denominator guard")
+    group.add_argument("--gpu-eps", type=float, help="[factorize,consensus,gpu] Replacement for exactly-zero MU denominators")
     group.add_argument("--gpu-check-every", type=int, help="[factorize,consensus,gpu] Eager-mode convergence check interval")
     group.add_argument("--gpu-compile-block", type=int, help="[factorize,consensus,gpu] Number of MU iterations per compiled block")
     group.add_argument("--gpu-batch", type=int, help="[factorize] Replicates run per GPU launch (batched MU); 1 = single-replicate")
@@ -310,16 +313,22 @@ def _mu_step(W, H, Xg, eps):
     tensors with shared `Xg[1,n,g]`. Operations are out-of-place for compile.
     """
     Ht = H.transpose(-2, -1)                               # [g,k] or [R,g,k]
-    W = W * ((Xg @ Ht) / (W @ (H @ Ht) + eps))            # W *= XHᵀ / (W·HHᵀ)   (uses old H)
+    denominator = W @ (H @ Ht)
+    denominator = denominator.where(denominator != 0, eps)
+    W = W * ((Xg @ Ht) / denominator)                      # W *= XHᵀ / (W·HHᵀ)   (uses old H)
     Wt = W.transpose(-2, -1)                               # [k,n] or [R,k,n]
-    H = H * ((Wt @ Xg) / ((Wt @ W) @ H + eps))            # H *= WᵀX / (WᵀW·H)   (uses new W)
+    denominator = (Wt @ W) @ H
+    denominator = denominator.where(denominator != 0, eps)
+    H = H * ((Wt @ Xg) / denominator)                      # H *= WᵀX / (WᵀW·H)   (uses new W)
     return W, H
 
 
 def _mu_step_fixed_h(W, H, Xg, eps):
     """One fixed-H MU update; only W changes. Supports 2D or stacked W."""
     Ht = H.transpose(-2, -1)                               # [g,k] or [1,g,k]
-    return W * ((Xg @ Ht) / (W @ (H @ Ht) + eps))
+    denominator = W @ (H @ Ht)
+    denominator = denominator.where(denominator != 0, eps)
+    return W * ((Xg @ Ht) / denominator)
 
 
 # ---------------------------------------------------------------------
@@ -438,7 +447,7 @@ def _to_device_factors(torch, W0, H0, dtype, device):
 
 
 def _to_device_eps(torch, eps, dtype, device):
-    """Create the MU denominator guard on runtime dtype/device."""
+    """Create the exact-zero denominator replacement on runtime dtype/device."""
     return torch.tensor(eps, dtype=dtype, device=device)
 
 

tests/test_nmf_gpu.py

@@ -39,11 +39,20 @@
     Uses fake backends for portable device/dtype policy checks, verifies loud
     dependency errors, covers the current sparse densify path, and keeps
     CUDA-only fp32/bf16/TF32 behavior behind device-gated tests.
+
+sklearn MU parity
+    Compares final aligned W/H factors against sklearn Frobenius MU for fp64
+    and fp32 with identical initialization, seed, and iteration count. Small
+    cases run routinely on the torch CPU backend and CUDA when available. An
+    opt-in, memory-gated CUDA stress case uses a 100,000 x 20,000 matrix.
 """
 
-from types import SimpleNamespace
 import argparse
 import builtins
+import gc
+import os
+import warnings
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
@@ -165,14 +174,42 @@ def test_mu_step_updates_w_first_using_old_h_then_h_using_new_w(kernel):
     H0 = torch.tensor([[0.6, 0.8], [1.0, 1.2]], dtype=torch.float64)
     eps = torch.tensor(1e-9, dtype=torch.float64)
 
-    expected_W = W0 * ((Xg @ H0.T) / (W0 @ (H0 @ H0.T) + eps))
-    expected_H = H0 * ((expected_W.T @ Xg) / ((expected_W.T @ expected_W) @ H0 + eps))
+    denominator = W0 @ (H0 @ H0.T)
+    denominator = denominator.where(denominator != 0, eps)
+    expected_W = W0 * ((Xg @ H0.T) / denominator)
+    denominator = (expected_W.T @ expected_W) @ H0
+    denominator = denominator.where(denominator != 0, eps)
+    expected_H = H0 * ((expected_W.T @ Xg) / denominator)
     W, H = kernel._mu_step(W0, H0, Xg, eps)
 
     assert torch.allclose(W, expected_W)
     assert torch.allclose(H, expected_H)
 
 
+def test_mu_step_fixed_h_matches_sklearn_exact_zero_protection(kernel):
+    """Replace exact zeros while leaving tiny nonzero denominators untouched."""
+    torch = require_nmf_runtime()
+    eps = torch.tensor(np.finfo(np.float32).eps, dtype=torch.float64)
+
+    zero_result = kernel._mu_step_fixed_h(
+        torch.ones((1, 1), dtype=torch.float64),
+        torch.zeros((1, 1), dtype=torch.float64),
+        torch.ones((1, 1), dtype=torch.float64),
+        eps,
+    )
+    assert torch.equal(zero_result, torch.zeros_like(zero_result))
+    assert torch.isfinite(zero_result).all()
+
+    tiny = torch.tensor(1e-12, dtype=torch.float64)
+    tiny_result = kernel._mu_step_fixed_h(
+        torch.ones((1, 1), dtype=torch.float64),
+        tiny.reshape(1, 1),
+        torch.ones((1, 1), dtype=torch.float64),
+        eps,
+    )
+    assert torch.equal(tiny_result, (1 / tiny).reshape(1, 1))
+
+
 def test_fit_mu_early_stops_when_relative_error_drop_is_below_tol(kernel):
     """Confirm the MU loop stops after a flat relative-error drop crosses the tolerance."""
     torch = require_nmf_runtime()
@@ -340,6 +377,191 @@ def test_nmf_gpu_custom_init_raises(kernel):
         kernel._init_wh(small_nonnegative_matrix(), 2, 0, "custom")
 
 
+# ---------------------------------------------------------------------
+# sklearn Frobenius-MU parity: same input, initialization, seed, iterations
+# ---------------------------------------------------------------------
+PARITY_CASES = [
+    pytest.param("fp64", np.float64, 2e-8, 1e-9, id="fp64"),
+    pytest.param("fp32", np.float32, 2e-5, 2e-6, id="fp32"),
+]
+
+LARGE_PARITY_ENV = "CNMF_RUN_LARGE_GPU_PARITY"
+LARGE_PARITY_ROWS = 100_000
+LARGE_PARITY_COLUMNS = 20_000
+LARGE_PARITY_COMPONENTS = 2
+
+
+def _parity_nmf_kwargs(n_components, seed, max_iter):
+    """Return the common, unregularized sklearn/GPU Frobenius-MU contract."""
+    return {
+        "n_components": n_components,
+        "init": "random",
+        "random_state": seed,
+        "solver": "mu",
+        "beta_loss": "frobenius",
+        "tol": 0.0,
+        "max_iter": max_iter,
+        "alpha_W": 0.0,
+        "alpha_H": 0.0,
+        "l1_ratio": 0.0,
+    }
+
+
+def _sklearn_mu_reference(X, nmf_kwargs):
+    """Run sklearn MU and require the requested fixed iteration count."""
+    pytest.importorskip("sklearn")
+    from sklearn.decomposition import non_negative_factorization
+    from sklearn.exceptions import ConvergenceWarning
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", ConvergenceWarning)
+        W, H, n_iter = non_negative_factorization(X, **nmf_kwargs)
+    assert n_iter == nmf_kwargs["max_iter"]
+    return H, W
+
+
+def _gpu_parity_result(kernel, X, nmf_kwargs, dtype_name, device):
+    """Run the PyTorch kernel without TF32, compilation, or early stopping."""
+    return kernel._nmf_gpu(
+        X,
+        nmf_kwargs,
+        {
+            "device": device,
+            "dtype": dtype_name,
+            "allow_tf32": False,
+            "compile": False,
+            # One convergence block means every requested MU iteration runs
+            # before the first possible tolerance check.
+            "check_every": nmf_kwargs["max_iter"],
+        },
+    )
+
+
+def _relative_reconstruction_error(X, H, W):
+    """Compute reconstruction error in float64, independent of compute dtype."""
+    X64 = np.asarray(X, dtype=np.float64)
+    H64 = np.asarray(H, dtype=np.float64)
+    W64 = np.asarray(W, dtype=np.float64)
+    return np.linalg.norm(X64 - W64 @ H64) / np.linalg.norm(X64)
+
+
+@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", PARITY_CASES)
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("seed", [0, 17, 101])
+@pytest.mark.parametrize("max_iter", [1, 25])
+def test_sklearn_mu_matches_gpu_kernel_on_small_matrix(
+    kernel, dtype_name, np_dtype, rtol, atol, device, seed, max_iter
+):
+    """Match sklearn numerically on CPU and CUDA across aligned seeds."""
+    torch = require_nmf_runtime()
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    X = np.random.default_rng(42).random((64, 48), dtype=np_dtype)
+    X += np_dtype(0.1)  # strictly positive denominators avoid degenerate MU behavior
+    nmf_kwargs = _parity_nmf_kwargs(
+        n_components=4,
+        seed=seed,
+        max_iter=max_iter,
+    )
+
+    expected_H, expected_W = _sklearn_mu_reference(X, nmf_kwargs)
+    actual_H, actual_W = _gpu_parity_result(
+        kernel, X, nmf_kwargs, dtype_name, device
+    )
+
+    # Bitwise CPU parity is not portable: torch (ATen) and numpy/sklearn (BLAS)
+    # reduce the MU matmuls in different orders on some builds (e.g. Linux x86,
+    # where they link different BLAS), so the aligned factors differ by ~ULP.
+    # Assert tight numerical parity instead of exact equality on every device.
+    np.testing.assert_allclose(actual_H, expected_H, rtol=rtol, atol=atol)
+    np.testing.assert_allclose(actual_W, expected_W, rtol=rtol, atol=atol)
+    np.testing.assert_allclose(
+        _relative_reconstruction_error(X, actual_H, actual_W),
+        _relative_reconstruction_error(X, expected_H, expected_W),
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+def _available_host_memory_bytes():
+    """Return Linux available host memory when sysconf exposes it."""
+    try:
+        return int(os.sysconf("SC_AVPHYS_PAGES")) * int(os.sysconf("SC_PAGE_SIZE"))
+    except (AttributeError, OSError, ValueError):
+        return None
+
+
+def _require_large_cuda_parity_capacity(torch, np_dtype):
+    """Skip the destructive-size stress case unless explicitly enabled and safe."""
+    if os.environ.get(LARGE_PARITY_ENV) != "1":
+        pytest.skip(f"set {LARGE_PARITY_ENV}=1 to run the 100K x 20K parity stress test")
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    x_bytes = (
+        LARGE_PARITY_ROWS
+        * LARGE_PARITY_COLUMNS
+        * np.dtype(np_dtype).itemsize
+    )
+    available_host = _available_host_memory_bytes()
+    required_host = 3 * x_bytes
+    if available_host is not None and available_host < required_host:
+        pytest.skip(
+            f"large parity needs about {required_host / 2**30:.1f} GiB available host "
+            f"memory; found {available_host / 2**30:.1f} GiB"
+        )
+
+    properties = torch.cuda.get_device_properties(torch.cuda.current_device())
+    required_device = int(1.25 * x_bytes)
+    if properties.total_memory < required_device:
+        pytest.skip(
+            f"large parity needs a device with at least {required_device / 2**30:.1f} "
+            f"GiB addressable memory; found {properties.total_memory / 2**30:.1f} GiB"
+        )
+    if not getattr(properties, "is_integrated", False):
+        free_device, _ = torch.cuda.mem_get_info()
+        if free_device < required_device:
+            pytest.skip(
+                f"large parity needs about {required_device / 2**30:.1f} GiB free GPU "
+                f"memory; found {free_device / 2**30:.1f} GiB"
+            )
+
+
+@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", PARITY_CASES)
+def test_sklearn_mu_matches_cuda_on_100k_by_20k_matrix(
+    kernel, dtype_name, np_dtype, rtol, atol
+):
+    """Stress sklearn/CUDA parity on a 100K x 20K dense matrix for one MU iteration."""
+    torch = require_nmf_runtime()
+    _require_large_cuda_parity_capacity(torch, np_dtype)
+
+    X = np.random.default_rng(42).random(
+        (LARGE_PARITY_ROWS, LARGE_PARITY_COLUMNS),
+        dtype=np_dtype,
+    )
+    X += np_dtype(0.1)
+    nmf_kwargs = _parity_nmf_kwargs(
+        n_components=LARGE_PARITY_COMPONENTS,
+        seed=17,
+        max_iter=1,
+    )
+
+    try:
+        expected_H, expected_W = _sklearn_mu_reference(X, nmf_kwargs)
+        actual_H, actual_W = _gpu_parity_result(
+            kernel, X, nmf_kwargs, dtype_name, "cuda"
+        )
+
+        np.testing.assert_allclose(actual_H, expected_H, rtol=rtol, atol=atol)
+        np.testing.assert_allclose(actual_W, expected_W, rtol=rtol, atol=atol)
+    finally:
+        # Do not make a following dtype inherit this case's multi-GiB CUDA cache.
+        del X
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
 # ---------------------------------------------------------------------
 # Input validation and degenerate shapes
 # ---------------------------------------------------------------------
@@ -439,6 +661,11 @@ def test_resolve_gpu_opts_uses_defaults_when_gpu_kwargs_is_missing(kernel):
     assert kernel._resolve_gpu_opts(None) == kernel.DEFAULT_GPU
 
 
+def test_default_gpu_epsilon_matches_sklearn_exact_value(kernel):
+    """Pin sklearn's float32 epsilon without importing its private EPSILON symbol."""
+    assert kernel.DEFAULT_GPU["eps"] == float(np.finfo(np.float32).eps)
+
+
 def test_resolve_gpu_opts_dict_values_override_defaults(kernel):
     """Explicit gpu_kwargs values override defaults and are normalized to typed options."""
     opts = kernel._resolve_gpu_opts(
@@ -1180,7 +1407,9 @@ def test_mu_step_fixed_h_matches_manual_w_only_update(kernel):
     eps = torch.tensor(1e-9, dtype=torch.float64)
     H_before = H.clone()
 
-    expected_W = W0 * ((Xg @ H.T) / (W0 @ (H @ H.T) + eps))
+    denominator = W0 @ (H @ H.T)
+    denominator = denominator.where(denominator != 0, eps)
+    expected_W = W0 * ((Xg @ H.T) / denominator)
     W = kernel._mu_step_fixed_h(W0, H, Xg, eps)
 
     assert torch.allclose(W, expected_W)