设备

主线位置：张量主线的第一步和最后一步：建张量时选设备，计算时两边必须在同一个设备上，读回主机时等 GPU 做完。

版本：PyTorch 2.15.0（CPU、MPS）、candle 0.11.0（CPU，accelerate；Metal）

flowchart TB
  host["主机上的数据"] --> mk["建张量时选设备"]
  mk --> chk{"两边同一设备"}
  chk -->|否| err["device mismatch"]
  chk -->|是| enc["Metal 只编码"]
  enc --> more["继续排队"]
  more --> rb["读回或 synchronize"]
  rb --> wait["这时才等 GPU"]

实验代码是一个 candle 程序 dev（release）和一个 PyTorch 脚本 dev.py；平台 M2 Max（一块 GPU）、macOS 27、rustc 1.99.0。

设备是一个值

mps = torch.device("mps")
torch.device("mps") == torch.device("mps")
torch.device("mps") == torch.device("mps:0")

let cpu = Device::Cpu;
let metal = Device::new_metal(0)?;
let again = Device::new_metal(0)?;
metal.same_device(&again);
metal.same_device(&metal.clone());
Device::metal_if_available(0)?.location();

# PyTorch
device('mps') == device('mps') True
device('mps') == 'mps:0'     False
# candle
metal.same_device(again)     false
metal.same_device(clone)     true
metal_if_available(0)        Metal { gpu_id: 4294968354 }

torch.device 只是「类型 + 序号」两个字段，按值比较（"mps" 的序号是 None，所以和 "mps:0" 不相等，但张量照样能一起算）。candle 的 Device 是一个枚举，Metal 变体里带着设备句柄、命令队列和缓冲池。Device::new_metal 每调一次就建一套新的，same_device 比的是进程内的编号，所以两次 new_metal(0) 是两个设备，互相不能运算。clone() 出来的是同一个设备。机制和报错见源码卷的 device mismatch in add。

对应的写法：PyTorch 里到处写 "mps" 字符串没问题；candle 里在 main 里建一次 Device，以 &Device 传给加载函数、生成函数。要「有 GPU 就用、没有就退回 CPU」，用 Device::metal_if_available，对应 torch.backends.mps.is_available() 加一个 if。

不做隐式搬运

a = torch.ones(2, 2)
m = torch.ones(2, 2, device=mps)
a + m
m + torch.tensor(1.0)      # 0 维 CPU 张量
m + torch.tensor([1.0])    # 1 个元素的 1 维 CPU 张量
a @ m
emb = torch.nn.Embedding(4, 3).to(mps)
emb(torch.tensor([1, 2]))  # 下标留在 CPU

let a = Tensor::ones((2, 2), DType::F32, &cpu)?;
let m = Tensor::ones((2, 2), DType::F32, &metal)?;
show("cpu + metal", &a + &m);
let s = Tensor::new(1f32, &cpu)?;
show("metal + 0-dim cpu (bcast)", m.broadcast_add(&s));
show("metal + 1.0 (f64 scalar)", &m + 1.0);
show("cpu.matmul(metal)", a.matmul(&m));
let emb = candle_nn::Embedding::new(Tensor::ones((4, 3), DType::F32, &metal)?, 3);
show("metal Embedding(cpu ids)", emb.forward(&Tensor::new(&[1u32, 2], &cpu)?));

# PyTorch
cpu + mps                    ERR RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!
mps + 0-dim cpu              ok torch.float32 mps:0 [2.0, 2.0, 2.0, 2.0]
mps + 1-elem cpu [1]         ERR RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!
cpu.matmul(mps)              ERR RuntimeError: Tensor for argument #1 'mat1' is on CPU, but expected it to be on GPU (while checking arguments for mm)
mps Embedding(cpu ids)       ERR RuntimeError: Passed CPU tensor to MPS op
# candle
cpu + metal                  ERR device mismatch in add, lhs: Cpu, rhs: Metal { gpu_id: 4294968354 }
metal + 0-dim cpu (bcast)    ERR device mismatch in add, lhs: Metal { gpu_id: 4294968354 }, rhs: Cpu
metal + 1.0 (f64 scalar)     ok F32 Metal { gpu_id: 4294968354 } [2.0, 2.0, 2.0, 2.0]
cpu.matmul(metal)            ERR device mismatch in matmul, lhs: Cpu, rhs: Metal { gpu_id: 4294968354 }
metal Embedding(cpu ids)     ERR device mismatch in index-select, lhs: Metal { gpu_id: 4294968354 }, rhs: Cpu

