Skip to content

Create the CUDA context only when a GPU tensor is actually used - #847

Open
DenisDrobyshev wants to merge 2 commits into
PythonOT:masterfrom
DenisDrobyshev:lazy-cuda-context
Open

Create the CUDA context only when a GPU tensor is actually used#847
DenisDrobyshev wants to merge 2 commits into
PythonOT:masterfrom
DenisDrobyshev:lazy-cuda-context

Conversation

@DenisDrobyshev

Copy link
Copy Markdown

Types of changes

  • Bug fix (non-breaking change which fixes an issue)

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 CUDA
context, 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:

import torch, ot
torch.cuda.is_initialized()        # False
x = torch.randn(2000, 2)           # CPU
ot.dist(x, x)
torch.cuda.is_initialized()        # True
torch.cuda.memory_allocated()      # 1024

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 as
    before; the CUDA entries are appended the first time the list is read, so
    anything that enumerates dtypes and devices (the tests, _bench) sees exactly
    what it saw before.
  • rng_cuda_ becomes a property with a setter, so seed(torch.Generator) keeps
    working.
  • seed(int) records the seed instead of forcing the CUDA generator into
    existence, 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:

nx.seed(42); nx.rand(3, type_as=torch.zeros(1, device="cuda"))
# [0.61296, 0.010059, 0.398414]   identical on both

len(nx.__type_list__) is still 4 when CUDA is available.

How has this been tested

New test_no_cuda_context_for_cpu_only_work in test/test_backend.py, a
non-regression test for #612. It asserts that CPU-only work leaves
torch.cuda.is_initialized() false and memory_allocated() at zero. The check
needs an interpreter that has not touched CUDA yet, so it runs in a subprocess,
following test_torch_optimizer_after_tensorflow_import in the same file. It
fails on master and passes here, and is skipped without CUDA.

test/test_backend.py is 23 passed, 1 skipped (22 passed, 1 skipped before this
PR — the new test is the difference). pre-commit run on the changed files is
clean.

Ran on Windows with torch 2.9.1+cu126 and an RTX 4060 Laptop GPU.

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

codecov Bot commented Aug 31, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 56.25000% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.82%. Comparing base (9439b8a) to head (c2a90b2).

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:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ot.solve uses GPU even though tensors are on CPU?

1 participant