Transformers 中的 Cohere Command-R：RAG 与工具调用场景下的因果语言模型使用指南

Cohere Command-R 是由 Cohere 训练并贡献进 Transformers 的一个 35B 参数多语言大语言模型，针对检索增强生成（RAG）与外部 API / 工具调用做了专门训练，原生支持单步与多步工具使用，上下文长度可达 128K tokens。本文以 Cohere 模型文档 为主线，结合仓库中 configuration_cohere.py、modeling_cohere.py 与 tokenization_cohere.py 的实现细节，系统讲解如何在当前 Transformers 仓库中加载、推理、量化并诊断 Command-R 系列模型，读完你即可在本地用 Pipeline、AutoModel 与 CLI 三种方式跑通 Command-R，并掌握 RAG / 工具调用专用提示模板与内存优化要点。

Command-R 是什么：面向 grounded generation 与工具调用的 35B 模型

Command-R（见原文档对 Command-R 博客 的引用）是一个 35B 参数的多语言 LLM，其设计目标并非通用闲聊，而是长上下文工作负载：

• 检索增强生成（RAG）：能够依据检索到的文档片段生成带引用的答案，即文档所述的"grounded generation"；
• 调用外部 API 与工具：支持单步（single-step）与多步（multi-step）工具调用；
• 128K tokens 上下文窗口：可一次性纳入长文档或长对话历史。

在 Transformers 中，Command-R 的 checkpoint 可直接通过自动模型类加载，cohere 已注册进 modeling_auto.py（CohereModel、CohereForCausalLM）对应的映射表。除 Command-R 本体外，同一系列的其他 checkpoint 可参见 Command Models 集合。注意本页只覆盖第一代 Command-R（CohereForAI/c4ai-command-r-v01）；仓库中另有独立的 cohere2 文档 介绍采用滑动窗口注意力的 7B 开源版本 Command R7B，两者实现分属 cohere 与 cohere2 两个模块。

三种方式快速上手 Command-R 生成

原文档用同一段"以问答为例"的代码展示了三种入口，下面逐一给出可完整运行的形态。

方式一：Pipeline（一行完成任务）

from transformers import pipeline

pipeline = pipeline(
    task="text-generation",
    model="CohereForAI/c4ai-command-r-v01",
    device=0  # 指定 GPU 设备
)
pipeline("Plants create energy through a process known as")

Pipeline 内部会自动组装 tokenizer 与模型，适合快速验证与脚本化推理。

方式二：AutoModelForCausalLM + chat template（推荐）

这是最贴近生产实践的写法——它通过 apply_chat_template 走 Command-R 专属聊天模板，而不是把用户句子裸塞给模型：

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
model = AutoModelForCausalLM.from_pretrained(
    "CohereForAI/c4ai-command-r-v01",
    device_map="auto",
    attn_implementation="sdpa",
)

