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

Parameters:

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

Raises:

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