Universal CUDA Tools: GPU-Safe Execution Made Simple for PyTorch

username: Tunahanyrd
cooked: <p><strong>Note: This is not an ad, but a suggestion to add to the Torch library.</strong></p>
<h2><a name="p-470108-what-is-universal-cuda-tools-1" class="anchor" href="#p-470108-what-is-universal-cuda-tools-1"></a>What is <em>universal-cuda-tools</em>?</h2>
<p><strong>universal-cuda-tools</strong> is a lightweight Python utility package designed to simplify and harden GPU/CPU execution workflows in PyTorch (and optionally TensorFlow).</p>
<p>It provides:</p>
<ul>
<li>A <code>DeviceContext</code> context manager for clean device, AMP, and cache handling</li>
<li>Simple and advanced decorators (<code>@cuda</code>, <code>@cuda_advanced</code>) to make any function run safely on GPU or CPU</li>
<li>Optional automatic tensorization (e.g. <code>int</code>, <code>list</code>, <code>np.ndarray</code> → <code>torch.Tensor</code>)</li>
<li>Built-in support for:
<ul>
<li>Memory profiling</li>
<li>Timeout handling</li>
<li>Retry on failure</li>
<li>Fallback to CPU when GPU out-of-memory (OOM)</li>
<li>Mixed precision (AMP)</li>
<li>Multi-GPU dispatch</li>
<li>Dry-run and telemetry logging</li>
</ul>
</li>
<li>Utilities for converting and moving arbitrary Python, NumPy, or even TensorFlow inputs to the correct device</li>
</ul>
<p>It’s built for <strong>real-world training pipelines</strong>, especially on constrained hardware, and aims to reduce the boilerplate and fragility in device management, AMP setup, and error handling.</p>
<hr>
<h2><a name="p-470108-why-does-this-matter-2" class="anchor" href="#p-470108-why-does-this-matter-2"></a>Why does this matter?</h2>
<p>Working with GPU-accelerated code in PyTorch often requires repetitive setup:</p>
<ul>
<li>Explicit device selection</li>
<li>Manual <code>.to(device)</code> calls for all inputs</li>
<li>Handling CUDA memory errors (OOM)</li>
<li>Managing AMP (<code>autocast</code>) scopes</li>
<li>Clearing cache and profiling memory</li>
<li>Converting raw Python or NumPy types into tensors</li>
</ul>
<p><strong>universal-cuda-tools</strong> abstracts all of that into a minimal, composable, and Pythonic interface.</p>
<hr>
<h2><a name="p-470108-installation-3" class="anchor" href="#p-470108-installation-3"></a>Installation</h2>
<pre data-code-wrap="bash"><code class="lang-bash">pip install universal-cuda-tools
</code></pre>
<p>Or install from source:</p>
<pre data-code-wrap="bash"><code class="lang-bash">git clone https://github.com/Tunahanyrd/universal-cuda-tools.git
cd universal-cuda-tools
python -m build
pip install dist/universal_cuda_tools-*.whl
</code></pre>
<hr>
<h2><a name="p-470108-api-overview-4" class="anchor" href="#p-470108-api-overview-4"></a>API Overview</h2>
<h3><a name="p-470108-cuda-decorator-5" class="anchor" href="#p-470108-cuda-decorator-5"></a><code>@cuda</code> Decorator</h3>
<p>A simple device-aware wrapper for lightweight functions. Supports retry, cache clearing, auto tensorization, and CPU fallback.</p>
<h3><a name="p-470108-cuda_advanced-decorator-6" class="anchor" href="#p-470108-cuda_advanced-decorator-6"></a><code>@cuda_advanced</code> Decorator</h3>
<p>All of the above, plus timeout, AMP, multi-GPU dispatching, live dashboard, and dry-run support.</p>
<h3><a name="p-470108-devicecontext-7" class="anchor" href="#p-470108-devicecontext-7"></a><code>DeviceContext</code></h3>
<p>A scoped context manager that handles device selection, AMP enablement, and auto tensorization for all code within the block.</p>
<h3><a name="p-470108-utilities-8" class="anchor" href="#p-470108-utilities-8"></a>Utilities</h3>
<ul>
<li><code>tensorize_for_universal(obj, device)</code></li>
<li><code>move_to_torch(device, obj)</code></li>
<li><code>patch_numpy_with_cupy()</code></li>
</ul>
<hr>
<h2><a name="p-470108-example-9" class="anchor" href="#p-470108-example-9"></a>Example</h2>
<pre data-code-wrap="python"><code class="lang-python">from cuda_tools import cuda

@cuda(device="cuda", auto_tensorize=True, to_list=True)
def add(a, b):
    return a + b

print(add([1, 2], [3, 4]))  # → [4, 6]
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">from cuda_tools import cuda_advanced

@cuda_advanced(timeout=0.5, use_amp=True, telemetry=True)
def train_step(model, x, y):
    pred = model(x)
    loss = (pred - y).square().mean()
    loss.backward()
    return loss.item()
</code></pre>
<hr>
<h2><a name="p-470108-documentation-demo-10" class="anchor" href="#p-470108-documentation-demo-10"></a>Documentation &amp; Demo</h2>
<ul>
<li><a href="https://www.funeralcs.com/posts/cuda_tools_dokuman/" rel="noopener nofollow ugc">Full Documentation</a></li>
<li><a href="https://github.com/pytorch/pytorch/issues/152679" rel="noopener nofollow ugc">PyTorch RFC Issue #152679</a></li>
<li>Pypi +1.4 downloads</li>
</ul>
<hr>
<h2><a name="p-470108-project-structure-11" class="anchor" href="#p-470108-project-structure-11"></a>Project Structure</h2>
<pre><code class="lang-plaintext">cuda_tools/
├── __init__.py        # Exports decorators and context
├── decorators.py      # @cuda and @cuda_advanced
├── context.py         # DeviceContext class
├── utils.py           # tensorize, device tools, CuPy support
</code></pre>
<hr>
<h2><a name="p-470108-license-12" class="anchor" href="#p-470108-license-12"></a>License</h2>
<p>MIT License © 2025 – Tunahan Yardımcı</p>

username: ptrblck
cooked: <p>Thanks for sharing! While some utilities sound interesting for PyTorch Core the majority of the proposal sounds like a feature for higher-level APIs. Did you already talk to these teams, e.g. <code>pytorch-lightning</code>?</p>

username: Tunahanyrd
cooked: <p>Thank you so much for taking the time to respond!</p>
<p>You’re absolutely right — most of the logic can indeed live in higher-level libraries like PyTorch Lightning or Accelerate. However, the internal structure of this tool is highly modular and granular.</p>
<p>For example, even without using the decorators, core functions such as <code>safe_to_device()</code>, <code>try_batch_size()</code>, or <code>run_with_amp()</code> can be used individually — making them potential candidates for inclusion in <code>torch.utils</code> or <code>torch.cuda</code>.</p>
<p>To be clear: this is not a proposal to merge the whole package as-is. Rather, I wanted to offer a small set of reusable utilities that address common patterns like:</p>
<ul>
<li>clean device transfers</li>
<li>OOM-safe function calls</li>
<li>automatic AMP context</li>
<li>fallback handling</li>
<li>safe tensorization of inputs</li>
</ul>
<p>These tools came out of repeated pain points in long training runs on low-memory devices, and I believe parts of them could help simplify day-to-day PyTorch code, even in core-level examples or docs.</p>
<p>If there’s any interest from the core team, I’d be happy to isolate and PR one or two focused components (e.g., <code>safe_to_device()</code> or a simple <code>cuda_guard()</code> context manager).</p>
<p>Thanks again!</p>