多卡推理时，如果显存仍然不足，可以把模型的一部分层固定放到 CPU 上，只让剩余层占用 GPU 显存。这种做法的本质是模型并行中的“层级切分”，不是把所有参数都赶进显存，而是按层的执行顺序把一部分层拆到 CPU 内存里。它和单纯依赖 PyTorch 的 swap 机制不同，后者是页面调度，层卸载则是提前确定每个层所在的设备，推理时按需执行。

层卸载到 CPU 能解决显存容量不够，但代价是推理延迟明显增加。它适合单次请求延迟不敏感、批量小、模型很大且必须跑通的场景；不适合要求高吞吐或低延迟的在线服务。实现上可以用 accelerate 或 transformers 的 device_map 自动分配，也可以手动把层搬到 CPU。需要先评估可接受的响应时间，再决定卸载多少层。

先判断卸载 CPU 是否真的解决问题

显存不足时，先确认是“峰值显存超限”还是“整体参数放不下”。如果只是上下文长度导致激活值暴涨，优先用梯度检查点（gradient checkpointing）、减小 batch size 或截断输入长度，这些不会引入 CPU 传输开销。只有模型权重本身超过 GPU 总显存时，才适合把部分层放到 CPU。

可以在每张卡上逐个加载层，观察各卡显存占用比例，找出需要腾出的空间。卸载的层通常是 Transformer 中间层，而不是 embedding 或 lm_head，因为后两者在推理每次请求时几乎都要访问，放在 CPU 会导致每次前向都拷贝特征，延迟增加更明显。

用 accelerate 的 device_map 快速实现分层卸载

transformers 在加载模型时支持 device_map 参数，可以同时用 multi-GPU 和 cpu。accelerate 会自动划分层，你也可以手动指定某几层放到 CPU。下面是一个通用骨架，替换成你自己的模型路径和层名即可。

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("your/model")
model = AutoModelForCausalLM.from_pretrained(
    "your/model",
    device_map="auto",
    max_memory={0: "20GiB", 1: "20GiB", "cpu": "64GiB"},
)

max_memory 中的 0 和 1 是 GPU 编号，按实际卡号修改。cpu 上限需要留足系统内存余量，避免 OOM。accelerate 会尽量把层分配到满足条件的设备上，放不下的层自动落到 CPU。

如果自动分配结果不合意，可以手动指定 layers 所属设备：

device_map = {
    "model.embed_tokens": 0,
    "model.layers.0": 0,
    "model.layers.1": 0,
    "model.layers.2": "cpu",
    "model.layers.3": "cpu",
    "model.layers.4": 1,
    "model.layers.5": 1,
    "model.layers.norm": 1,
    "lm_head": 1,
}
model = AutoModelForCausalLM.from_pretrained(
    "your/model", device_map=device_map
)

这里需要先确认模型的层名结构，通常用 model.layers.0 这种格式。手动指定后，accelerate 会跳过自动编排，直接按你的映射放置张量。

验证卸载效果和性能是否可接受

加载模型后，先跑一个短输入，确认前向和增量生成都能执行，没有设备不匹配报错。然后观察各卡显存占用是否降到目标范围内，CPU 内存占用是否还在安全线以下。

推理前记下开始时间，生成固定长度输出后计算耗时。多做几组，对比全 GPU 模式和 CPU 卸载模式的延迟差异。如果延迟从原本的 100 毫秒变成 10 秒，说明卸载层太多，需要减少 CPU 层数，或者把部分层分配到第二张卡。

同时用 nvidia-smi 看 GPU 显存变化，用 free -h 看系统内存。CPU 层在被访问时会持续占用内存，如果系统内存也紧张，卸载层数必须继续调低。

风险和边界：不是所有层都适合放 CPU

embedding、lm_head、norm 层在每轮前向都会被高频率访问，放在 CPU 会迫使 CPU 和 GPU 在每次生成 step 之间频繁交换数据。建议只卸载中间的 transformer 层，而且尽量把相邻层切到同一个设备，减少跨设备拼接次数。

CPU 层在推理时仍然参与计算，没有 CUDA 加速，速度会明显慢于 GPU。如果 CPU 是瓶颈，可以考虑用 torch.compile 优化 CPU 上的计算，或者把部分层分配到第二张 GPU 而不是 CPU。

另外，不同框架对 device_map 的支持程度不同。上面的写法适用于 transformers 4.30 以上版本和 accelerate 0.20 以上版本。如果你的代码里自己用 nn.Module.to('cpu') 搬层，需要手动确保前向传播时把输入张量也搬到对应设备，否则会报 device mismatch。

手动切分的替代做法：自定义 forward

如果不想引入 accelerate，也可以手动拆。把模型按层列表切分成两个子列表，前一段放到 GPU:0，后一段放到 CPU。前向时，先把输入送到 GPU 上的层，得到中间张量后搬到 CPU 继续，最后结果再搬回 GPU。这个做法需要你自己管理每层输入输出的设备，适合对模型结构非常熟悉的情况。

def forward_with_cpu_offload(model, input_ids, gpu_layers, cpu_layers, device_gpu, device_cpu):
    hidden = model.embed_tokens(input_ids).to(device_gpu)
    for layer in gpu_layers:
        hidden = layer(hidden)[0]
    hidden = hidden.to(device_cpu)
    for layer in cpu_layers:
        hidden = layer(hidden)[0]
    hidden = hidden.to(device_gpu)
    return model.lm_head(hidden)

这段代码是示意，实际还要处理 KV cache、位置编码、输出 logits 的形状。手动切分的好处是可控，坏处是每个模型的结构不一样，迁移成本高。建议优先用 device_map 自动方案，手动切分只在自动方案无法满足特定设备布局时才用。