merge attn_q etc. into attn_qkv

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

src/llmcompressor/modifiers/transform/spinquant/base.py

@@ -177,9 +177,7 @@ def _create_r1_scheme(self) -> TransformScheme:
                 ),
                 TransformArgs(
                     targets=[
-                        self.mappings.attn_q,
-                        self.mappings.attn_k,
-                        self.mappings.attn_v,
+                        *self.mappings.attn_qkv,
                         *self.mappings.mlp_in,
                         self.mappings.lm_head,
                     ],

src/llmcompressor/modifiers/transform/spinquant/mappings.py

@@ -10,9 +10,7 @@
 class SpinQuantMapping(BaseModel):
     embedding: str
 
-    attn_q: str
-    attn_k: str
-    attn_v: str
+    attn_qkv: List[str] # q_proj, k_proj, v_proj
     attn_o: str
     attn_head_dim: Optional[int] = Field(default=None)
 
@@ -21,7 +19,7 @@ class SpinQuantMapping(BaseModel):
 
     lm_head: str
 
-    @field_validator("mlp_in", "mlp_out", mode="before")
+    @field_validator("attn_qkv", "mlp_in", "mlp_out", mode="before")
     def cast_to_list(cls, value):
         if isinstance(value, str):
             return [value]
@@ -31,9 +29,7 @@ def cast_to_list(cls, value):
 
 _default_mappings = SpinQuantMapping(
     embedding="re:.*embed_tokens$",
-    attn_q="re:.*q_proj$",
-    attn_k="re:.*k_proj$",
-    attn_v="re:.*v_proj$",
+    attn_qkv=["re:.*q_proj$","re:.*k_proj$","re:.*v_proj$"],
     attn_o="re:.*o_proj$",
     mlp_in=["re:.*up_proj$", "re:.*gate_proj$"],
     mlp_out="re:.*down_proj$",