Using string_list for model.handlers argument.

Author: balancap

docs/float8.md

@@ -9,7 +9,7 @@ Launch training job with the following command (or alternatively set configs in
 ```
 CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh --float8.enable_float8_linear --float8.enable_fsdp_float8_all_gather --float8.precompute_float8_dynamic_scale_for_fsdp
 ```
-<!-- * `--float8.enable_float8_linear`: swap `nn.Linear` with `Float8Linear` to perform float8 matmul. -->
+* `--float8.enable_float8_linear`: swap `nn.Linear` with `Float8Linear` to perform float8 matmul.
 * `--float8.enable_fsdp_float8_all_gather`: cast `Float8Linear.weight` from high precision to float8 before FSDP all-gather so we can communicate in float8 to save bandwidth.
 * `--float8.precompute_float8_dynamic_scale_for_fsdp` (optional): communicate AMAX/scales efficiently in a single all-reduce for all parameters instead of doing many small all-reduce for each parameter.
 

tests/unit_tests/test_job_config.py

@@ -116,6 +116,15 @@ 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_converters_split(self):
+        config = JobConfig()
+        config.parse_args([])
+        assert config.model.converters == []
+
+        config = JobConfig()
+        config.parse_args(["--model.converters", "float8,mxfp"])
+        assert config.model.converters == ["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.
@@ -183,13 +196,14 @@ def __init__(self):
             help="Tokenizer path",
         )
         self.parser.add_argument(
-            "--model.handlers",
-            type=str,
-            default="",
+            "--model.converters",
+            type=string_list,
+            nargs="+",
+            default=[],
             help="""
-                Comma separated list of handlers to apply to the model.
+                Comma separated list of converters to apply to the model.
 
-                For instance, the `float8` handler swaps `torch.nn.Linear`
+                For instance, the `float8` converter swaps `torch.nn.Linear`
                 with `Float8Linear`. This feature requires you to install 'torchao'
                 which can be found here: https://github.com/pytorch/ao
             """,
@@ -541,15 +555,15 @@ def __init__(self):
         )
 
         # float8 configs
-        # self.parser.add_argument(
-        #     "--float8.enable_float8_linear",
-        #     action="store_true",
-        #     help="""
-        #         If true, swaps `torch.nn.Linear` with `Float8Linear`.
-        #         This feature requires you to install 'torchao' which can be found
-        #         here: https://github.com/pytorch/ao
-        #     """,
-        # )
+        self.parser.add_argument(
+            "--float8.enable_float8_linear",
+            action="store_true",
+            help="""
+                If true, swaps `torch.nn.Linear` with `Float8Linear`.
+                This feature requires you to install 'torchao' which can be found
+                here: https://github.com/pytorch/ao
+            """,
+        )
         self.parser.add_argument(
             "--float8.enable_fsdp_float8_all_gather",
             action="store_true",
@@ -618,18 +632,11 @@ 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/float8.py

@@ -20,7 +20,7 @@
 
 from torchtitan.config_manager import JobConfig
 from torchtitan.logging import logger
-from torchtitan.model_handler import ModelHandler, register_model_handler
+from torchtitan.model_converter import ModelConverter, register_model_converter
 from torchtitan.parallelisms import ParallelDims
 
 
@@ -29,7 +29,7 @@ def _is_sm89_or_later():
     return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 9)
 
 
-class Float8Handler(ModelHandler):
+class Float8Handler(ModelConverter):
     def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
         self.enabled = False
 
@@ -109,4 +109,4 @@ def precompute_float8_dynamic_scale_for_fsdp(
             precompute_float8_dynamic_scale_for_fsdp(m)
 
 
-register_model_handler(Float8Handler, "float8")
+register_model_converter(Float8Handler, "float8")

torchtitan/model_converter.py

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+from typing import Dict, List, Protocol, Union
+
+import torch.nn as nn
+
+from torchtitan.config_manager import JobConfig
+from torchtitan.parallelisms import ParallelDims
+
+
+class ModelConverter(Protocol):
+    """General model converter interface.
+
+    A model converter is applying a modification to PyTorch model.
+    Typical use cases are:
+        - Quantization: using QAT, FP8, ... specialized linear layers;
+        - Fused optimized layers (e.g. flash-attention, norms, ...)
+    """
+
+    def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
+        ...
+
+    def convert(self, model: nn.Module):
+        """Inplace convertion of the model."""
+        ...
+
+    def post_optimizer_hook(self, model: Union[nn.Module, List[nn.Module]]):
+        """Post-optimizer (optional) hook (e.g. compute weights statistics)."""
+        ...
+
+
+_registry_model_converter_cls: Dict[str, type[ModelConverter]] = {}
+"""Registry of model converter classes.
+"""
+
+
+def register_model_converter(converter_cls: type[ModelConverter], name: str):
+    """Register a model converter class.
+
+    A registered model converter can be applied on any model
+    using the `model.converters` config parameter.
+    """
+    assert (
+        name not in _registry_model_converter_cls
+    ), f"A model converter '{name}' is already registered."
+    _registry_model_converter_cls[name] = converter_cls
+
+
+class ModelConvertersContainer(ModelConverter):
+    """Model converters sequential container.
+
+    The class build the sequence of model converters defined in `model.converters`
+    job config, and apply them to the model sequentially.
+    """
+
+    def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
+        converter_classes = [
+            _registry_model_converter_cls[name] for name in job_config.model.converters
+        ]
+        self.converters = [
+            mh_cls(job_config, parallel_dims) for mh_cls in converter_classes
+        ]
+
+    def convert(self, model: nn.Module):
+        for mh in self.converters:
+            mh.convert(model)
+
+    def post_optimizer_hook(self, model: Union[nn.Module, List[nn.Module]]):
+        for mh in self.converters:
+            mh.post_optimizer_hook(model)
+
+
+def build_model_converters(
+    job_config: JobConfig, parallel_dims: ParallelDims
+) -> ModelConvertersContainer:
+    """Build the collection of model converters to apply to the model."""
+    return ModelConvertersContainer(job_config, parallel_dims)

torchtitan/model_handler.py

@@ -33,8 +33,9 @@
 
 from torchtitan.config_manager import JobConfig, TORCH_DTYPE_MAP
 from torchtitan.logging import logger
+from torchtitan.model_converter import parse_model_converters
 from torchtitan.parallelisms.parallel_dims import ParallelDims
-from torchtitan.model_handler import parse_model_handlers
+
 
 def parallelize_llama(
     model: nn.Module,
@@ -56,7 +57,7 @@ def parallelize_llama(
             and not job_config.training.compile
         ):
             raise RuntimeError("Async TP requires --training.compile")
-        enable_float8 = "float8" in parse_model_handlers(job_config)
+        enable_float8 = "float8" in parse_model_converters(job_config)
         apply_tp(
             model,
             world_mesh["tp"],

torchtitan/parallelisms/parallelize_llama.py

@@ -18,7 +18,7 @@
 from torchtitan.datasets import build_hf_data_loader, build_tokenizer
 from torchtitan.logging import init_logger, logger
 from torchtitan.metrics import build_device_memory_monitor, build_metric_logger
-from torchtitan.model_handler import build_model_handlers_container
+from torchtitan.model_converter import build_model_converters
 from torchtitan.models import model_name_to_cls, model_name_to_tokenizer, models_config
 from torchtitan.optimizer import build_lr_schedulers, build_optimizers
 from torchtitan.parallelisms import (
@@ -110,9 +110,9 @@ def main(job_config: JobConfig):
     with torch.device("meta"):
         model = model_cls.from_model_args(model_config)
 
-    # Build the collection of model handlers. No-op if `model.handlers` empty
-    model_handlers = build_model_handlers_container(job_config, parallel_dims)
-    model_handlers.convert(model)
+    # Build the collection of model converters. No-op if `model.converters` empty
+    model_converters = build_model_converters(job_config, parallel_dims)
+    model_converters.convert(model)
 
     # log model size
     model_param_count = utils.get_num_params(model)
@@ -325,10 +325,10 @@ def loss_fn(pred, labels):
             optimizers.step()
             lr_schedulers.step()
 
-            # Post-optimizer model handlers hook.
+            # Post-optimizer model converters hook.
             # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
             # it issues a single all-reduce for all parameters at once for better performance
-            model_handlers.post_optimizer_hook(model_parts)
+            model_converters.post_optimizer_hook(model_parts)
 
             # log metrics
             if (