Allow CLI overrides

Author: jackzhxng

extension/llm/export/export_llm.py

@@ -38,14 +38,20 @@
 from executorch.examples.models.llama.config.llm_config import LlmConfig
 from executorch.examples.models.llama.export_llama_lib import export_llama
 from hydra.core.config_store import ConfigStore
+from hydra.core.hydra_config import HydraConfig
+from hydra.core.override_parser.overrides_parser import OverridesParser
 from omegaconf import OmegaConf
 
 cs = ConfigStore.instance()
 cs.store(name="llm_config", node=LlmConfig)
 
 
+# Need this global variable to pass an llm_config from yaml
+# into the hydra-wrapped main function.
+llm_config_from_yaml = None
+
+
 def parse_config_arg() -> Tuple[str, List[Any]]:
-    """First parse out the arg for whether to use Hydra or the old CLI."""
     parser = argparse.ArgumentParser(add_help=True)
     parser.add_argument("--config", type=str, help="Path to the LlmConfig file")
     args, remaining = parser.parse_known_args()
@@ -65,28 +71,34 @@ def pop_config_arg() -> str:
 
 @hydra.main(version_base=None, config_name="llm_config")
 def hydra_main(llm_config: LlmConfig) -> None:
-    export_llama(OmegaConf.to_object(llm_config))
+    global llm_config_from_yaml
+
+    # Override the LlmConfig constructed from the provide yaml config file
+    # with the CLI overrides.
+    if llm_config_from_yaml:
+        # Get CLI overrides (excluding defaults list).
+        overrides_list: List[str] = list(HydraConfig.get().overrides.get("task", []))
+        override_cfg = OmegaConf.from_dotlist(overrides_list)
+        merged_config = OmegaConf.merge(llm_config_from_yaml, override_cfg)
+        export_llama(merged_config)
+    else:
+        export_llama(OmegaConf.to_object(llm_config))
 
 
 def main() -> None:
+    # First parse out the arg for whether to use Hydra or the old CLI.
     config, remaining_args = parse_config_arg()
     if config:
-        # Check if there are any remaining hydra CLI args when --config is specified
-        # This might change in the future to allow overriding config file values
-        if remaining_args:
-            raise ValueError(
-                "Cannot specify additional CLI arguments when using --config. "
-                f"Found: {remaining_args}. Use either --config file or hydra CLI args, not both."
-            )
-
+        global llm_config_from_yaml
+        # Pop out --config and its value so that they are not parsed by
+        # Hyra's main.
         config_file_path = pop_config_arg()
         default_llm_config = LlmConfig()
-        llm_config_from_file = OmegaConf.load(config_file_path)
-        # Override defaults with values specified in the .yaml provided by --config.
-        merged_llm_config = OmegaConf.merge(default_llm_config, llm_config_from_file)
-        export_llama(merged_llm_config)
-    else:
-        hydra_main()
+        # Construct the LlmConfig from the config yaml file.
+        default_llm_config = LlmConfig()
+        from_yaml = OmegaConf.load(config_file_path)
+        llm_config_from_yaml = OmegaConf.merge(default_llm_config, from_yaml)
+    hydra_main()
 
 
 if __name__ == "__main__":

extension/llm/export/test/test_export_llm.py

@@ -21,27 +21,37 @@ class TestExportLlm(unittest.TestCase):
     def test_parse_config_arg_with_config(self) -> None:
         """Test parse_config_arg when --config is provided."""
         # Mock sys.argv to include --config
-        test_argv = ["script.py", "--config", "test_config.yaml", "extra", "args"]
+        test_argv = ["export_llm.py", "--config", "test_config.yaml", "extra", "args"]
         with patch.object(sys, "argv", test_argv):
             config_path, remaining = parse_config_arg()
             self.assertEqual(config_path, "test_config.yaml")
             self.assertEqual(remaining, ["extra", "args"])
 
     def test_parse_config_arg_without_config(self) -> None:
         """Test parse_config_arg when --config is not provided."""
-        test_argv = ["script.py", "debug.verbose=True"]
+        test_argv = ["export_llm.py", "debug.verbose=True"]
         with patch.object(sys, "argv", test_argv):
             config_path, remaining = parse_config_arg()
             self.assertIsNone(config_path)
             self.assertEqual(remaining, ["debug.verbose=True"])
 
     def test_pop_config_arg(self) -> None:
         """Test pop_config_arg removes --config and its value from sys.argv."""
