本地部署 EmbeddingGemma 2——显存和批大小该怎么配才不 OOM？

本地跑 EmbeddingGemma 2 报 OOM，先别急着降精度或换设备。显存不足通常是几件事叠加：框架没看到目标设备、加载精度选得偏大、批大小超过当前卡在给定输入长度下能承受的量。可行的顺序是先确认设备可见与版本匹配，再测一次空载显存拿到基线，然后从批大小 1 阶梯往上测出 OOM 边界并回退一档，最后才在半精度、分块、设备回退之间做取舍。下面的命令和脚本都以“能被日志验证”为准，具体数值请以你自己机器上的输出为准。

本地部署 EmbeddingGemma 2 出现 OOM，多数情况是设备未被框架识别、精度偏大、批大小超限三者叠加。建议先确认设备与版本，再测空载占用，然后从批大小 1 阶梯上调记录峰值显存，找到边界后回退一档使用。验证方式是可复现的冷启动加固定长度样例；风险边界是输入长度分布一变、并发一上来，边界就会移动，不要照搬别人截图里的数字。

确认运行环境与可用加速设备

这一步的目的是把“设备不可见”“驱动与运行时版本不匹配”这类前置问题排除掉，否则后面测出来的批大小没有意义。先在目标机器上跑一段短脚本，确认框架视角看到的设备：

python - <<'PY'
import torch
print('torch:', torch.__version__)
print('cuda available:', torch.cuda.is_available())
print('device count:', torch.cuda.device_count())
if torch.cuda.is_available():
    print('device name:', torch.cuda.get_device_name(0))
    print('capability:', torch.cuda.get_device_capability(0))
PY

再对照驱动与设备视角的读数：

nvidia-smi

输出形如（版本号与占用请以实际输出为准）：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 5xx.xx      Driver Version: 5xx.xx      CUDA Version: 12.x       |
|-------------------------------+----------------------+----------------------+
|   0  NVIDIA ...          On   | 00000000:01:00.0  On |                  N/A |
| 30%   45C    P0    70W / 250W |   1024MiB / 24576MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

常见“看不见设备”的原因，按排查顺序大致是：容器启动时没有把设备透传进去；驱动版本与框架编译时使用的 CUDA 运行时差得较远；在 WSL 或某些虚拟化环境里设备被宿主占用；机器上有多张卡但环境变量只暴露了一张；显存已被其他进程或上一次没退干净的进程占住。AMD、Apple 等平台对应的查询命令不同，需要按各自后端确认。确认完这一步，再进入加载测试。

用最小脚本加载模型并打印参数与占用

目标是拿到“空载基线”：只有模型权重时占多少，加上一条输入又占多少。这样后面出现 OOM 时，能判断是模型本身就放不下，还是输入侧撑爆的。加载骨架如下，模型类与加载方式需要结合你手上的模型仓库说明确认，这里只给可替换的结构：

import torch
from transformers import AutoTokenizer, AutoModel

MODEL = '/models/embeddinggemma-2'   # 替换为本地路径或模型标识
DEVICE = 'cuda:0'

tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModel.from_pretrained(MODEL, torch_dtype=torch.float16).to(DEVICE).eval()

print('params(M):', sum(p.numel() for p in model.parameters()) / 1e6)
print('dtype:', next(model.parameters()).dtype)
torch.cuda.synchronize()
print('allocated MB:', torch.cuda.memory_allocated() / 1024**2)
print('reserved  MB:', torch.cuda.memory_reserved() / 1024**2)

然后单独跑一条输入，观察增量：

texts = ['一段用于测试的中文文本']
enc = tok(texts, padding=True, truncation=True, max_length=512, return_tensors='pt').to(DEVICE)
torch.cuda.reset_peak_memory_stats()
with torch.no_grad():
    out = model(**enc)
print('tokens:', enc['input_ids'].shape)
print('peak MB:', torch.cuda.max_memory_allocated() / 1024**2)

读法上注意两点：memory_allocated 是当前进程视角的分配量，memory_reserved 包含框架留着的缓存，跟 nvidia-smi 看到的进程占用通常不完全一致，差值主要来自缓存分配器和上下文开销。日志里同时记下这两个值和 nvidia-smi 的占用，后续对比才不容易误判。

阶梯测试批大小找到 OOM 边界

批大小的边界不是一个固定数字，它随输入长度、padding 策略、是否截断而变。所以测试时要把输入长度固定住，再让批大小从 1 往上走：

本地部署 EmbeddingGemma 2——显存和批大小该怎么配才不 OOM？

