Rust in Everything
基于 candle 0.11.0

同样的取样参数，分布不一样

主线位置：推理主线的「带 KV 缓存生成」一步里，从 logits 到下一个 token：把 generate 的取样参数原样搬到 LogitsProcessor。

版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate）
flowchart TB
  hf["generate 参数"] --> d1["do_sample 温度"]
  hf --> d2["top_k 加 top_p"]
  d1 -->|"隐含 top_k 50"| c1["TopK k 50"]
  d1 -->|"照字面"| w1["All 多出候选"]
  d2 -->|"照字面"| w2["TopKThenTopP"]
  d2 -->|"先 top_k 再归一化"| c2["sample_f 闭包加 TopP"]

实验代码：Python 侧 sampling_ref.py、edge_py.py，candle 侧 crate pitb-exp 的 bin p28_chat（生成上下文）、p26_samp、p26_fix、p26_hfsem、p26_edge（release）。模型是本地的 SmolLM2-135M-Instruct 权重，f32。平台 M2 Max、macOS 27、rustc 1.99.0，uptime 负载 10 左右，只影响耗时。比较方法（两边拿同一组 logits，比每个 token 的概率，不比抽出来的 token）和公式的逐个核对，见取样。

PyTorch 的写法

transformers 里开取样通常只写两三个参数：

SETTINGS = {
  "t07_default": dict(do_sample=True, temperature=0.7),
  "t07": dict(do_sample=True, temperature=0.7, top_k=0),
  "k40": dict(do_sample=True, temperature=0.7, top_k=40),
  "p09": dict(do_sample=True, temperature=0.7, top_k=0, top_p=0.9),
  "k40p09": dict(do_sample=True, temperature=0.7, top_k=40, top_p=0.9),
}
g = model.generate(ids, attention_mask=torch.ones_like(ids), max_new_tokens=1,
                   output_scores=True, output_logits=True, return_dict_in_generate=True, **kw)

上下文是对话模板包着 How do you say thank you in Chinese?（39 个 token），加上回答的前 7 个 token In Chinese, you can say "，下一个 token 的分布比较平。scores 里有限值的个数就是留下的候选数：

t07_default  kept=50
t07          kept=49152
k40          kept=40
p09          kept=31
k40p09       kept=14

只写了 do_sample=True, temperature=0.7 的那一行留下 50 个，不是整个词表。

搬到 candle 的现象

照字面搬：temperature=0.7 对 Sampling::All { temperature: 0.7 }，top_k=40, top_p=0.9 对 Sampling::TopKThenTopP { k: 40, p: 0.9, temperature: 0.7 }。p26_samp 把 transformers 的 logits 读进来，算出 candle 每种取法下的分布（照抄 LogitsProcessor 的截断逻辑，再用真的 LogitsProcessor 抽 10 万次核对），和 transformers 的 scores 做 softmax 后比：

t07      kept hf=49152 candle=49152  max|p_hf-p_candle|=6.98e-9  TV=2.14e-8
    LogitsProcessor x100000: max|freq-p_candle|=0.0017  freq outside candle set=0.0000  outside hf set=0.0000
t07_def  kept hf=   50 candle=49152  max|p_hf-p_candle|=1.34e-2  TV=6.06e-2
k40      kept hf=   40 candle=   40  max|p_hf-p_candle|=1.49e-8  TV=5.32e-8
    LogitsProcessor x100000: max|freq-p_candle|=0.0016  freq outside candle set=0.0000  outside hf set=0.0000
p09      kept hf=   31 candle=   31  max|p_hf-p_candle|=1.49e-8  TV=2.68e-8
    LogitsProcessor x100000: max|freq-p_candle|=0.0016  freq outside candle set=0.0000  outside hf set=0.0000
    top-40 original mass = 0.9208, p = 0.9
k40p09   kept hf=   14 candle=   31  max|p_hf-p_candle|=1.85e-2  TV=7.45e-2
    LogitsProcessor x100000: max|freq-p_candle|=0.0016  freq outside candle set=0.0000  outside hf set=0.0748

TV 是总变差距离 ½Σ|p_hf − p_candle|。单独的温度、top-k、top-p 两边一致（1e-8 量级，只差 softmax 的实现）；不一致的正好是照字面搬的两行：

  • t07_def：candle 留整个词表，transformers 留 50 个，TV 0.06。
  • k40p09：candle 留 31 个，transformers 留 14 个，10 万次里有 7.48% 落在 transformers 删掉的 token 上。

没有报错，生成的文字也通顺，只是「更随机」一些。第三处差别在并列的 logits 上（edge_py.py、p26_edge）：

top_k=2 on [3,2,2,1]: [True, True, True, False]

TopK k=2 on [3,2,2,1], 10000 draws: [7349, 2651, 0, 0]

transformers 的 top-2 留下 3 个，candle 正好 2 个。

原因

隐含的 top_k。 transformers 5.22.0 的 GenerationConfig._get_default_generation_params() 给没设的字段填默认值，其中 "top_k": 50；_get_logits_processor 只要 top_k 不是 None 也不是 0 就加一个 TopKLogitsWarper。所以 do_sample=True, temperature=0.7 的意思是「温度 0.7 加 top-50」。candle 的 Sampling::All 就是字面上的「整个词表」。

