get_targets interface

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

Author: kylesayrs

src/llmcompressor/args/dataset_arguments.py

@@ -183,12 +183,7 @@ class DatasetArguments(CustomDatasetArguments):
             ),
         },
     )
-    batch_size: int = field(
-        default=1,
-        metadata={
-            "help": "TODO"
-        }
-    )
+    batch_size: int = field(default=1, metadata={"help": "TODO"})
     # --- pipeline arguments --- #
     pipeline: str | None = field(
         default="independent",

src/llmcompressor/modifiers/autoround/base.py

@@ -21,7 +21,8 @@
 from llmcompressor.modifiers.quantization.calibration import apply_calibration_status
 from llmcompressor.modifiers.quantization.quantization import QuantizationMixin
 from llmcompressor.transformers.compression.compressed_tensors_utils import (
-    untie_if_target_shared_embedding,
+    targets_embeddings,
+    untie_word_embeddings,
 )
 from llmcompressor.utils.pytorch.module import get_no_split_params
 
@@ -109,7 +110,6 @@ class AutoRoundModifier(Modifier, QuantizationMixin):
     enable_torch_compile: bool = True
 
     # private variables
-    _module_names: Dict[torch.nn.Module, str] = PrivateAttr(default_factory=dict)
     _all_module_input: Dict[str, List[Tuple]] = PrivateAttr(default_factory=dict)
     _q_input: Optional[torch.Tensor] = PrivateAttr(default=None)
 
@@ -124,10 +124,6 @@ def on_initialize(self, state: State, **kwargs) -> bool:
             QuantizationMixin.initialize_quantization(self, state.model)
 
         # prepare module names
-        self._module_names = {
-            m: name
-            for name, m in match_named_modules(state.model, self.targets, self.ignore)
-        }
         self._add_temporary_names(state.model)
         # freeze all model parameters
         for _, param in state.model.named_parameters():
@@ -142,7 +138,8 @@ def start_calibration(self, model: torch.nn.Module):
 
         :param model: model to prepare for calibration
         """
-        untie_if_target_shared_embedding(model, self._module_names.keys())
+        if targets_embeddings(model, self.get_targets(model)):
+            untie_word_embeddings(model)
 
         for _, module in match_named_modules(model, self.targets, self.ignore):
             # Note: No need to register observers for auto-round

src/llmcompressor/modifiers/awq/base.py

@@ -7,6 +7,7 @@
 from compressed_tensors.utils import (
     align_modules,
     get_execution_device,
+    match_modules_set,
     match_named_modules,
     update_offload_parameter,
 )
@@ -26,6 +27,7 @@
 from llmcompressor.modifiers.quantization.quantization import QuantizationMixin
 from llmcompressor.modifiers.utils.hooks import HooksMixin
 from llmcompressor.pipelines.cache import IntermediatesCache
+from llmcompressor.typing import NamedModules
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
 from llmcompressor.utils.helpers import calibration_forward_context
 from llmcompressor.utils.pytorch.module import get_layer_by_name
@@ -306,6 +308,12 @@ def on_finalize(self, state: State, **kwargs) -> bool:
 
         return True
 
+    def get_targets(self, model: torch.nn.Module) -> NamedModules:
+        for mapping in self.mappings:
+            yield from match_modules_set(
+                model, (*mapping.balance_layers, mapping.smooth_layer)
+            )
+
     def _set_resolved_mappings(self, model: Module) -> None:
         """
         Transforms the list of activations to smooth and their corresponding weights

src/llmcompressor/modifiers/distillation/output/base.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Tuple, Union
 
+import torch
 from torch.nn import Module
 
 from llmcompressor.core import Event, EventType, State
@@ -9,6 +10,7 @@
     KDModelWrapper,
     KDModuleWrapper,
 )
+from llmcompressor.typing import NamedModules
 from llmcompressor.utils.fsdp.context import summon_full_params_context
 from llmcompressor.utils.fsdp.helpers import maybe_get_wrapped, set_wrapped_model
 from llmcompressor.utils.pytorch.module import get_layers, set_layer
@@ -138,6 +140,16 @@ def on_end(self, state: State, event: Event, **kwargs):
             teacher_wrapper.kd_enabled = False
         self.wrapped_kd_model_.kd_enabled = False
 