两边都不替你搬数据。差别在两处：

PyTorch 对 0 维的 CPU 张量网开一面，把它当标量用；1 个元素的 1 维张量就不行了。candle 没有这个例外，0 维张量也要在同一个设备上。和 Rust 的 f64 字面量运算（&m + 1.0）没有设备问题，它走 affine，标量作为参数传给内核。
candle 的报错格式统一：device mismatch in {op}, lhs: .., rhs: ..，op 是内部算子名。Embedding 报的是 index-select，因为它的 forward 就是一次 index_select。检查在 Storage::same_device，每个二元的 Storage 方法开头都调用。

transformers 里常见的 inputs = tokenizer(..., return_tensors="pt").to(model.device)，在 candle 里对应的是建 token 张量时就传模型的设备：Tensor::new(tokens.as_slice(), &device)?.unsqueeze(0)?。

模型在加载时定设备

PyTorch 的 model.to("mps") 把已经建好的模块的参数原地搬过去。candle 的模块（candle_nn::Linear、candle_transformers::models::llama::Llama 等）没有这一步：candle-nn 和 candle-transformers 的 src 里一个 fn to_device 都没有。设备在建模块时由 VarBuilder 决定：

let vb = unsafe { VarBuilder::from_mmaped_safetensors(&files, DType::F32, &device)? };
let model = Llama::load(vb, &config)?;

VarBuilder 每次取一个权重时，按自己记着的 dtype 和设备把它读出来。VarBuilder::set_device 只是换一个记着的设备，返回新的 VarBuilder，已经建好的模块不受影响。想把模型换到另一个设备，就用新设备的 VarBuilder 再建一次。权重文件怎么读、选什么 dtype，见safetensors。

to_device

m.to(mps) is m
m.to("cpu") is m

let same = m.to_device(&metal)?;
println!("{}", shares_storage(&m, &same));
let back = m.to_device(&cpu)?;
println!("{}", shares_storage(&m, &back));
show("m.to_device(again)", m.to_device(&again));

# PyTorch
m.to(mps) is m               True
m.to('cpu') is m             False
# candle
m.to_device(metal) shares    true
m.to_device(cpu) shares      false
m.to_device(again)           ERR not implemented yet, self.device: Metal(MetalDevice(DeviceId(1))), device: Metal(MetalDevice(DeviceId(2)))

（shares_storage 用公开的 storage_and_layout 比较两个张量底层的 Storage 是不是同一个。）已经在目标设备上时，两边都原样返回：PyTorch 返回同一个对象，candle 的 to_device 返回 self.clone()，共享 storage。换设备时整块复制。两个 Metal 设备之间没有实现，只能经 CPU 中转，这又是「设备只建一次」的理由。

dtype 和设备要一起选

PyTorch 的 MPS 不支持 float64，在搬过去的那一刻就拒绝。candle 的 Metal 能存 F64，却不能拿它计算，连转成 F32 都不行：

torch.tensor([1.0, 2.0], dtype=torch.float64).to(mps)
torch.tensor([1.0, 2.0], device=mps)

let lit = Tensor::new(&[1.0, 2.0], &metal)?;  // 浮点字面量是 F64
show("literal [1.0, 2.0] on metal", Ok(lit.clone()));
show("  ... + itself", &lit + &lit);
show("  ... to_dtype(F32)", lit.to_dtype(DType::F32));
show("F64 cpu.to_device(metal)", Tensor::new(&[1.0, 2.0], &cpu)?.to_device(&metal));
show("  ... then to F32 on cpu", Tensor::new(&[1.0, 2.0], &cpu)?.to_dtype(DType::F32)?.to_device(&metal));

# PyTorch
float64 to mps               ERR TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
literal [1.0, 2.0] on mps    ok torch.float32 mps:0 [1.0, 2.0]
# candle
literal [1.0, 2.0] on metal  ok F64 Metal { gpu_id: 4294968354 } [1.0, 2.0]
  ... + itself               ERR Metal error Error while loading function: badd_f64
  ... to_dtype(F32)          ERR Metal contiguous to_dtype F64 F32 not implemented
F64 cpu.to_device(metal)     ok F64 Metal { gpu_id: 4294968354 } [1.0, 2.0]
  ... then to F32 on cpu     ok F32 Metal { gpu_id: 4294968354 } [1.0, 2.0]

