Skip to content

Refactor KTO coordinated with DPO [b/N]: Simplify truncation logic#4808

Merged
albertvillanova merged 3 commits intohuggingface:mainfrom
albertvillanova:refactor-kto-coordinated-with-dpo-b
Jan 13, 2026
Merged

Refactor KTO coordinated with DPO [b/N]: Simplify truncation logic#4808
albertvillanova merged 3 commits intohuggingface:mainfrom
albertvillanova:refactor-kto-coordinated-with-dpo-b

Conversation

@albertvillanova
Copy link
Member

@albertvillanova albertvillanova commented Jan 12, 2026

Refactor KTO coordinated with DPO [b/N]: Simplify truncation logic.

This PR simplifies the truncation logic by removing max_prompt_length and truncation_mode parameters, keeping only max_length.

This simplification makes the KTO trainer significantly easier to use and understand, while maintaining correct truncation behavior.

Follow-up to:

Part of:

Coordinated with DPO refactoring, as discussed with @qgallouedec :

Key Changes

  1. KTOConfig

    • Removed max_prompt_length parameter and documentation
    • Removed truncation_mode parameter and documentation
    • Now only uses max_length for sequence length control
  2. KTOTrainer

    Simplified _process_tokens() function:

    • Removed complex two-stage truncation logic
    • Old logic: First truncate prompt (with keep_start/keep_end modes), then truncate completion
    • New logic: Simple single-stage truncation from the end (completion only)
    • Updated docstring to reflect simpler behavior

    Trainer Initialization:

    • Removed max_prompt_length setup logic
    • Removed self.max_prompt_length attribute
    • Removed self.truncation_mode attribute
    • Removed both parameters from fn_kwargs passed to _process_tokens()
  3. Tests

    • Updated test to remove truncation_mode and max_prompt_length from fn_kwargs

@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Copy link
Member

@qgallouedec qgallouedec left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

FYI, answer, response, completion all refer to the same thing. We used to mix these terms, but now we try have a unified use of completion only

@albertvillanova albertvillanova mentioned this pull request Jan 12, 2026
6 tasks
@albertvillanova albertvillanova merged commit 6402e70 into huggingface:main Jan 13, 2026
2 of 3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

Comments