Create the CUDA context only when a GPU tensor is actually used - #847
Open
DenisDrobyshev wants to merge 2 commits into
Open
Create the CUDA context only when a GPU tensor is actually used#847DenisDrobyshev wants to merge 2 commits into
DenisDrobyshev wants to merge 2 commits into
Conversation
TorchBackend.__init__ built a CUDA generator and CUDA entries for the type
list as soon as the backend was instantiated. Both initialise a CUDA context,
so the first POT call on CPU-only torch tensors claimed device memory and woke
the GPU:
import torch, ot
x = torch.randn(2000, 2)
ot.dist(x, x)
torch.cuda.is_initialized() # True
torch.cuda.memory_allocated() # 1024
Both are now built on first access. A seed set before any GPU work is
remembered and applied when the CUDA generator is created, so seeded results
are unchanged.
Closes PythonOT#612
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #847 +/- ##
==========================================
- Coverage 96.86% 96.82% -0.04%
==========================================
Files 128 128
Lines 25947 25970 +23
==========================================
+ Hits 25133 25146 +13
- Misses 814 824 +10 🚀 New features to boost your workflow:
|
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.
Types of changes
Motivation and context / Related issue
Closes #612.
TorchBackend.__init__builds a CUDA generator and the CUDA entries of__type_list__as soon as the backend is instantiated. Both initialise a CUDAcontext, so the first POT call on CPU-only torch tensors claims device memory
and wakes the GPU — which is what the issue reports as rising power draw with
CPU tensors.
On an RTX 4060 with POT 0.9.7.post1:
The backend is instantiated lazily, which is why this happens at the first call
rather than at
import ot.Description
The CUDA generator and the CUDA half of the type list are now built on first
access instead of in
__init__:__type_list__becomes a property. The CPU entries are built eagerly asbefore; the CUDA entries are appended the first time the list is read, so
anything that enumerates dtypes and devices (the tests,
_bench) sees exactlywhat it saw before.
rng_cuda_becomes a property with a setter, soseed(torch.Generator)keepsworking.
seed(int)records the seed instead of forcing the CUDA generator intoexistence, and applies it when that generator is first created. A seed set
before any GPU work therefore still governs it.
After the change the same script leaves CUDA untouched, and a later GPU call
initialises it as usual.
Reproducibility
Seeded GPU sampling is unchanged. With the same seed, before and after:
len(nx.__type_list__)is still 4 when CUDA is available.How has this been tested
New
test_no_cuda_context_for_cpu_only_workintest/test_backend.py, anon-regression test for #612. It asserts that CPU-only work leaves
torch.cuda.is_initialized()false andmemory_allocated()at zero. The checkneeds an interpreter that has not touched CUDA yet, so it runs in a subprocess,
following
test_torch_optimizer_after_tensorflow_importin the same file. Itfails on
masterand passes here, and is skipped without CUDA.test/test_backend.pyis 23 passed, 1 skipped (22 passed, 1 skipped before thisPR — the new test is the difference).
pre-commit runon the changed files isclean.
Ran on Windows with torch 2.9.1+cu126 and an RTX 4060 Laptop GPU.