Skip to content

scbutterfly: replace dense TF-IDF and cut redundant copies - #71

Merged
VladimirShitov merged 1 commit into
openproblems-bio:mainfrom
benjaminfreyuu:fix/scbutterfly-memory
Sep 23, 2026
Merged

VladimirShitov merged 1 commit into
openproblems-bio:mainfrom
benjaminfreyuu:fix/scbutterfly-memory

Conversation

@benjaminfreyuu

Copy link
Copy Markdown
Contributor

scButterfly's TFIDF tiles the per-cell and per-peak totals into dense (n_peaks, n_cells) float64 matrices and divides densely, so it costs O(n_peaks * n_cells) regardless of how sparse the input is -- 37 GiB per array at NeurIPS2021 Multiome scale, with three or four alive at once. That is what likely drove the ~265 GB peak and the OOM on the 2022 Multiome datasets.

Patch it with a sparse O(nnz) equivalent, and trim the wrapper's own avoidable allocations:

  • sparse TF-IDF monkeypatch: same arithmetic in the same order, on the CSR .data array. Bit-identical, since scButterfly binarizes first and the totals are exact integers in float32.
  • _placeholder_block gathers rows on the CSR instead of densifying the whole training block first (>12 GiB of ATAC at Multiome scale).
  • _to_counts_X builds a minimal AnnData instead of a full .copy(), which was duplicating the normalized layer and obsm['gene_activity'] only for sc.concat to discard them.
  • skip the no-op var reindex whose .copy() duplicated the matrix.
  • scatter_to_target copies in column blocks rather than one numpy call per feature (up to ~116k of them for ATAC).
  • suppress test_model's unconditional PCA + neighbour graph over both predicted matrices; neither is read back.
  • log torch version, CUDA build and device at startup, so a silent CPU fallback shows up in the run log instead of only as a slow run.
  • request highsharedmem, matching the other torch methods.

Verified bit-identical to the previous implementation: TF-IDF in float32 and float64 over binarized and raw counts, every touched helper against the real test h5ads, and build_butterfly end to end in both directions. Both component tests pass, and the predict test loads weights trained by the previous code, so the reconstructed architecture -- and therefore the preprocessing -- is confirmed unchanged.

Open:
Local testing showed CUDA: Not available; but local machine does not have CUDA gpu support, so proper testing on gpu needed.

Describe your changes

Checklist before requesting a review

  • [x ] I have performed a self-review of my code

  • Check the correct box. Does this PR contain:

    • Breaking changes
    • New functionality
    • Major changes
    • Minor changes
    • [ x] Bug fixes
  • Proposed changes are described in the CHANGELOG.md

  • [x ] CI Tests succeed and look good!

scButterfly's TFIDF tiles the per-cell and per-peak totals into dense
(n_peaks, n_cells) float64 matrices and divides densely, so it costs
O(n_peaks * n_cells) regardless of how sparse the input is -- 37 GiB per
array at NeurIPS2021 Multiome scale, with three or four alive at once.
That is what drove the ~265 GB peak and the OOM on the 2022 Multiome
datasets.

Patch it with a sparse O(nnz) equivalent, and trim the wrapper's own
avoidable allocations:

- sparse TF-IDF monkeypatch: same arithmetic in the same order, on the
  CSR .data array. Bit-identical, since scButterfly binarizes first and
  the totals are exact integers in float32.
- _placeholder_block gathers rows on the CSR instead of densifying the
  whole training block first (>12 GiB of ATAC at Multiome scale).
- _to_counts_X builds a minimal AnnData instead of a full .copy(), which
  was duplicating the normalized layer and obsm['gene_activity'] only
  for sc.concat to discard them.
- skip the no-op var reindex whose .copy() duplicated the matrix.
- scatter_to_target copies in column blocks rather than one numpy call
  per feature (up to ~116k of them for ATAC).
- suppress test_model's unconditional PCA + neighbour graph over both
  predicted matrices; neither is read back.
- log torch version, CUDA build and device at startup, so a silent CPU
  fallback shows up in the run log instead of only as a slow run.
- request highsharedmem, matching the other torch methods.

Verified bit-identical to the previous implementation: TF-IDF in float32
and float64 over binarized and raw counts, every touched helper against
the real test h5ads, and build_butterfly end to end in both directions.
Both component tests pass, and the predict test loads weights trained by
the previous code, so the reconstructed architecture -- and therefore the
preprocessing -- is confirmed unchanged.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@benjaminfreyuu
benjaminfreyuu marked this pull request as draft September 22, 2026 22:09
@benjaminfreyuu
benjaminfreyuu marked this pull request as ready for review September 22, 2026 22:09
@benjaminfreyuu

Copy link
Copy Markdown
Contributor Author

@VladimirShitov could you take a look at the PR?

@VladimirShitov VladimirShitov left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks good to me. Funny solution with turning off PCA and neighbors computation

@VladimirShitov
VladimirShitov merged commit f48faee into openproblems-bio:main Sep 23, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants