ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Candle 从 Hugging Face Hub 加载预训练权重:下载、safetensors 解析与张量并行分片实战

Candle 从 Hugging Face Hub 加载预训练权重:下载、safetensors 解析与张量并行分片实战 人工智能大模型机器学习深度学习本地部署模型推理服务【免费下载链接】candleMinimalist ML framework for Rust项目地址https://gitcode.com/GitHub_Trending/ca/candle点击查看免费下载本篇指南围绕 Candle 推理链路中的第一步——模型权重获取展开讲解如何借助hf-hub从 Hugging Face Hub 下载预训练权重以bert-base-uncased为例将safetensors文件解析为candle-core的Tensor并进一步打通内存映射加载mmap与多 GPU 张量并行Tensor Parallel场景下的按分片加载方案。读完本文你将掌握一套可复用的下载 → 加载 → 接入模型 → 分片完整流程并了解其底层 API 与测试验证。为什么需要 hubCandle 中的权重获取方式Candle 自身不托管模型文件绝大多数预训练模型的权重以safetensors或老旧的pytorch_model.bin格式存放在 Hugging Face Hub 上。因此任何推理或微调流程的第一步都是从 Hub 拉取模型文件。Candle 官方在 candle-book/src/inference/hub.md 中给出了标准做法使用官方维护的hf-hubRust crate它可以处理模型库与数据集库的鉴权、缓存、断点续传与 revision 固定。仓库中candle-examples的大量示例都依赖这一套流程例如 bert 示例 会先下载config.json、tokenizer.json与model.safetensors再构建VarBuilder加载模型candle-examples甚至把 hub 的常用封装收敛到了 hub.rs统一了缓存目录、revision 与下载进度条的处理。安装依赖并下载第一个模型文件在Cargo.toml中加入hf-hubcargo add hf-hub随后用下面的代码下载bert-base-uncased仓库中的model.safetensorsuse hf_hub::{split_id, HFClientSync}; use candle_core::Device; let api HFClientSync::new().unwrap(); let (owner, name) split_id(bert-base-uncased); let repo api.model(owner, name); let weights repo.download_file().filename(model.safetensors).send().unwrap();这里有几个关键点split_id(bert-base-uncased)把owner/name形式的仓库 ID 拆成(bert-base-uncased, )——对于没有组织前缀的仓库owner 与 name 相同对sentence-transformers/all-MiniLM-L6-v2这类 ID 则会正确拆为 owner 与 name 两部分。api.model(owner, name)返回一个针对模型仓库的句柄HFRepositorySyncRepoTypeModel与之对应还有api.dataset(...)用于数据集仓库。download_file().filename(model.safetensors)是可链式配置的请求构造器send()返回下载到本地缓存的文件路径PathBuf默认缓存目录遵循hf-hub约定也支持通过环境变量如HF_TOKEN、HF_HUB_CACHE配置鉴权与缓存位置。如果你使用的是异步运行时hf-hub同样提供HFClientasync 版本candle-book 的测试代码 candle-book/src/lib.rs#L13-L27book_hub_1中即采用use candle::Device; use hf_hub::{split_id, HFClient}; let api HFClient::new().unwrap(); let (owner, name) split_id(bert-base-uncased); let repo api.model(owner, name); let weights_filename repo.download_file().filename(model.safetensors).send().await.unwrap(); let weights candle::safetensors::load(weights_filename, Device::Cpu).unwrap();两者的区别仅在同步/异步HFClientSync直接阻塞返回PathBufHFClient返回Future需要.await按你的运行时环境二选一即可。将 safetensors 解析为 Tensor 集合下载完成后用candle_core::safetensors::load一次性把文件读入内存并解析let weights candle_core::safetensors::load(weights, Device::Cpu);load的返回值是HashMapString, Tensorkey 为张量名例如bert.encoder.layer.0.attention.self.query.weightvalue 为对应Tensor。其底层实现位于 candle-core/src/safetensors.rs#L408-L419先用std::fs::read读入完整字节再调用load_buffer通过SafeTensors::deserialize解析头部并逐个张量转换为 Candle 的Tensor。转换过程convert/convert_slice会按safetensors的 dtype 与形状直接构造存储因此F32、F16、BF16等常见权重格式都开箱即用。值得注意candle-book 的对应测试book_hub_1中有一条断言assert_eq!(weights.len(), 206)即bert-base-uncased的model.safetensors一共包含 206 个张量——你可以通过这个数字快速验证自己的加载流程是否完整。把权重接入真实模型以 BERT 的一个 Linear 层为例拿到HashMapString, Tensor之后就可以按张量名取用参数。文档中给出了最直接的用法——取 BERT 第一层 self-attention 中 query 投影的权重和偏置构造一个Linear并前向计算use candle_core::{Device, Tensor, DType}; use candle_nn::{Linear, Module}; let weights candle_core::safetensors::load(weights, Device::Cpu).unwrap(); let weight weights.get(bert.encoder.layer.0.attention.self.query.weight).unwrap(); let bias weights.get(bert.encoder.layer.0.attention.self.query.bias).unwrap(); let linear Linear::new(weight.clone(), Some(bias.clone())); let input_ids Tensor::zeros((3, 768), DType::F32, Device::Cpu).unwrap(); let output linear.forward(input_ids).unwrap();Linear::new来自 candle-nn输入形状(3, 768)对应bert-base-uncased的 hidden size 768。前向得到(3, 768)的输出该层输出维度与输入一致因为 query 投影保持隐层维度。如果要在生产代码里完整加载 BERT 而非手写张量名推荐直接复用仓库中的完整实现。可以参考 bert 示例其build_model_and_tokenizer展示了真实工程的做法用Api::new()构造客户端默认从环境变量读取配置用.with_revision(revision)固定 commit/revision如refs/pr/21分别get(config.json)、get(tokenizer.json)、get(model.safetensors)通过VarBuilder::from_mmaped_safetensors([weights_filename], DTYPE, device)构建变量加载器再BertModel::load(vb, config)组装模型。VarBuilder::from_mmaped_safetensors定义在 candle-nn/src/var_builder.rs#L642-L647它内部调用candle::safetensors::MmapedSafetensors::multi支持同时映射多个权重文件这也正好衔接到下一节的内存映射方案。内存映射加载mmap更高效的大模型启动方式对于动辄数 GB 的权重文件整文件读入内存会产生不必要的分配与拷贝。Candle 支持借助memmap2将文件映射到虚拟内存只在实际访问时按页调入cargo add memmap2use candle::Device; use hf_hub::{split_id, HFClientSync}; use memmap2::Mmap; use std::fs; let api HFClientSync::new().unwrap(); let (owner, name) split_id(bert-base-uncased); let repo api.model(owner, name); let weights_filename repo.download_file().filename(model.safetensors).send().unwrap(); let file fs::File::open(weights_filename).unwrap(); let mmap unsafe { Mmap::map(file).unwrap() }; let weights candle::safetensors::load_buffer(mmap[..], Device::Cpu).unwrap();与load的区别在于这里不再std::fs::read而是直接把Mmap的字节切片交给load_buffer解析。这一用法在 candle-book 测试 candle-book/src/lib.rs#L31-L49book_hub_2中被验证同样断言可解析出 206 个张量。需要注意的是文档与源码中的三重提醒unsafeMmap::map是 unsafe 操作语义上要求映射期间底层文件不被截断或改写。参见memmap2的 Safety 说明。实际上模型文件在推理期间不会被修改且映射通常保持只读因此该风险在常规场景下基本不触发但仍应时刻留意。Windows / WSL 兼容性内存映射在 Windows 与 WSL 环境下可能出现问题社区已有相关 issue 反馈跨平台项目需要评估。网络挂载盘如果权重文件位于网络挂载的磁盘NFS 等mmap 会触发更多小粒度读调用性能反而明显变慢此时整文件读入更合适。进一步地candle-core 在 safetensors.rs 中把 mmap 封装成了多种加载器按需选用MmapedSafetensors对单文件new或多文件multi做 mmap 惰性解析load(name, dev)按需取张量这正是VarBuilder::from_mmaped_safetensors与ShardedVarBuilder的底层SliceSafetensors/BufferedSafetensors分别面向外部借用的字节切片与自持有的Vecu8缓冲MmapedFile仅映射文件、按需deserialize获取SafeTensors。张量并行分片加载每个 GPU 只读自己那份权重在多 GPU 做张量并行Tensor Parallel以降低延迟时每个 rank 其实只需要权重的一个切片。此时应直接使用safetensorscrate按张量维度切出本卡需要的区间而不是把整个张量都加载进来cargo add safetensorsuse candle::{DType, Device, Tensor}; use hf_hub::{split_id, HFClientSync}; use memmap2::Mmap; use safetensors::slice::IndexOp; use safetensors::SafeTensors; use std::fs; let api HFClientSync::new().unwrap(); let (owner, name) split_id(bert-base-uncased); let repo api.model(owner, name); let weights_filename repo.download_file().filename(model.safetensors).send().unwrap(); let file fs::File::open(weights_filename).unwrap(); let mmap unsafe { Mmap::map(file).unwrap() }; // 直接使用 safetensors 反序列化拿到张量视图 let tensors SafeTensors::deserialize(mmap[..]).unwrap(); let view tensors .tensor(bert.encoder.layer.0.attention.self.query.weight) .unwrap(); // 以 rank1、world_size4 为例沿第 0 维切出本卡所需分片 VIEW[start..stop, :] let rank 1; let world_size 4; let dim 0; let dtype view.dtype(); let mut tp_shape view.shape().to_vec(); let size tp_shape[0]; if size % world_size ! 0 { panic!(The dimension is not divisible by world_size); } let block_size size / world_size; let start rank * block_size; let stop (rank 1) * block_size; // 一切按张量维度表达字节偏移由 safetensors 自动处理 let iterator view.slice(start..stop).unwrap(); tp_shape[dim] block_size; // 将 safetensors 的 Dtype 转换为 candle 的 DType let dtype: DType dtype.try_into().unwrap(); // 收集该分片的原始字节 let raw: Vecu8 iterator.into_iter().flatten().cloned().collect(); let tp_tensor Tensor::from_raw_buffer(raw, dtype, tp_shape, Device::Cpu).unwrap();这段代码的要点view.slice(start..stop)来自safetensors::slice::IndexOp在张量维度上切分safetensors会自动把维度区间换算成字节偏移无需手工计算data_offsets每个 rank 只需设置自己的rank与全局world_size即可只把[start, stop)区间的数据收集到Vecu8Tensor::from_raw_buffer定义于 candle-core/src/safetensors.rs#L208-L288按原始字节 dtype 形状直接构造张量避免了一次中间分配源码注释中的 TODO 也指出未来可进一步实现from_buffer_iterator以省去这段 CPU 拷贝safetensors::Dtype到candle::DType的转换通过try_into()完成映射关系可在 candle-core/src/safetensors.rs#L43-L64 中核对。以bert-base-uncased的 query 权重形状[768, 768]为例world_size4、rank1时切出的分片形状为[192, 768]——candle-book 测试 candle-book/src/lib.rs#L107-L108book_hub_3正是用这两条断言验证了整个分片逻辑的正确性。工程落地缓存、revision 与进度条的参考实现如果要在自己的项目中复刻 candle-examples 的完整下载体验可以借鉴 candle-examples/src/hub.rs 的封装Api::with_cache_dir(cache_dir)通过hf_hub::HFClient::builder().cache_dir(...)自定义缓存目录Repo::with_revision(revision)把后续所有下载固定到指定 revision保证结果可复现Repo::get(filename)先以local_files_only(true)尝试命中本地缓存未命中才真正发起下载并挂载StderrProgress进度处理器——在终端上以\r原地刷新百分比重定向时按行输出文件名: 42% (42.0/100.0 MiB)形式的日志该格式化逻辑与单位换算B/KiB/MiB/GiB均有单测覆盖于 hub.rs#L198-L234。小结至此一条完整的 Candle 权重接入链路已经打通hf-hub下载同步/异步、revision 固定、本地缓存→candle::safetensors::load整文件解析为HashMapString, Tensor→ 按张量名接入Linear等candle-nn模块或经VarBuilder::from_mmaped_safetensors直接加载进 BERT 等完整模型 → 面向大模型场景改用memmap2load_buffer减少拷贝 → 面向多 GPU 张量并行场景用safetensors的slice只加载本 rank 分片。每一步都在仓库源码与 candle-book 测试candle-book/src/lib.rs中留下了可验证的实现与断言你可以直接参考这些测试代码把流程复刻到自己的项目里。赞分享人工智能大模型机器学习深度学习本地部署模型推理服务【免费下载链接】candleMinimalist ML framework for Rust项目地址https://gitcode.com/GitHub_Trending/ca/candle点击查看免费下载相关推荐Shortcircuit XT 免费采样器三平台支持全景指南Windows、macOS、LinuxShortcircuit XT 免费采样器三平台支持全景指南Windows、macOS、Linux Shortcircuit XT 是一款由 Surge Sy人工智能大模型机器学习深度学习本地部署模型推理服务Text Generation Inference 中的 Safetensors 权重格式安全加载、张量并行分片与自动转换机制Text Generation Inference 中的 Safetensors 权重格式安全加载、张量并行分片与自动转换机制 Safetensors 是 T模型推理服务大模型后端从 PyTorch、Transformers 与 Safetensors 三种途径加载 GPT-2 预训练权重LLMs-from-scratch 实战指南从 PyTorch、Transformers 与 Safetensors 三种途径加载 GPT 2 预训练权重LLMs from scratch 实战指南示例工程大模型人工智能上一篇URH终极指南无线协议分析与信号处理的完整教程下一篇CTF-Wiki 深入解析 ptmalloc2 堆检查机制malloc / free / unlink 全链路校验与漏洞利用对抗创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表