+    def get_targets(self, model: torch.nn.Module) -> NamedModules:
+        module_targets = dict()
+        targets = self.targets if isinstance(self.targets, list) else [self.targets]
+        for target in targets:
+            # only return targets of student model, not teacher model
+            target = target[0] if isinstance(target, tuple) else target
+            module_targets.update(get_layers(target, model))
+
+        return module_targets.items()
+
     def _create_model_wrapper(
         self, student_model: Module, teacher_model: Module, state: State
     ) -> KDModelWrapper:

src/llmcompressor/modifiers/modifier.py

@@ -1,6 +1,7 @@
 from abc import abstractmethod
-from typing import Optional
+from typing import Iterable, Optional
 
+import torch
 from pydantic import ConfigDict
 
 from llmcompressor.core.events import Event, EventType
@@ -238,3 +239,12 @@ def on_event(self, state: State, event: Event, **kwargs):
         :param kwargs: Additional arguments for updating the model
         """
         pass
+
+    def get_targets(
+        self, model: torch.nn.Module
+    ) -> Iterable[tuple[str, torch.nn.Module]]:
+        """
+        Return all of the named modules which will be updated by this modifier. This
+        function can only be called after the modifier has been initialized.
+        """
+        raise NotImplementedError()

src/llmcompressor/modifiers/pruning/sparsegpt/sgpt_base.py

@@ -12,6 +12,7 @@
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers.modifier import Modifier
 from llmcompressor.modifiers.utils.hooks import HooksMixin
+from llmcompressor.typing import NamedModules
 from llmcompressor.utils.pytorch.module import (
     get_layers,
     get_no_split_params,
@@ -192,6 +193,9 @@ def on_end(self, state: State, event: Event, **kwargs):
         self.ended_ = True
         self.remove_hooks()
 
+    def get_targets(self, model: torch.nn.Module) -> NamedModules:
+        return get_layers(self.targets, model).items()
+
     def _infer_sequential_targets(self, model: torch.nn.Module) -> str | list[str]:
         match self.sequential_targets:
             case None:

src/llmcompressor/modifiers/quantization/gptq/base.py

@@ -32,7 +32,7 @@
 __all__ = ["GPTQModifier"]
 
 
-class GPTQModifier(Modifier, QuantizationMixin):
+class GPTQModifier(QuantizationMixin, Modifier):
     """
     Implements the GPTQ algorithm from https://arxiv.org/abs/2210.17323. This modifier
     uses activations to calibrate a hessian matrix, which is then used to determine

src/llmcompressor/modifiers/quantization/quantization/mixin.py

@@ -35,8 +35,10 @@
 )
 from llmcompressor.modifiers.utils.hooks import HooksMixin
 from llmcompressor.transformers.compression.compressed_tensors_utils import (
-    untie_if_target_shared_embedding,
+    targets_embeddings,
+    untie_word_embeddings,
 )
+from llmcompressor.typing import NamedModules
 
 __all__ = ["QuantizationMixin"]
 
