Optimized rotation matrix #1073
Labels
Comments
Giuseppe5
added
enhancement
vishwamartur
added a commit
to vishwamartur/brevitas
that referenced
this issue
Nov 3, 2024
Related to Xilinx#1073 Add `RotationEqualizedLayer` class to `equalized_layer.py` to hold two learnable matrices for rotation. * Implement `forward` method in `RotationEqualizedLayer` to apply rotation and call the wrapped layer. * Add method to fuse rotation matrices into the wrapped layer. * Add unit tests for `RotationEqualizedLayer` class in `test_equalized_layer.py`. * Test forward pass with rotation matrices. * Test fusing rotation matrices into the wrapped layer.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Is your feature request related to a problem? Please describe.
Current rotation equalization method is based on orthogonal and hadamard matrices. We can extend this to learned rotational matrices
Describe the solution you'd like
We need to have an extra rotation class in
equalized_layers.py, that holds the two learnable matrices.After the rotation is applied and the layers are wrapped with this new equalized layer, we can apply the learning process, e.g., the one described in SpinQuant.
Post learning, the rotation matrices can be fused in their respective wrapped layers and remove the wrapped layers.
Additional context
The infrastructure to pick up the layers to rotate is already in place. The most effort is to put together the training process.
@i-colbert
The text was updated successfully, but these errors were encountered: