Scale batched GPU NMF factorize to millions of cells#7
Merged
OlivierBakker merged 3 commits intoJul 7, 2026
Merged
Conversation
_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).
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Three scalability fixes to the batched GPU NMF factorize (
--engine gpu --gpu-batch,added in #6), all in
src/cnmf/nmf_gpu.py. They were surfaced by running cNMFfactorize on a 2.4M-cell Perturb-seq dataset on an A100-80GB, where the batched GPU
path hit three separate scaling walls.
Motivation
At
batch=4on 2.4M cells × 2000 genes,factorize_gpufailed in three independentways: an OOM-scale convergence check, a
torch.dotinteger overflow, and repeated re-densification of the sparse matrix. This PR makes the batched GPU factorize usable at
millions of cells.
Changes
Identity reconstruction error (
_recon_err). The convergence checktorch.linalg.norm(X − W@H)materialized the full[R, n, g]residual (R replicates ×2.4M cells × 2000 genes ≈ tens of GB). Replaced with the algebraic identity ‖X−WH‖² =
‖X‖² − 2⟨X,WH⟩ + ‖WH‖², computed from small
[R,k,g]/[R,k,k]matmuls (WᵀX, WᵀW, HHᵀ) —it never forms the large product.
int32-safe ‖X‖² (
_sq_norm). Computing ‖X‖² as a flattorch.dot(X.reshape(-1), X.reshape(-1))overflows the BLAS int32 vector-length cap (2³¹ − 1): 2.4M × 2000 = 4.8Belements crashed with
bound <= 2147483647. Now reduced in ~2²⁸-element row chunks,index-safe at any size.
Densify once in
factorize_gpu. The sparsenorm_counts.Xwas passed straight to_nmf_gpu_mu, whose_gpu_setupcalls.toarray()— so every batch re-densified thewhole n×g matrix (O(#batches) redundant conversions + RAM churn). Densify once up front;
downstream sees a dense ndarray.
Validation
(K = 2,4,6,12,16,18,20,22,24,36,48,60 × 100 iters) in ~9.8 h with no OOM and flat
~55 GB VRAM across all K (peak 55,945 MiB at k=2 through k=60), GPU at ~99% util.