Metal: Tensor::copy() shares the source's buffer, so in-place updates (Var::set, zero_set) show through the "copy" #4027

Version / platform: candle-core 0.9.1 and 0.11.0 (the code is unchanged on main); macOS 27, Apple M1 Max, metal feature; rustc 1.97.1.

Description

Tensor::copy is documented as "Compared to clone, this copies the actual storage". It builds the new storage with BackendStorage::try_clone. On Metal that is:

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

and MetalStorage holds buffer: Arc<Buffer>. So the "copy" shares the original's GPU buffer, and any later in-place write to the original (Var::set, which the candle-nn optimizers use, or zero_set / const_set / slice_set) also changes the copy. The CPU backend (clones the Vec) and the CUDA backend (device copy) make independent copies, so the same code behaves differently on Metal, with no error.

One practical effect: snapshotting a parameter with .copy() before an optimizer step (an EMA or teacher copy, or a saved checkpoint) silently tracks the live weights on Metal. We hit this when a published snapshot of model weights changed after the next training step.

Minimal repro (Cargo.toml: candle-core = { version = "=0.11.0", features = ["metal"] })

use candle_core::{Device, Tensor, Var};

fn check(dev: &Device) -> candle_core::Result<bool> {
    let v = Var::new(&[1f32, 2., 3.], dev)?;
    let c = v.as_tensor().copy()?;
    v.set(&Tensor::new(&[10f32, 20., 30.], dev)?)?;
    let c1 = c.to_vec1::<f32>()?;
    let t = Tensor::new(&[1f32, 2., 3.], dev)?;
    let d = t.copy()?;
    t.zero_set()?;
    let c2 = d.to_vec1::<f32>()?;
    println!("{:?}: copy after Var::set = {:?}, copy after zero_set = {:?}", dev.location(), c1, c2);
    Ok(c1 == [1., 2., 3.] && c2 == [1., 2., 3.])
}

fn main() -> candle_core::Result<()> {
    let cpu_ok = check(&Device::Cpu)?;
    let metal_ok = check(&Device::new_metal(0)?)?;
    println!("cpu ok: {cpu_ok}, metal ok: {metal_ok}");
    Ok(())
}

Observed (0.9.1 and 0.11.0):

Cpu: copy after Var::set = [1.0, 2.0, 3.0], copy after zero_set = [1.0, 2.0, 3.0]
Metal { .. }: copy after Var::set = [10.0, 20.0, 30.0], copy after zero_set = [0.0, 0.0, 0.0]
cpu ok: true, metal ok: false

Expected: the copy keeps [1.0, 2.0, 3.0] on Metal, as it does on CPU and CUDA.

Suggested fix: in MetalStorage::try_clone, allocate a new buffer of the same length and blit the source into it, rather than cloning the Arc<Buffer>.

Workaround: produce the copy with any kernel that writes a new buffer, for example t.affine(1.0, 0.0)?.