fix: add repetition penalty to mitigate multi-turn repetition (fixes #1125)

[modimihir07]

Problem

Fixes #1125

The sample() function in generate.py only supports temperature scaling. There is no mechanism to penalize tokens that have already been generated, which causes a self-reinforcing repetition loop in multi-turn dialogues:

1. Model generates a pattern (e.g., "Based on the reasoning above...")
2. Pattern enters conversation history
3. Model sees the pattern in context and is more likely to repeat it
4. Loop escalates with each turn

Fix

Added repetition_penalty parameter using the approach from the CTRL paper (Keskar et al., 2019) — the same method used by HuggingFace Transformers and vLLM:

• For each previously generated token, scale down its logit:
  • If logit > 0: divide by penalty factor
  • If logit < 0: multiply by penalty factor

Changes:

• sample() — new repetition_penalty and generated_tokens parameters
• generate() — passes token history to sample() for penalty
• main() — threads repetition_penalty through
• CLI — new --repetition-penalty flag (default: 1.0 = backward compatible)

Usage

# Default (no change from before)
torchrun ... generate.py --ckpt-path ... --config ... --interactive

# With repetition penalty (recommended 1.1–1.3 for multi-turn)
torchrun ... generate.py --ckpt-path ... --config ... --interactive --repetition-penalty 1.2

[ai-nurmamat]

Great fix for the repetition issue! Consider adding a configurable repetition penalty range in the inference config.

[modimihir07]

Thanks for the approval and the great suggestion, @Financier-Nuri! I've added --repetition-penalty-min (default: 1.0) and --repetition-penalty-max (default: 2.0) CLI arguments. The penalty value is now validated against these bounds at startup, raising a clear ValueError if it's out of range.

…

On Wed, 11 Mar 2026 at 05:40, Nurmamat.Omar ***@***.***> wrote: ***@***.**** approved this pull request. Great fix for the repetition issue! Consider adding a configurable repetition penalty range in the inference config. — Reply to this email directly, view it on GitHub <#1129 (review)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/BK54AHPY6GPNT56TBFKIY6L4QCVINAVCNFSM6AAAAACWDOJ6DOVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZTSMRWGE2TIOBZGY> . You are receiving this because you authored the thread.Message ID: ***@***.***>