Using string_list for model.handlers argument.

Author: balancap

tests/unit_tests/test_job_config.py

@@ -116,6 +116,11 @@ def test_parse_pp_split_points(self):
                 config.experimental.pipeline_parallel_split_points == cmdline_splits
             ), config.experimental.pipeline_parallel_split_points
 
+    def test_job_config_model_handlers_split(self):
+        config = JobConfig()
+        config.parse_args(["--model.handlers", "float8,mxfp"])
+        assert config.model.handlers == ["float8", "mxfp"]
+
     def test_print_help(self):
         config = JobConfig()
         parser = config.parser

torchtitan/config_manager.py

@@ -26,9 +26,22 @@
 
 
 def string_list(raw_arg):
+    """Comma-separated string list argument."""
     return raw_arg.split(",")
 
 
+def check_string_list_argument(args_dict: dict[str, any], fullargname: str):
+    section, name = fullargname.split(".")
+    # Split string list which are still raw strings.
+    if (
+        section in args_dict
+        and name in args_dict[section]
+        and isinstance(args_dict[section][name], str)
+    ):
+        sec = args_dict[section]
+        sec[name] = string_list(sec[name])
+
+
 class JobConfig:
     """
     A helper class to manage the train configuration.
@@ -184,8 +197,9 @@ def __init__(self):
         )
         self.parser.add_argument(
             "--model.handlers",
-            type=str,
-            default="",
+            type=string_list,
+            nargs="+",
+            default=[],
             help="""
                 Comma separated list of handlers to apply to the model.
 
@@ -617,19 +631,12 @@ def parse_args(self, args_list: list = sys.argv[1:]):
                 )
                 logger.exception(f"Error details: {str(e)}")
                 raise e
-
+        
+        # Checking string-list arguments are properly split into a list
         # if split-points came from 'args' (from cmd line) it would have already been parsed into a list by that parser
-        if (
-            "experimental" in args_dict
-            and "pipeline_parallel_split_points" in args_dict["experimental"]
-            and isinstance(
-                args_dict["experimental"]["pipeline_parallel_split_points"], str
-            )
-        ):
-            exp = args_dict["experimental"]
-            exp["pipeline_parallel_split_points"] = string_list(
-                exp["pipeline_parallel_split_points"]
-            )
+        string_list_argnames = self._get_string_list_argument_names()
+        for n in string_list_argnames:
+            check_string_list_argument(args_dict, n)
 
         # override args dict with cmd_args
         cmd_args_dict = self._args_to_two_level_dict(cmd_args)
@@ -657,13 +664,21 @@ def _validate_config(self) -> None:
         assert self.model.flavor
         assert self.model.tokenizer_path
 
+    def _get_string_list_argument_names(self) -> list[str]:
+        """Get the parser argument names of type `string_list`."""
+        string_list_args = [
+            v.dest for v in self.parser._actions if v.type is string_list
+        ]
+        return string_list_args
+
     def parse_args_from_command_line(
         self, args_list
     ) -> Tuple[argparse.Namespace, argparse.Namespace]:
         """
         Parse command line arguments and return the parsed args and the command line only args
         """
         args = self.parser.parse_args(args_list)
+        string_list_argnames = set(self._get_string_list_argument_names())
 
         # aux parser to parse the command line only args, with no defaults from main parser
         aux_parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
@@ -672,7 +687,7 @@ def parse_args_from_command_line(
                 aux_parser.add_argument(
                     "--" + arg, action="store_true" if val else "store_false"
                 )
-            elif arg == "experimental.pipeline_parallel_split_points":
+            elif arg in string_list_argnames:
                 # without this special case, type inference breaks here,
                 # since the inferred type is just 'list' and it ends up flattening
                 # e.g. from ["layers.0", "layers.1"] into ["l", "a", "y", "e", "r", "s", ".0", ...]

torchtitan/parallelisms/parallelize_llama.py

@@ -33,8 +33,9 @@
 
 from torchtitan.config_manager import JobConfig, TORCH_DTYPE_MAP
 from torchtitan.logging import logger
-from torchtitan.parallelisms.parallel_dims import ParallelDims
 from torchtitan.model_handler import parse_model_handlers
+from torchtitan.parallelisms.parallel_dims import ParallelDims
+
 
 def parallelize_llama(
     model: nn.Module,