[transformers] Delegate creating GradSampler to PruningModifiers (#831)

* Add customizable grad sampler

Add customizable grad sampler

Change default arg to None

* Update MFAC with GradSampler changes

Update mfac tests to reflect GradSampler changes

Add grad_sampler_kwargs as ModifierProp

Fix style

* Remove apostrophe from docstring

* Fix typos

* quality fix

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

Author: eldarkurtic

src/sparseml/pytorch/sparsification/pruning/modifier_pruning_mfac.py

@@ -21,7 +21,7 @@
 import os
 from abc import ABC, abstractmethod
 from functools import wraps
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -83,6 +83,8 @@ class MFACPruningModifier(BaseGradualPruningModifier):
     |       num_grads: {0.0: 64, 0.5: 128, 0.75: 256, 0.85: 512}
     |       fisher_block_size: 10000
     |       available_devices: ["cuda:0"]
+    |       grad_sampler_kwargs:
+    |           batch_size: 8
 
     :param init_sparsity: the initial sparsity for the param to start with at
         start_epoch
@@ -130,6 +132,8 @@ class MFACPruningModifier(BaseGradualPruningModifier):
     :param mask_type: String to define type of sparsity to apply. May be 'unstructred'
         for unstructured pruning or 'block4' for four block pruning or a list of two
         integers for a custom block shape. Default is 'unstructured'
+    :param grad_sampler_kwargs: kwargs to override default train dataloader config
+        for gradient sampling.
     """
 
     def __init__(
@@ -151,6 +155,7 @@ def __init__(
         num_pages: int = 1,  # break computation into pages when block size is None
         available_devices: Optional[List[str]] = None,
         mask_type: str = "unstructured",
+        grad_sampler_kwargs: Optional[Dict[str, Any]] = None,
     ):
         super().__init__(
             params=params,
@@ -172,6 +177,7 @@ def __init__(
         self._fisher_block_size = fisher_block_size
         self._num_pages = num_pages
         self._mask_type = mask_type
+        self._grad_sampler_kwargs = grad_sampler_kwargs
         if available_devices is None:
             if torch.cuda.device_count() > 0:
                 self._available_devices = ["cuda:0"]
@@ -229,6 +235,13 @@ def available_devices(self) -> Optional[List[str]]:
         """
         return self._available_devices
 
+    @ModifierProp(serializable=True)
+    def grad_sampler_kwargs(self) -> Optional[Dict[str, Any]]:
+        """
+        Return dict of training dataloader configs overridden for gradient sampling
+        """
+        return self._grad_sampler_kwargs
+
     @ModifierProp()
     def mask_type(self) -> str:
         """
@@ -259,14 +272,21 @@ def initialize(
         if "grad_sampler" in kwargs and self._use_gradient_buffering is not True:
             # set grad sampler, must be done before initialize in case pruning step
             # occurs on initialize epoch
-            grad_sampler = kwargs["grad_sampler"]
-            if not isinstance(grad_sampler, GradSampler):
-                raise ValueError(
-                    "grad_sampler must be an instance of the GradSampler class"
+            if (
+                "data_loader_builder" not in kwargs["grad_sampler"]
+                or "loss_function" not in kwargs["grad_sampler"]
+            ):
+                raise RuntimeError(
+                    "grad_sampler dict with data_loader_builder and loss_function "
+                    "must be provided to initialize GradSampler"
                 )
-            self._grad_sampler = grad_sampler
+            self._grad_sampler = GradSampler(
+                kwargs["grad_sampler"]["data_loader_builder"](
+                    self._grad_sampler_kwargs
+                ),
+                kwargs["grad_sampler"]["loss_function"],
+            )
             self.log_string("Using provided GradSampler")
-
         elif self._use_gradient_buffering is False:
             raise RuntimeError(
                 "grad_sampler must be provided when use_gradient_buffering is set"

src/sparseml/transformers/sparsification/trainer.py

@@ -28,7 +28,6 @@
 import torch
 from torch import distributed as dist
 from torch.nn import Module
-from torch.utils.data import RandomSampler
 from transformers import Trainer as TransformersTrainer
 from transformers import TrainerCallback, TrainerControl, TrainingArguments
 from transformers.file_utils import WEIGHTS_NAME
@@ -38,7 +37,6 @@
 
 from sparseml.pytorch.optim import ScheduledModifierManager, ScheduledOptimizer
 from sparseml.pytorch.utils import (
-    GradSampler,
     LoggerManager,
     ModuleSparsificationInfo,
     TensorBoardLogger,
@@ -157,9 +155,6 @@ def __init__(
         self.callback_disable_fp16 = DisableHalfPrecisionCallback(self)
         self.callback_handler.add_callback(self.callback_disable_fp16)
         self._add_tensorboard_logger_if_available()
-        self.grad_sampler = GradSampler(
-            self._mfac_data_loader(), self._mfac_loss_function
-        )
 
         model_signature = inspect.signature(self.model.forward)
         self._model_signature_columns = list(model_signature.parameters.keys())
@@ -275,7 +270,10 @@ def create_optimizer(self):
                 wrap_optim=self.scaler,
                 loggers=self.logger_manager,
                 distillation_teacher=self.teacher,
-                grad_sampler=self.grad_sampler,
+                grad_sampler={
+                    "data_loader_builder": self._data_loader_builder,
+                    "loss_function": self._loss_function,
+                },
             )
         else:
             wrap_optim_key = "optimizer"
@@ -286,16 +284,22 @@ def create_optimizer(self):
                 steps_per_epoch=self.manager_steps_per_epoch,
                 loggers=self.logger_manager,
                 initialize_kwargs={
-                    "grad_sampler": self.grad_sampler,
                     "distillation_teacher": self.teacher,
+                    "grad_sampler": {
+                        "data_loader_builder": self._data_loader_builder,
+                        "loss_function": self._loss_function,
+                    },
                 },
             )
             if not self.manager.initialized:
                 self.manager.initialize(
                     self.model,
                     loggers=self.logger_manager,
                     distillation_teacher=self.teacher,
-                    grad_sampler=self.grad_sampler,
+                    grad_sampler={
+                        "data_loader_builder": self._data_loader_builder,
+                        "loss_function": self._loss_function,
+                    },
                 )
         self.manager_initialized = True
         _LOGGER.info(
@@ -663,25 +667,21 @@ def _reload_model_state(self, load_path: str, orig_state_dict: Dict[str, Any]):
             delayed_load=False,
         )
 
-    def _mfac_data_loader(self):
-        def dataloader():
-            data_loader_template = self.get_train_dataloader()
-
-            data_loader = torch.utils.data.DataLoader(
-                dataset=data_loader_template.dataset,
-                batch_size=data_loader_template.batch_size // 2,
-                sampler=RandomSampler(data_loader_template.dataset, replacement=False),
-                num_workers=data_loader_template.num_workers,
-                collate_fn=data_loader_template.collate_fn,
-                pin_memory=data_loader_template.pin_memory,
-                drop_last=data_loader_template.drop_last,
-                timeout=data_loader_template.timeout,
-                worker_init_fn=data_loader_template.worker_init_fn,
-                generator=data_loader_template.generator,
-                prefetch_factor=data_loader_template.prefetch_factor,
-                persistent_workers=data_loader_template.persistent_workers,
-            )
+    def _data_loader_builder(self, kwargs: Optional[Dict[str, Any]] = None):
+        default_loader = self.get_train_dataloader()
+        template = dict(default_loader.__dict__)
 
+        # drop attributes that will be auto-initialized
+        to_drop = [k for k in template if k.startswith("_") or k == "batch_sampler"]
+        for item in to_drop:
+            template.pop(item)
+
+        # override defaults if kwargs are given, for example via recipe
+        if kwargs:
+            template.update(kwargs)
+        data_loader = type(default_loader)(**template)
+
+        while True:  # infinite dataloading
             for sample in data_loader:
                 if self.label_smoother is not None and "labels" in sample:
                     label = sample.pop("labels")
@@ -690,9 +690,7 @@ def dataloader():
                 sample = self._prepare_inputs(sample)
                 yield [], sample, label
 
-        return dataloader
-
-    def _mfac_loss_function(self, model_outputs, loss_target):
+    def _loss_function(self, model_outputs, loss_target):
         if loss_target is not None:
             loss = self.label_smoother(model_outputs, loss_target)
         else:

tests/sparseml/pytorch/sparsification/pruning/test_modifier_pruning_mfac.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from typing import Any, Dict, Optional
 
 import pytest
 import torch
 from torch.utils.data import DataLoader
 
 from flaky import flaky
 from sparseml.pytorch.sparsification.pruning import MFACPruningModifier
-from sparseml.pytorch.utils import GradSampler, tensor_sparsity
+from sparseml.pytorch.utils import tensor_sparsity
 from sparseml.utils import FROM_PARAM_TOKEN
 from tests.sparseml.pytorch.helpers import MLPDataset, MLPNet
 from tests.sparseml.pytorch.sparsification.pruning.helpers import (
@@ -38,29 +39,32 @@
 )
 
 
-def _device_data_loader(data_loader):
-    for sample in data_loader:
-        img, target = [t for t in sample]
-        yield [img], {}, target
-
-
-def _mfac_loss_function(model_outputs, loss_target):
-    return torch.nn.functional.mse_loss(model_outputs[0], loss_target)
+def _get_loss_function():
+    return lambda model_outputs, loss_target: torch.nn.functional.mse_loss(
+        model_outputs[0], loss_target
+    )
 
 
-def _build_gradient_sampler(
+def _get_dataloader_builder(
     dataset_lambda,
-    loss_function,
-    data_generator,
-    batch_size,
+    mfac_batch_size,
     num_grads,
     num_epochs,
     update_frequency,
 ):
-    data_length = int(batch_size * num_grads * num_epochs * (1 / update_frequency) * 2)
-    dataset = dataset_lambda(length=data_length)
-    data_loader = DataLoader(dataset, batch_size=batch_size)
-    return GradSampler(data_generator(data_loader), loss_function)
+    def dataloader_builder(kwargs: Optional[Dict[str, Any]] = None):
+        batch_size = kwargs["batch_size"] if kwargs else mfac_batch_size
+        data_length = int(
+            mfac_batch_size * num_grads * num_epochs * (1 / update_frequency) * 2
+        )
+        dataset = dataset_lambda(length=data_length)
+        data_loader = DataLoader(dataset, batch_size=batch_size)
+
+        for sample in data_loader:
+            img, target = [t for t in sample]
+            yield [img], {}, target
+
+    return dataloader_builder
 
 
 @flaky(max_runs=3, min_passes=2)
@@ -106,6 +110,9 @@ def _build_gradient_sampler(
             inter_func="cubic",
             num_grads=8,
             global_sparsity=True,
+            grad_sampler_kwargs={
+                "batch_size": 4,
+            },
         ),
     ],
     scope="function",
@@ -121,8 +128,8 @@ def _build_gradient_sampler(
 )
 class TestMFACPruningModifier(ScheduledUpdateModifierTest):
     @pytest.mark.parametrize(
-        "dataset_lambda,loss,mfac_batch_size",
-        [(MLPDataset, _mfac_loss_function, 4)],
+        "dataset_lambda,mfac_batch_size",
+        [(MLPDataset, 4)],
     )
     def test_lifecycle(
         self,
@@ -131,21 +138,21 @@ def test_lifecycle(
         optim_lambda,
         test_steps_per_epoch,  # noqa: F811
         dataset_lambda,
-        loss,
         mfac_batch_size,
     ):
         modifier = modifier_lambda()
         model = model_lambda()
         optimizer = optim_lambda(model)
-        grad_sampler = _build_gradient_sampler(
-            dataset_lambda,
-            loss,
-            _device_data_loader,
-            mfac_batch_size,
-            modifier.num_grads,
-            modifier.end_epoch - modifier.start_epoch + 1,
-            modifier.update_frequency,
-        )
+        grad_sampler = {
+            "data_loader_builder": _get_dataloader_builder(
+                dataset_lambda,
+                mfac_batch_size,
+                modifier.num_grads,
+                modifier.end_epoch - modifier.start_epoch + 1,
+                modifier.update_frequency,
+            ),
+            "loss_function": _get_loss_function(),
+        }
 
         self.initialize_helper(modifier, model, grad_sampler=grad_sampler)
         if modifier.start_epoch > 0:
@@ -222,8 +229,8 @@ def _test_final_sparsity_applied():
             _test_final_sparsity_applied()
 
     @pytest.mark.parametrize(
-        "dataset_lambda,loss,mfac_batch_size",
-        [(MLPDataset, _mfac_loss_function, 4)],
+        "dataset_lambda,mfac_batch_size",
+        [(MLPDataset, 4)],
     )
     def test_scheduled_update(
         self,
@@ -233,19 +240,20 @@ def test_scheduled_update(
         test_epoch,  # noqa: F811
         test_steps_per_epoch,  # noqa: F811
         dataset_lambda,
-        loss,
         mfac_batch_size,
     ):
         modifier = modifier_lambda()
-        grad_sampler = _build_gradient_sampler(
-            dataset_lambda,
-            loss,
-            _device_data_loader,
-            mfac_batch_size,
-            modifier.num_grads,
-            modifier.end_epoch - modifier.start_epoch + 1,
-            modifier.update_frequency,
-        )
+        grad_sampler = {
+            "data_loader_builder": _get_dataloader_builder(
+                dataset_lambda,
+                mfac_batch_size,
+                modifier.num_grads,
+                modifier.end_epoch - modifier.start_epoch + 1,
+                modifier.update_frequency,
+            ),
+            "loss_function": _get_loss_function(),
+        }
+
         super().test_scheduled_update(
             modifier_lambda,
             model_lambda,
@@ -290,6 +298,7 @@ def test_mfac_pruning_yaml(params, init_sparsity, final_sparsity):
     num_pages = 1
     available_devices = ["cpu"]
     mask_type = "block4"
+    batch_size = 4
     yaml_str = f"""
     !MFACPruningModifier
         init_sparsity: {init_sparsity}
@@ -307,6 +316,8 @@ def test_mfac_pruning_yaml(params, init_sparsity, final_sparsity):
         num_pages: {num_pages}
         available_devices: {available_devices}
         mask_type: {mask_type}
+        grad_sampler_kwargs:
+            batch_size: {batch_size}
     """
     yaml_modifier = MFACPruningModifier.load_obj(yaml_str)
     serialized_modifier = MFACPruningModifier.load_obj(
@@ -328,6 +339,9 @@ def test_mfac_pruning_yaml(params, init_sparsity, final_sparsity):
         num_pages=num_pages,
         available_devices=available_devices,
         mask_type=mask_type,
+        grad_sampler_kwargs={
+            "batch_size": batch_size,
+        },
     )
     assert isinstance(yaml_modifier, MFACPruningModifier)
     pruning_modifier_serialization_vals_test(
@@ -373,3 +387,8 @@ def test_mfac_pruning_yaml(params, init_sparsity, final_sparsity):
         == str(serialized_modifier.mask_type)
         == str(obj_modifier.mask_type)
     )
+    assert (
+        str(yaml_modifier._grad_sampler_kwargs)
+        == str(serialized_modifier._grad_sampler_kwargs)
+        == str(obj_modifier._grad_sampler_kwargs)
+    )