Automatically choose batch size based on available memory#84
Merged
Conversation
Contributor
☂️ Python Coverage
Overall Coverage
New FilesNo new covered files... Modified Files
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves sampler UX by making batch_size default to an automatic value derived from an estimated per-sample memory footprint and currently available memory, reducing the chance of OOM crashes when requesting many shots.
Changes:
- Add automatic batch size selection when
batch_sizeis not provided. - Add unit test coverage for the new auto-batching behavior.
- Add
psutildependency and update docs/README to reflect the new default behavior.
Reviewed changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/tsim/sampler.py |
Implements auto batch sizing based on available memory and an estimated per-sample footprint; updates sample() defaults/docs. |
test/unit/test_sampler.py |
Adds a parametrized test to ensure auto-batching splits sampling into equal-sized batches. |
pyproject.toml |
Adds psutil as a runtime dependency to query system available RAM on non-GPU platforms. |
docs/demos/from_stim_to_tsim.ipynb |
Updates demo text to describe the new batch_size behavior. |
README.md |
Adds documentation/coverage/arXiv/stars badges. |
uv.lock |
Locks psutil and includes additional lockfile churn, including a registry/source change for griffe-kirin. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
|
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
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.
This should be a UX improvement. Previously, batch_size was always chosen as shots when not specified, and I have heard of multiple users crashing their kernel when requesting many shots.
Now, a simple memory model is used and Tsim chooses the batch size to approximately use half of the available memory.