support log completions

axolotl-ai-cloud · Feb 7, 2025 · 9b05aac · 9b05aac
1 parent 8957a9a
commit 9b05aac
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/axolotl/core/trainers/grpo/__init__.py b/src/axolotl/core/trainers/grpo/__init__.py
@@ -54,6 +54,7 @@ def set_training_args_kwargs(cls, cfg):
             if cfg.trl and cfg.trl.ref_model_sync_steps:
                 grpo_args_kwargs["ref_model_sync_steps"] = cfg.trl.ref_model_sync_steps
         grpo_args_kwargs["max_completion_length"] = cfg.trl.max_completion_length
+        grpo_args_kwargs["log_completions"] = cfg.trl.log_completions
         return grpo_args_kwargs
 
     @classmethod

diff --git a/src/axolotl/utils/config/models/input/v0_4_1/trl.py b/src/axolotl/utils/config/models/input/v0_4_1/trl.py
@@ -25,8 +25,11 @@ class TrlConfig(BaseModel):
     vllm_gpu_memory_utilization: Optional[float] = 0.9
     vllm_max_model_len: Optional[int] = None
     vllm_dtype: Optional[str] = "auto"
+
     reward_funcs: Optional[List[str]] = None
     num_generations: Optional[int] = None
+    log_completions: Optional[bool] = False
+
     sync_ref_model: Optional[bool] = False
     ref_model_mixup_alpha: Optional[float] = 0.9
     ref_model_sync_steps: Optional[int] = 64