Feat (equalize): enable parametrized rotations #1148

pablomlago · 2025-01-09T20:06:38Z

Reason for this PR

Changes Made in this PR

Testing Summary

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.

src/brevitas/graph/base.py

Giuseppe5 · 2025-01-13T09:46:02Z

src/brevitas/graph/base.py

+        return model
+
+
+class RotationWeightParametrization(torch.nn.Module):


Not sure if this should live here, maybe pytorch utils?

It seems to me something very specific to live in pytorch utils. Maybe worth having a graph/rotation_utils.py?

Let's do utils/rotation_utils

Maybe at some point we will also have a utils/equalize_utils and we'll leave only a few core classes in graph/equalize.py

Giuseppe5 · 2025-01-13T09:47:43Z

src/brevitas/graph/base.py

+            if old_module is self.old_module_instance:
+                # register the parametrization in the old_module
+                parametrize.register_parametrization(
+                    old_module, self.tensor_name, self.parametrization_module, unsafe=True)


What's the unsafe flag for?

It checks that the dtype and shape of the parametrized tensor is the same as the original one. As it computes parametrization(self.weight) and then the checks are done on the result, the parametrizations take some time to register. I had disabled it for faster experimentation, but probably it's worth doing those checks

Giuseppe5 · 2025-01-13T09:50:03Z

src/brevitas/graph/base.py

+    def apply(self, model: GraphModule) -> GraphModule:
+        for old_module in model.modules():
+            if old_module is self.old_module_instance:
+                if hasattr(old_module, 'allocate_params'):


Add a comment, maybe TODO, whether this should live here or outside the apply function. I'm not sure about it

Giuseppe5 · 2025-01-13T09:50:26Z

src/brevitas/graph/base.py

+        return weight
+
+
+class ModuleInstanceFuseRotationWeights(Transform):


Seems a very specific functions, add some comments

Is this equivalent to the old behaviour?

Yes, the point about adding this function is to avoid duplicating the code between the parametrization module and the in-place transformation. Also, by having an specific rewriter for the in-place fusing of the rotation, we can avoid doing any modifications at all in the FX model, thus preventing potential inconsistencies in fused_no_fx

Giuseppe5 · 2025-01-13T11:27:00Z

src/brevitas/graph/base.py

+    def __init__(
+            self, old_module_instance: Module, tensor_name: str,
+            parametrization_module: Module) -> None:
+        self.old_module_instance = old_module_instance


I believe names are not correctly representing what variables do

Giuseppe5 · 2025-01-13T11:27:50Z

src/brevitas/graph/base.py

+        self.axis = axis
+        self.K = K
+
+    def forward(self, weight: torch.Tensor) -> torch.Tensor:


Change weight to tensor

Giuseppe5 · 2025-01-13T11:33:57Z

src/brevitas/graph/equalize.py

-                raise RuntimeError("Not supported yet")
-            module.weight.data = weight
+            if fuse_rotations:
+                rewriter = ModuleInstanceFuseRotationWeights(


If this is equivalent to the old behavior, it seems a bit over-complicated. Do we need it?

I did it mainly for preventing duplication. Before, we had the same logic for sinks/sources, and when we do not fuse the rotations, I had to duplicate that logic again in the parametrization module

Giuseppe5 · 2025-01-13T11:36:17Z

tests/brevitas/graph/test_equalization.py

+
+
+@pytest_cases.parametrize('N', [1, 2, 3], ids=lambda x: f"N={x}")
+def test_composition_unfused_rotations(N):


Thanks for the tests!

Giuseppe5 · 2025-01-13T21:33:07Z

src/brevitas/graph/equalize.py

+        tied_param_name_split = tied_param_name.split(".")
+        # Check if the tied parameter is the original parameter in the module
+        if len(tied_param_name_split) >= 3 and tied_param_name_split[
+                -3] == "parametrizations" and tied_param_name_split[-1] == "original":


why [-3]?
Seems pretty arbitrary. What if the hierarchy is smaller than 3?

I'll add a comment to explain so it's clearer

pablomlago requested a review from Giuseppe5 January 9, 2025 20:06

pablomlago changed the title ~~[DO NOT MERGE] Enable parametrized rotations~~ Enable parametrized rotations Jan 10, 2025

pablomlago force-pushed the feat-unfused-rotations branch from 7d27050 to 3888df4 Compare January 11, 2025 10:39

Giuseppe5 reviewed Jan 13, 2025

View reviewed changes

src/brevitas/graph/base.py Show resolved Hide resolved

Giuseppe5 reviewed Jan 13, 2025

View reviewed changes

pablomlago force-pushed the feat-unfused-rotations branch from 20cffcc to e6fb34f Compare January 13, 2025 14:37

Giuseppe5 reviewed Jan 13, 2025

View reviewed changes

pablomlago force-pushed the feat-unfused-rotations branch from f9ae9ae to 82287b0 Compare January 13, 2025 22:21

Giuseppe5 self-requested a review January 13, 2025 22:46

pablomlago added 14 commits January 14, 2025 11:21

Enable parametrized rotations

bf50a3a

Fix tests

baf41e0

Run precommit

f1c9ae8

Layerwise fix

f4c75bc

Fix tests

b899e66

Force PyTorch version in test

727a576

Lower PyTorch version to ensure test is run

b8264ff

Revert change to generator

c175f80

Remove call to clone state

54555a2

Fix for rotations with tied weights

5cad364

Address comments

04e6f92

Unify interfaces rewriters

6819521

Address final comments

ebb7bf5

Cast rotation matrix to the weight dtype

565c3dc

pablomlago added 4 commits January 14, 2025 11:21

Fix tests

a363804

Move import to top

0575f0e

Minor refactor

18fff54

Fix tests

b1b59a2

pablomlago force-pushed the feat-unfused-rotations branch from beb2cbc to b1b59a2 Compare January 14, 2025 11:22

Add function source

2d5ef36

Giuseppe5 requested review from Giuseppe5 and removed request for Giuseppe5 January 14, 2025 12:23

Skip test on windows

26857df

Giuseppe5 requested review from Giuseppe5 and removed request for Giuseppe5 January 14, 2025 14:05

Giuseppe5 merged commit 52cfffd into Xilinx:dev Jan 14, 2025
393 of 396 checks passed

pablomlago changed the title ~~Enable parametrized rotations~~ Feat (equalize): enable parametrized rotations Jan 15, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feat (equalize): enable parametrized rotations #1148

Feat (equalize): enable parametrized rotations #1148

pablomlago commented Jan 9, 2025

Giuseppe5 Jan 13, 2025

pablomlago Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

pablomlago Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

pablomlago Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

pablomlago Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

Giuseppe5 Jan 13, 2025

pablomlago Jan 14, 2025

		return model


		class RotationWeightParametrization(torch.nn.Module):

		return weight


		class ModuleInstanceFuseRotationWeights(Transform):



		@pytest_cases.parametrize('N', [1, 2, 3], ids=lambda x: f"N={x}")
		def test_composition_unfused_rotations(N):

Feat (equalize): enable parametrized rotations #1148

Feat (equalize): enable parametrized rotations #1148

Conversation

pablomlago commented Jan 9, 2025

Reason for this PR

Changes Made in this PR

Testing Summary

Risk Highlight

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment