feat: add input validation for targets_per_subgraph

Signed-off-by: Ayush <aayush7511@gmail.com>

Author: aayush7511

src/llmcompressor/pipelines/sequential/helpers.py

@@ -100,7 +100,6 @@ def trace_subgraphs(
     :param sequential_targets: list of patterns matching sequential targets
     :param ignore: function and method names to skip during tracing
     :param targets_per_subgraph: number of targets to include per subgraph
-    
     :return: a list of Subgraphs in order of execution
     """
     # find modules
@@ -277,15 +276,18 @@ def topological_partition(graph: GraphModule, targets: set[Module], targets_per_
     """
     assert graph_is_well_formed(graph.graph)
     target_nodes = find_target_nodes(graph, targets)
+    
+    if(targets_per_subgraph <= 0):
+        raise ValueError("targets_per_subgraph is required to be greater than or equal to one")
 
     partitions: list[list[Node]] = [[]]
     remaining_indegrees = {
         node: len([node for node in node.all_input_nodes if node.op != "get_attr"])
         for node in graph.graph.nodes
     }
     partition_index = 0  # global counter
-    targets_seen = 0 # global counter
-
+    targets_seen = 0 # number of targets encountered so far
+    
     # start with graph input nodes,
     # but delay the `get_attr` nodes as long as possible
     queue = deque(
@@ -303,7 +305,7 @@ def topological_partition(graph: GraphModule, targets: set[Module], targets_per_
         if node in target_nodes:
             targets_seen += 1
 
-            if(targets_seen % targets_per_subgraph == 0):
+            if targets_seen % targets_per_subgraph == 0:
                 partition_index += 1
                 partitions.append([])
 

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -103,7 +103,8 @@ def __call__(
             sample_input,
             sequential_targets,
             ignore,
-            dataset_args.sequential_targets_per_subgraph)
+            dataset_args.sequential_targets_per_subgraph,
+        )
         num_subgraphs = len(subgraphs)
 
         LifecycleCallbacks.calibration_epoch_start()

src/llmcompressor/transformers/tracing/debug.py

@@ -27,7 +27,7 @@ def parse_args():
     parser.add_argument("--trust_remote_code", type=bool, default=False, help="Whether to trust model remote code")  # noqa: E501
     parser.add_argument("--skip_weights", type=bool, default=True, help="Whether to load the model with dummy weights")  # noqa: E501
     parser.add_argument("--device_map", type=str, default="cpu", help="Device to load model and inputs onto")  # noqa: E501
-    parser.add_argument("--targets_per_subgraph", type=int, default=1, help="Number of sequential targets to include per subgraph") # noqa: E501
+    parser.add_argument("--targets_per_subgraph", type=int, default=1, help="Number of sequential targets to include per subgraph")  # noqa: E501
     return parser.parse_args()
 
 
@@ -106,7 +106,11 @@ def trace(
         f"    ignore={dataset_args.tracing_ignore}\n"
     )
     subgraphs = trace_subgraphs(
-        model, sample, sequential_targets, dataset_args.tracing_ignore, targets_per_subgraph
+        model,
+        sample,
+        sequential_targets,
+        dataset_args.tracing_ignore,
+        targets_per_subgraph
     )
     print(f"Successfully traced model into {len(subgraphs)} subgraphs!\n")
 
@@ -168,7 +172,7 @@ def main():
         trust_remote_code=args.trust_remote_code,
         skip_weights=args.skip_weights,
         device_map=args.device_map,
-        targets_per_subgraph=DatasetArguments.sequential_targets_per_subgraph
+        targets_per_subgraph=args.sequential_targets_per_subgraph
     )