use pipeline registry

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/args/dataset_arguments.py

@@ -3,8 +3,6 @@
 
 from transformers import DefaultDataCollator
 
-from llmcompressor.pipelines.registry import PIPELINES
-
 
 @dataclass
 class DVCDatasetArguments:
@@ -176,7 +174,8 @@ class DatasetArguments(CustomDatasetArguments):
     pipeline: Optional[str] = field(
         default="independent",
         metadata={
-            "help": "Calibration pipeline used to calibrate model. "
-            f"Options: {PIPELINES.keys()}"
+            "help": "Calibration pipeline used to calibrate model"
+            "Options: ['basic', 'datafree', 'sequential', 'layer_sequential', "
+            "independent]"
         },
     )

src/llmcompressor/entrypoints/oneshot.py

@@ -7,7 +7,7 @@
 from llmcompressor.core.session_functions import active_session
 from llmcompressor.datasets import get_calibration_dataloader
 from llmcompressor.entrypoints.utils import post_process, pre_process
-from llmcompressor.pipelines.registry import get_pipeline_fn
+from llmcompressor.pipelines.registry import CalibrationPipeline
 
 __all__ = ["Oneshot", "oneshot"]
 
@@ -168,9 +168,10 @@ def apply_recipe_modifiers(
             recipe_args=self.recipe_args.recipe_args,
         )
 
+        user_pipeline = self.dataset_args.pipeline
         modifiers = session.get_modifiers()
-        _, pipeline_fn = get_pipeline_fn(self.dataset_args.pipeline, modifiers)
-        pipeline_fn(self.model, calibration_dataloader, self.dataset_args)
+        pipeline = CalibrationPipeline.from_modifiers(modifiers, user=user_pipeline)
+        pipeline(self.model, calibration_dataloader, self.dataset_args)
 
         session.finalize()
 

src/llmcompressor/modifiers/obcq/sgpt_mixin.py

@@ -12,7 +12,6 @@
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers.modifier import Modifier
 from llmcompressor.modifiers.utils.hooks import HooksMixin