-        test_argv = ["script.py", "--config", "test_config.yaml", "other", "args"]
+        test_argv = ["export_llm.py", "--config", "test_config.yaml", "other", "args"]
         with patch.object(sys, "argv", test_argv):
             config_path = pop_config_arg()
             self.assertEqual(config_path, "test_config.yaml")
-            self.assertEqual(sys.argv, ["script.py", "other", "args"])
+            self.assertEqual(sys.argv, ["export_llm.py", "other", "args"])
+
+    def test_with_cli_args(self) -> None:
+        """Test main function with only hydra CLI args."""
+        test_argv = ["export_llm.py", "debug.verbose=True"]
+        with patch.object(sys, "argv", test_argv):
+            with patch(
+                "executorch.extension.llm.export.export_llm.hydra_main"
+            ) as mock_hydra:
+                main()
+                mock_hydra.assert_called_once()
 
     @patch("executorch.extension.llm.export.export_llm.export_llama")
     def test_with_config(self, mock_export_llama: MagicMock) -> None:
@@ -70,7 +80,7 @@ def test_with_config(self, mock_export_llama: MagicMock) -> None:
             config_file = f.name
 
         try:
-            test_argv = ["script.py", "--config", config_file]
+            test_argv = ["export_llm.py", "--config", config_file]
             with patch.object(sys, "argv", test_argv):
                 main()
 
@@ -99,54 +109,48 @@ def test_with_config(self, mock_export_llama: MagicMock) -> None:
         finally:
             os.unlink(config_file)
 
-    def test_with_cli_args(self) -> None:
-        """Test main function with only hydra CLI args."""
-        test_argv = ["script.py", "debug.verbose=True"]
-        with patch.object(sys, "argv", test_argv):
-            with patch(
-                "executorch.extension.llm.export.export_llm.hydra_main"
-            ) as mock_hydra:
-                main()
-                mock_hydra.assert_called_once()
-
-    def test_config_with_cli_args_error(self) -> None:
-        """Test that --config rejects additional CLI arguments to prevent mixing approaches."""
+    @patch("executorch.extension.llm.export.export_llm.export_llama")
+    def test_with_config(self, mock_export_llama: MagicMock) -> None:
+        """Test main function with --config file and no hydra args."""
         # Create a temporary config file
         with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-            f.write("base:\n  checkpoint: /path/to/checkpoint.pth")
-            config_file = f.name
-
-        try:
-            test_argv = ["script.py", "--config", config_file, "debug.verbose=True"]
-            with patch.object(sys, "argv", test_argv):
-                with self.assertRaises(ValueError) as cm:
-                    main()
-
-                error_msg = str(cm.exception)
-                self.assertIn(
-                    "Cannot specify additional CLI arguments when using --config",
-                    error_msg,
-                )
-        finally:
-            os.unlink(config_file)
-
-    def test_config_rejects_multiple_cli_args(self) -> None:
-        """Test that --config rejects multiple CLI arguments (not just single ones)."""
-        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-            f.write("export:\n  max_seq_length: 128")
+            f.write(
+                """
+base:
+  model_class: llama2
+model:
+  dtype_override: fp16
+backend:
+  xnnpack:
+    enabled: False
+"""
+            )
             config_file = f.name
 
         try:
             test_argv = [
-                "script.py",
+                "export_llm.py",
                 "--config",
                 config_file,
-                "debug.verbose=True",
-                "export.output_dir=/tmp",
+                "base.model_class=stories110m",
+                "backend.xnnpack.enabled=True",
             ]
             with patch.object(sys, "argv", test_argv):
-                with self.assertRaises(ValueError):
-                    main()
+                main()
+
+            # Verify export_llama was called with config
+            mock_export_llama.assert_called_once()
+            called_config = mock_export_llama.call_args[0][0]
+            self.assertEqual(
+                called_config["base"]["model_class"], "stories110m"
+            )  # Override from CLI.
+            self.assertEqual(
+                called_config["model"]["dtype_override"].value, "fp16"
+            )  # From yaml.
+            self.assertEqual(
+                called_config["backend"]["xnnpack"]["enabled"],
+                True,  # Override from CLI.
+            )
         finally:
             os.unlink(config_file)