Rotation based equalization

[Giuseppe5]

Reason for this PR

Implement rotation based equalization for weights and activation.

Highlights:

• Graph based
• Layerwise
• Merge Affine parameters of RMSNorm to make it compatible with rotation equalization
• Extend LLM entrypoint to support
• Tests

Changes Made in this PR

The graph based region algorithm has been extended and generalized to support different set of supported layers based on what type of equalization is applied.

Important considerations:

• This PR depends on external library for fast implementation (I believe it should be made optional)
• This PR reuses code from other repositories (credited and with the appropriate license)

Testing Summary

Added one test that checks:

• Affine RMS norm parameter merging
• Rotation equalization being applied
• Mathematical invariance of the floating point output after rotation.

Risk Highlight

• This PR includes code from another work (please detail).
• This PR contains API-breaking changes.
• This PR depends on work in another PR (please provide links/details).
• This PR introduces new dependencies (please detail).
• There are coverage gaps not covered by tests.
• Documentation updates required in subsequent PR.

Checklist

• Code comments added to any hard-to-understand areas, if applicable.
• Changes generate no new warnings.
• Updated any relevant tests, if applicable.
• No conflicts with destination dev branch.
• I reviewed my own code changes.
• Initial CI/CD passing.
• 1+ reviews given, and any review issues addressed and approved.
• Post-review full CI/CD passing.

[Giuseppe5]

Current limitations:

• No support for scaled_dot_product in a FX graph
• No support for MultiheadAttention
• Missing conversion from LayerNorm to RMSNorm to ensure rotation compatibility

[nickfraser]

Could you add some basic tests for the LLM entry-point as well? Both for RMSNorm & rotations?

[nickfraser]

Other than the hadamard matrix instantiation issue, this LGTM (once above changes are applied).