Rust in Everything
基于 candle 0.11.0

重复惩罚罚的范围不同

主线位置:推理主线的「带 KV 缓存生成」一步里,从 logits 到下一个 token 之前的重复惩罚。
版本:PyTorch 2.15.0 / transformers 5.22.0(CPU)、candle 0.11.0(CPU,accelerate)
flowchart TB
  seq["系统提示 用户问题 已生成"] --> hf["transformers 罚全部"]
  seq --> win["candle 例子 取最后 N 个"]
  win --> drop["前面的不再罚"]
  hf --> s1["贪心 120 个"]
  drop --> s2["第 58 个分叉"]
  seq -->|"context 传全部"| s3["120 个全同"]

实验代码:Python 侧 rep.py、sampling_ref.py,candle 侧 crate pitb-exp 的 bin p27_rep、p26_samp(release)。模型是本地的 SmolLM2-135M-Instruct 权重,f32,transformers 用 eager、torch.set_num_threads(8)。平台 M2 Max、macOS 27、rustc 1.99.0。公式两边逐位相同的核对见取样的「重复惩罚」一节。

PyTorch 的写法

贪心加一个重复惩罚:

enc = tok.apply_chat_template([{"role": "user", "content": "How do you say thank you in Chinese?"}],
                              add_generation_prompt=True, return_tensors="pt", return_dict=True)
out = model.generate(**enc, max_new_tokens=max_new, do_sample=False, repetition_penalty=1.3)
In Mandarin Chinese: "您好,我们都爱你!" (Wǒu wèi yīn shì), which translates to "Thank You, I'm really sorry for your trouble." 

Alternatively, if the person is from China or Taiwan and has not spoken English fluently yet but can speak some basic phrases like 'hello' ('嗎'), they could also use this phrase. For example: "Hello! How have you been?" This will show that they appreciate your kindness without needing an explanation of their

(120 个 token,到长度上限。)RepetitionPenaltyLogitsProcessor 拿到的是 generate 当前的整个 input_ids:prompt 的 39 个 token(对话模板、默认系统提示词、用户问题)加上已经生成的全部。

搬到 candle 的现象

candle 的 llama 例子这样调重复惩罚(--repeat-last-n 默认 128,quantized 例子默认 64):

            let start_at = tokens.len().saturating_sub(args.repeat_last_n);
            candle_transformers::utils::apply_repeat_penalty(

照着写一个贪心循环(p27_rep),惩罚 1.3,分别传最后 64 个、最后 128 个、全部:

let start = match scope.as_str() {
  "all" => 0,
  n => tokens.len().saturating_sub(n.parse::<usize>()?),
};
let logits = apply_repeat_penalty(&logits, 1.3, &tokens[start..])?;
let next = logits.argmax(0)?.to_scalar::<u32>()?;

和 transformers 的 120 个 token 逐个比:

pt 120 same prefix 120
c_all 120 same prefix 120
c_64 69 same prefix 58
c_128 120 same prefix 102
  • 最后 64 个:第 58 个 token 起不同,第 69 个 token 遇到 eos 停了:..., you can also say it in English as such.<|im_end|>。
  • 最后 128 个:第 102 个起不同( How have you been? 变成 How are you?)。
  • 全部:120 个逐个相同。

两边的惩罚公式一样,差的只是罚谁。不报错,文字也通顺。

原因

apply_repeat_penalty 罚的就是传进来的 context,正 logit 除以 penalty、负 logit 乘以 penalty,每个 id 只罚一次,和 transformers 的 RepetitionPenaltyLogitsProcessor 逐位相同。差别在调用方传什么。用 p26_samp 在同一组 logits 上量(上下文是 prompt 39 个加回答的前 60 个,共 99 个):

rep13    context=99 changed candle=60 hf=60 max|diff|=0e0
rep13    last 64 of context: changed=39 max|diff vs hf|=4.663518e0

传全部时两边改动的 60 个 logit 逐位相同;只传最后 64 个,被罚的 token 从 60 个变成 39 个,单个 logit 最多差 4.66。没被罚的 21 个 token 只出现在前 35 个位置里:

[('system', 1), ('You', 3), ('Ġare', 4), ('Ġhelpful', 6), ('ĠAI', 7), ('Ġassistant', 8), ('Ġnamed', 9), ('ĠSm', 10), ('ol', 11), ('LM', 12), ('Ġtrained', 14), ('Ġby', 15), ('ĠH', 16), ('ugging', 17), ('ĠFace', 18), ('<|im_end|>', 19), ('user', 22), ('How', 24), ('Ġdo', 25), ('Ġthank', 28), ('?', 32)]

(token 和它在上下文里第一次出现的位置。)大部分是系统提示词,还有 <|im_end|>(eos)。transformers 一直罚着 prompt 里的 <|im_end|>;c_64 在第 69 个 token 选了 <|im_end|> 停下,这时序列长 107,最后 64 个从位置 43 开始,prompt 里最后一个 <|im_end|> 在位置 33,已经不在惩罚范围里。窗口为 128 时,上下文要超过 128 个(生成到第 90 个左右)prompt 才开始滑出,所以分叉晚得多。

transformers 的这个处理器有一个 prompt_ignore_length 参数可以不罚 prompt,但 generate 构造它时不传,默认就是罚全部。

改法

对齐 transformers:把整个序列(prompt 加已生成)传进去,就是上面的 all,120 个 token 逐个相同。

let logits = apply_repeat_penalty(&logits, 1.3, &tokens)?;

apply_repeat_penalty 每步把 logits 读回主机、用 HashSet 去重,传整个序列的开销随长度增加;序列很长时可以自己维护一个「已出现 id」的集合,只把新 token 加进去。反过来,如果想要「只罚最近的」,那是另一种语义,transformers 里要自己写处理器,两边就不该再拿来逐 token 对照。

怎么检查

  • 搬取样参数时把 repetition_penalty 和它的范围写在一起:「1.3,作用于整个 input_ids」。抄 candle 例子的 repeat_last_n 时知道它和 transformers 不是一回事。
  • 对照测试用贪心加惩罚跑足够长(超过窗口长度),逐 token 比;只跑 20 个 token 时窗口还没滑,两种写法看不出差别。
  • 同一组 logits 上比惩罚后的结果,max|diff| 应该是 0。

小结

  • transformers 的 repetition_penalty 罚整个 input_ids(含对话模板和系统提示词),candle 的 apply_repeat_penalty 罚你传的 context;candle 例子默认只传最后 128(llama)或 64(quantized)个。
  • 公式逐位相同。范围不同时,SmolLM2 贪心 120 个 token:窗口 64 在第 58 个分叉、第 69 个就停了,窗口 128 在第 102 个分叉,传全部 120/120。
  • 对齐 transformers 就传整个序列。

💬 关联讨论

发起讨论
加载中...