disable_kv_cache

torchtune.modules.common_utils.disable_kv_cache(model: Module) → Generator[None, None, None]

此上下文管理器临时禁用给定模型上的 KV 缓存，该模型必须已设置 KV 缓存。在此上下文管理器中使用该模型的所有前向传递将不使用 KV 缓存。

KV 缓存将在进入上下文管理器时禁用，并在退出时启用，而不会被修改。

这在某些情况下很有用，例如我们可能希望在使用 KV 缓存的模型调用和不使用 KV 缓存的模型调用之间交替，而无需每次都删除和设置缓存的额外开销。

示例

>>> from torchtune.models.llama3_2 import llama3_2_1b
>>> from torchtune.modules import disable_kv_cache
>>> import torch
>>> model = llama3_2_1b()
>>> # setup caches
>>> 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()
>>> # now temporarily disable caches
>>> with disable_kv_cache(model):
>>>     print(model.caches_are_setup())
True
>>>     print(model.caches_are_enabled())
False
>>>     print(model.layers[0].attn.kv_cache)
KVCache()
>>> # caches are now re-enabled, and their state is untouched
>>> print(model.caches_are_setup())
True
>>> print(model.caches_are_enabled())
True
>>> print(model.layers[0].attn.kv_cache)
KVCache()

参数:

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

Yields:

None – 将控制权返回给调用者，并在给定模型上禁用 KV 缓存。

Raises:

ValueError – 如果模型未设置缓存。请先使用 setup_caches() 设置缓存。