[Canonical LoRA] fix: use correct q_out_features for linear_q

src/megatron/bridge/peft/canonical_lora.py

@@ -257,8 +257,9 @@ def transform(self, m: nn.Module, name: Optional[str] = None, prefix: Optional[s
             if name == "linear_qkv":
                 adapter_q, adapter_k, adapter_v = None, None, None
                 kv_out_features = m.config.kv_channels * m.config.num_query_groups
+                q_out_features = m.config.kv_channels * m.config.num_attention_heads
                 if "linear_q" in canonical_submodules:
-                    adapter_q = ParallelLinearAdapter(in_features, in_features, **adapter_kwargs)
+                    adapter_q = ParallelLinearAdapter(in_features, q_out_features, **adapter_kwargs)
                 if "linear_k" in canonical_submodules:
                     adapter_k = ParallelLinearAdapter(in_features, kv_out_features, **adapter_kwargs)
                 if "linear_v" in canonical_submodules:

tests/unit_tests/peft/test_canonical_lora.py

@@ -61,6 +61,7 @@ class MockConfig:
             def __init__(self):
                 self.kv_channels = kv_channels or 64
                 self.num_query_groups = num_query_groups or 8
+                self.num_attention_heads = self.num_query_groups
                 self.sequence_parallel = False
 
         self.config = MockConfig()