[Transform] Implement multi-headed transforms

[kylesayrs]

Purpose

• Support applying transforms to modules involved in multi-headed attention
  • Specifically, the Spinquant R2 transform applied to v_proj and o_proj

Prerequisites

• [Transform] apply_transform_config #348

Implementation

In order to better support low-memory representations transform weight as well as size-specific features such as permutations, we led head_dim override the size of the generated transform. Then, at apply-time, we reshape (unflatten) the value to reflect the head_dim, apply the weight, then reshape (flatten) back to the original shape.

Changes

• Add head_dim field to TransformScheme
  • This field specifies the size of heads in mha. This can be inferred from most model configs
• Rename get_matrix_size to get_transform_size
  • Support overriding the size of the transform weight with head_dim
• Rename transform/utils/utils.py to transform/utils/matrix.py

Testing

• Add tests for head_dim as well as a mock attention operation