Gemma-4-E4B: the PLE table holds ~5.6 GB of VRAM that can live on CPU bit-identically – is there a supported device_map path?

快速结论：该问题发生在加载 google/gemma-4-E4B-it 模型时，模型内部的 per-layer embedding（PLE）表占用大量显存。优先排查方向是：在 from_pretrained 时是否使用了细粒度（per-module）的 device_map 字典，而不是只把顶层模块指定为 "cpu"。

适用环境：Transformers 5.8.1、Accelerate 1.13.0、PyTorch 2.6.0+cu124、NVIDIA RTX A5000 Laptop（16 GB）、模型 google/gemma-4-E4B-it。加载方式为 4-bit NF4 量化，且 PLE 相关模块保持 bf16。

最快修复方案：在 from_pretrained 中使用细粒度 device_map 字典，按模块全路径逐个指定设备。例如将某一层（如第 31 层）放到 CPU 并存储 lm_head 到磁盘，语法为 device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}。Issue 作者未在修复后重新验证加载流程，因此建议按此方式重新测试。

注意事项：目前 Transformers 维护者未确认是否会将 PLE 默认放到 CPU 作为官方加载方式；且作者实测 CPU 卸载后生成速度略快（而非变慢），这一点值得进一步验证。后加载（post-load）手动 offload 方案虽然有效，但并非官方支持路径，PEFT 包装等场景需按名称后缀匹配。

问题场景

用户在加载 google/gemma-4-E4B-it 模型时发现，单张 16 GB 显卡（RTX A5000 Laptop）在 8.5k token prefill 阶段直接 OOM。定位后发现 embed_tokens_per_layer（PLE 表）占了约 5.6 GB 显存，将这一模块手动移到 CPU 后显存占用从 10.11 GB 降到 4.48 GB，峰值从 OOM（15.5 GB）降到 11.84 GB。用户询问是否有官方支持的 device_map 路径可以让 PLE 表在加载时就直接放到 CPU。

报错原文

Gemma-4-E4B: the PLE table holds ~5.6 GB of VRAM that can live on CPU bit-identically - is there a supported device_map path?

Measured on RTX A5000 16 GB, bnb-4bit NF4:
| PLE on GPU  | resident 10.11 GB | peak 8.5k prefill: OOM at 15.5 GB |
| PLE on CPU  | resident 4.48 GB  | peak 8.5k prefill: 11.84 GB       |

We first tried the supported route — passing a `device_map` that pins `embed_tokens_per_layer` to `"cpu"` — and it did not hold:
the module appeared to be back on CUDA by the time init finished, and we OOMed.

原因分析

可能原因：顶层 device_map（如 {"": "cpu"} 或只指定模型主模块）无法精确控制嵌套的 per-layer embedding 表，模型初始化后该模块被重新放回 CUDA。Transformers 维护者指出，要正确将 PLE 层放到 CPU，必须使用 细粒度 device_map 字典，为每个子模块分别指定设备，而不是只指定顶层模块。用户之前尝试的方式可能没有覆盖到 embed_tokens_per_layer 这个具体子路径。

环境排查

Transformers 版本：5.8.1（需确认是否包含 PLE 相关的模块命名与调度逻辑）
Accelerate 版本：1.13.0（device_map 的解析和执行依赖于此库）
PyTorch 版本：2.6.0+cu124
显卡显存：16 GB（RTX A5000 Laptop）
模型：google/gemma-4-E4B-it，4-bit NF4 量化，PLE 相关模块保持 bf16
在 Prompt 中指定 device_map 时，注意是传 dict 还是字符串，dict 需要覆盖到 embed_tokens_per_layer 所在模块路径

解决步骤

确认模型的模块结构：先打印 model.named_modules()，找到 embed_tokens_per_layer 的完整路径（例如 model.layers.31.embed_tokens_per_layer）。
为这些 PLE 模块单独指定 "cpu"，其余层按需分配到 GPU 或 disk。参考维护者给出的格式：device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}。
在 from_pretrained(..., device_map=device_map) 中传入该细粒度字典，然后用 model.hf_device_map 确认所有 PLE 模块最终所在设备是否为 CPU。
如果细粒度 device_map 仍不稳定，可优先尝试 Issue 作者提供的后加载方案（按名称后缀匹配 embed_tokens_per_layer，将 forward 改为 CPU gather 后回传 GPU），但要注意这只是绕过手段。

验证方法

加载完成后打印 model.hf_device_map，确认 embed_tokens_per_layer 对应条目为 "cpu"。用 torch.cuda.memory_reserved() 或 nvidia-smi 核对显存占用是否从约 10 GB 下降到约 4.5 GB；然后用 8.5k token 的 prefill 测试是否不再 OOM。最后用 greedy decoding 生成一段长文本，与 PLE 在 GPU 时的输出做逐 token 对比，确认结果 bit-identical。