-from llmcompressor.pipelines.basic import run_pipeline as run_basic
 from llmcompressor.utils.pytorch.module import (
     get_layers,
     get_no_split_params,
@@ -247,7 +246,7 @@ def _infer_owl_layer_sparsity(
         return sparsities
 
     def _get_activations(self, model, dataloader, nsamples=128) -> Dict[str, int]:
-        from llmcompressor.args import DatasetArguments
+        from llmcompressor.pipelines.basic import run_calibration
 
         acts = defaultdict(int)
 
@@ -263,7 +262,7 @@ def save_acts(_module, input: Union[Tuple[Any, ...], torch.Tensor], name: str):
             if isinstance(mod, torch.nn.Linear) and "lm_head" not in name
         )
         with HooksMixin.disable_hooks(keep=hooks):
-            run_basic(model, dataloader, DatasetArguments())
+            run_calibration(model, dataloader)
         self.remove_hooks(hooks)
 
         return acts

src/llmcompressor/pipelines/__init__.py

@@ -0,0 +1,8 @@
+# flake8: noqa
+# populate registry
+from .basic import *
+from .data_free import *
+from .independent import *
+from .layer_sequential import *
+from .registry import *
+from .sequential import *

src/llmcompressor/pipelines/basic/__init__.py

@@ -1,2 +1,2 @@
 # flake8: noqa
-from .pipeline import run_pipeline
+from .pipeline import *

src/llmcompressor/pipelines/basic/pipeline.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 import torch
 import tqdm
@@ -7,40 +7,49 @@
 
 from llmcompressor.core import LifecycleCallbacks
 from llmcompressor.modifiers.utils.pytorch_helpers import apply_pad_mask_to_batch
+from llmcompressor.pipelines.registry import CalibrationPipeline
 from llmcompressor.pytorch.utils.helpers import tensors_to_device
 from llmcompressor.utils.helpers import calibration_forward_context
 
 if TYPE_CHECKING:
     from llmcompressor.args.dataset_arguments import DatasetArguments
 
-__all__ = ["run_pipeline"]
+__all__ = ["BasicPipeline", "run_calibration"]
 
 
-def run_pipeline(
-    model: torch.nn.Module,
-    dataloader: DataLoader,
-    dataset_args: "DatasetArguments",
-):
-    """
-    Run a basic data pipeline.
+@CalibrationPipeline.register("basic")
+class BasicPipeline(CalibrationPipeline):
+    @staticmethod
+    def __call__(
+        model: torch.nn.Module,
+        dataloader: DataLoader,
+        dataset_args: Union["DatasetArguments", None],
+    ):
+        """
+        Run a basic data pipeline.
 
-    Batches are fetched from the data loader and are used to perform forward passes
-    through the model. This pipeline is typically used for basic model calibration
-    and, unlike the sequential pipelines, does not propagate compression error when
-    used to calibrate model compression
+        Batches are fetched from the data loader and are used to perform forward passes
+        through the model. This pipeline is typically used for basic model calibration
+        and, unlike the sequential pipelines, does not propagate compression error when
+        used to calibrate model compression
 
-    :param model: model being calibrated
-    :param dataloader: loads data for calibration
-    :param dataset_args: dataset arguments relevant to pipelines
-    """
-    model_device = get_execution_device(model)
+        :param model: model being calibrated
+        :param dataloader: loads data for calibration
+        :param dataset_args: dataset arguments relevant to pipelines
+        """
+        model_device = get_execution_device(model)
 
-    LifecycleCallbacks.calibration_epoch_start()
+        LifecycleCallbacks.calibration_epoch_start()
 
-    with calibration_forward_context(model):
-        for batch in tqdm.tqdm(dataloader, desc="Calibrating"):
-            batch = apply_pad_mask_to_batch(batch)
-            batch = tensors_to_device(batch, model_device)
-            model(**batch)
+        with calibration_forward_context(model):
+            for batch in tqdm.tqdm(dataloader, desc="Calibrating"):
+                batch = apply_pad_mask_to_batch(batch)
+                batch = tensors_to_device(batch, model_device)
+                model(**batch)
 
-    LifecycleCallbacks.calibration_epoch_end()
+        LifecycleCallbacks.calibration_epoch_end()
+
+
+def run_calibration(model: torch.nn.Module, dataloader: DataLoader):
+    pipeline = BasicPipeline()
+    pipeline(model, dataloader, None)

src/llmcompressor/pipelines/data_free/__init__.py

@@ -1,2 +1,2 @@
 # flake8: noqa
-from .pipeline import run_pipeline
+from .pipeline import *

src/llmcompressor/pipelines/data_free/pipeline.py

@@ -1,27 +1,31 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import torch
 from torch.utils.data.dataloader import DataLoader
 
 from llmcompressor.core.session_functions import LifecycleCallbacks
+from llmcompressor.pipelines.registry import CalibrationPipeline
 
 if TYPE_CHECKING:
     from llmcompressor.args.dataset_arguments import DatasetArguments
 
-__all__ = ["run_pipeline"]
+__all__ = ["DataFreePipeline"]
 
 
-def run_pipeline(
-    model: torch.nn.Module,
-    dataloader: DataLoader,
-    dataset_args: "DatasetArguments",
-):
-    """
-    A pipeline for data-free calibration
+@CalibrationPipeline.register("datafree")
+class DataFreePipeline(CalibrationPipeline):
+    @staticmethod
+    def __call__(
+        model: torch.nn.Module,
+        dataloader: Optional[DataLoader],
+        dataset_args: "DatasetArguments",
+    ):
+        """
+        A pipeline for data-free calibration
 
-    :param model: model being calibrated
-    :param dataloader: loads data for calibration
-    :param dataset_args: dataset arguments relevant to pipelines
-    """
-    LifecycleCallbacks.calibration_epoch_start()
-    LifecycleCallbacks.calibration_epoch_end()
+        :param model: model being calibrated
+        :param dataloader: loads data for calibration
+        :param dataset_args: dataset arguments relevant to pipelines
+        """
+        LifecycleCallbacks.calibration_epoch_start()
+        LifecycleCallbacks.calibration_epoch_end()

src/llmcompressor/pipelines/independent/__init__.py

@@ -1,2 +1,2 @@
 # flake8: noqa
-from .pipeline import run_pipeline
+from .pipeline import *

src/llmcompressor/pipelines/independent/pipeline.py

@@ -6,43 +6,46 @@
 
 from llmcompressor.core import active_session
 from llmcompressor.modifiers.stage import StageModifiers
+from llmcompressor.pipelines.registry import CalibrationPipeline
 from llmcompressor.utils.helpers import patch_attr
 
 if TYPE_CHECKING:
     from llmcompressor.args.dataset_arguments import DatasetArguments
 
-__all__ = ["run_pipeline"]
-
-
-def run_pipeline(
-    model: torch.nn.Module,
-    dataloader: DataLoader,
-    dataset_args: "DatasetArguments",
-):
-    """
-    Data pipeline where each modifier is assigned its own calibration epoch and data
-    pipeline
-
-    :param model: model being calibrated
-    :param dataloader: loads data for calibration
-    :param dataset_args: dataset arguments relevant to pipelines
-    """
-    # avoid circular import
-    from llmcompressor.pipelines.registry import get_pipeline_fn
-
-    session = active_session()
-
-    modifiers = session.get_modifiers()
-    with patch_attr(session.lifecycle, "modifiers", None):
-        for index, modifier in enumerate(modifiers):
-            mod_type = str(type(modifier).__name__)
-            session.lifecycle.modifiers = [
-                StageModifiers(modifiers=[modifier], group=mod_type, index=index)
-            ]
-
-            pipeline, pipeline_fn = get_pipeline_fn(user=None, modifiers=[modifier])
-            logger.info(f"Inferred `{pipeline}` calibration pipeline for `{mod_type}`")
-
-            pipeline_fn(model, dataloader, dataset_args)
-
-        # restore modifiers on exit for proper model compression inference from recipe
+__all__ = ["IndependentPipeline"]
+
+
+@CalibrationPipeline.register("independent")
+class IndependentPipeline(CalibrationPipeline):
+    @staticmethod
+    def __call__(
+        model: torch.nn.Module,
+        dataloader: DataLoader,
+        dataset_args: "DatasetArguments",
+    ):
+        """
+        Data pipeline where each modifier is assigned its own calibration epoch and data
+        pipeline
+
+        :param model: model being calibrated
+        :param dataloader: loads data for calibration
+        :param dataset_args: dataset arguments relevant to pipelines
+        """
+        _logger = logger.patch(lambda r: r.update(function="IndependentPipeline"))
+
+        session = active_session()
+        modifiers = session.get_modifiers()
+        with patch_attr(session.lifecycle, "modifiers", None):
+            for index, modifier in enumerate(modifiers):
+                mod_type = str(type(modifier).__name__)
+                session.lifecycle.modifiers = [
+                    StageModifiers(modifiers=[modifier], group=mod_type, index=index)
+                ]
+
+                pipeline = CalibrationPipeline.from_modifiers([modifier])
+                pipeline_name = pipeline.__class__.__name__
+                _logger.info(f"Inferred `{pipeline_name}` for `{mod_type}`")
+
+                pipeline(model, dataloader, dataset_args)
+
+            # restore modifiers on exit so model can be compressed based on recipe