PyTorch如何处理原地操作以保留反向传播所需信息？

我看到不少关于原地操作效率的讨论，但更关心PyTorch处理这类操作的内部机制，这里通过几个示例来拆解问题：

常规非原地操作的梯度逻辑 #

先看一个简单的非原地操作示例：

import torch
a = torch.randn(10, requires_grad=True)
b = torch.randn(10, requires_grad=True)
c = torch.randn(10, requires_grad=True)

x1 = a * b
x2 = x1 * c 

这种情况下反向传播的逻辑非常清晰：

x2.grad <- 1
c.grad <- x2.grad * x1 = a * b
x1.grad <- x2.grad * c = c
b.grad <- x1.grad * a = c * a
a.grad <- x1.grad * b = c * b

这个场景会分配x1、x2两个中间缓冲区，用于保存计算过程的中间结果，为反向传播提供所需的信息。

原地覆盖操作的梯度问题 #

但如果改成如下的原地覆盖操作：

x = a * b
x = x * c 

从正向计算结果看和之前完全一致，但直接按常规逻辑计算梯度会出错：

x.grad <- 1
c.grad <- x.grad * x = a * b * c

问题出在第二次赋值时覆盖了原本存储a*b的缓冲区，导致计算c的梯度时，无法获取到x在乘法前的原始值（也就是a*b），进而丢失了计算上游梯度（a.grad和b.grad）的必要信息。

可能的解决方案猜测 #

针对这个问题，我有两种猜测：

• 框架将代码编译优化为x = a * b * c，但这种优化在复杂表达式场景下很可能失效；
• 框架内部实际仍会创建类似x1的中间缓冲区，保存被覆盖前的中间结果。

复杂原地操作的疑问 #

那对于更复杂的链式原地操作：

x = a * b
x *= c
x *= d
x *= e
x *= f 

PyTorch会创建x1、x2、x3这类临时缓冲区吗？现代深度学习框架到底是如何解决这类原地操作的梯度计算问题的？

内容的提问来源于stack exchange，提问作者Victor Chavauty