[Sequential Pipeline] Return subgraph on sequential_epoch_end, remove layer_sequential pipeline (#1998)

## Purpose ##
* Enable better targeting of modules by modifiers such as
[AutoRound](#1994)
* Remove legacy pipeline (which is incompatible with this change)

## Changes ##
* Pass subgraph to `sequential_epoch_end`, allowing modifiers to view
all of the module that were called in the subgraph
* Implement `submodules` method on `Subgraph` which returns all the
modules called by this subgraph
* Remove `LayerSequentialPipeline`, which does not use the `Subgraph`
API and has been superseded by the sequential pipeline

---------

Signed-off-by: Kyle Sayers <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Dipika Sikka <[email protected]>

Author: 3 people

src/llmcompressor/args/dataset_arguments.py

@@ -188,8 +188,7 @@ class DatasetArguments(CustomDatasetArguments):
         default="independent",
         metadata={
             "help": "Calibration pipeline used to calibrate model"
-            "Options: ['basic', 'datafree', 'sequential', 'layer_sequential', "
-            "independent]"
+            "Options: ['basic', 'datafree', 'sequential', independent]"
         },
     )
     tracing_ignore: list[str] = field(

src/llmcompressor/core/session_functions.py

@@ -7,12 +7,16 @@
 
 import threading
 from contextlib import contextmanager
-from typing import Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from llmcompressor.core.events import EventType
 from llmcompressor.core.session import CompressionSession
 from llmcompressor.core.state import ModifiedState
 
+if TYPE_CHECKING:
+    from llmcompressor.pipelines.sequential import Subgraph
+
+
 __all__ = [
     "create_session",
     "active_session",
@@ -150,15 +154,15 @@ def calibration_epoch_start(cls, **kwargs) -> ModifiedState:
         return cls.event(EventType.CALIBRATION_EPOCH_START, **kwargs)
 
     @classmethod
-    def sequential_epoch_end(cls, **kwargs) -> ModifiedState:
+    def sequential_epoch_end(cls, subgraph: "Subgraph", **kwargs) -> ModifiedState:
         """
         Invoke a sequential epoch end event for the active session. This event should be
         called after one sequential layer has been calibrated/trained for one epoch
 
         This is called after a sequential layer has been calibrated with one batch, see
         `src/llmcompressor/pipelines/sequential/pipeline.py` for usage example
         """
-        return cls.event(EventType.SEQUENTIAL_EPOCH_END, **kwargs)
+        return cls.event(EventType.SEQUENTIAL_EPOCH_END, subgraph=subgraph, **kwargs)
 
     @classmethod
     def calibration_epoch_end(cls, **kwargs) -> ModifiedState:

src/llmcompressor/modifiers/pruning/wanda/base.py

@@ -36,7 +36,7 @@ class WandaPruningModifier(SparsityModifierBase):
     Lifecycle:
         - on_initialize
             - register_hook(module, calibrate_module, "forward")
-            - run_sequential / run_layer_sequential / run_basic
+            - run_sequential / run_basic
                 - make_empty_row_scalars
                 - accumulate_row_scalars
         - on_sequential_batch_end

src/llmcompressor/pipelines/__init__.py

@@ -13,6 +13,5 @@
 from .basic import *
 from .data_free import *
 from .independent import *
-from .layer_sequential import *
 from .registry import *
 from .sequential import *

src/llmcompressor/pipelines/layer_sequential/__init__.py

@@ -1,2 +1,3 @@
 # ruff: noqa
 from .pipeline import *
+from .helpers import *

src/llmcompressor/pipelines/layer_sequential/helpers.py

@@ -79,6 +79,14 @@ def forward(self, *args, **kwargs) -> Dict[str, Any]:
 
         return outputs
 
+    def submodules(self, model: Module, recurse: bool = False) -> Set[Module]:
+        nodes = self.graph.find_nodes(op="call_module")
+        modules = set(model.get_submodule(node.target) for node in nodes)
+        if recurse:
+            modules = set(m for module in modules for m in module.modules())
+
+        return modules
+
 
 def trace_subgraphs(
     model: PreTrainedModel,

src/llmcompressor/pipelines/layer_sequential/pipeline.py

@@ -99,7 +99,7 @@ def __call__(
                         inputs = activations.fetch(batch_idx, subgraph.input_names)
                         subgraph.forward(model, **inputs)
 
-                    LifecycleCallbacks.sequential_epoch_end()
+                    LifecycleCallbacks.sequential_epoch_end(subgraph)
 
                     # this pass does not trigger modifier hooks
                     # and is only used for capturing outputs of newly compressed modules