Add num_completions_to_print for trl and grpo (#2604)

2025-04-30 19:00:30 -06:00
parent 6ba5c0ed2c
commit c3f2b1c5c2
2 changed files with 7 additions and 0 deletions
--- a/src/axolotl/utils/schemas/trl.py
+++ b/src/axolotl/utils/schemas/trl.py
@@ -67,6 +67,12 @@ class TRLConfig(BaseModel):
        default=False,
        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(
        default=False,
        json_schema_extra={