# 用 Command-R 的 chat template 格式化消息
messages = [{"role": "user", "content": "How do plants make energy?"}]
input_ids = tokenizer.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
).to(model.device)
output = model.generate(
    input_ids,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.3,
    cache_implementation="static",
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

几点参数解释（均可替换为其他采样参数）：

• attn_implementation="sdpa" 启用 PyTorch 原生 SDPA 内核（见下文"注意力后端"）；
• temperature=0.3 偏低，适合事实性问答；do_sample=True 开启随机采样；
• cache_implementation="static" 使用静态 KV cache 以提升长序列解码吞吐；
• add_generation_prompt=True 保证输入末尾带上"开始生成助手回复"的控制 token。

方式三：transformers CLI（无需写 Python）

# pip install -U flash-attn --no-build-isolation
transformers chat CohereForAI/c4ai-command-r-v01 --dtype auto --attn_implementation flash_attention_2

命令行中的 --dtype auto 让框架按设备自动选择权重精度；若本地已安装 FlashAttention，可用 flash_attention_2 后端获得高性能解码。CLI 需要较新版本仓库配套的 transformers 命令，可通过仓库 cli 目录了解 chat 子命令的其他开关。

从源码看 Command-R 的架构细节

原文档对本模型的代码说明非常克制（正文仅提示它"基于 EleutherAI 的 GPT-NeoX 代码改写"），但翻开 modeling_cohere.py 可以发现它与 Llama 的同源关系与四处关键差异，理解这些差异有助于你调参和做二次开发：

1. logit_scale 输出缩放（与 Llama 最显眼的区别）。配置中 logit_scale 默认为 0.0625。在 modeling_cohere.py 中，CohereForCausalLM.forward 计算完 lm_head 后执行 logits = logits * self.logit_scale，这是训练时就引入的固定缩放，推理与微调时不应移除，否则 logits 量纲错位会破坏采样分布。

2. QK Norm（query-key 归一化）。use_qk_norm 默认 False，开启后注意力会先对 Q、K 在每个 head_dim 上做一次 CohereLayerNorm 归一化再进 RoPE 与注意力。代码注释将其标注为"main diff from Llama"。

3. Rotary Embedding 的 interleave 拼接。RoPE 频率以 interleave 方式交织而非 Llama 式的 cat，见 CohereRotaryEmbedding.forward 中 torch.repeat_interleave(freqs, 2, dim=-1) 的注释。

4. pre-LayerNorm + 并行分支的 decoder layer。每个 CohereDecoderLayer 内部是 residual + attention + mlp：先对残差做一次 LayerNorm，attention 与 MLP（SwiGLU 结构 gate/up/down 三个线性层）读取同一份归一化结果后并行计算再相加，这与 Llama 顺序式（先 attn 后 mlp 各自带残差）不同，减少了 LayerNorm 次数。

此外，配置类 CohereConfig 本身是 @strict 数据类，默认超参数直接对应 Command-R 的官方规模，rope_theta 默认 500000.0 支撑 128K 长上下文的外推。模型类还声明了 base_model_tp_plan / base_model_pp_plan，表明支持张量并行与流水线并行的分层切分计划。

注意力后端支持

从 modeling_cohere.py 中 CoherePreTrainedModel 的能力标记看，cohere 模型类支持 _supports_flash_attn、_supports_sdpa、_supports_flex_attn，因此可以传入：

• flash_attention_2：需自行安装 flash-attn，速度最快，仅支持 fp16/bf16（见下文注意事项）；
• sdpa：PyTorch 内置，无需额外依赖，是文档示例中的默认推荐；
• eager：走纯 PyTorch 算子，便于调试；
• flex_attn 与分页注意力等能力在仓库 generation 与 tests/generation/test_flash_attention_parity.py 中有更完整的覆盖测试。

用 bitsandbytes 把 Command-R 压到 4-bit

Command-R 有 35B 参数，fp16 权重即需约 70GB 显存，单卡难以直接加载。文档给出的做法是用 bitsandbytes 量化到 4-bit：

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(load_in_4bit=True)
tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
model = AutoModelForCausalLM.from_pretrained(
    "CohereForAI/c4ai-command-r-v01",
    device_map="auto",
    quantization_config=bnb_config,
    attn_implementation="sdpa",
)

# 用 Command-R 的 chat template 格式化消息
messages = [{"role": "user", "content": "How do plants make energy?"}]
input_ids = tokenizer.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
).to(model.device)
output = model.generate(
    input_ids,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.3,
    cache_implementation="static",
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

要点：

• BitsAndBytesConfig(load_in_4bit=True) 也可按需追加 bnb_4bit_compute_dtype、bnb_4bit_quant_type 等字段细化精度与计算 dtype；
• quantization_config 与 device_map="auto" 配合，使各层在加载时就被分片到可用设备；
• 关于量化后如何微调（QLoRA）以及更多量化后端（AWQ、GPTQ、HQQ 等），参见仓库 bitsandbytes 指南 与 量化总览。

用 AttentionMaskVisualizer 理解注意力掩码

Command-R 面向长上下文、支持左填充批处理，理解"哪些 token 能 attend 到哪些 token"对排查生成质量很有价值。文档给出的诊断工具是注意力掩码可视化器，代码位于仓库 attention_visualizer.py（该文件中的 AttentionMaskVisualizer 类会基于模型的因果掩码把每个 token 的可 attend 关系渲染成矩阵图）：

from transformers.utils.attention_visualizer import AttentionMaskVisualizer

visualizer = AttentionMaskVisualizer("CohereForAI/c4ai-command-r-v01")
visualizer("Plants create energy through a process known as")

输出会以逐 token 的网格展示当前 prompt 下哪些位置相互可见。由于本仓库不含该输出图的副本（原文档中的示意图托管在文档仓库），运行时请以终端打印结果为准。

使用注意事项

原文档"Notes"一节专门提醒了一类易错点，此处完整保留并补充解释：

• 使用 FlashAttention-2 时不要在 AutoModel.from_pretrained 中传 dtype 参数。FlashAttention-2 只支持 fp16 或 bf16，直接指定 dtype=torch.float32 会报错或静默失效。正确姿势是：
  • 训练场景：在 Trainer 中开启 fp16=True 或 bf16=True（自动混合精度，AMP）；
  • 推理脚本：用 torch.autocast 上下文包裹前向传播，让算子以半精度执行；
  • 或用 --dtype auto（CLI）让框架按设备能力自动选择。

这一限制同样适用于所有依赖 flash-attn 的模型，属于使用 Transformers 的通用经验。

API 速查：配置、Tokenizer 与模型类

原文档以 autodoc 形式给出了四个公开类，本文将其整理为便于检索的速查表。

CohereConfig

定义在 configuration_cohere.py，继承 PreTrainedConfig，model_type = "cohere"。默认值即 Command-R 官方超参：

参数	默认值	说明
vocab_size	256000	词表大小
hidden_size	8192	隐藏层维度
intermediate_size	22528	SwiGLU 中间层维度
num_hidden_layers	40	解码器层数
num_attention_heads	64	注意力头数
num_key_value_heads	None（= attention heads）	GQA 的 KV 头数，缺省时与 Q 头相同
hidden_act	"silu"	MLP 激活函数
max_position_embeddings	8192	训练最大序列长度
initializer_range	0.02	参数初始化范围
layer_norm_eps	1e-5	LayerNorm 的 epsilon
logit_scale	0.0625	输出 logits 的固定缩放系数
use_cache	True	是否缓存 past_key_values
pad_token_id / bos_token_id / eos_token_id	0 / 5 / 255001	特殊 token id
tie_word_embeddings	True	是否捆绑输入/输出词嵌入
rope_parameters.rope_theta	500000.0	RoPE 基数，支撑长上下文
attention_bias / attention_dropout	False / 0.0	注意力线性层偏置与 dropout
use_qk_norm	False	是否启用 QK 归一化

__post_init__ 中会在 num_key_value_heads=None 时自动补齐为 num_attention_heads，因此你不需要手动设置 KV 头数也能得到与 Command-R 一致的配置。

CohereTokenizer

定义在 tokenization_cohere.py，底层是 byte-level BPE，并启用了 ByteFallback（未登录 UTF-8 字节回退）与 NFC 归一化（预处理阶段即执行，见其 normalizer 配置），其 padding_side 固定为 "left"（左侧填充，适合 decoder-only 长上下文批处理）。特殊 token 约定：<BOS_TOKEN>=bos、<|END_OF_TURN_TOKEN|>=eos、<PAD>=pad。

注意：如果自定义了 bos_token / eos_token，必须同步调用 tokenizer.update_post_processor() 重新生成 post-processor，否则编码结果的首/尾 token 值不准确；而 add_prefix_space=True 虽可绕过部分行为，但模型预训练并未采用该设置，可能造成精度下降。当配合 is_split_into_words=True 做逐词标注时，则必须设置 add_prefix_space=True。

除常规聊天模板外，该 tokenizer 还暴露两个 Command-R 特性方法，是 RAG / Agent 场景的杀手锏：

• apply_grounded_generation_template(conversation, documents, citation_mode="accurate")：把对话历史与检索文档（形如 {"title": ..., "text": ...} 的字典列表）渲染为带引用指令的 RAG prompt。citation_mode 取 "accurate"（先答后补引用，引用质量更高）或 "fast"（直接生成带引用答案，所需生成 token 更少）；
• apply_tool_use_template(conversation, tools)：把工具清单（name / description / parameter_definitions）与对话渲染为工具调用 prompt，模型会输出类似 [{"tool_name": "internet_search", "parameters": {...}}] 的结构化动作，多轮循环即可实现多步工具调用。

二者在实现上都会委托给 apply_chat_template（分别对应 chat_template="rag" 与 "tool_use"），具体调用示例见 tokenization_cohere.py 中两个方法的 doctest，可按其中给出的 internet_search / directly_answer 等工具 schema 直接改写复用。

CohereModel 与 CohereForCausalLM

• CohereModel（裸主干）：词嵌入 → 40 层 CohereDecoderLayer → 末层 LayerNorm，输出 BaseModelOutputWithPast，前向仅暴露 forward 方法。加载原始 checkpoint 权重时直接对应仓库中的 model 权重键。
• CohereForCausalLM：在主干之上叠加 lm_head（默认与 embed_tokens 权重捆绑，见 _tied_weights_keys），并在 forward 中执行 logits * logit_scale；它混入 GenerationMixin，可直接调用 generate()。该类的 doctest 演示了最基本的续写流程（model.generate(...) 后 tokenizer.batch_decode）。

两者的完整参数签名与返回值均可在 modeling_cohere.py 中查阅，常见入参包括 input_ids、attention_mask、position_ids、past_key_values（KV cache，类型为 cache_utils.py 中的 Cache）、inputs_embeds、use_cache 等。

相关阅读

• 同类后一代模型见 Cohere2 模型文档（Command R7B，滑动窗口注意力）；
• 量化总览 与 bitsandbytes 提供更多低比特加载后端；
• 模型与 tokenizer 的官方测试分别在 test_modeling_cohere.py 与 test_tokenization_cohere.py，其中包含与参考实现一致性校验、长上下文与注意力掩码等用例，可作为自建适配的回归基线。