How to Compare Custom CUDA Gradients with PyTorch's Autograd Gradients?

OmkarV23: <p>Hi PyTorch Community,</p>
<p>I’m working on a project where I’ve implemented a custom <code>autograd.Function</code> in PyTorch that wraps a CUDA operation. My custom CUDA functions (both forward and backward) are critical for performance and currently support only <code>float32</code> (<code>torch.float</code>) data types.</p>
<p>I want to verify that my custom backward implementation is correct by comparing the gradients computed by my custom backward function with the gradients computed by PyTorch’s autograd engine. However, I’m facing challenges due to the data type constraints and the use of custom CUDA code.</p>
<p><strong>Here’s what I’m trying to achieve:</strong></p>
<p><strong>1. Include my custom CUDA function in the computational graph.</strong><br>
<strong>2. Use PyTorch’s autograd to compute gradients based on the forward pass.</strong><br>
<strong>3. Compare my custom backward gradients with the autograd-computed gradients to ensure correctness.</strong></p>
<p><strong>My Challenges:</strong></p>
<ul>
<li><strong>Data Type Constraint:</strong> <code>torch.autograd.gradcheck</code> expects inputs of type <code>torch.double</code>, but my CUDA functions only support <code>torch.float</code> (<code>float32</code>). Modifying the CUDA code to support <code>double</code> would be a significant effort I’d like to avoid.</li>
<li><strong>Gradient Checking:</strong> I’m unsure how to perform gradient checking in this scenario where my custom functions are limited to <code>float32</code>.</li>
</ul>
<p><strong>My Questions:</strong></p>
<ol>
<li><strong>Can a custom CUDA function be part of the computational graph, allowing PyTorch to compute autograd gradients based on the forward pass?</strong></li>
<li><strong>If so, how can I extract PyTorch’s autograd gradients for comparison with my custom backward gradients?</strong></li>
<li><strong>Is there a recommended approach or tool for performing gradient checking with custom CUDA functions that only support <code>float32</code>?</strong></li>
</ol>
<p><strong>Illustrative Pseudo-Code:</strong></p>
<p>Here’s a simplified pseudo-code to illustrate my setup:</p>
<pre><code class="lang-auto">import torch

class MyCustomFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # Call my custom CUDA forward function (supports float32 only)
        output = my_cuda_forward(input)
        ctx.save_for_backward(input)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        # Call my custom CUDA backward function
        grad_input = my_cuda_backward(grad_output, input)
        return grad_input

# Example usage
input = torch.randn(10, dtype=torch.float32, requires_grad=True)
output = MyCustomFunction.apply(input)
loss = output.sum()
loss.backward()

# Now, I want to compare my custom backward gradients with PyTorch's autograd gradients
# How can I achieve this?

</code></pre>
<p><strong>Additional Context:</strong></p>
<ul>
<li><strong>PyTorch Version:</strong> (2.5.1)</li>
<li><strong>CUDA Version:</strong> (e.g., 12.1)</li>
<li><strong>Operating System:</strong> (e.g., Ubuntu 22.04)</li>
</ul>
<p><strong>What I’ve Tried:</strong></p>
<ul>
<li><strong>Using <code>gradcheck</code> with <code>float32</code> tensors:</strong> <code>gradcheck</code> requires <code>torch.double</code> tensors and doesn’t accept <code>float32</code> inputs.</li>
<li><strong>Implementing Custom Gradient Checking:</strong> Considering writing a custom finite differences gradient checker, but I’m wondering if there’s a better approach.</li>
</ul>
<p><strong>Any guidance or suggestions would be greatly appreciated!</strong></p>
<p>Thank you!</p>

soulitzer: <aside class="quote no-group" data-username="OmkarV23" data-post="1" data-topic="213431">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/omkarv23/48/74199_2.png" class="avatar"> OmkarV23:</div>
<blockquote>
<p><strong>Using <code>gradcheck</code> with <code>float32</code> tensors:</strong> <code>gradcheck</code> requires <code>torch.double</code> tensors and doesn’t accept <code>float32</code> inputs.</p>
</blockquote>
</aside>
<p>gradcheck does allow float32 inputs, though running finite differencing on float32 would result in greater loss in precision and false positives (you would want to increase the tolerance)</p>

soulitzer: <aside class="quote no-group" data-username="OmkarV23" data-post="1" data-topic="213431">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/omkarv23/48/74199_2.png" class="avatar"> OmkarV23:</div>
<blockquote>
<p>If so, how can I extract PyTorch’s autograd gradients for comparison with my custom backward gradients?</p>
</blockquote>
</aside>
<p>You must use plain torch operations for the autograd computational graph to be generated. If you can express your forward computation using torch operations you can compute the gradient using autograd and compare that with your custom Function’s gradients.</p>