docstrings

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Author: brian-dellabetta

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

@@ -41,6 +41,20 @@ class SpinQuantModifier(Modifier, use_enum_values=True):
     existing weights and therefore do not induce runtime cost. R3 and R4 are "online"
     rotations, meaning that they require additional computation at runtime.
 
+    Lifecycle:
+        - on_initialize
+            - infer SpinQuantMappings & NormMappings
+            - as needed, create transform schemes for R1, R2, R3, & R4
+        - on_start
+            - normalize embeddings
+            - fuse norm layers into subsequent Linear layers
+            - apply TransformConfig
+                - fuse transforms into weights for mergeable transforms
+                - add hooks for online transforms
+        - on sequential epoch end
+        - on_end
+        - on_finalize
+
     :param rotations: A list containing the names of rotations to apply to the model.
         Possible rotations include R1, R2, R3, and R4
     :param transform_type: The type of transform to apply to the model.

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

@@ -8,6 +8,25 @@
 
 
 class SpinQuantMapping(BaseModel):
+    """
+    SpinQuant needs to know the entire architecture of the model,
+    as R1, R2, R3, and R4 rotations need to be applied to specific
+    layers (https://arxiv.org/pdf/2405.16406 Fig. 1).
+
+    :param embedding: name or regex of embedding layer
+    :param attn_q: name or regex of q_proj layer in attention block
+    :param attn_k: name or regex of k_proj layer in attention block
+    :param attn_v: name or regex of v_proj layer in attention block
+    :param attn_o: name or regex of o_proj layer in attention block
+    :param attn_head_dim: head_dim of the attention module, needed
+        because R2 needs to be applied "head-wisely" to v_proj and
+        o_proj
+    :param mlp_in: list of names or regexes for the mlp blocks that
+        receive the input to the MLP block, usually up_proj and gate_proj
+    :param mlp_out: list of names or regexes for the mlp blocks that
+        consitute the output of the MLP block, usually down_proj
+    """
+
     embedding: str
 
     attn_q: str

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

@@ -8,6 +8,16 @@
 
 
 class NormMapping(BaseModel):
+    """
+    SpinQuant needs to know where every norm layer exists in the model,
+    as well as all the subsequent Linear layers the norm passes into.
+    This is because the norm layer weights need to normalized before
+    transforms can be fused into Linear layers.
+
+    :param norm: name or regex that matches norm layer in model
+    :param linears: list of names or regexes of Linear layers that
+    receive input from norm.
+    """
     norm: str
     linears: List[str]