Rust in Everything
基于 candle-core 0.11.0

改了 copy 出来的张量，原张量也变了

主线位置：求导主线的最后一步。优化器用 Var::set 把新值原地写进变量的存储；在这之前用 copy() 留的快照，在 Metal 上和变量共用同一块显存。Storage 的共享规则见 Tensor 的结构，MetalStorage::try_clone 和缓冲池见 Metal 后端。

flowchart TB
  w["Var w"] --> cp["w.copy()"]
  cp --> d{"设备"}
  d -->|CPU| own["新 Vec 独立"]
  d -->|Metal| alias["同一个 Buffer"]
  st["w.set 原地写"] --> alias
  alias --> bad["快照也变了"]
  fix["affine 1 0"] --> fresh["新缓冲区"]

实验 crate exp-candle/pa-pit，bin p11_copy，debug；M2 Max、macOS 27、rustc 1.99.0。

现象

训练前把初始权重留一份，用来事后比较参数变了多少，或者出问题时回滚：

let w = Var::from_tensor(&Tensor::ones(4, DType::F32, &dev)?)?;
let snapshot = w.as_tensor().copy()?;
let by_affine = w.as_tensor().affine(1.0, 0.0)?;
let on_cpu = w.as_tensor().to_device(&Device::Cpu)?;
w.set(&(w.as_tensor() * 0.5)?)?;

另外对一个普通张量 orig 做 copy()，再往副本里 slice_set：

== Cpu
w after set:      [0.5, 0.5, 0.5, 0.5]
copy() snapshot:  [1.0, 1.0, 1.0, 1.0]
affine snapshot:  [1.0, 1.0, 1.0, 1.0]
to_device(Cpu):   [0.5, 0.5, 0.5, 0.5]
orig after c.slice_set: [0.0, 0.0, 0.0, 0.0]
== Metal { gpu_id: 4294968354 }
w after set:      [0.5, 0.5, 0.5, 0.5]
copy() snapshot:  [0.5, 0.5, 0.5, 0.5]
affine snapshot:  [1.0, 1.0, 1.0, 1.0]
to_device(Cpu):   [1.0, 1.0, 1.0, 1.0]
orig after c.slice_set: [1.0, 1.0, 0.0, 0.0]

CPU 上 copy() 是独立的副本；Metal 上它跟着 w 变成了 0.5，往副本里写也改到了原张量。

这个坑在写本卷时被实际踩到过：RmsNorm 之后梯度是 None 一页的训练程序最初用 w1.as_tensor().copy()? 记初始权重，Metal 上算出「1000 步后 w1 变化量为 0」，CPU 上同一段代码是 0.0427。

出处

Tensor::copy 靠 Storage::try_clone 复制数据：

    pub fn copy(&self) -> Result<Tensor> {
        let op = BackpropOp::new1(self, Op::Copy);
        let tensor_ = Tensor_ {
            id: TensorId::new(),
            storage: Arc::new(RwLock::new(self.storage().try_clone(self.layout())?)),
            // ...
        };
        Ok(Tensor(Arc::new(tensor_)))
    }

CPU 的 CpuStorage 是一组 Vec，clone 是深拷贝。Metal 的 MetalStorage::try_clone 是：

    fn try_clone(&self, _: &Layout) -> Result<Self> {
        Ok(self.clone())
    }

MetalStorage 里的缓冲区是 Arc<Buffer>，clone 只多一个引用。新张量有自己的 TensorId 和 RwLock，底下是同一块显存。这一行从 Metal 后端最初的实现起就是这样。

Var::set 不换存储，而是把新值 copy_strided_src 写进变量现有的存储；slice_set 同样原地写。于是共用缓冲区的快照一起变。只读的运算（加法、matmul）总是分配新的输出缓冲区，不受影响，这也是这个 bug 平时不显形的原因。这是 0.11.0 的上游 bug：copy 的文档注释写着「Compared to clone, this copies the actual storage」，Metal 上没有做到。

to_device(&Device::Cpu) 的两行结果相反，原因是 to_device 第一行：目标设备和当前设备相同时直接 Ok(self.clone())。在 Metal 上它是一次真正的下载，在 CPU 上它什么都不做，快照和变量共用存储。

修法

要独立副本，用一个总会分配新输出的运算：

let snapshot = w.as_tensor().affine(1.0, 0.0)?;

上面的输出里 affine 的快照在两个设备上都保持 1.0。不要用这几个看起来像复制的方法：

  • to_dtype(同一个 dtype)、contiguous()（已连续时）、clone()：都只是 clone，和原张量共用存储。
  • to_device(&Device::Cpu)：Metal 上可以，CPU 上不行，同一段代码换设备行为就变。
  • detach()：只丢掉计算图，存储共用。

要把快照放在主机内存里，to_vec* 读成 Vec 最明确。

怎么确认没中招

  • 在 Metal 上搜 .copy()：只要之后对原张量或副本有 Var::set、slice_set、zero_set、inplace_op*，就会互相影响。candle-core、candle-nn、candle-transformers 的 src/ 里没有调用 .copy()，问题只出在用户代码。
  • 写一个最小检查：let c = t.copy()?; c.zero_set()?; 后读 t，看它有没有变成 0。
  • 比较「训练前后参数变化」得到 0 时，先怀疑快照，而不是优化器。

小结

  • Metal 上 Tensor::copy() 不复制：MetalStorage::try_clone 只克隆 Arc<Buffer>（0.11.0 上游 bug）。
  • 原地写（Var::set、slice_set、zero_set）会同时改到原张量和副本；只读运算不受影响。
  • 要副本用 affine(1.0, 0.0)；to_device(&Cpu)、clone、detach、同 dtype 的 to_dtype 都不复制（CPU 上的 to_device(&Cpu) 也不）。

💬 关联讨论

发起讨论
加载中...