Address review: move deep_equal, make prefetch a global setting via State

Signed-off-by: Avishek Goswami <avishek.goswami@ibm.com>

Author: Avishek Goswami

src/llmcompressor/core/state.py

@@ -114,6 +114,7 @@ class State:
     _last_log_step: float | int | None = None
     loss_masks: list[torch.Tensor] | None = None
     current_batch_idx: int = -1
+    sequential_prefetch: bool = False
 
     @property
     def compression_ready(self) -> bool:

src/llmcompressor/modifiers/awq/base.py

@@ -138,10 +138,6 @@ class AWQModifier(Modifier, QuantizationMixin):
         requirements but requires more time to move data between cpu and execution
         device. Defaults to None, so cached args are not offloaded. Consider setting
         to torch.device("cpu") if you are encountering OOM errors
-    :param prefetch: when offloading, prefetch the next batch in a background thread
-        to overlap CPU-to-device onload with the forward pass, reducing wall-clock
-        time. Default False; set True when offload_device is set and GPU memory
-        allows two batches on device simultaneously
     :param duo_scaling: whether to use duo scaling, which uses both input activations
         and weights to determine the scaling factor. Defaults to True
         If True, both activations and weights are used.
@@ -161,7 +157,6 @@ class AWQModifier(Modifier, QuantizationMixin):
     sequential_targets: str | list[str] | None = None
     mappings: list[AWQMapping] | None = None
     offload_device: torch.device | None | Sentinel = Sentinel("not_provided")
-    prefetch: bool = False
     duo_scaling: bool | Literal["both"] = True
     n_grid: int = 20
 
@@ -613,8 +608,8 @@ def _smooth(
     @torch.no_grad()
     def _run_samples(self, module: Module) -> list[torch.Tensor]:
         cache = self._parent_args_cache[module]
-        # Prefetch overlaps CPU->device onload with forward pass when offloading.
-        batch_iter = cache.iter_prefetch() if self.prefetch else cache
+        use_prefetch = active_session().state.sequential_prefetch
+        batch_iter = cache.iter_prefetch() if use_prefetch else cache
         outputs = [module(**batch_kwargs) for batch_kwargs in batch_iter]
         return [
             # If tuple, assume that first argument is the input

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -131,14 +131,16 @@ def __call__(
             else:
                 session.state.loss_masks = None
 
+            use_prefetch = getattr(dataset_args, "sequential_prefetch", False)
+            session.state.sequential_prefetch = use_prefetch
+
             for subgraph_index, subgraph in enumerate(subgraphs):
                 # prepare tqdm description texts
                 calib_desc = f"({subgraph_index + 1}/{num_subgraphs}): Calibrating"
                 prop_desc = f"({subgraph_index + 1}/{num_subgraphs}): Propagating"
 
                 # reduce memory movement by keeping modules onloaded
                 num_batches = len(dataloader)
-                use_prefetch = getattr(dataset_args, "sequential_prefetch", False)
                 with disable_offloading():
                     # do a preliminary pass to trigger modifier hooks
                     for batch_idx, inputs in _get_batches(

tests/llmcompressor/pipelines/test_cache.py

@@ -41,6 +41,30 @@ def sample_cache(sample_dataloader):
 ]
 
 
+def deep_equal(a, b) -> bool:
+    if type(a) is not type(b):
+        return False
+
+    match a:
+        case torch.Tensor():
+            return torch.equal(a, b)
+        case list() | tuple():
+            if len(a) != len(b):
+                return False
+            return all(deep_equal(_a, _b) for _a, _b in zip(a, b))
+        case dict():
+            if a.keys() != b.keys():
+                return False
+            return all(deep_equal(a[key], b[key]) for key in a.keys())
+        case _ if is_dataclass(a):
+            a_dict = {field.name: getattr(a, field.name) for field in fields(a)}
+            b_dict = {field.name: getattr(b, field.name) for field in fields(b)}
+
+            return deep_equal(a_dict, b_dict)
+        case _:
+            return a == b
+
+
 @pytest.mark.unit
 def test_initialization(sample_dataloader):
     cache = IntermediatesCache.from_dataloader(
@@ -76,30 +100,6 @@ def batch_dicts_equal(a: dict, b: dict) -> bool:
         assert batch_dicts_equal(b_iter, b_prefetch), f"batch {i} differs"
 
 
-def deep_equal(a, b) -> bool:
-    if type(a) is not type(b):
-        return False
-
-    match a:
-        case torch.Tensor():
-            return torch.equal(a, b)
-        case list() | tuple():
-            if len(a) != len(b):
-                return False
-            return all(deep_equal(_a, _b) for _a, _b in zip(a, b))
-        case dict():
-            if a.keys() != b.keys():
-                return False
-            return all(deep_equal(a[key], b[key]) for key in a.keys())
-        case _ if is_dataclass(a):
-            a_dict = {field.name: getattr(a, field.name) for field in fields(a)}
-            b_dict = {field.name: getattr(b, field.name) for field in fields(b)}
-
-            return deep_equal(a_dict, b_dict)
-        case _:
-            return a == b
-
-
 @pytest.mark.unit
 def test_fetch_inputs(sample_cache):
     fetched = sample_cache.fetch(0, ["input_ids", "attention_mask"])