@@ -182,11 +184,8 @@ def start_calibration(self, model: torch.nn.Module):
 
         :param model: model to prepare for calibration
         """
-
-        matched_module_generator = (
-            x[1] for x in match_named_modules(model, self.resolved_targets, self.ignore)
-        )
-        untie_if_target_shared_embedding(model, matched_module_generator)
+        if targets_embeddings(model, self.get_targets(model)):
+            untie_word_embeddings(model)
 
         for _, module in match_named_modules(model, self.resolved_targets, self.ignore):
             self._initialize_observers(module)
@@ -263,6 +262,9 @@ def resolve_quantization_config(self) -> QuantizationConfig:
             ignore=ignore,
         )
 
+    def get_targets(self, model: torch.nn.Module) -> NamedModules:
+        return match_named_modules(model, self.resolved_targets, self.ignore)
+
     def _initialize_observers(self, module: torch.nn.Module):
         if not hasattr(module, "quantization_scheme"):
             return

src/llmcompressor/modifiers/smoothquant/base.py

@@ -2,7 +2,7 @@
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
-from compressed_tensors.utils import align_module_device
+from compressed_tensors.utils import align_module_device, match_modules_set
 from loguru import logger
 from pydantic import ConfigDict, Field
 from torch.nn import Module
@@ -13,6 +13,7 @@
     get_layer_mappings_from_architecture,
     handle_mapping_resolution_errors,
 )
+from llmcompressor.typing import NamedModules
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
 from llmcompressor.utils.pytorch.module import (
     get_layers,
@@ -54,6 +55,7 @@ class SmoothQuantMapping:
 
     smooth_name: str
     smooth_layer: Module
+    balance_names: List[str]
     balance_layers: List[Module]
 
 
@@ -178,6 +180,13 @@ def on_finalize(self, state: State, **kwargs) -> bool:
 
         return True
 
+    def get_targets(self, model: torch.nn.Module) -> NamedModules:
+        if not self.initialized_:
+            raise ValueError("Cannot get targets before modifier has been initialized")
+
+        for balance_targets, smooth_target in self.mappings:
+            yield from match_modules_set(model, (*balance_targets, smooth_target))
+
     def _infer_mappings_from_model(
         self,
         model: Module,
@@ -207,18 +216,20 @@ def _resolve_mappings(self, model: Module) -> List[SmoothQuantMapping]:
             to_smooth_layers = get_layers(to_smooth, model)
             for layer_name, smooth_layer in to_smooth_layers.items():
                 if not match_targets(layer_name, self.ignore)[0]:
+                    balance_names = []
                     balance_layers = []
                     for balance_suffix in to_balance:
                         # find the submodule that matches the activation layer
-                        _, balance_layer = get_matching_layer(
+                        balance_name, balance_layer = get_matching_layer(
                             balance_suffix, layer_name, model
                         )
                         if balance_layer:
+                            balance_names.append(balance_name)
                             balance_layers.append(balance_layer)
                     # each mapping can contain multiple layers to balance, but only
                     # one layer to smooth
                     mapping = SmoothQuantMapping(
-                        layer_name, smooth_layer, balance_layers
+                        layer_name, smooth_layer, balance_names, balance_layers
                     )
                     resolved_mappings.append(mapping)
         return resolved_mappings

src/llmcompressor/modifiers/transform/quip/base.py

@@ -13,8 +13,10 @@
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
 from llmcompressor.transformers.compression.compressed_tensors_utils import (
-    untie_if_target_shared_embedding,
+    targets_embeddings,
+    untie_word_embeddings,
 )
+from llmcompressor.typing import NamedModules
 
 __all__ = ["QuIPModifier"]
 
@@ -102,18 +104,13 @@ def on_initialize(self, state: State, **kwargs) -> bool:
 
     def on_start(self, state: State, event: Event, **kwargs):
         self.started_ = True
+        model = state.model
 
-        def matched_module_generator():
-            for scheme in self.transform_config.config_groups.values():
-                for arg in scheme.apply:
-                    gen = match_named_modules(state.model, arg.targets, arg.ignore)
-                    for _, module in gen:
-                        yield module
+        # untie embeddings if they will be targeted by transforms
+        if targets_embeddings(model, self.get_targets(model)):
+            untie_word_embeddings(model)
 
-        # Untie embeddings if they will be targeted by transforms
-        untie_if_target_shared_embedding(state.model, matched_module_generator())
-
-        apply_transform_config(state.model, self.transform_config)
+        apply_transform_config(model, self.transform_config)
 
     def on_event(self, state: State, event: Event, **kwargs):
         if event.type_ == EventType.CALIBRATION_EPOCH_START:
@@ -136,6 +133,17 @@ def on_finalize(self, state: State, **kwargs) -> bool:
 
         return True
 
+    def get_targets(self, model: torch.nn.Module) -> NamedModules:
+        if not self.initialized_:
+            raise ValueError("Cannot get targets before modifier has been initialized")
+
+        return [
+            (name, module)
+            for scheme in self.transform_config.config_groups.values()
+            for arg in scheme.apply
+            for name, module in match_named_modules(model, arg.targets, arg.ignore)
+        ]
+
     def _create_config(self) -> TransformConfig:
         config_groups = dict()
         if "v" in self.rotations: