[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

@@ -3,6 +3,7 @@
 import pytest
 from transformers import (
     AutoModelForCausalLM,
+    Cohere2VisionForConditionalGeneration,
     Gemma3ForConditionalGeneration,
     Idefics3ForConditionalGeneration,
     Llama4ForConditionalGeneration,
@@ -88,6 +89,12 @@
             ["torchvision"],
         ),
         (
+            "CohereLabs/command-a-vision-07-2025",
+            Cohere2VisionForConditionalGeneration,
+            ["Cohere2DecoderLayer"],
+            "vision",
+            [],
+        )(
             "Qwen/Qwen2-VL-2B-Instruct",
             Qwen2VLForConditionalGeneration,
             ["Qwen2VLDecoderLayer"],