MuseAI 是由阿里集团爱橙科技研发的面向阿里内部的 AIGC 创作工作台，同时通过与阿里云旗下魔搭社区合作共建的形式，将主体能力通过魔搭社区的 AIGC 专区对公众开放。本文主要分析了平台由于频繁切换 Diffusion Pipeline 引起的用户体验与资源浪费问题，并从网络传输、内存管理、Host-to-Device、模型量化等方面着手优化。
考虑到 MuseAI 平台本身是公司内部服务，下文通过底层技术同源的魔搭社区 AIGC 专区来做说明与介绍，为避免混淆，下文如未特意提及，“MuseAI”和“魔搭社区 AIGC 专区”指代同一事物。
魔搭社区 AIGC 专区链接：https://modelscope.cn/aigc/home
MuseAI 是一款专为设计专业人士量身定制的先进 AI 绘图工具，旨在提供卓越的绘画体验，并为设计团队打造一个既稳定又易于管理的创作平台。基于尖端的扩散模型（Diffusion Model）技术，MuseAI 不仅提供了强大的推理与训练服务，还允许用户通过简化的文本输入和参数设置（包括选择模型版本及设定分辨率等），轻松实现心中构想的视觉作品。
在 MuseAI 的核心——Diffusion Pipeline 中，多个模型协同工作以生成高质量图像。这一流程包括以下关键组件：
基础模型：包括 SD1.5、SDXL、SD3 和 FLUX 算法的基础模型，文件大小从 1GB 到 20GB 不等。
LoRA 微调模型：采用轻量级调整技术，通过添加少量可训练层来对基础模型进行个性化调整，适应多样化的风格需求，如卡通或写实效果，体积通常介于 100MB 至 1GB 之间。
ControlNet 控制模型：用于根据具体条件（如人物姿态、面部表情或背景设定）精确指导图像合成，其文件大小大约在 500MB 至 10GB 范围内。
辅助性模型：包含 VAE 编码器、CLIP 特征提取器、T5 语言模型以及其他注解工具，它们虽功能独立但不可或缺，多数情况下预先加载以便即时使用，单个文件大小约在 100MB-10GB 左右。
MuseAI 集成了丰富的资源库，涵盖数百款 Checkpoint 模型、数千种 LoRA 模型以及数十种适用于不同场景的 ControlNet 方案，并支持用户上传自定义模型，既促进了个人创意表达，也鼓励了社区内的资源共享。然而，面对如此庞大的模型库，我们面临的主要挑战之一是如何有效地管理这些资源，在不影响用户体验的前提下最小化模型切换时间。
为了应对这一挑战，MuseAI 不断探索并实践创新方法，致力于减少请求间的流水线切换时间。尽管学术界和工业界已有诸多关于提升模型推理速度的研究成果，但在减少模型切换时间方面的工作仍然相对稀缺。因此，我们将分享一些 MuseAI 在过去一段时间内积累的经验，希望能为相关研究和技术发展提供有价值的参考。
在深入探讨之前，有必要对一些关键概念进行明确定义，以确保讨论的准确性和一致性。以下是与 MuseAI 平台性能优化相关的几个重要时间度量：
端到端生成时间：从用户提交请求开始，直到推理集群完成图像生成并将结果返回给用户为止的总时间。
模型下载时间：从 MuseAI 的远程存储中下载模型参数至推理集群所在机器磁盘所需的时间。
模型读取时间：将模型参数从磁盘加载到内存所需的时间。
模型切换时间：从模型参数加载到内存后，到其在 GPU 上准备就绪、能够执行推理任务的这段时间。
模型推理时间：模型在 GPU 上实际执行推理计算所需的时间。
请注意，虽然“请求排队时间”（即从用户发起请求到请求进入推理集群的时间）是用户体验的一个重要因素，但它主要受用户请求数量和可用集群资源的影响，因此在本文后续讨论中不会被考虑。这使得我们可以集中精力研究其他影响端到端生成时间的因素，即：
上图仅形式化地展示了各个时间阶段对应的流程内容，并不反映实际的时间占比。在实际情况中，每个阶段的时间消耗会受到多种因素的影响，包括但不限于模型类型、模型大小、模型存储方式、缓存命中情况以及硬件性能等。
为了更好地理解这些因素如何共同作用于端到端生成时间，以下是一组基于 MuseAI 平台真实请求的数据，展示了各阶段时间消耗的分布情况：
表 1. MuseAI 真实请求下端到端生图时间的耗时分布
在上述简易测试中，我们通过清空 Linux 的 PageCache（使用命令 sudo sh -c "echo 1 > /proc/sys/vm/drop_caches"），模拟了首次推理的冷启动场景。此时，系统完全没有任何缓存支持，因此模型下载时间、模型读取时间和模型切换时间都显著增加。具体表现如下：
第一次推理（冷启动）：所有数据必须从远程存储下载到本地磁盘，再加载到内存，并最终迁移到 GPU 上准备就绪。这个过程涉及大量的 I/O 操作和数据传输，导致整体耗时较长。
第二次推理（PageCache 命中）：这次推理是在重启 Worker 但不重启物理机的情况下进行的。虽然无法命中 Worker 内部的内存缓存，但由于 Linux PageCache 的存在，模型文件可以从缓存中快速加载，减少了模型下载时间和部分读取时间。然而，模型切换时间仍然相对较长，因为需要重新加载模型到 GPU 上。
第三次推理（内存缓存命中）：当模型已经被加载到 Worker 的内存缓存中时，模型下载时间和模型读取时间几乎为零。模型切换仅需执行 Host-to-Device (H2D) 的传输操作，并且由于之前的推理已经预热了 CUDA 层面的资源，模型推理时间也达到了最短。
从上述简易测试的数据可以看出，在没有缓存命中的情况下，模型下载、模型加载和模型切换时间占据了端到端生成时间的绝大部分。这表明这些阶段的时间消耗是不可忽视的，尤其是在 MuseAI 这样的多模型环境下，缓存未命中的情况更为普遍。
MuseAI 平台集成了大量不同类型的模型，使得将所有模型都缓存到磁盘或内存中变得极为困难。因此，平台不得不频繁面对缓存未命中的挑战，这对用户体验和服务效率产生了负面影响。
鉴于此，我们的研究重点将集中在以下几个方面：
1. 增强硬件与软件协同优化：
2. 提升模型构建与加载效率:
3. 内存管理与复用：
4. 模型量化:
5. 模块拆解并行：
模型加载 是指推理服务从存储介质加载模型数据到内存的时间，涵盖了前文所述的“模型下载时间”与“模型读取时间”。模型加载时间不仅依赖于存储介质的理论性能，还需要通过最佳实践来充分发挥其性能。因此，本节将围绕“基于业务特性选择存储介质”和“如何充分发挥存储介质性能”两点展开讨论。
在 Diffusion 生图社区中，模型种类繁多且数量庞大，难以将所有模型都存储在服务端本地磁盘上。因此，选择适当的存储介质至关重要，以确保模型能够高效地保存并快速读取。
截至撰写本文时，最大的生图模型网站 Civitai AI 上常用的模型属性如下表所示。考虑到算法版本的迭代，SD1.5 的模型将逐渐被淘汰，同时高质量的生成效果模型相对较少。综合评估后，400 TB 的存储容量对于 MuseAI 来说是充足的。
表 2. 模型总量存储统计
根据第二章提供的“MuseAI 真实请求下端到端生图时间的耗时分布”表格，当前 MuseAI 第一次生图的实际推理时间约为 10 秒，但“模型下载时间”与“模型读取时间”的总和却接近其 10 倍。这不仅严重影响了用户体验，还造成了 GPU 资源的极大浪费。目前，模型下载速度仅为 100 MB/s，而服务器的网卡带宽为 2 GB/s，显然，提升带宽利用率以减少非推理时间占比是当务之急。
简而言之，MuseAI 的业务特性对存储方案提出了以下要求：
在选择适合 MuseAI 平台的存储介质时，我们考虑了多种常见的文件存储方案，包括对象存储服务（OSS）、网络附加存储（NAS）以及阿里内部的分布式集群存储服务“盘古”，三者的特点简单归纳如下：
对象存储服务 (OSS)：广泛应用于互联网业务中，具备高可用性和无限扩展性，适用于大规模数据存储。然而，从 OSS 读取模型时需要先下载到本地磁盘再加载到内存，增加了两次数据拷贝的时间开销。
网络附加存储 (NAS)：提供 POSIX 兼容的文件系统接口，易于集成和使用。NAS 的主要挑战在于单个实例的带宽瓶颈，当达到极限时，需要通过复制数据到多个 NAS 实例并分组绑定服务器来扩展，这增加了管理和维护的复杂度。
盘古： 是阿里内部的分布式集群存储系统，以高性能和稳定性著称。盘古是一个复杂的分布式服务，由 Client、ChunkServer 和 Master 组成，当上层应用通过 Client 提出读取请求时，Client 需要首先向 Master 通信查找相关 metadata，从而得知数据在哪几个 Chunkserver 上，然后拿到实际的数据。Client 提供的接口相对底层，且配置较为复杂，为降低使用难度，阿里内部为其封装了自定义 fuse（Filesystem in Userspace） 文件系统 fsfuse，并配套 dcache 服务进一步提高读取性能：
meta 缓存，保证本地能快速获取 meta 而降低高延时的 rpc 请求量；
block 缓存，read 系统调用通常单次读取 128K，而单次 RPC 耗时较长，每个 128K 都走 RPC 调用代价过大，所以 fsfuse 会以 8M 数据请求数据块，并将多个 8M block 组成较大的 block 缓存池，以让业务的读取请求尽可能命中本地缓存；
block 预读，业务连续顺序读取 block 时，会触发预读更多 block。
每种方案都有其独特的优势和适用场景，如下表所示。
表 3. 存储方案特点
三种存储方案都能满足 MuseAI 的容量需求，但在读取性能和扩展性方面存在显著差异：
综上所述，基于 MuseAI 的业务需求和技术特点，我们决定采用以下存储策略：
在 MuseAI 平台中，选择合适的 NAS 类型和优化其配置对于确保高效的模型加载至关重要。阿里云 NAS 提供了通用型和极速型两种主要类型，前者适合存储大量数据且对总吞吐量有需求的场景，后者则适用于处理大量的读写请求和对响应延迟要求较高的场景。鉴于 MuseAI 的业务特性——需要存储大量模型数据并保证高吞吐量，我们选择了通用型 NAS。
表 4. 通用型和极速型 NAS
为了验证 ECS 实例与 NAS 之间的连接性能，我们在 ecs.g7ne.12xlarge 规格（48 核 192G 内存 40 Gbit/s 带宽）的 ECS 上进行了性能测试。由于通用型 NAS 的读带宽与容量有关，我们首先创建了一个 33TB 的文件以确保 NAS 能够达到 20 GB/s 的读带宽。
通过阿里云控制台默认挂载参数进行测试，采用 nfsv3 协议并通过 TCP 访问，具体挂载参数如下。
vers=3,nolock,proto=tcp,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport
使用 fio 测速工具测试 NAS 的随机读取吞吐量，命令如下。
fio -numjobs=2 -iodepth=128 -direct=1 -ioengine=libaio -sync=1 -rw=randread -bs=4M -size=1G -time_based -runtime=60 -name=Fio -directory=/mnt/muse
初次测试结果显示，读取速度仅为 500 MB/s，远低于机器网卡上限（40 Gbit/s）和 NAS 读带宽上限（160 Gbit/s），显然存在性能瓶颈。
经过排查，发现 NFS 客户端和服务器之间默认仅通过一个 TCP 连接通信，并且阿里云 NAS 前端机限制了每条连接的带宽上限。为了解决这一问题，我们调整了 Linux 内核参数，特别是 nconnect 参数，该参数控制 NAS 与客户端之间的连接数，默认值为 1，最大值为 16。
vers=3,nolock,proto=tcp,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport,nconnect=16
增加 TCP 连接数后，读取带宽显著提升至 4500 MB/s，达到了网卡带宽的大约 90%，效果明显。此外，我们还采用了多线程并发读取模型文件的方法，进一步充分利用网络带宽，确保了高性能的数据传输。
inline int read_file(int fd, void *buf, size_t nbytes, off_t offset) { std::stringstream oss; uint64_t total_read_bytes = 0; while (total_read_bytes < nbytes) { void *buf = (char*)buf + total_read_bytes; ssize_t n = pread(fd, buf, nbytes - total_read_bytes, offset + total_read_bytes); if (n == 0) { oss << "cannot read expected " << nbytes << " bytes, total read bytes: " << total_read_bytes; throw SafetensorsException(oss.str()); } if (n < 0) { oss << "pread failed. fd: " << fd << ", bytes: " << nbytes << ", errorno: " << errno; throw SafetensorsException(oss.str()); } total_read_bytes += n; } return 0;}
void multi_thread_read_file(int fd, void *buf, size_t start_offset, size_t end_offset, int num_threads) { uint64_t file_offset = start_offset; uint64_t buf_offset = 0; uint64_t data_size = end_offset - start_offset; uint64_t block_size = (data_size + num_threads - 1) / num_threads; ThreadPool pool(num_threads); std::vector<std::future<int>> futures;
char *buf = reinterpret_cast<char*>(buf); for (int i = 0; i < num_threads; ++i) { size_t read_bytes = std::min(end_offset - file_offset, block_size); void *store = reinterpret_cast<void*>(buf + buf_offset); futures.emplace_back( pool.enqueue(read_file, fd, store, read_bytes, file_offset));
file_offset += read_bytes; buf_offset += read_bytes; } for (auto &future : futures) { future.get(); }}
在 MuseAI 平台中，盘古结合 fsfuse 提供了一种高效、透明的数据访问方式，极大地提升了模型加载的性能。为了充分发挥这一组合的优势，我们总结了以下最佳实践：
挂载与目录管理：fsfuse 通过接口将盘古分布式文件系统挂载到本地，使得用户可以像访问本地路径一样读取数据。然而，fuse 的最大挂载上限数为 1024，对于 MuseAI 这样的服务，如果每次使用不同模型时都挂载不同的目录，很容易超过该上限，导致数据读取失败。最佳做法是将所有模型的父目录统一挂载到本地，这样只需要一个挂载点即可访问所有模型，避免了频繁挂载带来的问题。
缓存机制优化：fsfuse 在读取数据时会从盘古中预读取比用户请求更多的数据，并将其存储在缓存中。这种机制特别适用于顺序读取场景，因为预读取的数据能够显著提高缓存命中率，从而加快数据访问速度。推荐尽量采用顺序读取模式，以充分利用 fsfuse 的预读取和缓存机制，确保高效的数据访问。
Direct I/O 技术应用： 默认情况下，不使用 Direct I/O 技术时，每次 read 系统调用的大小为 128 KB。当读取大文件时，这会导致大量的系统调用，成为性能瓶颈。为了减少系统调用次数并提高读取效率，fsfuse 推荐使用 Direct I/O 技术。使用 Direct I/O 每次读取的 block 大小设置为 2 MB，这样可以将系统调用开销降低至原来的 1/16，极大减少了系统资源的消耗，在 fsfuse 缓存全命中的情况下，读取速度可达到 9 GB/s。
在 MuseAI 平台中，模型切换时间是影响整体性能的关键因素之一。本章将详细介绍如何优化模型从内存加载到 GPU 的过程，包括 state dict 的传输、nn.Module 的构造和装载，以及进一步的性能提升策略。
内存中的模型切换时间主要包括以下几个部分：
构造 nn.Module：创建模型实例。
nn.Module 装载 state dict：使用 load_state_dict 方法将 state dict 应用到模型上。
state dict 从 CPU 传输到 GPU：通过 to("cuda:0") 将内存中的 state dict 转换为 GPU 上的 state dict。
为了达成上述目标，一是可以让 nn.Module 先装载 state dict，再让 nn.Module 调用 to("cuda:0")，二是先让 state dict 中的 tensor 执行 to("cuda:0")，再让 nn.Module 装载 state dict。这两者之间是否有性能差异呢？
我们构造了一个 4.1 GB 的 safetensors 文件（忽略其读取时间），并通过实验对比了两种常见方案的性能差异。为了消除磁盘性能对实验的影响，我们在 profile 之前对 state dict 进行了 clone，确保每个 tensor 都已读取到内存中。
from safetensors.torch import load_file, save_fileimport torch.nn as nnclass MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.linear1 = nn.Linear(10, 1024) self.linear2 = nn.Linear(1024, 1024 * 1024) self.linear3 = nn.Linear(1024 * 1024, 10)
def forward(self, x): x = self.linear1(x) x = self.linear2(x) x = self.linear3(x) return x
state_dict = load_file("test_model.safetensors")
def func1(state_dict): model = MyModel() model.load_state_dict(state_dict) model = model.to("cuda:0") torch.cuda.synchronize()
def func2(state_dict): model = MyModel() model = model.to("cuda:0") model.load_state_dict(state_dict) torch.cuda.synchronize()
if __name__ == "__main__": state_dict = load_file("test_model.safetensors") new_state_dict = {} for key, value in state_dict.items(): new_state_dict[key] = value.clone() func1(new_state_dict)
先 load_state_dict 再 to：
先 to 再 load_state_dict：
实验结果显示，先 load_state_dict 再 to("cuda:0") 的方法显著优于先 to("cuda:0") 再 load_state_dict。
问题关键点在于，MyModel 其实在构造完成后就已经持有了许多 tensor，而 load_state_dict 操作本质做的事情相当于：
model.load_state_dict(state_dict) 相当于for key, param in model.name_parameters(): param.data.copy_(state_dict[key])
如果模型遵循 cpu-> gpu -> load_state_dict 的流程，并且 state_dict 本身的 tensor 在 cpu 内存上，那么就需要进行一批额外的 cpu to gpu 操作。因为 gpu tensor 只能拷贝 gpu tensor 的值，所以 state dict 被隐式的进行了 h2d 传输。
如果模型遵循 cpu -> load_state_dict -> gpu 的流程，那么就相当于把 gpu 上的 tensor copy 变成了 cpu 上的 tensor copy，节省了一大批的 h2d 操作，因此降低了整个流程的耗时。
简单总结两种执行顺序涉及的内存操作：
因此，推荐采用 cpu->load_state_dict->gpu 的流程，以节省大量 H2D 操作，从而降低整个流程的耗时。
为了进一步缩短时间，可以在 load_state_dict 中传入参数 assign=True。这会直接用 state dict 的 tensor 充当 parameter 的 data，省去了参数数据拷贝的过程，从而加速模型装载。
Tensor 从 CPU 到 GPU 的传输较为缓慢，主要原因在于通常持有的 tensor 内存是由 Linux 分配的 pagable memory，在传输到 GPU 前 CUDA API 会开辟临时的 pinned memory 空间，将 pagable memory 中的内容复制到该空间后再通过 PCI-E 总线传输到 GPU。
这里有一个简易的传输时间的对比，tensor 的大小是 16GB：
我们可以采取措施减少内存分配和拷贝的次数：
这部分优化细节放在下一章展开叙述。
实际场景中的 state dict 往往包含成百上千个 tensor，顺序进行 H2D 操作效率较低。参考 PyTorch 官方文档的建议，我们可以使用 tensordict 包装 state dict，并结合 pin_memory、non_blocking 和多线程操作，实现并发异步传输，从而大幅提升性能。
我们可以看到 2.2 章节实验中的 MyModel 的构造时间其实也占了很大一块时间（5s 以上），但 2.3 章节实验中分配 16GB 的 tensor 只需要 500ms，说明 MyMode l 构造的时间大头并不是构造 tensor，这并不符合我们的预期。
我们对 nn.Linear 的 __init__ 函数进行单独 profile 可以发现，在 nn.Linear 构造过程中，绝大多数时间都集中在 reset_parameters 上：
翻看 torch 的代码可以发现：
reset_parameters 本质上做的事情是在 nn.Parameters 构造后为 weight 和 bias 遵循 kaiming 初始化方法赋予初始值。这个步骤在模型训练中非常重要，但是在推理中显得很冗余，因为不管初始化 Parameter 的值是什么，我们都需要 load_state_dict 去覆盖这个值。PyTorch 提供了一种 skip init 技术，允许我们将模型构造在虚无的 meta 设备上，然后再迁移到 CPU 或 GPU，从而跳过初始化过程。
在 MuseAI 平台中，内存管理和复用是优化模型加载的关键环节。本章将详细介绍现有方案的性能问题，并提出相应的优化措施，以实现更高效的内存利用和更快的模型切换。
MuseAI 原有的模型加载和推理链路大致如下：
load_file 加载 safetensors：使用 Huggingface Safetensors 库 load_file 将 safetensors 文件加载为 state dict。
load_state_dict 初始化权重：模型调用 load_state_dict 方法初始化权重。
H2D 数据拷贝至显存：执行 model.to("cuda:0") 将数据从 CPU 拷贝到 GPU 显存。
推理执行：最终进行推理操作。
import torchimport torch.nn as nnfrom safetensors.torch import load_file, save_file
class MyModel(nn.Module): def __init__(self, device): super(MyModel, self).__init__() self.linear1 = nn.Linear(10, 1024, device=device) self.linear2 = nn.Linear(1024, 1024 * 1024, device=device) self.linear3 = nn.Linear(1024 * 1024, 1024, device=device) self.linear4 = nn.Linear(1024, 1024 * 1024, device=device) self.linear5 = nn.Linear(1024 * 1024, 1024, device=device) self.linear6 = nn.Linear(1024, 10, device=device)
def forward(self, x): x = self.linear1(x) x = self.linear2(x) x = self.linear3(x) x = self.linear4(x) x = self.linear5(x) x = self.linear6(x) return x
@profiledef load_baseline(model_path): state_dict = load_file(model_path) model = torch.nn.utils.skip_init(MyModel, device="cuda:0") model.load_state_dict(state_dict, assign=True) model = model.to("cuda:0") torch.cuda.synchronize() return model
if __name__ == "__main__": model_path = "test_model.safetensors" load_baseline(model_path)
通过对 infer 函数进行 line profiler 分析，我们观察到了以下几个现象：
load_file 耗时 12 ms：原因是 Huggingface Safetensors 库利用 mmap 加载文件，此时模型数据实际上还未加载到内存。
model.to 耗时 14.7 秒：单线程遍历 state dict 中的所有数据，触发缺页异常实际把数据从文件读取到内存中，并且将 tensor 从内存拷贝至显存。此过程包括 malloc pinned_memory、copy pagable_memory to pinned_memory 和 copy pinned_memory to device_memory 三个步骤。
原方案存在以下性能问题：
单线程读取：无法充分利用分布式块存储服务的带宽。
非顺序访问：空间局部性差，对存储缓存不友好。
H2D 时的内存开销：
针对上述性能问题，我们采取了以下优化措施，并已集成至文首提及的代码仓中：
改进模型加载方式： 根据前文所述的不同存储介质（NAS 和 盘古 +fsfuse），采用优化后的方案分别读取模型数据，充分利用多线程和分布式存储的优势，提升读取速度。
内存分配与拷贝优化：
维护 pinned_memory 内存池：避免每次推理重复 malloc 内存，减少不必要的内存分配和释放操作。
直接读取到 pinned_memory：消除一次内存拷贝，直接将文件数据读取到预分配的 pinned_memory 中，从而简化传输路径并提高效率。
两级内存池设计
虽然 MuseAI 涉及的模型类型较多，但每种模型大小在 [10MB, 24GB] 范围内。我们提前分配并维护几个固定大小的内存块，形成两级内存池。当有需求时从中拿一个适当的空闲块给用户，示意图如下所示 。
举个例子，先加载一个 2.8G 的模型，再加载一个 8.5G 的模型，紧接着卸载最先加载的 2.8G 模型，最后加载一个 2.9G 的模型，内存池会进行以下行为：
圆整模型大小：将 2.8G 模型大小圆整至 3G，从 free_lists[2] 中拨出第一块内存块来存储模型，并更新 free_lists[2] 指向第二块内存。
处理大模型：将 8.5G 模型大小圆整至 9G，从 free_lists[8] 中拨出第一块内存块来存储模型，并更新 free_lists[8] 指向第二块内存。
回收内存：根据 2.8G 模型大小，我们知道这块内存来自 free_lists[2]，让 free_lists[2] 重新管理这块空闲内存，即更新 free_lists[2] 重新指向第一块内存。
再次分配：重新把第一块 3G 内存从 free_lists[2] 拨出用于存放 2.9G 模型数据。
在 pinned_memory 上构造 state_dict
PyTorch 中的 state_dict 是一个简单的字典对象，其中 key 类型为 str，代表 tensor 名称；value 类型为 tensor，记录了模型权重信息，是消耗内存的主要因素。为了控制在指定的内存区域构造 tensor，我们首先理解 tensor 的数据结构：
因此，我们可以首先用预分配好的 pinned_memory 构造 Storage，再设置 tensor 的 storageOffset 为适当的值。具体步骤如下：
解析 safetensors 文件格式：
前 8 个字节记录 header 大小 header_size。
[8, 8 + header_size) 记录 tensors 元信息，包括数据类型、tensor shape 和数据在文件中的位置 offsets。
[8 + header_size, file_size) 记录 tensor 实际数据。
构建 tensor：将文件中的 tensor 数据段拷贝到预分配的 pinned_memory 中，根据 safetensors header 中的 offsets 构建 tensor。需要注意的是，storage_offset 是指在 storage 中元素个数的偏移量，而非字节偏移量。若同个 storage 中不同 tensor 的 sizeof(dtype) 不一致，可能导致字节偏移量计算错误，这时需要将内存按每个 tensor 占用量切分为多块连续内存，在每块内存上单独构造 storage，并绑定到对应 tensor 上。
读取性能评估
通过比较以下三个函数的耗时，评价模型读取所带来的性能提升：
@profiledef load_baseline(model_path): state_dict = load_file(model_path) model = torch.nn.utils.skip_init(MyModel, device="cuda:0") model.load_state_dict(state_dict, assign=True) model = model.to("cuda:0") torch.cuda.synchronize() return model
@profiledef fast_safetensors_fuse(model_path): state_dict = fast_safetensors.load_safetensors(model_path, num_threads=1, direct_io=True) model = torch.nn.utils.skip_init(MyModel, device="cuda:0") model.load_state_dict(state_dict, assign=True) model = model.to("cuda:0") torch.cuda.synchronize() return
