[Tracing] Support Cohere Vision, Decouple vision tower from first layer

src/llmcompressor/pipelines/sequential/helpers.py

@@ -277,14 +277,14 @@ def topological_partition(graph: GraphModule, targets: Set[Module]) -> List[List
     while len(queue) > 0:
         node = queue.popleft()
 
-        # assign to partition
-        partitions[partition_index].append(node)
-
         # guarantee targets are assigned to disjoint partitions
-        if node in target_nodes:
+        if node in target_nodes and len(partitions[partition_index]) > 0:
             partition_index += 1
             partitions.append([])
 
+        # assign to partition
+        partitions[partition_index].append(node)
+
         # recurse on last indegree only in order to guarantee that
         # the node is assigned to maximal partition
         for user in node.users:

tests/llmcompressor/transformers/tracing/test_models.py

@@ -12,6 +12,7 @@
     Qwen2_5_VLForConditionalGeneration,
     Qwen2VLForConditionalGeneration,
     WhisperForConditionalGeneration,
+    Cohere2VisionForConditionalGeneration,
 )
 
 from llmcompressor.pipelines.sequential.helpers import match_modules
@@ -87,6 +88,13 @@
             "vision",
             ["torchvision"],
         ),
+        (
+            "CohereLabs/command-a-vision-07-2025",
+            Cohere2VisionForConditionalGeneration,
+            ["Cohere2DecoderLayer"],
+            "vision",
+            [],
+        )
         (
             "Qwen/Qwen2-VL-2B-Instruct",
             Qwen2VLForConditionalGeneration,