Merge branch 'kylesayrs/transform_factory' into kylesayrs/transform_permutations

Author: kylesayrs

src/compressed_tensors/transform/transform_args.py

@@ -13,15 +13,31 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Any, List
+from typing import List
 
 from pydantic import BaseModel, Field, field_validator
 
 
-__all__ = ["TransformLocation", "TransformArgs"]
+__all__ = ["TransformArgs", "TransformLocation"]
 
 
 class TransformLocation(str, Enum):
+    """
+    Enum representing which parameters/activations a transform weight should be applied
+    to on a given module.
+
+    | -------------------------------------------------------------------------------------------------------- |  # noqa: E501
+    | Name            | Runtime     | Values        | Locations Where Inverse Could Be Applied                 |  # noqa: E501
+    | --------------- | ----------- | ------------- | -------------------------------------------------------- |  # noqa: E501
+    | `INPUT`         | online      | activations   | `prev.WEIGHT_OUTPUT`, `prev.OUTPUT`, `this.WEIGHT_INPUT` |  # noqa: E501
+    | `WEIGHT_INPUT`  | offline     | weight        | `prev.WEIGHT_OUTPUT`, `prev.OUTPUT`, `this.INPUT`        |  # noqa: E501
+    | `WEIGHT_OUTPUT` | offline     | weight        | `this.OUTPUT`, `next.INPUT`, `next.WEIGHT_INPUT`         |  # noqa: E501
+    | `OUTPUT`        | online      | activations   | `this.WEIGHT_OUTPUT`, `next.INPUT`, `next.WEIGHT_INPUT`  |  # noqa: E501
+    | `K_CACHE`       | online      | key_values    | `q_proj.Q_ATTN`                                          |  # noqa: E501
+    | `Q_ATTN`        | online      | query_values  | `k_proj.K_CACHE`                                         |  # noqa: E501
+    | -------------------------------------------------------------------------------------------------------- |  # noqa: E501
+    """
+
     INPUT = "input"
     WEIGHT_INPUT = "weight_input"
     WEIGHT_OUTPUT = "weight_output"

src/compressed_tensors/transform/utils/utils.py

@@ -41,25 +41,28 @@ def apply_transform_weight(
 ) -> torch.Tensor:
     """
     Using the transform location, determine how to apply the transform weight to the
-    given value
+    given value. For more info on input and output transforms, see `TransformLocation`
+
+    The following explains how weights should be applied to values according to location
 
     let  x          be input activation
          W          be weight,
          yh, xh, Wh be transformed output, input, weight
 
     note that
          y  = (x W.T)        // torch.nn.Linear
-         yh = (xh) (Wh).T    // transformed
+
+    Choose values for yh, xh, and Wh which incorporate matrix transforms
 
     let  V, Vi      be transform matrices on input side
          U, Ui      be transform matrices on output side
 
-    show that the following values for yh, xh, and Wh are consistent
-
     pick xh = (x V)
          Wh = (U.T W Vi.T)
          yh = (y U)
 
+    The following shows that `yh = (xh) (Wh).T` for the chosen values of yh, xh, and Wh
+
     (xh) (Wh).T = (x V) (U.T W Vi.T).T
                 = (x V) (Vi W.T U)        // transpose matrix product identity
                 = (x W.T) U
@@ -83,3 +86,6 @@ def apply_transform_weight(
 
     elif location == TransformLocation.OUTPUT:
         return value @ weight
+
+    else:
+        raise NotImplementedError(f"{location} has not been implemented yet")