Does PyTorch `.to(device)` propagate gradients back to original device?

gammag4: <p>I am trying to understand how the gradients are propagated back when using <code>loss.backward()</code>. From what I’ve seen, when you call <code>tensor.to('cuda')</code> the returned tensor is not a leaf anymore, so it has no <code>.grad</code> and, theoretically, the gradients should be propagated back to the original tensor in the CPU, which does not seem efficient.</p>
<p>So I want to know what exactly happens in this case when you call <code>loss.backward()</code> and why it is not inefficient, since I see a lot of people using <code>tensor.to(device) instead of specifying </code>device=device<code>in tensor construction. I also want to know if using</code>requires_grad = True` in the GPU tensor to make it a leaf tensor would be better in this case.</p>

Naming-isDifficult: <p>Worth a try <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<pre data-code-wrap="python"><code class="lang-python">import torch

# torch.Tensor.to(device) approach
x_cpu = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
x_gpu = x_cpu.to('cuda')
y_gpu = x_gpu * 2
loss = y_gpu.sum()
loss.backward()

print("Gradient of x_cpu (original tensor on CPU):", x_cpu.grad)
print("Gradient of x_gpu (tensor on GPU, not a leaf):", x_gpu.grad)

# specify device approach
x_gpu_leaf = torch.tensor([1.0, 2.0, 3.0], requires_grad=True, device='cuda')
y_gpu_leaf = x_gpu_leaf * 2
loss_leaf = y_gpu_leaf.sum()
loss_leaf.backward()

print("Gradient of x_gpu_leaf (leaf tensor on GPU):", x_gpu_leaf.grad)
</code></pre>
<hr>
<p><em><strong>DO NOT READ THE FOLLOWING COMMENTS BEFORE TRYING THE CODE ABOVE</strong></em></p>
<hr>
<p>You are correct ---- gradients <em>will</em> propagates back to CPU if you use <code>.to(device)</code> approach. I’ve also tried to test whether there is a significant gap between these two approaches via the following code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import time
import numpy as np

size = (10000, 10000)

# assuming that there is a CUDA device
tensor_cpu = torch.randn(size, requires_grad=True)
tensor_gpu_transfer = tensor_cpu.to("cuda")
tensor_gpu_direct = torch.randn(size, requires_grad=True, device="cuda")

transfer_time_list = []
for _ in range(1000):
    output_transfer = tensor_gpu_transfer.sum()
    start_transfer = time.time()    # this is not the proper way though, I'm using it for convenience
    output_transfer.backward()
    torch.cuda.synchronize()
    backward_time_transfer = time.time() - start_transfer
    transfer_time_list.append(backward_time_transfer)
print(
    "Transfer approach:\n"
    f"Mean: {np.mean(transfer_time_list):.6f}\n"
    f"std: {np.std(transfer_time_list):.6f}\n"
)

direct_time_list = []
for _ in range(1000):
    output_direct = tensor_gpu_direct.sum()
    start_direct = time.time()
    output_direct.backward()
    torch.cuda.synchronize()
    backward_time_direct = time.time() - start_direct
    direct_time_list.append(backward_time_direct)
print(
    "Direct approach:\n"
    f"Mean: {np.mean(direct_time_list):.6f}\n"
    f"std: {np.std(direct_time_list):.6f}\n"
)
</code></pre>
<p>Tyr it by yourself, but I will post my results here:</p>
<ol>
<li>There <strong>is</strong> a significant gap between these two methods</li>
<li>The larger the <code>size</code> is, the larger the gap is</li>
</ol>
<p>When <code>size</code> is <code>(1000, 1000)</code>, I have:</p>
<pre><code class="lang-auto">Transfer approach:
Mean: 0.000594
std: 0.000239

Direct approach:
Mean: 0.000037
std: 0.000187
</code></pre>
<p>But if I set <code>size</code> to be <code>(10000, 10000)</code>, I have:</p>
<pre><code class="lang-auto">Transfer approach:
Mean: 0.134293
std: 0.001979

Direct approach:
Mean: 0.001331
std: 0.000157
</code></pre>
<p>So, back to your question: “why it is not inefficient?”<br>
Unfortunately, most likely it is <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"><br>
I think there might be a tradeoff between efficiency and convenience/flexibility or simply common practice. Generally speaking, if you use <code>.to(device)</code>, it’s basically a one-liner and will, in theory, cause the least bugs.</p>
<hr>
<p>Note:<br>
My initial idea is that the transfer process should be an asynchronized process so the difference should be small. However, my test results are telling a different story. I would be glad if anyone could dive deeper and explain how PyTorch handles these things. I’m not familiar with this part…</p>

ptrblck: <aside class="quote no-group" data-username="Naming-isDifficult" data-post="2" data-topic="217641">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/naming-isdifficult/48/58644_2.png" class="avatar"> Naming-isDifficult:</div>
<blockquote>
<p><code>start_transfer = time.time()    # this is not the proper way though, I'm using it for convenience</code></p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="Naming-isDifficult" data-post="2" data-topic="217641">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/naming-isdifficult/48/58644_2.png" class="avatar"> Naming-isDifficult:</div>
<blockquote>
<p>I would be glad if anyone could dive deeper and explain how PyTorch handles these things. I’m not familiar with this part…</p>
</blockquote>
</aside>
<p>Based on your code you know that your profiling might be invalid as it accumulates previously launched kernels on the first iterations so you could fix this.<br>
A full profiler timeline could also give you more information about the launched operations and their overhead.</p>

ptrblck: <aside class="quote no-group" data-username="gammag4" data-post="1" data-topic="217641">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/gammag4/48/76063_2.png" class="avatar"> gammag4:</div>
<blockquote>
<p>and why it is not inefficient, since I see a lot of people using <code>tensor.to(device) instead of specifying </code>device=device<code>in tensor construction</code></p>
</blockquote>
</aside>
<p>The majority of the use cases don’t compute gradients in the input tensor. Moving the input batch inside the <code>DataLoader</code> loop to the GPU is fine as usually only the model parameters will receive gradients, not the input.</p>

gammag4: <p>Thank you both for the answers, you both helped me clear out some ideas.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="3" data-topic="217641">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>as it accumulates previously launched kernels on the first iterations</p>
</blockquote>
</aside>
<p>What do you mean specifically by this? do you mean the fact that it sends the kernels in the first iteration, making it take longer than the other iterations? How could one profile the part of sending the kernels separately?</p>
<p>I would also like to know if there is anything in the documentation or somewhere else explaining this pipeline, since it is still kind of a black box to me. For example, I don’t know exactly how the process of scheduling the functions to the GPU work and the best practices for preventing bottleneck from blocking it by needing to send data back and forth between CPU and GPU.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="217641">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>The majority of the use cases don’t compute gradients in the input tensor.</p>
</blockquote>
</aside>
<p>This actually sounds pretty obvious now, since the variables are just the parameters, not the data. I guess my confusion came from the fact that the examples I saw propagated back to all input tensors. Thank you.</p>