From e1e3fc421453bb80dae38228f59bd99a887acd6b Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:38:20 +0000 Subject: [PATCH] Match GPU NMF denominator handling to sklearn --- README.md | 19 ++++ src/cnmf/nmf_gpu.py | 21 ++-- tests/test_nmf_gpu.py | 237 +++++++++++++++++++++++++++++++++++++++++- 3 files changed, 267 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 5050f36..43a1e02 100644 --- a/README.md +++ b/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: diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py index 617170c..1b92cbe 100644 --- a/src/cnmf/nmf_gpu.py +++ b/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) diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index a66ce7e..4c31b5d 100644 --- a/tests/test_nmf_gpu.py +++ b/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)