optimization

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -95,26 +95,27 @@ def __call__(
                 # reduce memory movement by keeping modules onloaded
                 with disable_offloading():
                     # do a preliminary pass to trigger modifier hooks
-                    for batch_idx in tqdm(range(len(dataloader)), desc=calib_desc):
-                        inputs = activations.fetch(batch_idx, subgraph.input_names)
+                    for b_idx in tqdm(range(len(dataloader)), desc=calib_desc):
+                        inputs = activations.fetch(b_idx, subgraph.input_names)
                         outputs = subgraph.forward(model, **inputs)
 
                         if not dataset_args.propagate_error:
-                            activations.update(batch_idx, outputs)
-                            activations.delete(batch_idx, subgraph.consumed_names)
+                            activations.update(b_idx, outputs)
+                            activations.delete(b_idx, subgraph.consumed_names)
 
                     LifecycleCallbacks.sequential_epoch_end(subgraph)
 
-                    # this pass does not trigger modifier hooks
-                    # and is only used for capturing outputs of newly compressed modules
-                    with HooksMixin.disable_hooks():
-                        for batch_idx in tqdm(range(len(dataloader)), desc=prop_desc):
-                            inputs = activations.fetch(batch_idx, subgraph.input_names)
-                            outputs = subgraph.forward(model, **inputs)
-
-                            if dataset_args.propagate_error:
-                                activations.update(batch_idx, outputs)
-                                activations.delete(batch_idx, subgraph.consumed_names)
+                    if not dataset_args.propagate_error:
+                        # this pass does not trigger modifier hooks
+                        # and is only used for capturing outputs of compressed modules
+                        with HooksMixin.disable_hooks():
+                            for b_idx in tqdm(range(len(dataloader)), desc=prop_desc):
+                                inputs = activations.fetch(b_idx, subgraph.input_names)
+                                outputs = subgraph.forward(model, **inputs)
+
+                                if dataset_args.propagate_error:
+                                    activations.update(b_idx, outputs)
+                                    activations.delete(b_idx, subgraph.consumed_names)
 
             # redundant, finish any remaining compression
             LifecycleCallbacks.calibration_epoch_end()