Rust in Everything
基于 candle 0.11.0

clone 出来的快照跟着参数变

主线位置：训练主线的「优化器一步」。训练中途留一份「目前最好的」参数，之后优化器原地改写参数。
版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）
flowchart TB
  w["Var 参数"] --> c["clone 或 detach"]
  w --> cp["copy"]
  w --> mul["乘 1.0"]
  st["opt.step 原地写"] --> w
  c --> follow["跟着变"]
  cp -->|CPU| keep["留住旧值"]
  cp -->|Metal| follow
  mul --> keep

本页的对照：PyTorch 侧是实验代码里的 p11.py，candle 侧是实验 crate pita-t 的 p11_snapshot（features metal accelerate，release）。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

state_dict() 返回的张量和参数共享存储，所以留快照要 clone() 或 deepcopy：

lin = torch.nn.Linear(2, 1, bias=False)            # weight = [[1.0, 2.0]]
opt = torch.optim.SGD(lin.parameters(), lr=0.5)
sd = lin.state_dict()
best_clone = {k: v.clone() for k, v in lin.state_dict().items()}
best_deep = copy.deepcopy(lin.state_dict())
lin(torch.tensor([[1.0, 1.0]])).sum().backward(); opt.step()
buf = torch.zeros(2, 2); snap = buf.clone(); buf[0] = torch.tensor([7.0, 7.0])

weight after step        [[0.5, 1.5]]
state_dict() (no copy)   [[0.5, 1.5]]
{k: v.clone()}           [[1.0, 2.0]]
copy.deepcopy(state_dict) [[1.0, 2.0]]
buf.clone() after buf[0] = 7 [[0.0, 0.0], [0.0, 0.0]]

PyTorch 用户的习惯是：clone() 就是一份独立的数据。

搬到 candle 的现象

同一个线性层，参数放在 VarMap 里，五种「快照」：

let w = &vars["lin.weight"];
let snap_clone = w.as_tensor().clone();
let snap_detach = w.as_tensor().detach();
let snap_copy = w.as_tensor().copy()?;
let snap_mul = (w.as_tensor() * 1.0)?;
let snap_cpu = w.as_tensor().to_device(&Device::Cpu)?;
let mut opt = SGD::new(varmap.all_vars(), 0.5)?;
opt.backward_step(&lin.forward(&x)?.sum_all()?)?;
// a preallocated buffer written in place, as a KV cache does
let buf = Tensor::zeros((2, 2), DType::F32, &dev)?;
let (snap, snap_copy) = (buf.clone(), buf.copy()?);
buf.slice_set(&Tensor::new(&[[7f32, 7.]], &dev)?, 0, 0)?;

== cpu
weight after step      [[0.5, 1.5]]
as_tensor().clone()    [[0.5, 1.5]]
as_tensor().detach()   [[0.5, 1.5]]
as_tensor().copy()     [[1.0, 2.0]]
as_tensor() * 1.0      [[1.0, 2.0]]
to_device(Cpu)         [[0.5, 1.5]]
buf.clone() after slice_set [[7.0, 7.0], [0.0, 0.0]]
buf.copy() after slice_set  [[0.0, 0.0], [0.0, 0.0]]
== metal
weight after step      [[0.5, 1.5]]
as_tensor().clone()    [[0.5, 1.5]]
as_tensor().detach()   [[0.5, 1.5]]
as_tensor().copy()     [[0.5, 1.5]]
as_tensor() * 1.0      [[1.0, 2.0]]
to_device(Cpu)         [[1.0, 2.0]]
buf.clone() after slice_set [[7.0, 7.0], [0.0, 0.0]]
buf.copy() after slice_set  [[7.0, 7.0], [0.0, 0.0]]

clone() 和 detach() 两个设备上都跟着变。copy() 在 CPU 上留住了旧值，在 Metal 上也跟着变。to_device(&Device::Cpu) 正好相反：Metal 上是一次真的下载，留住了旧值；CPU 上它什么也不做，跟着变。只有 * 1.0 在两个设备上都对。

原因

  • candle 的 Tensor 是 Arc 包着的结构，clone() 只多一个引用，和 PyTorch 的 clone() 正好相反，见引用与不可变。
  • detach 新建一个 id，但 storage: self.storage.clone()，和 PyTorch 的 detach() 一样共用存储。
  • 优化器用 Var::set 把新值写进参数自己的存储，slice_set 同样是原地写，所有共用这块存储的张量都看到新值。
  • copy 调后端的 try_clone。CPU 上它复制数据；Metal 的 MetalStorage::try_clone 是 Ok(self.clone())，只多一个缓冲区引用。这是 0.11.0 的上游问题，origin/main 上没有改，现象和来历见源码卷的改了 copy 出来的张量，原张量也变了。
  • to_device 遇到同一个设备时直接 Ok(self.clone())。

几种快照在 opt.step 前后的变化，没有 no_grad的「detach 和参数快照」一节用一个两元素的 Var 讲过；本页换成 VarMap 和 slice_set，结论相同。

改法

用一次真正的计算产生新存储：

let snap = (w.as_tensor() * 1.0)?;       // or w.as_tensor().affine(1.0, 0.0)?

as_tensor() * 1.0      [[1.0, 2.0]]

* 1.0 走 affine，没有「乘 1 加 0 就原样返回」的捷径，后端分配新的输出缓冲区再写进去，CPU 和 Metal 都留住了旧值。结果仍带着一个指向参数的反向节点，快照如果之后还要参与计算，再接一个 .detach()：这时存储已经独立，detach 只是切断图。对整个 VarMap 做快照，就对 all_vars() 逐个这样做；要把快照留到进程之外，VarMap::save 写 safetensors 文件。

怎么检查

  • 搜 .clone()、.detach()、.copy() 后面紧跟着优化器一步或 slice_set 的地方；在 candle 里这些都不保证留住旧值。
  • 写一个两步的小测试：留快照、改参数、断言快照没变。要在 CPU 和 Metal 上各跑一次，copy() 只在 Metal 上出错。
  • 快照用来做 EMA、早停、「最好的权重」时，错误不会报出来，只会表现为「回滚之后效果没变」。

小结

  • candle 的 clone() 是引用，detach() 共用存储；opt.step（Var::set）和 slice_set 之后它们跟着变。
  • copy() 在 CPU 上复制，在 Metal 上也共用缓冲区（上游问题）；to_device(&Device::Cpu) 在 CPU 上什么也不做。
  • 改法：(p * 1.0)? 或 p.affine(1.0, 0.0)?，两个设备都对。

💬 关联讨论

发起讨论
加载中...