delete_kv_caches

torchtune.modules.common_utils.delete_kv_caches(model: Module)

从模型中所有注意力层删除 KV 缓存，并确保 cache_enabled 设置为 False。

示例

>>> from torchtune.models.llama3_2 import llama3_2_1b
>>> from torchtune.modules import delete_kv_caches
>>> import torch
>>> model = llama3_2_1b()
>>> model.setup_caches(batch_size=1,
>>>                     dtype=torch.float32,
>>>                     decoder_max_seq_len=1024)
>>> print(model.caches_are_setup())
True
>>> print(model.caches_are_enabled())
True
>>> print(model.layers[0].attn.kv_cache)
KVCache()
>>> delete_kv_caches(model)
>>> print(model.caches_are_setup())
False
>>> print(model.caches_are_enabled())
False
>>> print(model.layers[0].attn.kv_cache)
None

参数:

model (nn.Module) – 要为其启用 KV 缓存的模型。

引发:

ValueError – 如果在未设置缓存的模型上调用此函数。请先使用 setup_caches() 设置缓存。