PyTorch 的 .bin / .pt

model = AutoModelForCausalLM.from_pretrained(model_dir)  # 默认 dtype，bf16
sd = model.state_dict()
torch.save(sd, f"{out}/sd_zip.pt")
torch.save(sd, f"{out}/sd_legacy.pt", _use_new_zipfile_serialization=False)
torch.save(model, f"{out}/full.pt")
torch.save({"model": sd, "step": 100}, f"{out}/ckpt.pt")

let vb = VarBuilder::from_pth(format!("{dir}/sd_zip.pt"), DType::F32, &dev)?;
let model = Llama::load(vb, &cfg)?;

load 258.0 ms
logits                   n=49152  max|diff|=2.98e-05
worst=2.98e-05 tol=1e-04 OK

/// Lazy tensor loader.
pub struct PthTensors {
    tensor_infos: HashMap<String, TensorInfo>,
    path: std::path::PathBuf,
    // We do not store a zip reader as it needs mutable access to extract data. Instead we
    // re-create a zip reader for each tensor.
}

for f in ["sd_zip.pt", "sd_legacy.pt", "ckpt.pt", "full.pt", "odd.pt", "evil.pt"]:
  try:
    obj = torch.load(f"{d}/{f}")
    keys = list(obj)[:2] if isinstance(obj, dict) else type(obj).__name__
    print(f"{f:13} weights_only=True  ok, {type(obj).__name__}, {keys}")
  except Exception as e:
    print(f"{f:13} weights_only=True  {type(e).__name__}: {str(e).splitlines()[0][:110]}")

for f in ["sd_zip.pt", "sd_legacy.pt", "ckpt.pt", "full.pt"] {
  let path = format!("{dir}/{f}");
  match pickle::read_pth_tensor_info(&path, false, None) {
    Ok(infos) => println!("{f}: read_pth_tensor_info -> {} tensors", infos.len()),
    Err(e) => println!("{f}: read_pth_tensor_info -> Err: {e}"),
  }
  let r = VarBuilder::from_pth(&path, DType::F32, &dev).and_then(|vb| Llama::load(vb, &cfg));
  // ... 打印 ok 或 Err
}

# PyTorch
sd_zip.pt     weights_only=True  ok, OrderedDict, ['model.embed_tokens.weight', 'model.layers.0.self_attn.q_proj.weight']
sd_legacy.pt  weights_only=True  ok, OrderedDict, ['model.embed_tokens.weight', 'model.layers.0.self_attn.q_proj.weight']
ckpt.pt       weights_only=True  ok, dict, ['model', 'step']
full.pt       weights_only=True  UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps onl
# candle
sd_zip.pt: read_pth_tensor_info -> 273 tensors
sd_zip.pt: from_pth + Llama::load -> ok
sd_legacy.pt: read_pth_tensor_info -> Err: invalid Zip archive: Could not find EOCD
sd_legacy.pt: from_pth + Llama::load -> Err: invalid Zip archive: Could not find EOCD
ckpt.pt: read_pth_tensor_info -> 0 tensors
ckpt.pt: from_pth + Llama::load -> Err: cannot find tensor model.embed_tokens.weight
full.pt: read_pth_tensor_info -> 0 tensors
full.pt: from_pth + Llama::load -> Err: cannot find tensor model.embed_tokens.weight

let vb = VarBuilder::from_pth_with_state(format!("{dir}/ckpt.pt"), DType::F32, "model", &dev)?;

load 209.4 ms
logits                   n=49152  max|diff|=2.98e-05
worst=2.98e-05 tol=1e-04 OK

UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint. 
	(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
	(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
	WeightsUnpickler error: Unsupported global: GLOBAL transformers.models.llama.modeling_llama.LlamaForCausalLM was not an allowed global by default. ...

        let obj = match obj {
            Object::Build { callable, args } => match *callable {
                Object::Reduce { callable, args: _ } => match *callable {
                    Object::Class {
                        module_name,
                        class_name,
                    } if module_name == "__torch__" && class_name == "Module" => *args,
                    _ => continue,
                },
                _ => continue,
            },
            obj => obj,
        };

lin_sd.pt: read_all -> ["bias [2]", "weight [2, 3]"]
lin_full.pt: read_all -> []
jit.pt: read_all -> []

base = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)
odd = {
  "f32": torch.arange(6, dtype=torch.float32).reshape(2, 3),
  "f64": torch.arange(3, dtype=torch.float64),
  "i64": torch.arange(3, dtype=torch.int64),
  "i32": torch.arange(3, dtype=torch.int32),
  "u8": torch.arange(3, dtype=torch.uint8),
  "bool": torch.tensor([True, False, True]),
  "transposed": torch.arange(6, dtype=torch.float32).reshape(2, 3).t(),
  "sliced": torch.arange(10, dtype=torch.float32)[4:7],
  "permuted": base.permute(1, 0, 2),
}
torch.save(odd, f"{out}/odd.pt")

transposed torch.float32  shape=[3, 2] stride=(1, 3) offset=0 contiguous=False
sliced     torch.float32  shape=[3] stride=(1,) offset=4 contiguous=True
permuted   torch.float32  shape=[3, 2, 4] stride=(4, 12, 1) offset=0 contiguous=False

let pth = pickle::PthTensors::new(format!("{dir}/odd.pt"), None)?;
for name in ["f32", "f64", "i64", "i32", "u8", "bool", "transposed", "sliced", "permuted"] {
  match pth.get(name) {
    Ok(Some(t)) => {
      let v: Vec<f64> = t.flatten_all()?.to_dtype(DType::F64)?.to_vec1()?;
      println!("{name:10} {:?} {:?} {v:?}", t.dtype(), t.dims())
    }
    Ok(None) => println!("{name:10} None"),
    Err(e) => println!("{name:10} Err: {e}"),
  }
}

skipping: unsupported storage type BoolStorage
skipping: unsupported storage type IntStorage
tensor infos: ["f32", "f64", "i64", "permuted", "sliced", "transposed", "u8"]
f32        F32 [2, 3] [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
f64        F64 [3] [0.0, 1.0, 2.0]
i64        I64 [3] [0.0, 1.0, 2.0]
i32        None
u8         U8 [3] [0.0, 1.0, 2.0]
bool       None
transposed F32 [3, 2] [0.0, 3.0, 1.0, 4.0, 2.0, 5.0]
sliced     F32 [3] [4.0, 5.0, 6.0]
permuted   Err: cannot retrieve non-contiguous tensors Layout { shape: [3, 2, 4], stride: [4, 12, 1], start_offset: 0 }

class Shout:
  # torch.load(weights_only=False) 时，pickle 会调用 __reduce__ 返回的函数
  def __reduce__(self):
    return (os.system, (f"echo pickle ran code > {out}/pwned.txt",))


torch.save({"w": torch.ones(2), "hook": Shout()}, f"{out}/evil.pt")

evil.pt       weights_only=True  UnpicklingError: Weights only load failed. In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torc
evil.pt weights_only=False: ['w', 'hook'] pwned.txt exists: True

for (n, t) in pickle::read_all(format!("{dir}/evil.pt"))? {
  println!("{n}: {:?} {:?}", t.dims(), t.to_vec1::<f32>()?);
}
println!("pwned.txt exists: {}", std::path::Path::new(&format!("{dir}/pwned.txt")).exists());

w: [2] [1.0, 1.0]
pwned.txt exists: false