tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Add num_completions_to_print for trl and grpo (#2604)

Browse Source
This commit is contained in:
Dhruv Mullick
2025-04-30 19:00:30 -06:00
committed by GitHub
parent 6ba5c0ed2c
commit c3f2b1c5c2
2 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/axolotl/core/trainers/grpo/__init__.py
View File

@@ -63,6 +63,7 @@ class GRPOStrategy:
grpo_args_kwargs["max_completion_length"] = trl.max_completion_length grpo_args_kwargs["max_completion_length"] = trl.max_completion_length
grpo_args_kwargs["log_completions"] = trl.log_completions grpo_args_kwargs["log_completions"] = trl.log_completions
grpo_args_kwargs["num_completions_to_print"] = trl.num_completions_to_print
if trl.reward_weights: if trl.reward_weights:
grpo_args_kwargs["reward_weights"] = trl.reward_weights grpo_args_kwargs["reward_weights"] = trl.reward_weights

6
src/axolotl/utils/schemas/trl.py
View File

@@ -67,6 +67,12 @@ class TRLConfig(BaseModel):
default=False, default=False,
json_schema_extra={"description": "Whether to log completions"}, json_schema_extra={"description": "Whether to log completions"},
) )
num_completions_to_print: int | None = Field(
default=None,
json_schema_extra={
"description": "Number of completions to print. If `log_completions` is `True`, this will be the number of completions logged."
},
)
sync_ref_model: bool | None = Field( sync_ref_model: bool | None = Field(
default=False, default=False,
json_schema_extra={ json_schema_extra={