top-k 加 top-p。 transformers 的处理器依次作用在 scores 上：TopKLogitsWarper 把第 40 名之后填成 -inf，TopPLogitsWarper 接着在这个结果上做 softmax、累加，p = 0.9 量的是前 40 个里的 90%。candle 的 sample_topk_topp 取出前 k 个后不归一化：

            let mut prs = indices.iter().map(|&i| prs[i]).collect::<Vec<_>>();
            let sum_p = prs.iter().sum::<f32>();
            let index = if top_p <= 0.0 || top_p >= sum_p {
                self.sample_multinomial(&prs)?
            } else {
                self.sample_topp(&mut prs, top_p)?
            };

p 量的是原分布的质量。这一步前 40 个原始质量 0.9208，刚过 0.9，candle 的截断点就是「整个词表的 90%」，和单独 TopP { p: 0.9 } 一样留 31 个；前 40 个原始质量不到 0.9 时，candle 走 top_p >= sum_p 那一支，p 完全不起作用。

并列。 transformers 的 TopKLogitsWarper 删的是「严格小于第 k 名」的，和第 k 名相等的都留下；candle 的 sample_topk 用 select_nth_unstable_by 正好取 k 个位置，并列的哪个进来由排序算法决定。

三处都不是 bug，是两边对同一组参数名的定义不同。origin/main 上 generation/mod.rs 只给温度那一行加了一条 clippy 注释，取样逻辑不变。

改法

do_sample=True 且没写 top_k 时，搬成 TopK { k: 50, .. }（p26_fix，各抽 10 万次，和 transformers 的分布比）：

for (name, s) in [
  ("All { temperature: 0.7 }", Sampling::All { temperature: 0.7 }),
  ("TopK { k: 50, temperature: 0.7 }", Sampling::TopK { k: 50, temperature: 0.7 }),
] {
  let mut lp = LogitsProcessor::from_sampling(42, s);
  // ... 抽 draws 次，统计频率
}

All { temperature: 0.7 }           x100000: tokens drawn= 596 max|freq-p_hf|=0.0123 outside hf top-50=0.0614
TopK { k: 50, temperature: 0.7 }   x100000: tokens drawn=  50 max|freq-p_hf|=0.0012 outside hf top-50=0.0000

All 抽到了 596 个不同的 token，6.14% 落在 transformers 的前 50 之外；TopK { k: 50 } 只抽到那 50 个，频率和 transformers 的概率最多差 0.0012，在 10 万次抽样的统计误差之内。

top-k 加 top-p 要 transformers 的语义：用 TopP，在 sample_f 的闭包里先做 top-k 并重新归一化（闭包拿到的是 softmax 之后、截断之前的概率）：

fn hf_topk_renorm(prs: &mut [f32], k: usize) {
  let mut idx: Vec<usize> = (0..prs.len()).collect();
  idx.select_nth_unstable_by(k, |&i, &j| prs[j].total_cmp(&prs[i]));
  for &i in &idx[k..] {
    prs[i] = 0.0;
  }
  let s: f32 = prs.iter().sum();
  prs.iter_mut().for_each(|p| *p /= s);
}

let mut lp = LogitsProcessor::from_sampling(42, Sampling::TopP { p: 0.9, temperature: 0.7 });
let next = lp.sample_f(&logits, |prs| hf_topk_renorm(prs, 40))?;

TopP + top-k renorm in sample_f x100000: tokens drawn=14 max|freq-p_hf|=0.0015 outside hf set=0.0000

只抽到那 14 个，和 transformers 最多差 0.0015。这个闭包没有处理并列；真实的 f32 logits 上并列很少见，量化模型、被 -inf 屏蔽过的 logits 上会有，那时要在闭包里按「不小于第 k 名」保留。

怎么检查

  • 把 transformers 那边真正生效的参数打出来，不要看自己写了什么：generate(..., max_new_tokens=1, output_scores=True, return_dict_in_generate=True)，数 scores 里有几个有限值。
  • 模型目录的 generation_config.json 也可能写了 top_k、top_p、temperature，transformers 会用它，candle 不读。SmolLM2 的这个文件里没有取样参数。
  • 同一组 logits 交给两边，比分布（TV 或逐 token 的最大差），而不是比抽出来的 token：同一个种子两边的随机数流不同（取样的最后一节）。

小结

  • generate(do_sample=True, temperature=0.7) 隐含 top_k=50，对应 Sampling::TopK { k: 50, temperature: 0.7 }；搬成 All 时 TV 0.06，6% 的抽样落在 transformers 永远不选的 token 上。
  • top_k 加 top_p：transformers 在前 k 个里重新归一化再截 p，candle 的 TopKThenTopP 用原分布的质量，这一步留 31 个对 14 个。用 TopP 加 sample_f 闭包里的 top-k 归一化就对上了。
  • 并列时 transformers 留下所有等于第 k 名的，candle 正好 k 个。

💬 关联讨论

发起讨论
加载中...