From 7a093aed52fea130ab5131f40385ee66d576d9ff Mon Sep 17 00:00:00 2001
From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com>
Date: Fri, 3 Jul 2026 02:16:54 +0900
Subject: [PATCH 1/3] Scale batched-fit convergence check to large n (identity
 recon error)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_fit_mu / _fit_mu_fixed_h computed the reconstruction error as norm(X - W@H),
which materializes the full [R, n_cells, n_genes] product at every convergence
check — tens of GB, and the OOM ceiling for batched fits at 1e6+ cells (e.g.
batch=2 at 2.4M x 2000 needs ~2x38 GB just for that residual).

Replace it with _recon_err, the identity ||X-WH||^2 = ||X||^2 - 2<X,WH> + ||WH||^2
computed from the small WtX [R,k,g], WtW [R,k,k], HHt [R,k,k] matmuls (||X||^2
precomputed once) — no [R,n,g] tensor is ever formed. The value is identical to
the direct norm (verified 2D and batched), so convergence / early-stop behavior
is unchanged; only peak memory drops (batch>1 at 2.4M now fits on an 80 GB GPU).
---
 src/cnmf/nmf_gpu.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py
index fed8f50..b8c6bf1 100644
--- a/src/cnmf/nmf_gpu.py
+++ b/src/cnmf/nmf_gpu.py
@@ -327,10 +327,21 @@ def _mu_step_fixed_h(W, H, Xg, eps):
 # ---------------------------------------------------------------------
 
 
+def _recon_err(Xg, W, H, xnorm2):
+    """Per-replicate ‖X − WH‖_F via the identity ‖X‖² − 2⟨X,WH⟩ + ‖WH‖², using only small
+    [R,k,g]/[R,k,k] matmuls (WᵀX, WᵀW, HHᵀ) — it NEVER materializes the [R,n,g] product, so batched
+    fits scale to millions of cells (a full [R,n,g] residual would be ~R·n·g floats, tens of GB).
+    Returns a 0-dim tensor (unbatched) or shape-[R] (batched); xnorm2 = ‖X‖² is precomputed once."""
+    Wt = W.transpose(-2, -1)
+    cross = ((Wt @ Xg) * H).sum(dim=(-2, -1))                          # ⟨X, WH⟩ per replicate
+    whnorm = ((Wt @ W) * (H @ H.transpose(-2, -1))).sum(dim=(-2, -1))  # ‖WH‖² per replicate
+    return (xnorm2 - 2.0 * cross + whnorm).clamp_min(0).sqrt()
+
+
 def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
     """Run full MU until `max_iter` or all replicate slices meet `tol`."""
     _check_runtime_tensors(Xg, W, H, eps)
-    lead = W.shape[:-2]                                     # () unbatched, (R,) batched
+    xnorm2 = torch.dot(Xg.reshape(-1), Xg.reshape(-1))     # ‖X‖² once; error check avoids [R,n,g]
     err_init = prev_err = None
     with torch.no_grad(), _cuda_tf32(torch, tf32, device):
         it = 0
@@ -339,7 +350,7 @@ def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
             for _ in range(n):                             # MU updates run inside the (compiled) step
                 W, H = step(W, H, Xg, eps)
             it += n
-            err = torch.linalg.norm((Xg - W @ H).reshape(*lead, -1), dim=-1)
+            err = _recon_err(Xg, W, H, xnorm2)
             if err_init is None:
                 err_init = err.clamp_min(1e-30)            # avoid 0/0 on a degenerate (all-zero) slice
             elif prev_err is not None and bool((((prev_err - err) / err_init) < tol).all()):