PyTorch 里浮点字面量默认是 float32，torch.tensor([1.0, 2.0], device=mps) 没问题。candle 里 Tensor::new(&[1.0, 2.0], ..) 是 F64（见创建与 dtype），直接建在 Metal 上能成功、能读回，到第一次运算才报错；想在 Metal 上把它转成 F32 也不行。只能在 CPU 上转好再搬，或者一开始就写 1f32。CPU 上 BF16 不能 matmul 是同一类问题，各 dtype 在两个设备上能做什么，见源码卷的 unsupported dtype BF16 for op matmul。

📝NOTE

没有蠢问题

问： F64 张量能上传到 Metal、也能读回来，为什么在 Metal 上连 to_dtype(F32) 都做不了？

答：上传和读回是整块拷贝字节，不看 dtype。to_dtype 要跑一个转换内核，而 cast.metal 只给 f32、f16、bf16、i64、u32、u8 六种类型两两生成了 cast_* 内核，没有 f64。MetalStorage::to_dtype 按（源，目标）配对找内核名，找不到就 bail!("Metal contiguous to_dtype {left:?} {right:?} not implemented")。加法则是拼出了内核名 badd_f64，到内核库里找不到这个函数，报 Error while loading function。

异步：读回才等

MPS 和 candle 的 Metal 后端都是异步的：一个 op 编码进命令缓冲区就返回，GPU 稍后执行。两个 (2048, 2048) 的 f32 随机矩阵，热身一次后连做 20 次 matmul，分别在循环刚结束时和同步之后读时钟：

t0 = time.perf_counter()
outs = [x @ y for _ in range(reps)]
naive = time.perf_counter() - t0
torch.mps.synchronize()
synced = time.perf_counter() - t0

let t0 = Instant::now();
let mut outs = Vec::new();
for _ in 0..reps {
  outs.push(x.matmul(&y)?);
}
let naive = t0.elapsed().as_secs_f64();
metal.synchronize()?;
let synced = t0.elapsed().as_secs_f64();

# PyTorch MPS
round 0: 20 matmuls 2048x2048 f32: returned after 0.60 ms, synchronized after 48.8 ms -> 2.438 ms each
round 1: 20 matmuls 2048x2048 f32: returned after 0.43 ms, synchronized after 39.4 ms -> 1.971 ms each
round 2: 20 matmuls 2048x2048 f32: returned after 0.41 ms, synchronized after 39.4 ms -> 1.971 ms each
# candle Metal
round 0: 20 matmuls 2048x2048 f32: returned after 7.75 ms, synchronized after 53.1 ms -> 2.654 ms each
round 1: 20 matmuls 2048x2048 f32: returned after 0.06 ms, synchronized after 38.5 ms -> 1.927 ms each
round 2: 20 matmuls 2048x2048 f32: returned after 0.06 ms, synchronized after 38.6 ms -> 1.932 ms each

（release；三轮取最小，每次 1.93 ms 对 1.97 ms；跑之前 uptime 的 1 分钟负载 3.35，其他实验在并行；Metal 着色器缓存是热的。MPS 那边线程数固定为 4，对 GPU 计时没有影响。）

习惯是一样的：PyTorch 里计时前写 torch.mps.synchronize()，candle 里写 device.synchronize()，或者读回一个值（to_vec*、to_scalar、to_device(&Device::Cpu) 都会等之前排着的所有 op 做完）。不同步时循环「只花了」0.06 ms，那是编码的时间。candle 第 0 轮返回时已经花了 7.75 ms：这一轮的 20 个 16 MiB 输出缓冲区都要新分配，后两轮能复用上一轮释放回缓冲池的那些（缓冲池见源码卷 Metal 后端）。提交和等待的细节、第一次运行的着色器编译，见源码卷的 Metal 上计时几乎为零和 Metal 第一次运行要好几秒。

小结

torch.device 按值比较；candle 的 Device 带着命令队列，两次 new_metal(0) 是两个互不相通的设备。建一次，&Device 传下去。
两边都不隐式搬运；PyTorch 允许 0 维 CPU 张量参与 GPU 运算，candle 不允许，f64 标量可以。报错都是 device mismatch in {op}。
candle 的模块没有 .to(device)，设备在 VarBuilder 加载时定，换设备就重建模块；输入张量建在同一个设备上。
to_device 在同一设备上原样返回、共享 storage；Metal 到另一个 Metal 设备没实现。
F64 能上传到 Metal 但不能计算、不能在 Metal 上转换；浮点字面量是 F64，先在 CPU 上转 F32。
Metal 是异步的，计时前 device.synchronize()，和 torch.mps.synchronize() 一样。

设备