torch-max-mem
This package provides decorators for memory utilization maximization with PyTorch and CUDA by starting with a maximum parameter size and applying successive halving until no more out-of-memory exception occurs.
💪 Getting Started
Assume you have a function for batched computation of nearest neighbors using brute-force distance calculation.
import torch
def knn(x, y, batch_size, k: int = 3):
return torch.cat(
[
torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
for start in range(0, x.shape[0], batch_size)
],
dim=0,
)
With torch_max_mem you can decorate this function to reduce the batch size
until no more out-of-memory error occurs.
import torch
from torch_max_mem import maximize_memory_utilization
@maximize_memory_utilization()
def knn(x, y, batch_size, k: int = 3):
return torch.cat(
[
torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
for start in range(0, x.shape[0], batch_size)
],
dim=0,
)
In the code, you can now always pass the largest sensible batch size, e.g.,
x = torch.rand(100, 100, device="cuda")
y = torch.rand(200, 100, device="cuda")
knn(x, y, batch_size=x.shape[0])
If you would rather not have to pass the batch size explicitly at all, you can
additionally decorate the function with infer_maximum_batch_size, which infers
it from the length of another parameter (x by default) whenever batch_size
is not given explicitly
(source).
import torch
from torch_max_mem import infer_maximum_batch_size, maximize_memory_utilization
@infer_maximum_batch_size()
@maximize_memory_utilization()
def knn(x, y, batch_size, k: int = 3):
return torch.cat(
[
torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
for start in range(0, x.shape[0], batch_size)
],
dim=0,
)
x = torch.rand(100, 100, device="cuda")
y = torch.rand(200, 100, device="cuda")
knn(x, y)
Sharing the GPU with a display
Memory utilization maximization grows the batch size until an allocation fails. If the GPU also drives your display, or is shared with other processes, that means growing until somebody else fails: the compositor can stutter, freeze, or trigger a driver reset well before PyTorch raises anything catchable.
set_memory_budget caps how much the caching allocator may hand out, so running
out of budget surfaces as an ordinary, catchable torch.cuda.OutOfMemoryError
that the decorators already handle.
from torch_max_mem import set_memory_budget
# leave everything beyond 2 GiB to the desktop
set_memory_budget(2 * 1024**3)
Note that the budget only bounds PyTorch's caching allocator. The CUDA context, as well as cuBLAS and cuDNN workspaces, live outside of it and add a few hundred MiB on top, so leave some slack.
🚀 Installation
The most recent release can be installed from PyPI with uv:
uv pip install torch_max_mem
or with pip:
python3 -m pip install torch_max_mem
The most recent code and data can be installed directly from GitHub with uv:
uv pip install git+https://github.com/mberr/torch-max-mem.git
or with pip:
python3 -m pip install git+https://github.com/mberr/torch-max-mem.git
👐 Contributing
Contributions, whether filing an issue, making a pull request, or forking, are appreciated. See CONTRIBUTING.md for more information on getting involved.
👋 Attribution
Parts of the logic have been developed with Laurent Vermue for PyKEEN.
⚖️ License
The code in this package is licensed under the MIT License.
🍪 Cookiecutter
This package was created with @audreyfeldroy's cookiecutter package using @cthoyt's cookiecutter-snekpack template.
🛠️ For Developers
See developer instructions
The final section of the README is for if you want to get involved by making a code contribution.
Development Installation
To install in development mode, use the following:
git clone git+https://github.com/mberr/torch-max-mem.git
cd torch-max-mem
uv pip install -e .
Alternatively, install using pip:
python3 -m pip install -e .
Updating Package Boilerplate
This project uses cruft to keep boilerplate (i.e., configuration, contribution
guidelines, documentation configuration) up-to-date with the upstream
cookiecutter package. Install cruft with either uv tool install cruft or
python3 -m pip install cruft then run:
cruft update
More info on Cruft's update command is available here.
🥼 Testing
After cloning the repository and installing tox with
uv tool install tox --with tox-uv or python3 -m pip install tox tox-uv, the
unit tests in the tests/ folder can be run reproducibly with:
tox -e py
Additionally, these tests are automatically re-run with each commit in a GitHub Action.
📖 Building the Documentation
The documentation can be built locally using the following:
git clone git+https://github.com/mberr/torch-max-mem.git
cd torch-max-mem
tox -e docs
open docs/build/html/index.html
The documentation automatically installs the package as well as the docs extra
specified in the pyproject.toml. sphinx plugins like
texext can be added there. Additionally, they need to be added to the
extensions list in docs/source/conf.py.
The documentation can be deployed to ReadTheDocs using
this guide. The
.readthedocs.yml YAML file contains all the configuration
you'll need. You can also set up continuous integration on GitHub to check not
only that Sphinx can build the documentation in an isolated environment (i.e.,
with tox -e docs-test) but also that
ReadTheDocs can build it too.
🧑💻 For Maintainers
See maintainer instructions
Initial Configuration
Configuring ReadTheDocs
ReadTheDocs is an external documentation hosting service that integrates with GitHub's CI/CD. Do the following for each repository:
- Log in to ReadTheDocs with your GitHub account to install the integration at https://readthedocs.org/accounts/login/?next=/dashboard/
- Import your project by navigating to https://readthedocs.org/dashboard/import then clicking the plus icon next to your repository
- You can rename the repository on the next screen using a more stylized name (i.e., with spaces and capital letters)
- Click next, and you're good to go!
Configuring Archival on Zenodo
Zenodo is a long-term archival system that assigns a DOI to each release of your package. Do the following for each repository:
- Log in to Zenodo via GitHub with this link: https://zenodo.org/oauth/login/github/?next=%2F. This brings you to a page that lists all of your organizations and asks you to approve installing the Zenodo app on GitHub. Click "grant" next to any organizations you want to enable the integration for, then click the big green "approve" button. This step only needs to be done once.
- Navigate to https://zenodo.org/account/settings/github/, which lists all of your GitHub repositories (both in your username and any organizations you enabled). Click the on/off toggle for any relevant repositories. When you make a new repository, you'll have to come back to this
After these steps, you're ready to go! After you make "release" on GitHub (steps for this are below), you can navigate to https://zenodo.org/account/settings/github/repository/mberr/torch-max-mem to see the DOI for the release and link to the Zenodo record for it.
Registering with the Python Package Index (PyPI)
The Python Package Index (PyPI) hosts packages so they can
be easily installed with pip, uv, and equivalent tools.
- Register for an account here
- Navigate to https://pypi.org/manage/account and make sure you have verified your email address. A verification email might not have been sent by default, so you might have to click the "options" dropdown next to your address to get to the "re-send verification email" button
- 2-Factor authentication is required for PyPI since the end of 2023 (see this blog post from PyPI). This means you have to first issue account recovery codes, then set up 2-factor authentication
- Issue an API token from https://pypi.org/manage/account/token
This only needs to be done once per developer.
Configuring your machine's connection to PyPI
This needs to be done once per machine.
uv tool install keyring
keyring set https://upload.pypi.org/legacy/ __token__
keyring set https://test.pypi.org/legacy/ __token__
Note that this deprecates previous workflows using .pypirc.
📦 Making a Release
Publishing to PyPI happens automatically via the
release.yml GitHub Actions workflow, using
PyPI's Trusted Publishing (OIDC).
No API token needs to be stored anywhere - the workflow only needs to be
triggered by tagging and publishing a GitHub release.
Tagging the release
After installing the package in development mode and installing tox with
uv tool install tox --with tox-uv or python3 -m pip install tox tox-uv, run
the following from the console:
tox -e tag-release
This does the following:
- Uses bump-my-version to
switch the version number in the
pyproject.toml,CITATION.cff,src/torch_max_mem/version.py, anddocs/source/conf.pyto not have the-devsuffix - Pushes the resulting commit and tag to GitHub
- Bumps the version to the next
-devpatch and pushes that too. If you made big changes and want to bump the version by minor instead, usetox -e bumpversion -- minorafter.
Releasing on GitHub
- Navigate to https://github.com/mberr/torch-max-mem/releases/new to draft a new release
- Click the "Choose a Tag" dropdown and select the tag corresponding to the release you just made
- Click the "Generate Release Notes" button to get a quick outline of recent changes. Modify the title and description as you see fit
- Click the big green "Publish Release" button
Publishing the release triggers the release.yml workflow, which builds the
package and uploads it to PyPI. It also triggers Zenodo to assign a DOI to your
release.
One-time setup: registering the trusted publisher
Before the first automated release, a repository maintainer needs to register this workflow as a trusted publisher on PyPI:
- On the
PyPI project page
(or, for a brand new project, https://pypi.org/manage/account/publishing/),
add a new GitHub publisher with:
- Owner:
mberr - Repository name:
torch-max-mem - Workflow name:
release.yml - Environment name:
pypi
- Owner:
- Optionally, create a
pypienvironment under the repository's Settings -> Environments with required reviewers, so publishing needs manual approval.
The manual, credential-based release commands (tox -e finish,
tox -e release, tox -e release-via-env) are still available below as a
fallback if needed.
Uploading to PyPI manually
If you need to publish without going through GitHub Actions, run:
tox -e finish
which performs the same version bump as tox -e tag-release, but also builds
and uploads the package directly to PyPI using
uv publish
with credentials from keyring, before pushing the tag and bumping to the next
-dev version.
Updating Package Boilerplate
This project uses cruft to keep boilerplate (i.e., configuration, contribution
guidelines, documentation configuration) up-to-date with the upstream
cookiecutter package. Install cruft with either uv tool install cruft or
python3 -m pip install cruft then run:
$ cruft update
More info on Cruft's update command is available here.
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file torch_max_mem-0.1.6.tar.gz.
File metadata
- Download URL: torch_max_mem-0.1.6.tar.gz
- Upload date:
- Size: 19.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
uv/0.12.6 {"installer":{"name":"uv","version":"0.12.6","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
8cde6e5991324e4f04aa34ddd1cd8b7af8cd42a6c9a6a0b952c56b1a68c3675a
|
|
| MD5 |
09ea3d93c6845e8f65f3723cec432f72
|
|
| BLAKE2b-256 |
abc6d2a299ae4cdadf768b96fae681bfe6c5dc120d610be5df5e028df522c06d
|
File details
Details for the file torch_max_mem-0.1.6-py3-none-any.whl.
File metadata
- Download URL: torch_max_mem-0.1.6-py3-none-any.whl
- Upload date:
- Size: 18.5 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
uv/0.12.6 {"installer":{"name":"uv","version":"0.12.6","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
814a5fb2048c5c1abbf724b372a952bcea0d828d89e36b08b76e33c4e4526dd2
|
|
| MD5 |
23c13aa26b1b08b9802dc49730f3669a
|
|
| BLAKE2b-256 |
8dca5520562e312dd48da467f07d7652dc44eee3b82af70030728aeee77875a7
|