📋 Allow calling trl cli in sft mode with config file (#3380)

Co-authored-by: Quentin Gallouédec <[email protected]>

Author: CloseChoice

tests/test_cli.py

@@ -13,12 +13,15 @@
 # limitations under the License.
 
 
+import os
 import sys
 import tempfile
 import unittest
 from io import StringIO
 from unittest.mock import patch
 
+import yaml
+
 
 @unittest.skipIf(
     sys.version_info < (3, 10),
@@ -67,6 +70,33 @@ def test_sft(self):
             with patch("sys.argv", command.split(" ")):
                 main()
 
+    def test_sft_config_file(self):
+        from trl.cli import main
+
+        with tempfile.TemporaryDirectory() as tmp_dir:  # Create a temporary directory
+            output_dir = os.path.join(tmp_dir, "output")
+
+            # Create a temporary config file
+            config_path = os.path.join(tmp_dir, "config.yaml")
+            config_content = {
+                "model_name_or_path": "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
+                "dataset_name": "trl-internal-testing/zen",
+                "dataset_config": "standard_language_modeling",
+                "report_to": "none",
+                "output_dir": output_dir,
+                "lr_scheduler_type": "cosine_with_restarts",
+            }
+            with open(config_path, "w") as config_file:
+                yaml.dump(config_content, config_file)
+
+            # Test the CLI with config file
+            command = f"trl sft --config {config_path}"
+            with patch("sys.argv", command.split(" ")):
+                main()
+
+            # Verify that output directory was created
+            self.assertTrue(os.path.exists(output_dir))
+
 
 if __name__ == "__main__":
     unittest.main()

trl/cli.py

@@ -46,7 +46,7 @@ def main():
     make_vllm_serve_parser(subparsers)
 
     # Parse the arguments
-    args = parser.parse_args()
+    args = parser.parse_args_and_config()[0]
 
     if args.command == "chat":
         (chat_args,) = parser.parse_args_and_config()

trl/scripts/utils.py

@@ -51,7 +51,7 @@ class ScriptArguments:
             type, inplace operation. See https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992.
     """
 
-    dataset_name: str = field(metadata={"help": "Dataset name."})
+    dataset_name: Optional[str] = field(default=None, metadata={"help": "Dataset name."})
     dataset_config: Optional[str] = field(
         default=None,
         metadata={