Cartridges

Cartridges 是一种提示学习方法，它将压缩的长上下文表示存储为参数化的 KV-cache 前缀。其核心思想源自论文 Cartridges: Lightweight and general-purpose long context representations via self-study。

有关高层概述和动机，请参阅博客文章 Cartridges: Storing long contexts in tiny caches with self-study。

Cartridges 与 Prefix Tuning 的区别

Prefix Tuning 和 Cartridges 均通过向基础模型注入 past_key_values（前缀 KV 缓存）来实现服务。

Prefix Tuning 学习虚拟 token 嵌入（以及可选的 MLP 投影），并生成 KV 前缀。
Cartridges 直接学习 KV 前缀本身（即 p 个虚拟 token 的逐层键/值向量），并且被设计为从真实的 prefill KV 进行初始化（例如语料库/系统提示词的前 p 个 token）。

该论文还建议冻结第一个 token 作为注意力汇聚点（attention sink）以提升稳定性（num_frozen_tokens=1 为默认值）。

用法（推理）

加载训练好的 CARTRIDGE 适配器并运行生成：

from transformers import AutoModelForCausalLM, AutoTokenizer

from peft import PeftModel

model_id = "Qwen/Qwen2.5-0.5B-Instruct"
adapter_path = "path/to/cartridge_adapter"

base = AutoModelForCausalLM.from_pretrained(model_id)
model = PeftModel.from_pretrained(base, adapter_path)

tok = AutoTokenizer.from_pretrained(model_id)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

out = model.generate(**tok("Question about the corpus:", return_tensors="pt"), max_new_tokens=64)
print(tok.decode(out[0], skip_special_tokens=True))

如果需要在训练之前创建并初始化 cartridge，请参阅下方的初始化选项。

初始化选项

该论文讨论了几种实用的初始化策略：

随机 KV（默认）：创建一个 CartridgeConfig 并开始训练。这会随机初始化 KV 前缀。
来自提示/语料库前几个 token 的 KV：使用 initialize_kv_prefix_from_text(model, tokenizer, text=...)。这会对 text 执行 prefill，并将前 num_virtual_tokens 个 token 对应的 KV 缓存结果复制到适配器中。
来自现有缓存的 KV：如果您已经从基础模型的预填充中获得了 past_key_values 对象，请使用 initialize_kv_prefix_from_past_key_values(model, past_key_values=...)。

训练

Cartridges 论文提出了一种 SELF-STUDY 蒸馏目标（一个冻结的基础模型提供教师 logits；CARTRIDGE 适配器经过训练，使学生模型在目标片段上的下一个 token 分布与教师模型相匹配）。 PEFT 将训练逻辑与核心库分离；参考工作流请参见 https://github.com/huggingface/peft/tree/main/examples/cartridge_self_study。示例脚本使用冻结的基础模型作为教师，使用适配后的模型作为学生，因此两者共享相同的基础检查点。

组合

要将独立训练的 cartridge 拼接为单个 adapter，请使用 compose_cartridge_adapters(...)。

API

CartridgeConfig

类 peft.CartridgeConfig

< source >

( task_type: Optional[Union[str, TaskType]] = Nonepeft_type: Optional[Union[str, PeftType]] = Noneauto_mapping: Optional[dict] = Nonepeft_version: Optional[str] = Nonebase_model_name_or_path: Optional[str] = Nonerevision: Optional[str] = Noneinference_mode: bool = Falsenum_virtual_tokens: int = Nonetoken_dim: int = Nonenum_transformer_submodules: Optional[int] = Nonenum_attention_heads: Optional[int] = Nonenum_layers: Optional[int] = Nonemodules_to_save: Optional[list[str]] = Nonenum_frozen_tokens: int = 1 )

参数

num_frozen_tokens (int, 默认为 1) — 在 cartridge 开头保持冻结（不计算梯度）的前缀 token 数量。Cartridges 论文建议冻结第一个 token 作为注意力汇聚点（attention sink）以增强稳定性（将此值设为 1），因为许多 LLM 使用早期 token 作为注意力汇聚点，更改它们可能会损害训练。

CARTRIDGE 的配置，这是一种基于 KV 缓存参数化的前缀适配器。

这与前缀微调（prefix-tuning）的提供方式类似（作为 past_key_values），但它直接将 KV 缓存存储为可训练参数，而不是通过 MLP 投影来学习它。

初始化： Cartridges 论文讨论了多种初始化选项。在 PEFT 中，初始化是独立于构建适配器配置的单独步骤：

随机 KV 初始化（论文选项 2）：通过 get_peft_model(...) 创建适配器。CARTRIDGE 提示编码器的参数由 PyTorch 随机初始化。
从提示/语料库的前几个 token 派生的 KV（论文选项 3）：在 基础模型 上运行无梯度 prefill，并将缓存的前 num_virtual_tokens 个 KV token 复制到适配器中。PEFT 为此提供了实用工具（可从 peft 或 peft.tuners.cartridge.utils 导入）：
- initialize_kv_prefix_from_text(model, tokenizer, text=...)
- initialize_kv_prefix_from_past_key_values(model, past_key_values=...)
如果您已经拥有展平的 KV 前缀张量，可以直接通过提示编码器的 load_prompt_embeddings(...) 方法加载它。

CartridgeEncoder

class peft.CartridgeEncoder

< source >

( config )

参数化的前缀 KV 缓存。

参数以与 PrefixEncoder 输出相同的扁平化布局存储：`[num_virtual_tokens, num_layers

2 * token_dim], wheretoken_dimis per-head hidden size times number of heads (after any GQA adjustment performed by_prepare_prompt_learning_config`).

如果 num_frozen_tokens > 0，则前 num_frozen_tokens 个虚拟令牌将作为不可训练参数存储，其余令牌可训练。

load_prompt_embeddings

< source >

( prompt_embeddings: torch.Tensor )

加载由 PEFT 保存的展平后的提示嵌入（prompt_embeddings）。

PEFT 将提示学习适配器保存为单个 prompt_embeddings 张量。对于 CARTRIDGE，我们根据 self.num_frozen_tokens 将该张量拆分为冻结段和可训练段。

initialize_kv_prefix_from_past_key_values

peft.initialize_kv_prefix_from_past_key_values

< source >

( modeladapter_name: Optional[str] = Nonepast_key_values: Anynum_virtual_tokens: Optional[int] = None )

从现有的缓存前缀（past_key_values）初始化 KV 前缀提示学习适配器。

返回已加载到适配器中的提示嵌入张量。

prompt_embeddings_from_past_key_values

peft.prompt_embeddings_from_past_key_values

< source >

( past_key_values: Anynum_virtual_tokens: int )

将（旧版）past_key_values 缓存转换为 PEFT 保存的展平提示嵌入张量。

输出符合 PeftModel.get_prompt() 对前缀式提示学习所期望的布局：形状为 [num_virtual_tokens, num_layers * 2 * token_dim]。

Update on GitHub