@@ -351,7 +362,7 @@ def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
 def _fit_mu_fixed_h(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
     """Run fixed-H MU until `max_iter` or all replicate slices meet `tol`."""
     _check_runtime_tensors(Xg, W, H, eps)
-    lead = W.shape[:-2]                                     # () unbatched, (R,) batched
+    xnorm2 = torch.dot(Xg.reshape(-1), Xg.reshape(-1))     # ‖X‖² once; error check avoids [R,n,g]
     err_init = prev_err = None
     with torch.no_grad(), _cuda_tf32(torch, tf32, device):
         it = 0
@@ -360,7 +371,7 @@ def _fit_mu_fixed_h(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, devi
             for _ in range(n):
                 W = step(W, H, Xg, eps)
             it += n
-            err = torch.linalg.norm((Xg - W @ H).reshape(*lead, -1), dim=-1)
+            err = _recon_err(Xg, W, H, xnorm2)
             if err_init is None:
                 err_init = err.clamp_min(1e-30)            # avoid 0/0 on a degenerate (all-zero) slice
             elif prev_err is not None and bool((((prev_err - err) / err_init) < tol).all()):

From 7c1829236089b385131deb2fb8f8c642a7a9490a Mon Sep 17 00:00:00 2001
From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com>
Date: Fri, 3 Jul 2026 05:09:09 +0900
Subject: [PATCH 2/3] =?UTF-8?q?Compute=20=E2=80=96X=E2=80=96=C2=B2=20in=20?=
 =?UTF-8?q?row=20chunks=20(torch.dot=20caps=20vector=20length=20at=202?=
 =?UTF-8?q?=C2=B3=C2=B9)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The identity recon-error check added in a59acd7 precomputed ‖X‖² with
torch.dot(X.reshape(-1), X.reshape(-1)). torch.dot is BLAS-backed and caps
its vector length at 2³¹-1 elements, so it raises RuntimeError on X with
more than ~2.1B entries (2.4M cells × 2000 genes = 4.8B) before the first
MU step — defeating the scaling goal of a59acd7.

Replace with _sq_norm(): a chunked square-sum reduction that is index-safe
at any size and also bounds the squaring temporary. Verified on CUDA to
match a float64 reference to 6e-08 on a 2.2B-element tensor (where torch.dot
raises); CD4 2.4M×2000 batch=4 factorize now runs at ~55 GB on an 80 GB A100.
---
 src/cnmf/nmf_gpu.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py
index b8c6bf1..a9201ad 100644
--- a/src/cnmf/nmf_gpu.py
+++ b/src/cnmf/nmf_gpu.py
@@ -338,10 +338,23 @@ def _recon_err(Xg, W, H, xnorm2):
     return (xnorm2 - 2.0 * cross + whnorm).clamp_min(0).sqrt()
 
 
+def _sq_norm(X):
+    """‖X‖² (every element squared, then summed) in row chunks. torch.dot is BLAS-backed and caps
+    its vector length at 2³¹ elements, so a flat dot overflows past ~2.1B entries (e.g. 2.4M cells
+    × 2000 genes = 4.8B); this chunked reduction is index-safe at any size and also bounds the
+    squaring temporary. Returns a 0-dim tensor."""
+    rows = X.shape[0]
+    step = max(1, (1 << 28) // max(1, X.numel() // max(1, rows)))      # ~2²⁸ elems/chunk (~1GB fp32)
+    total = X.new_zeros(())
+    for i in range(0, rows, step):
+        total = total + X[i:i + step].square().sum()
+    return total
+
+
 def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
     """Run full MU until `max_iter` or all replicate slices meet `tol`."""
     _check_runtime_tensors(Xg, W, H, eps)
-    xnorm2 = torch.dot(Xg.reshape(-1), Xg.reshape(-1))     # ‖X‖² once; error check avoids [R,n,g]
+    xnorm2 = _sq_norm(Xg)                                  # ‖X‖² once; error check avoids [R,n,g]
     err_init = prev_err = None
     with torch.no_grad(), _cuda_tf32(torch, tf32, device):
         it = 0
@@ -362,7 +375,7 @@ def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
 def _fit_mu_fixed_h(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device):
     """Run fixed-H MU until `max_iter` or all replicate slices meet `tol`."""
     _check_runtime_tensors(Xg, W, H, eps)
-    xnorm2 = torch.dot(Xg.reshape(-1), Xg.reshape(-1))     # ‖X‖² once; error check avoids [R,n,g]
+    xnorm2 = _sq_norm(Xg)                                  # ‖X‖² once; error check avoids [R,n,g]
     err_init = prev_err = None
     with torch.no_grad(), _cuda_tf32(torch, tf32, device):
         it = 0

From bfe6402567a1b4f30e04ed9d1a2ded95cb42c3e4 Mon Sep 17 00:00:00 2001
From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com>
Date: Fri, 3 Jul 2026 08:23:23 +0900
Subject: [PATCH 3/3] Densify norm_counts once in factorize_gpu (avoid
 per-batch re-densify)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

factorize_gpu read norm_counts once but passed the sparse X into every
_nmf_gpu_mu batch call; _gpu_setup then ran X.toarray() (+ isfinite/min/
ascontiguousarray) on each launch. At millions of cells with many k-values
that re-densifies the full n×g matrix O(#batches) times — dominating runtime
(~4 reps/min observed) and churning tens of GB of RAM per batch, which risks
OOM in the eager path. Matches the module's own "sparse X is still densified
by this prototype" TODO.

Densify X once before the batch loop and pass the dense ndarray downstream
(toarray then becomes a no-op). NOTE: root-caused from the CD4 2.4M run but
not yet runtime-validated — the pod died before I could re-run; needs a CUDA
smoke test before merge.
---
 src/cnmf/nmf_gpu.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py
index a9201ad..617170c 100644
--- a/src/cnmf/nmf_gpu.py
+++ b/src/cnmf/nmf_gpu.py
@@ -576,6 +576,10 @@ def factorize_gpu(cnmf_obj, gpu_kwargs, worker_i=0, total_workers=1, skip_comple
         job_idx = worker_filter(run_params.index[run_params['completed'] == False], worker_i, total_workers)
 
     genes = norm_counts.var.index
+    # Densify X once. _nmf_gpu_mu -> _gpu_setup calls X.toarray(), so passing the sparse matrix
+    # would re-densify the full n×g matrix on *every* batch (O(#batches) redundant conversions plus
+    # RAM churn at millions of cells). Convert once here; downstream then sees a dense ndarray.
+    X_dense = norm_counts.X.toarray() if hasattr(norm_counts.X, "toarray") else np.asarray(norm_counts.X)
     by_k = OrderedDict()
     for idx in job_idx:
         p = run_params.iloc[idx, :]
@@ -589,7 +593,7 @@ def factorize_gpu(cnmf_obj, gpu_kwargs, worker_i=0, total_workers=1, skip_comple
             seeds = [s for _, s in chunk]
             print('[Worker %d]. k=%d: launching %d replicate(s), iters=%s.'
                   % (worker_i, k, len(chunk), iters))
-            results = _nmf_gpu_mu(norm_counts.X, seeds, run_kwargs, gpu_kwargs)
+            results = _nmf_gpu_mu(X_dense, seeds, run_kwargs, gpu_kwargs)
             for (spectra, _usages), it in zip(results, iters):
                 spectra = pd.DataFrame(spectra, index=np.arange(1, k + 1), columns=genes)
                 save_df_to_npz(spectra, cnmf_obj.paths['iter_spectra'] % (k, it))