12.3.2 并行计算与通信

非常棒的问题！

你提到的这段代码中：

def copy_to_cpu(x, non_blocking=False):
    return [y.to('cpu', non_blocking=non_blocking) for y in x]

其中的 non_blocking 参数是 PyTorch 中用于控制 数据传输是否异步（非阻塞） 的一个关键选项。

---

✅ 简单回答：

non_blocking=True 表示：在 GPU 和 CPU 之间复制张量时，不等待当前的 GPU 操作完成，直接返回控制权给 Python。

也就是说，这是一种异步（asynchronous）数据传输方式，可以提高整体性能。

---

🔍 详细解释

🧠 背景知识：GPU 是异步执行的

PyTorch 默认情况下，所有 GPU 操作都是异步的。这意味着：

• 当你调用 x + y、model(x) 等操作时，Python 只是将任务提交给 GPU；
• 实际计算是在后台进行的；
• Python 主线程继续往下运行，不会卡住。

但如果你要做一些“需要结果”的操作（如打印、保存、转成 NumPy 数组等），就需要确保这些操作的数据已经真正准备好 —— 这时候你就需要同步（torch.cuda.synchronize()）。

---

📌 回到 non_blocking

当你使用 .to('cpu') 或 .cpu() 把一个 GPU 张量复制到 CPU 上时，默认行为是：

阻塞（blocking）：Python 线程会一直等待，直到这个张量从 GPU 内存复制完成，才继续执行后续代码。

但如果你设置：

y = x.to('cpu', non_blocking=True)

它的意思是：

异步（non-blocking）：Python 不会等待复制完成，而是立即返回，继续执行其他任务；
后续如果访问 y 的内容（比如打印、.numpy()），会自动等待复制完成。

---

⚠️ 使用 non_blocking=True 的前提条件：

你必须确保：

1. 数据在使用前已经完全复制完成；
2. 当前流（CUDA stream）上没有正在进行的操作会修改该张量；

否则可能会导致数据竞争（data race）或读取未定义值。

---

🔄 阻塞 vs 非阻塞对比表

特性	non_blocking=False（默认）	non_blocking=True
是否等待复制完成	✅ 是	❌ 否
是否影响性能	❌ 更慢（因为等待）	✅ 更快（异步）
安全性	✅ 安全	⚠️ 需要手动管理同步
推荐使用场景	单次拷贝、调试	多个张量并行拷贝、高性能流水线

---

🧪 示例说明

import torch

# 创建一个 GPU 张量
x = torch.randn(10000, 10000, device='cuda')

# 阻塞式拷贝
y1 = x.to('cpu')  # Python 会在这里等待拷贝完成

# 非阻塞式拷贝
y2 = x.to('cpu', non_blocking=True)  # Python 不等待，继续执行下一行
print("Non-blocking copy started")  # 这行会立刻输出

# 如果你想确保拷贝完成
torch.cuda.synchronize()
print("Copy finished")

---

🧠 回到你的完整代码

def copy_to_cpu(x, non_blocking=False):
    return [y.to('cpu', non_blocking=non_blocking) for y in x]

with d2l.Benchmark('Run on GPU1'):
    y = run(x_gpu1)
    torch.cuda.synchronize()

with d2l.Benchmark('Copy to CPU'):
    y_cpu = copy_to_cpu(y)
    torch.cuda.synchronize()

分析：

• run(x_gpu1) 在 GPU 上执行；
• copy_to_cpu(...) 将结果复制回 CPU；
• 如果你设置 non_blocking=True，那么：
  • Python 不会等待每个张量复制完成；
  • 整体复制过程更高效；
  • 最后通过 torch.cuda.synchronize() 来强制等待所有拷贝完成，以保证计时准确。

---

✅ 总结一句话：

non_blocking=True 的作用是让张量从 GPU 到 CPU 的复制操作变成异步（非阻塞）方式，从而提升性能，但你需要确保在访问数据前调用 torch.cuda.synchronize() 来等待复制完成。

---

如果你想了解如何利用 non_blocking 和 CUDA streams 做更高级的性能优化，我也可以继续深入讲解 😊

对于代码

with d2l.Benchmark('在GPU1上运行并复制到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, True)
    torch.cuda.synchronize()

和

with d2l.Benchmark('在GPU1上运行并复制到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, False)
    torch.cuda.synchronize()

后者会更慢，因为copy_to_cpu(y, non_blocking=False)这句代码在执行的时候，有可能会因为列表中后面的y先计算完，但是Python线程在等待前面的y计算完后拷贝了前面的y才能拷贝后面的y