显存又不够了。这大概是过去两年里我听到最多的一句话。最近接了一个多模态项目模型从7B涨到15B单卡连加载都费劲更别提训练。换卡预算不够后来我们在一台普通工作站上用SSD卸载技术把训练规模硬生生扩了一大截。这篇文章就聊聊这条“用盘换显存”的路——SSD卸载到底解决什么问题、该怎么做、有哪些坑。无论你是正在被显存卡住的算法工程师还是做推理服务优化的同学这篇文章里的方案和踩坑经验应该都能直接用上。1. 为什么AI规模会被显存卡死先说清瓶颈在哪1.1 显存里到底装了什么参数、梯度、优化器状态和激活值很多人对显存的认识停留在“能装下模型参数就行”实际训练时显存消耗远远不止参数本身。一个可训练的深度学习模型显存里通常要同时放四样东西模型参数Weight前向和反向都要用占用与模型规模直接相关。梯度Gradient反向传播算出来的梯度更新参数前要保留。优化器状态Optimizer State比如Adam里的momentum和variance每个参数往往对应多份FP32副本。激活值Activation前向传播时每层算出来的中间结果反向传播要用来算梯度。以7B参数模型为例FP16参数本身约14GB但训练时Adam优化器状态可能需要56GB以上再加上梯度14GB、中间激活值若干一张80GB的A100光塞这些状态就捉襟见肘。如果换到30B、70B级别的模型单卡训练根本不可能。这就是“AI规模被显存卡死”的直接原因卡越来越贵显存扩容的速度远跟不上模型参数膨胀的速度。1.2 CPU Offload与SSD Offload的逻辑内存不够之后轮到盘显存放不下下一步自然是往内存放。CPU内存便宜、容量大于是出现了CPU Offload把优化器状态、梯度这些“不急着用”的数据挪到系统内存参数更新时再取回来。这是DeepSpeed ZeRO-Offload的经典思路也是很多人在单机训练时用来压缩显存的第一步。但CPU内存也有上限。一台工作站插满也就256GB到512GB面对上百B参数级别的模型依旧不够。再往外一层就是磁盘了——这就是SSD Offload的基本动机。存储层级里SSD是离CPU内存最近的一层延迟在微秒级带宽可以到数GB/s容量从几百GB到几十TB随便选。ZeRO-Infinity这类系统做的核心事情就是把卸载链路从GPU显存扩展到CPU内存再扩展到NVMe SSD把“放不下”的空间直接延伸到盘上。其实可以把这套逻辑理解成“仓库-货架-桌面”的关系。GPU显存是桌面CPU内存是手边的货架SSD是仓库。桌面放不下就放货架货架放不下就放仓库。每次干活先把要用的东西从仓库搬到桌面用完放回去。SSD卸载就是这个“仓库”。1.3 什么场景真正适合SSD卸载不是所有场景都该无脑上SSD卸载。我判断一个项目适不适合主要看三点训练场景单机单卡或双卡想跑一个参数超过显存承载能力的模型能接受训练速度下降。SSD卸载能把“跑不起来”变成“跑得慢”这已经是质变。推理场景模型权重一次性放不进显存可以在层间切换时从SSD按需加载也就是常说的weight streaming。这种场景对SSD顺序读带宽要求高。超大规模训练多机多卡显存全用完还不够SSD作为最后一层“溢出缓冲区”兜住偶发的显存尖峰。不适合的场景也有训练节奏极其敏感的模型、超低延迟在线推理、大量小数据块随机读写的任务。SSD虽然快但跟显存和内存比还是差几个数量级。硬上卸载方案性能会很难看。方案容量上限读写延迟带宽量级成本典型场景纯GPU训练受单卡/多卡显存限制纳秒级TB/s级很高小模型、大集群CPU Offload受内存容量限制百纳秒级数十GB/s中单机中等模型SSD Offload受盘容量限制微秒级数GB/s低超大模型、显存不足2. SSD卸载的工作原理与主流技术思路2.1 ZeRO-Infinity思路简述分而治之按生命周期分配SSD卸载尽管在工程实现上各有不同但底层的核心思想非常一致把模型状态按使用频率和生命周期拆分放到不同存储层级。以DeepSpeed ZeRO-Infinity的公开思路为例训练过程中模型状态被分成参数、梯度、优化器状态三类。优化器状态只在参数更新阶段被读写平时完全闲着优先丢到CPU内存或SSD梯度在反向传播结束后用于更新参数更新完就可以丢弃参数前向和反向都要用所以尽量留在显存。激活值则通过重计算机制处理需要时重新算一遍而不是一直占着显存。这就像团队协作最核心的人留在会议室不常参与的人在外面待命需要开会的瞬间再叫进来。SSD卸载不是把所有东西都倒到盘上而是把该留的留下该挪的挪走。2.2 卸载什么、保留什么四种状态怎么取舍实际做卸载方案时最关键的决策是“什么东西值得卸载”。我的经验是这样推理场景优先卸载权重。尤其量化后的权重按层加载计算完一层就释放再加载下一层。权重体积最大卸载收益最明显。训练场景优先卸载优化器状态和较老的梯度。更新完的参数梯度立刻释放优化器状态可以挪到内存或SSD参数本身尽量留在显存因为前向反向都要频繁访问。激活值最灵活。可以设置卸载阈值超过多少就压缩落盘也可以用重计算换显存——不是所有激活都要留留下反向传播最需要的那部分就够了。有一个常见误区是想把所有东西都卸载到SSD。比如我曾经见过有人把embedding层也卸载到盘上结果每步训练都要从头读一遍embedding表IO开销直接吃满。embedding这类每步都访问的模块再大也得想办法留显存真正该卸载的是那些访问频率低、生命周期短的中间状态。2.3 与量化、稀疏化组合叠加效果更明显SSD卸载不是孤立的它和量化、稀疏化的组合能产生112的效果。FP16权重直接落盘30B模型要占60GB量化为8bit只要30GB4bit只要15GB。同样的盘容量能装的模型直接翻倍。训练场景同样受益。优化器状态里momentum和variance经常用FP32保存可以压缩成FP16甚至更低的精度再写盘减少写放大也能明显降低SSD的磨损。稀疏化则是利用梯度里的稀疏结构只写入非零部分进一步压缩IO量。所以我在做方案设计时有个习惯先量化压缩再考虑卸载。压缩后的模型状态越小卸载需要消耗的带宽、容量和SSD寿命都更宽裕。3. 硬件选型AI卸载场景下的SSD该看哪些指标3.1 接口与协议NVMe是必须的SATA只能兜底SSD卸载说白了是在用带宽换容量。如果带宽本身不够那整个方案就是空中楼阁。接口协议直接决定了带宽上限SATA SSD6Gbps接口实际顺序读带宽约500-550MB/s。这个速度跑小模型凑合跑大模型卸载根本不够看。PCIe 3.0 NVMex4通道理论带宽约3.5GB/s属于“勉强能干活”的底线。PCIe 4.0 NVMe理论带宽约7GB/s这是目前性价比最高的区间。PCIe 5.0 NVMe理论带宽约14GB/s适合大规模训练卸载还有发热和稳定性的考量。AI卸载场景主要是大块连续读中等频率写NVMe是硬门槛。如果预算实在紧张至少也要PCIe 3.0 NVMeSATA盘只适合放不常用的checkpoint不适合参与高频卸载。3.2 容量、寿命与缓存策略不是越大越好是有多“抗造”很多人在选SSD时只盯着容量忽略了两个关键指标TBW总写入字节数和DWPD每日全盘写入次数。AI训练卸载的写入量非常惊人优化器状态每轮训练都要写盘一天写几TB很正常。我算过一笔账一块2TB的消费级SSD标称TBW 1200TB每天写3TB理论上能用400天。如果换一块1.6TB企业级盘TBW做到8.76PB每天写5TB也能用将近5年。训练任务跑得越重越不能省这块成本。选盘还要注意预留空间。SSD通常需要留20%左右做OPOver-Provisioning一方面是为了磨损均衡另一方面是保证写性能不崩。容量规划时我一般按“模型大小的3-5倍”来准备卸载空间既留足余量又避免盘被写满后性能断崖式下跌。另外一个小细节挂载文件系统时建议加noatime参数减少不必要的元数据写入系统swap也尽量别放在卸载盘上否则内存不足时SSD一边做模型卸载一边做swap换页两边IO互相挤占性能直接雪崩。3.3 我用过的选型参考表类型顺序读带宽4K随机读写TBW价格区间适合场景SATA SSD~550MB/s一般较低低checkpoint存储、冷数据PCIe 3.0 NVMe~3.5GB/s中等中等中低小规模卸载、个人实验PCIe 4.0 NVMe~7GB/s良好较高中常规训练卸载、模型流式加载企业级 NVMe~7GB/s以上很好很高高7x24训练、频繁卸载如果只让我选一块盘做AI卸载我会选PCIe 4.0或5.0的TLC企业级盘容量按上面说的3-5倍准备。企业级盘不仅寿命长更重要的是写性能不会像消费盘那样在长时间写入后缩水。4. 实操配置在PyTorch环境里把SSD卸载跑起来4.1 系统与驱动准备检查NVMe设备、文件系统、挂载参数动手之前先把底层的存储环境确认好。SSD卸载最怕的就是硬件链路有隐患比如M.2插槽走的是PCH共享带宽或者PCIe通道数不够都会让实际速度远低于标称值。推荐先做三件事用lspci或nvme list确认系统正确识别NVMe设备驱动加载正常。用lsblk确认盘挂在哪个控制器上避免多块盘抢同一条总线。挂载文件系统时设置好参数。我常用mkfs.xfs或mkfs.ext4挂载参数加noatime,nodiratime。# 格式化并挂载NVMe SSD sudo mkfs.xfs /dev/nvme0n1 sudo mkdir -p /mnt/nvme sudo mount -o noatime,nodiratime /dev/nvme0n1 /mnt/nvme值得注意的一点如果主板有多个M.2插槽最好把卸载盘插到直连CPU的插槽上。走PCH的插槽会和SATA、USB等设备共享带宽跑满时可能出现带宽减半的情况。4.2 一个最小可跑的PyTorch卸载示例很多人一听到SSD卸载就觉得要上DeepSpeed、要改分布式框架其实底层逻辑没那么玄乎。用PyTorch和mmap就能搭一个最小可跑的示例出来。后面这个方法我一直在用把模型权重写进一个或多个大文件运行时用mmap把文件映射进虚拟内存地址空间按需读取某一层。mmap的好处是内核会帮你做页缓存管理读过的数据短时间内在内存里还有备份多次访问不会反复打盘。import mmap import torch import numpy as np # 1. 在SSD上分配一个大文件作为权重仓库 f open(/mnt/nvme/weight_store.bin, wb) f.truncate(2 * 1024**3) # 预分配2GB空间 mm mmap.mmap(f.fileno(), 0) # 2. 将某一层权重写入文件 def write_tensor_to_mmap(tensor: torch.Tensor, offset: int) - int: arr tensor.cpu().numpy().tobytes() mm[offset:offset len(arr)] arr return offset len(arr) # 3. 按需读取某一层权重到GPU def read_tensor_from_mmap(offset: int, size: int, shape, dtype) - torch.Tensor: buf mm[offset:offset size] arr np.frombuffer(buf, dtypedtype).reshape(shape).copy() return torch.from_numpy(arr).cuda()生产环境里我通常不直接用裸mmap去管理权重文件而是基于safetensors格式做二次封装。safetensors本身支持懒加载meta设备上的模型可以先创建空tensor再用load_file按需把某个key对应的权重读到GPU。这种方案既安全又省事。from safetensors import safe_open from safetensors.torch import save_file # 一次性保存整个checkpoint到SSD save_file({layer.0.weight: layer0_weight, layer.1.weight: layer1_weight}, /mnt/nvme/model.safetensors) # 按需加载某一层权重 tensors {} with safe_open(/mnt/nvme/model.safetensors, frameworkpt, devicecpu) as f: layer0_weight f.get_tensor(layer.0.weight).cuda()4.3 把示例扩展到真实训练场景分块、预取、后台线程最小示例跑通后真正的痛点来了如果每层都同步等待SSD读取GPU利用率会掉到惨不忍睹。我自己的经验是卸载工程化最关键的一步是把IO和计算重叠起来。思路很简单计算第N层的同时后台线程预取第N1层的权重用双缓冲机制在计算和IO之间切换。下面是预取逻辑的骨架from concurrent.futures import ThreadPoolExecutor executor ThreadPoolExecutor(max_workers1) future None for layer in range(num_layers): # 计算结果之前先拿到下一层的权重 weights future.result() if future is not None else initial_weight # 提前发起下一次读取 future executor.submit(load_layer_weights, layer 1) # 用当前层权重做计算 output forward_one_layer(input_tensor, weights)在线程数上不用贪多1个后台IO线程配合足够大的预取队列就够用了。因为单块NVMe的顺序读带宽是有限的同时开几十个线程反而会增加调度开销和IO争抢。训练场景里同理。优化器状态更新完后立即写回SSD这步可以用异步IO做参数更新前先等上一轮写盘完成避免覆盖未落盘的数据。内核里其实有io_uring这样的异步IO接口Python生态里可以用libaio或aiofiles去接。我之前在项目里用ThreadPoolExecutor 大块连续读写也能达到差不多的效果关键是别在计算主路径上同步等IO。4.4 监控与调优看哪些指标判断卸载生效SSD卸载方案最怕的就是“看起来在跑实际都在等IO”。我通常用四个指标判断系统状态SSD利用率utiliostat -x 1看%util持续100%说明IO是瓶颈需要减少卸载数据量或提高带宽。GPU利用率nvidia-smi dmon看GPU utilization偏低且磁盘读写高说明IO和计算没有重叠。CPU等待观察vmstat里的wa值如果CPU大量时间在等IO说明同步IO路径太多。块大小同样数据量4KB块读和1MB块读的吞吐差距可能相差数十倍。卸载文件的设计要保证每次读取至少是256KB以上的大块。实际调优时我会把“SSD顺序读带宽”当作第一约束所有设计都围绕它展开优先大块连续读写、减少随机小IO、用预取掩盖延迟。这个思路贯穿了后面所有踩坑修复。5. 实测中的坑与工程化教训5.1 小文件读写的灾难把checkpoint拆成几万个bin盘直接废了我第一次做SSD卸载时犯过一个大错把优化器状态按每个参数的名字拆成了几千个小文件保存。加载时程序直接卡死iostat显示吞吐不到几百MB/s但IO次数高得吓人。原因在于SSD对顺序大块IO很友好但对随机小IO的吞吐远低于标称值。一个几百KB的文件要经过反复寻址、映射、元数据更新效率极低。这也是为什么上面示例里我反复强调大文件 偏移量 连续块的做法。文件越小元数据开销和IO次数占比就越高。解决办法很简单把成千上万个小文件合并成一个大文件用mmap或safetensors管理按偏移量读取。文件系统里的目录层级也尽量扁平化不要让一个目录下堆几万个文件否则读目录本身都会成为瓶颈。5.2 预取与计算重叠别让GPU闲等另一个高频问题是卸载本身实现了但训练/推理速度反而更慢。我看过不少人的实现卡就卡在“同步加载”上——GPU每算完一层都要干等CPU从SSD读下一层GPU利用率从90%直接掉到40%。修这个问题的核心思路不是优化单次读取速度而是让IO永不阻塞计算。双缓冲预取只是第一步更复杂的场景还要考虑卸载和重计算的顺序。比如训练时可以先算好某几个层的激活值并落盘反向传播时再读回来而不是临时算。这中间有个权衡重计算省带宽但费计算落盘省计算但费带宽。我通常的做法是先量化压缩激活值再结合预取和重计算做混合调度。没有绝对最优只能通过profiler测出实际瓶颈再调。5.3 断电与文件完整性加载更可靠比备份更可靠SSD卸载有个容易被忽视的问题如果训练过程中突然断电或进程被杀卸载到盘上的临时状态可能处于半写状态。注意这类临时数据本身不要求绝对一致训练重启后重来就行但checkpoint的落盘必须是原子的。我现在的规范是checkpoint先写到临时文件写完再rename覆盖旧文件。为什么因为rename在同一个文件系统里是原子操作要么成功要么失败不会出现“读到一半的checkpoint”这种最恶心的状况。import os def save_checkpoint_atomic(state, path): tmp_path path .tmp torch.save(state, tmp_path) os.replace(tmp_path, path) # 原子替换另外mmap映射文件如果被意外截断进程直接收到SIGBUS崩溃。所以mmap文件写完后要flush别指望操作系统最后帮你兜底。5.4 多进程读写同一块盘的竞争问题真实项目里SSD往往不是只伺候模型卸载还要同时处理dataloader的图片读取、checkpoint保存、日志写入等任务。多个进程同时读写同一块盘会造成严重的IO带宽互相挤占。我遇到过图像数据加载和模型卸载抢带宽两边吞吐一起掉到原来的三分之一。解决方式很直接分盘隔离。系统、代码放一块盘数据加载放一块盘卸载缓存和checkpoint放另一块NVMe。如果物理分盘做不到至少要分文件系统或限制IO优先级。Linux下可以用ionice和cgroup io.throttle限制某些进程的读写带宽让卸载路径始终有足够的余量。另一个容易踩的点是swap。如果系统内存吃紧内核会把部分内存页换到swap分区而swap如果恰好放在同一块卸载盘上模型卸载的读写会和swap换页互相放大延迟。我的建议是卸载专用盘上不要放swap必要的话干脆关掉swap或单独用一块低性能盘承载。5.5 文件系统与内核参数ext4、XFS怎么选文件系统选择看起来小事实际也会影响卸载稳定性。我试过ext4、XFS和btrfs最终在AI卸载场景里最常用的是XFS和ext4。btrfs在小文件场景有优势但在大文件高吞吐下容易因为写时复制CoW机制产生额外碎片和写放大卸载场景不太划算。XFS对大文件顺序读写的支持很成熟而且在掉电恢复上比ext4更稳。挂载参数上除了noatime还可以考虑logbsize调大日志缓存块大小提升高写入负载下的吞吐。对于内核参数有一个经验值把vm.dirty_ratio适当调低一点避免内核攒太多脏页一次性刷盘导致IO尖峰。这个可以根据实际写入量来调不建议照搬网上的配置。说到最后如果你现在正在被显存不够折腾我的建议不是先上框架而是先想清楚三个问题哪些状态可以卸载卸载盘用多大带宽IO和计算怎么重叠。框架只是工具这三个问题想清楚了SSD卸载这条路才走得通。