![[代码学习] vLLM 如何把大模型权重临时搬到 CPU：PR #51710 源码阅读](http://pic.xiahunao.cn/yaotu/[代码学习] vLLM 如何把大模型权重临时搬到 CPU：PR #51710 源码阅读)
最近读了 vLLM 的 PR #51710。这个 PR 解决的问题很直接MoE 模型太大专家权重放不进显存怎么办它给出的答案是先把一部分权重放在 CPU 内存里GPU 快用到它们时再提前拷贝回来。只要拷贝和计算重叠得足够好就能用一些 PCIe 流量换取几十 GB 显存。先弄清楚它在搬什么MoE 模型的每一层里通常有很多专家。虽然一次推理只会选中部分专家但所有专家的权重往往都要常驻显存。DeepSeek-V4 这类模型的专家权重非常大多卡切分以后每张卡仍然可能放不下。这里搬的是模型权重不是 KV Cache也不是请求数据。权重平时放在 pinned CPU memory 中某层快要执行时vLLM 把对应权重异步复制到 GPU 的临时缓冲区。CPU第2、5、8、11 层的专家权重 GPU固定大小的临时权重缓冲区 GPU 计算第2层时后台复制第5层 GPU 计算第5层时后台复制第8层哪些层会被卸载核心判断其实很朴素ifmodule_index%group_sizegroup_size-num_in_group:# 选择这一层进行 offload假设group_size3、num_in_group1代码会把每三个模块分成一组然后选择每组最后一个模块模块012|345|678卸载 √|√|√接着代码还会根据参数名做一次筛选。比如只卸载experts.w13_weight和experts.w2_weight而保留attention和其他小参数。匹配时在名字两边补上 .是为了让w2_weight不会误匹配到w2_weight_scale。GPU 缓冲区为什么要循环使用如果每次预取都重新申请显存开销会很大也很难放进 CUDA Graph。PR 使用StaticBufferPool提前申请固定缓冲区常驻 GPU 的权重不走这套逻辑slot_idxlayer_idx%prefetch_stepbufferbuffers[slot_idx]例如prefetch_step3GPU 只准备三组槽位offload unit0-slot0offload unit1-slot1offload unit2-slot2offload unit3-slot0前面的层用完后后面的层继续覆盖同一块显存。参数对象会提前指向这些静态 buffer因此 forward 时不需要反复替换参数也不需要动态申请 tensor。需要注意buffer 不能只按 tensor shape 复用。量化后的权重可能有不同的 stride 和 dtype所以实际 key 中包含参数名、shape、stride 和 dtype。预取是怎么插进 forward 的PR 给选中的模块包了一层 forward hook。简化以后大致是这样defforward(*args,**kwargs):wait_prefetch(current_layer)outputoriginal_forward(*args,**kwargs)start_prefetch(current_layerprefetch_step)returnoutput执行当前层之前计算 stream 等待该层权重复制完成。当前层算完后马上在单独的 copy stream 上预取后面的权重。理想情况下H2D copy 会被中间几层的计算时间盖住。这里的prefetch_step不是越大越好。太小权重可能来不及搬完GPU 只能停下来等太大又需要更多 GPU buffer节省的显存会变少。为什么还要处理 CUDA Graphcopy stream 和 compute stream 是两条不同的 CUDA stream。CUDA Graph capture 结束时如果 copy stream 启动了任务却没有重新汇合到当前 stream就会报cudaErrorStreamCaptureUnjoined所以 PR 增加了join_after_forward()。它会检查哪些预取任务是在 capture 期间启动、但还没有被等待然后通过 CUDA event 把这些任务重新接回当前 stream。这部分看起来只是同步细节但很关键。作者的 baseline 可以完成权重加载却会在 CUDA Graph capture 阶段失败补上这个生命周期以后服务才能真正启动。H2D copy 为什么会拖慢 TP在 PCIe 机器上CPU 到 GPU 的权重复制和 GPU 之间的 TP collective 可能争抢同一批链路。一笔很大的权重复制占着链路时all-reduce 或 all-gather 会变慢而 TP 中所有 rank 都在等它。PR 因此加入--offload-comm-aware把大 copy 切成小块collective 开始时暂缓复制collective 完成后再继续。这个逻辑只在 TP 1 时启用。作者在 TP8 的 DeepSeek-V4-Pro 上测得开启这项控制后 TTFT 降低约 9.7%吞吐提高约 10.9%。这说明预取不只是“尽量早搬”还要知道什么时候应该让路。Schedule Planner 是干什么的三个参数可以组合出很多方案。以前只能换一组参数、启动一次模型然后看是否 OOM。这个 PR 允许第一次启动时输出 manifest记录每个位置的权重大小、buffer layout 和实际卸载结果VLLM_PREFETCH_LOG_SCHEDULE1vllm serve...随后可以离线运行 plannerpython-m vllm.benchmarks.weight_offload.plannerplanner 会计算每个方案剩余多少常驻权重、需要多大的运行时 buffer以及每次 forward 要重新传输多少字节。它只计算内存不预测延迟因为真实延迟还取决于 PCIe 带宽、每层计算时间和 TP 通信竞争。实际效果和限制作者在 DeepSeek-V4-Pro、TP8、RTX PRO 6000 上每个 rank 卸载约 31.38 GiB 权重使用约 4.71 GiB GPU 临时 buffer最终净释放约 26.67 GiB 显存。每次 forward 需要重新搬运约 33.69 GB但超过 99.8% 的复制时间被计算覆盖。它也不是免费的CPU 必须有足够大的 pinned memoryPCIe 必须有足够带宽还要给 fused-MoE workspace 留出显存。
最近读了 vLLM 的 PR #51710。这个 PR 解决的问题很直接MoE 模型太大专家权重放不进显存怎么办它给出的答案是先把一部分权重放在 CPU 内存里GPU 快用到它们时再提前拷贝回来。只要拷贝和计算重叠得足够好就能用一些 PCIe 流量换取几十 GB 显存。先弄清楚它在搬什么MoE 模型的每一层里通常有很多专家。虽然一次推理只会选中部分专家但所有专家的权重往往都要常驻显存。DeepSeek-V4 这类模型的专家权重非常大多卡切分以后每张卡仍然可能放不下。这里搬的是模型权重不是 KV Cache也不是请求数据。权重平时放在 pinned CPU memory 中某层快要执行时vLLM 把对应权重异步复制到 GPU 的临时缓冲区。CPU第2、5、8、11 层的专家权重 GPU固定大小的临时权重缓冲区 GPU 计算第2层时后台复制第5层 GPU 计算第5层时后台复制第8层哪些层会被卸载核心判断其实很朴素ifmodule_index%group_sizegroup_size-num_in_group:# 选择这一层进行 offload假设group_size3、num_in_group1代码会把每三个模块分成一组然后选择每组最后一个模块模块012|345|678卸载 √|√|√接着代码还会根据参数名做一次筛选。比如只卸载experts.w13_weight和experts.w2_weight而保留attention和其他小参数。匹配时在名字两边补上 .是为了让w2_weight不会误匹配到w2_weight_scale。GPU 缓冲区为什么要循环使用如果每次预取都重新申请显存开销会很大也很难放进 CUDA Graph。PR 使用StaticBufferPool提前申请固定缓冲区常驻 GPU 的权重不走这套逻辑slot_idxlayer_idx%prefetch_stepbufferbuffers[slot_idx]例如prefetch_step3GPU 只准备三组槽位offload unit0-slot0offload unit1-slot1offload unit2-slot2offload unit3-slot0前面的层用完后后面的层继续覆盖同一块显存。参数对象会提前指向这些静态 buffer因此 forward 时不需要反复替换参数也不需要动态申请 tensor。需要注意buffer 不能只按 tensor shape 复用。量化后的权重可能有不同的 stride 和 dtype所以实际 key 中包含参数名、shape、stride 和 dtype。预取是怎么插进 forward 的PR 给选中的模块包了一层 forward hook。简化以后大致是这样defforward(*args,**kwargs):wait_prefetch(current_layer)outputoriginal_forward(*args,**kwargs)start_prefetch(current_layerprefetch_step)returnoutput执行当前层之前计算 stream 等待该层权重复制完成。当前层算完后马上在单独的 copy stream 上预取后面的权重。理想情况下H2D copy 会被中间几层的计算时间盖住。这里的prefetch_step不是越大越好。太小权重可能来不及搬完GPU 只能停下来等太大又需要更多 GPU buffer节省的显存会变少。为什么还要处理 CUDA Graphcopy stream 和 compute stream 是两条不同的 CUDA stream。CUDA Graph capture 结束时如果 copy stream 启动了任务却没有重新汇合到当前 stream就会报cudaErrorStreamCaptureUnjoined所以 PR 增加了join_after_forward()。它会检查哪些预取任务是在 capture 期间启动、但还没有被等待然后通过 CUDA event 把这些任务重新接回当前 stream。这部分看起来只是同步细节但很关键。作者的 baseline 可以完成权重加载却会在 CUDA Graph capture 阶段失败补上这个生命周期以后服务才能真正启动。H2D copy 为什么会拖慢 TP在 PCIe 机器上CPU 到 GPU 的权重复制和 GPU 之间的 TP collective 可能争抢同一批链路。一笔很大的权重复制占着链路时all-reduce 或 all-gather 会变慢而 TP 中所有 rank 都在等它。PR 因此加入--offload-comm-aware把大 copy 切成小块collective 开始时暂缓复制collective 完成后再继续。这个逻辑只在 TP 1 时启用。作者在 TP8 的 DeepSeek-V4-Pro 上测得开启这项控制后 TTFT 降低约 9.7%吞吐提高约 10.9%。这说明预取不只是“尽量早搬”还要知道什么时候应该让路。Schedule Planner 是干什么的三个参数可以组合出很多方案。以前只能换一组参数、启动一次模型然后看是否 OOM。这个 PR 允许第一次启动时输出 manifest记录每个位置的权重大小、buffer layout 和实际卸载结果VLLM_PREFETCH_LOG_SCHEDULE1vllm serve...随后可以离线运行 plannerpython-m vllm.benchmarks.weight_offload.plannerplanner 会计算每个方案剩余多少常驻权重、需要多大的运行时 buffer以及每次 forward 要重新传输多少字节。它只计算内存不预测延迟因为真实延迟还取决于 PCIe 带宽、每层计算时间和 TP 通信竞争。实际效果和限制作者在 DeepSeek-V4-Pro、TP8、RTX PRO 6000 上每个 rank 卸载约 31.38 GiB 权重使用约 4.71 GiB GPU 临时 buffer最终净释放约 26.67 GiB 显存。每次 forward 需要重新搬运约 33.69 GB但超过 99.8% 的复制时间被计算覆盖。它也不是免费的CPU 必须有足够大的 pinned memoryPCIe 必须有足够带宽还要给 fused-MoE workspace 留出显存。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
