Refactor KTO [2/N]: Improve config validation in KTOConfig by albertvillanova · Pull Request #4787 · huggingface/trl

albertvillanova · 2026-01-08T07:19:52Z

Refactor KTO [2/N]: Improve config validation in KTOConfig.

This PR moves validation logic from KTOTrainer.__init__() to KTOConfig.__post_init__() for earlier error detection, better user experience, and cleaner separation of concerns.

Principle: Fail-fast with clear, actionable error messages

Part of:

KTO refactoring #4786

Problem

Before:

Validation scattered between config and trainer
Errors discovered late (during trainer initialization)
Invalid configs could be created and passed around
generate_during_eval validated in trainer
No validation for loss_type, truncation_mode, beta, weights
No validation for max_length relationships

After:

All validation centralized in KTOConfig.__post_init__()
Errors discovered immediately at config creation
Invalid configs cannot be created
Clear, actionable error messages with guidance
Comprehensive validation for all critical parameters

HuggingFaceDocBuilderDev · 2026-01-08T07:22:48Z

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.

albertvillanova · 2026-01-08T07:29:10Z

-        if args.generate_during_eval and not (is_wandb_available() or is_comet_available()):
-            raise ValueError(
-                "`generate_during_eval=True` requires Weights and Biases or Comet to be installed."
-                " Please install `wandb` or `comet-ml` to resolve."
-            )
-


What do you think about this change? Before it was checked at trainer instantiation and now at config creation. This has an import time side effect: it imports wandb/comet checkers immediately.

In general, I prefer to keep all these kinds of checks in one place (in MyTrainer.__init__), so that the configuration remains minimal, and we avoid unintended duplication. That said, the codebase isn’t entirely consistent on this point, see for example:

trl/trl/trainer/grpo_config.py

Lines 845 to 888 in 1a93971

if self.generation_batch_size is None and self.steps_per_generation is None:

self.steps_per_generation = self.gradient_accumulation_steps

self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation

elif self.generation_batch_size is not None and self.steps_per_generation is None:

# Just ensure the value is divisible by the global batch size

if self.generation_batch_size % (self.per_device_train_batch_size * num_processes) != 0:

raise ValueError(

f"generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size "

f"({self.per_device_train_batch_size * num_processes})."

)

self.steps_per_generation = self.generation_batch_size // (

self.per_device_train_batch_size * num_processes

)

elif self.generation_batch_size is None and self.steps_per_generation is not None:

self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation

else:

raise ValueError(

"'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time"

)

if self.do_eval and self.eval_strategy != "no":

# Determine the number of generations to use for evaluation

num_generations = self.num_generations_eval or self.num_generations

# Just ensure the value is divisible by the global batch size

if (self.per_device_eval_batch_size * num_processes) % num_generations != 0:

raise ValueError(

f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be "

f"divisible by the number of generations used for evaluation ({num_generations})."

)

# The generation batch must contain full prompt groups (no partials), so it must be divisible by

# num_generations.

if self.generation_batch_size % self.num_generations != 0:

raise ValueError(

f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations "

f"({self.num_generations})."

)

if self.num_generations < 2:

raise ValueError(

"GRPO requires at least 2 generations per prompt to calculate the advantages. You provided "

f"{self.num_generations}, which is less than the minimum required."

)

For this specific argument, I think it could probably be removed. See point 7 here:
#3906 (comment)

qgallouedec · 2026-01-08T13:34:33Z

+        # Validate beta
+        if self.beta <= 0:
+            raise ValueError(
+                f"beta must be positive, got {self.beta}. Higher β means less deviation from the reference model."
+            )
+
+        # Validate weights
+        if self.desirable_weight <= 0:
+            raise ValueError(
+                f"desirable_weight must be positive, got {self.desirable_weight}. "
+                "This weight is used to balance desirable and undesirable examples."
+            )
+
+        if self.undesirable_weight <= 0:
+            raise ValueError(
+                f"undesirable_weight must be positive, got {self.undesirable_weight}. "
+                "This weight is used to balance desirable and undesirable examples."
+            )


I’m generally in favor of keeping validation as minimal as possible. Users experimenting with parameters are expected to understand the semantics, and overly defensive validation tends to grow without bound. I’d rather keep checks limited to hard invariants (i.e. impossible combinations), like here:

trl/trl/trainer/grpo_config.py

Lines 876 to 882 in 1a93971

# The generation batch must contain full prompt groups (no partials), so it must be divisible by

# num_generations.

if self.generation_batch_size % self.num_generations != 0:

raise ValueError(

f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations "

f"({self.num_generations})."

)

This one is required by the sampling logic: if you remove the check, you can end up in a failure mode that would be very hard to debug. Beyond that, I’d rather let mistakes fail naturally.

Also, strict value checks can unnecessarily block experimentation. For example, negative (un)desirable weights aren’t mathematically invalid: they change the objective (and are almost certainly not what we want), but I’m not sure we should hard-forbid them at the config level unless they actually cause breakage.

Improve config validation in KTOConfig

a25e640

albertvillanova commented Jan 8, 2026

View reviewed changes

Update test

73645a8

albertvillanova mentioned this pull request Jan 8, 2026

KTO refactoring #4786

Open

6 tasks

qgallouedec reviewed Jan 8, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor KTO [2/N]: Improve config validation in KTOConfig#4787

Refactor KTO [2/N]: Improve config validation in KTOConfig#4787
albertvillanova wants to merge 2 commits into
huggingface:mainfrom
albertvillanova:refactor-kto-1d

albertvillanova commented Jan 8, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jan 8, 2026

Uh oh!

albertvillanova Jan 8, 2026

Uh oh!

qgallouedec Jan 8, 2026

Uh oh!

qgallouedec Jan 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	if self.generation_batch_size is None and self.steps_per_generation is None:
	self.steps_per_generation = self.gradient_accumulation_steps
	self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation
	elif self.generation_batch_size is not None and self.steps_per_generation is None:
	# Just ensure the value is divisible by the global batch size
	if self.generation_batch_size % (self.per_device_train_batch_size * num_processes) != 0:
	raise ValueError(
	f"generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size "
	f"({self.per_device_train_batch_size * num_processes})."
	)
	self.steps_per_generation = self.generation_batch_size // (
	self.per_device_train_batch_size * num_processes
	)
	elif self.generation_batch_size is None and self.steps_per_generation is not None:
	self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation
	else:
	raise ValueError(
	"'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time"
	)

	if self.do_eval and self.eval_strategy != "no":
	# Determine the number of generations to use for evaluation
	num_generations = self.num_generations_eval or self.num_generations

	# Just ensure the value is divisible by the global batch size
	if (self.per_device_eval_batch_size * num_processes) % num_generations != 0:
	raise ValueError(
	f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be "
	f"divisible by the number of generations used for evaluation ({num_generations})."
	)

	# The generation batch must contain full prompt groups (no partials), so it must be divisible by
	# num_generations.
	if self.generation_batch_size % self.num_generations != 0:
	raise ValueError(
	f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations "
	f"({self.num_generations})."
	)

	if self.num_generations < 2:
	raise ValueError(
	"GRPO requires at least 2 generations per prompt to calculate the advantages. You provided "
	f"{self.num_generations}, which is less than the minimum required."
	)

Conversation

albertvillanova commented Jan 8, 2026

Problem

Uh oh!

HuggingFaceDocBuilderDev commented Jan 8, 2026

Uh oh!

albertvillanova Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants