PyTorch中切换GPU/CPU设备会影响反向传播的梯度计算吗？

核心结论 #

这种GPU+CPU混合运行前向传播的模式不会破坏PyTorch自动微分的正常逻辑，只要前向过程中没有主动断开计算图，反向传播可以正常执行，但会产生额外的性能开销，有几个注意事项需要关注。

具体影响说明 #

• 功能层面无影响：PyTorch的Autograd机制会自动跟踪张量所在设备，只要你在张量从GPU迁移到CPU时没有调用.detach()、也没有手动设置requires_grad=False，计算图就会保持完整。反向传播时梯度会自动匹配参数所在设备：GPU上的算子A参数会收到对应GPU存储的梯度，CPU上的算子B参数会收到对应CPU存储的梯度，优化器可以正常更新两类参数。
• 存在额外性能损耗：反向传播过程中梯度需要在CPU和GPU之间来回传输，会增加通信开销，训练/推理的速度会远低于全量在GPU运行的模式，中间迁移的张量尺寸越大，这部分开销越明显。
• 需要注意CPU内存容量：反向传播需要用到前向传播阶段的中间激活值，算子B的所有前向中间张量都会保留在CPU内存中，需要确保CPU内存足够容纳这些数据，否则会触发CPU侧的OOM报错。

正确实现示例 #

import torch
import torch.nn as nn

# 算子A部署在GPU
op_a = nn.Linear(1024, 2048).cuda()
# 算子B部署在CPU
op_b = nn.Linear(2048, 512).cpu()

# 输入数据在GPU，开启梯度跟踪
x = torch.randn(32, 1024, device='cuda', requires_grad=True)

# 前向传播：GPU侧计算
out_a = op_a(x)
# 张量迁移到CPU，不要调用detach()，保持计算图连接
out_a_cpu = out_a.cpu()
# CPU侧计算
out_b = op_b(out_a_cpu)
loss = out_b.sum()

# 反向传播正常执行
loss.backward()

# 各参数梯度自动生成在对应设备上
print(x.grad.device) # 输出cuda:0
print(op_a.weight.grad.device) # 输出cuda:0
print(op_b.weight.grad.device) # 输出cpu

注意：只有主动断开计算图的操作才会导致反向传播异常，迁移设备本身不会打断计算图的跟踪逻辑。

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