Skip to content

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

Merged
albertvillanova merged 3 commits into
huggingface: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 into
huggingface:mainfrom
albertvillanova:refactor-kto-coordinated-with-dpo-b

Conversation

@albertvillanova

@albertvillanova albertvillanova commented Jan 12, 2026

Copy link
Copy Markdown
Member

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

Copy link
Copy Markdown

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.

@qgallouedec qgallouedec left a comment

Copy link
Copy Markdown
Member

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
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