Repository navigation
Module.to should not be differentiable #1148
Description
Current Module.to implementation keeps the moved tensors in the same graph.
case Parameter param: {
var t = param.to(dtype, device); //in the same graph
t.retain_grad();
...
}
case Tensor tensor when (device.type != tensor.device_type || device.index != tensor.device_index): {
var t = tensor.to(dtype, device); //in the same graph
...
}While in the pytorch implementation, the moved tensors are detached.
for key, param in self._parameters.items():
if param is None:
continue
# Tensors stored in modules are graph leaves, and we don't want to
# track autograd history of `param_applied`, so we have to use
# `with torch.no_grad():`
with torch.no_grad(): #detached from the original tensor
param_applied = fn(param)
...
if param.grad is not None:
with torch.no_grad(): #detached from the original tensor
grad_applied = fn(param.grad)
...As the module tensors should be leaf tensors, it makes no sense to keep the original tensors. If a CUDA module is not detached from the original CPU tensors, there are unnecessary gradient traffics from the GPU to the CPU in very backpropagation.
Activity
NiklasGustafsson commentedon Dec 7, 2023
@shaltielshmid -- since you have been spending time in 'to' lately... Can you comment, please?
shaltielshmid commentedon Dec 7, 2023
It seems like @zgxnet has a good point.
As per the LibTorch documentation:
When we move the Parameters between devices, the graph is kept and then the new parameters are no longer from leaf tensors. The call to
Tensor.retain_grad()is a workaround which forces torch to have the new tensorsgrad()populated during the nextbackward()call.The use of
torch.no_grad()in PyTorch is a way of enforcing the moved Tensor to be a leaf tensor and therefore will accumulate gradients. In addition, it also disables the autograd tracking of moving a tensor across devices.[As a side note: the
retain_grad()call actually throws an error whenrequires_grad=False, which is a different bug].I'll prepare a PR with a fix for this.
@NiklasGustafsson - I see two possibilities for solving this:
1] Like PyTorch, move the tensors in a
no_grad()scope.2] After moving the tensor, call "Tensor.detach()" to make the new tensor a leaf tensor.
The advantage of 1 is that the graph doesn't have to track moving a tensor between devices. On the other hand, with 2 we will just dispose the old tensor and then that part of the graph will be thrown out regardless.
I'm leaning towards option 2 since that is a more direct way of solving the issue, what do you think?
NiklasGustafsson commentedon Dec 7, 2023
Whichever uses less memory. We don't want to hit the upper limit of GPU memory with temporary storage use.