import time

SAMPLE = '用于压测的中文句子。' * 20      # 固定一条样本的长度
MAXLEN = 512
POOL = [SAMPLE] * 256

records = []
for bs in range(1, 257):
    batch = POOL[:bs]
    enc = tok(batch, padding=True, truncation=True, max_length=MAXLEN,
              return_tensors='pt').to(DEVICE)
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    try:
        with torch.no_grad():
            out = model(**enc)
        torch.cuda.synchronize()
        peak = torch.cuda.max_memory_allocated() / 1024**2
        records.append((bs, enc['input_ids'].shape[1], round(peak), 'ok', round(time.time()-t0, 2)))
        print(f'bs={bs} peak={peak:.0f}MB time={time.time()-t0:.2f}s')
    except torch.cuda.OutOfMemoryError:
        records.append((bs, enc['input_ids'].shape[1], -1, 'OOM', -1))
        print(f'bs={bs} OOM')
        break

把每轮结果记成表，方便交接和复跑：

批大小 | 实际最大 token 数 | 峰值显存 MB | 结果(ok/OOM) | 单批耗时 s
1 / 8 / 16 / 32 / 64 … 逐档填写，遇到第一档 OOM 就停
备注列写清 padding 方式、是否截断、是否有其他进程占卡

使用时的取值建议是 OOM 边界的上一档，或者再留一档余量。原因是线上输入长度分布通常比测试样本更长、更不齐，按最长样本对齐 padding 时显存会明显上去。如果你的输入长度差异很大，可以考虑按长度分桶、或改用“每条单独编码”的方式控制单次 token 总量，而不是只盯条数。

显存不足时在半精度、分块、设备回退之间取舍

下面几个选项都能把显存压下来，但代价不同，需要按你的任务对结果一致性的敏感程度来选：

fp32 换 fp16/bf16：显存通常下降明显，速度一般也有改善；数值上与 fp32 有细微差别，用于检索召回通常影响不大，但用于逐位比对或对精度敏感的打分任务，需要自己用固定样例对比验证。bf16 需要设备支持，不支持的卡上会退回或报错。
限制 max_seq_length：显存和耗时一起下降，代价是超长文本被截断，向量语义会变。适合文档片段本身就不长的场景。
长文本分块再池化：单次显存可控，能覆盖任意长文档；但跨段信息丢失，结果与整段编码不一致，块与块之间的重叠策略会影响最终向量。
减小批大小、增加调用次数：对结果一致性影响最小，代价是整体吞吐下降、调度开销上升。
回退到 CPU：显存问题直接消失，适合低频、离线任务；速度会慢很多，而且部分后端在 CPU 上不支持半精度，通常需要改回 fp32 并确认算子可用。
量化或换更小模型：显存下降幅度通常最大，但速度是否提升要看后端是否真有对应加速，结果一致性也要单独验证，不能默认无损。

取舍顺序上，可以先在“减小批大小”和“限制长度”里找，这两项对结果影响最可控；实在放不下，再考虑半精度；设备回退和量化放到最后，因为它们的代价更难回退。

整理一份可复制的启动配置并复跑验证

测出边界后，把参数固化进配置文件，避免下次靠记忆复现。一个最小配置示例（字段名可按自己的加载代码调整）：

{
  "model_path": "/models/embeddinggemma-2",
  "device": "cuda:0",
  "dtype": "float16",
  "max_seq_length": 512,
  "batch_size": 16,
  "padding": "longest",
  "truncation": true,
  "normalize": true
}

其中 batch_size 填阶梯测试里稳定通过的那一档，max_seq_length 填你实际输入的上限而不是模型上限。然后按下面的步骤做一次冷启动验证：

确认没有其他进程占卡：先看 nvidia-smi 的显存占用，必要时结束残留进程。
新起一个进程按配置文件加载，打印参数量、dtype、空载 allocated 与 reserved，和之前记录的基线对比。
用固定样例（固定条数、固定长度）跑一遍，记录峰值显存、输出向量的 shape、以及归一化后的范数是否接近 1。
把 batch_size 临时调到边界上一档再跑一次，确认确实会 OOM，说明边界判断没有偏差。
把配置、设备信息、框架版本、这次日志一起放到同目录下，交接时直接照配置复跑。

验证时如果峰值显存和基线差得很多，先排查是不是有别的进程共用了这张卡、或者输入长度超出了配置里的上限被静默截断。配置能复跑、日志能对上，才算这套参数真的可以交给别人用。