Stop mckp from silently swallowing the requested device - #473
Open
BenjaminDEMAILLE wants to merge 1 commit into
Open
Stop mckp from silently swallowing the requested device#473BenjaminDEMAILLE wants to merge 1 commit into
BenjaminDEMAILLE wants to merge 1 commit into
Conversation
`MultipleChoiceKnapsack.estimate_noise` had no `device` parameter, so the `device=args.device` that run.py passes to every estimator was absorbed into `**kwargs` and dropped. mckp is the default estimator, so `--cuda` looked effective there and did nothing. The CPU choice itself is kept. It predates the 2026 refactor, and mckp is the memory-hungry estimator whose MAP step runs over a whole chunk at once, which is what broadinstitute#396 is about. Moving it onto the GPU is a memory trade-off for maintainers to make deliberately, not something to flip in passing. What changes is only that the choice is now visible: the parameter is accepted and documented, and a debug line says the MAP step runs on the CPU rather than on the requested device, and why. Nothing about the computation moves. Worth deciding separately: whether mckp should honour --cuda for this step, given the memory cost. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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
1 participant
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.
Noticed while measuring the effect of device selection on estimation for #467.
MultipleChoiceKnapsack.estimate_noisehas nodeviceparameter.run.pypassesdevice=args.deviceto every estimator, so for mckp it lands in**kwargsand is dropped. mckp is the default estimator, which means--cuda(and--mps) appear to apply to it and do not.What this PR does not do
It does not move the computation. The hard-coded CPU is deliberate: it predates the 2026 refactor, and mckp is the memory-hungry estimator whose MAP step runs over a whole chunk at once, which is what #396 (
OutOfMemory error when computing target noise counts per gene) is about. Flipping that to GPU is a memory trade-off you should make deliberately, and I cannot validate it against the datasets that hit #396.What it does
Makes the choice visible instead of invisible:
deviceis accepted and documented on the signature, rather than silently absorbed.device="cpu"call sites carry a comment explaining the intent, so the next reader does not take them for an oversight, as I initially did.Worth your decision
Should mckp honour
--cudafor that MAP step? There is a real speed-up available, at a memory cost that #396 suggests is not free. Happy to implement whichever way you want, including a flag to opt in.Verification
Full suite 176 passed, 69 skipped.
ruff check,ruff format --check,mypy cellbender testsclean. New test asserts that a caller passingdevice="cuda"gets told, rather than silence.