Feat (graph/equalize): implement permute regions #1380

Giuseppe5 · 2025-09-25T10:00:54Z

Reason for this PR

Refactor the graph equalization code, and include supports for permutation in regions.

Changes Made in this PR

Testing Summary

Missing tests for permutation yet.

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/equalize.py

Co-authored-by: Pablo Monteagudo Lago <44771380+pablomlago@users.noreply.github.com>

…itas into refactor_rotation

src/brevitas/graph/equalize.py

pablomlago · 2025-10-03T09:28:27Z

src/brevitas/graph/equalize.py

+        self.rewriters = None
+
+    def __enter__(self):
+        model, rewriters = self.rotation.apply(self.model)


Suggested change

model, rewriters = self.rotation.apply(self.model)

self.model, self.rewriters = self.rotation.apply(self.model)

and remove the two following lines.

src/brevitas/graph/equalize.py

pablomlago · 2025-10-03T10:00:52Z

src/brevitas/graph/equalize.py

This file is over 2K lines of code at the moment and it contains a lot of functionality, so maybe it is worth splitting it in a future refactor. One possibility would be to have a file with the base equalization functionality, and then other file with the code for scalar/rotation/permutation equalization, but other options could be explored.

One possibility would be to have a file with the base equalization functionality, and then other file with the code for scalar/rotation/permutation equalization, but other options could be explored.

Agreed, this would be better... Do we do that in this PR or in a future refactor of this section?

It will be done in the next release, as part of the re-org of our PTQ algorithms.
I understand it can be done without necessary breaking anything on the outside, but it will still be a half-job if we did it that way

pablomlago · 2025-10-03T10:45:24Z

src/brevitas/graph/equalize.py

            weight.size(0), -1))[self.equalization_indexes.start:self.equalization_indexes.end]

+    def permute(self, permute_index):
+        permutation_list = []


Maybe I would opt for doing a list comprehension:

permutation_tuple = tuple(permute_index if i == dim else slice(size)for dim, size in enumerate(self.module.weight.shape))

Also, for indexing the tensor I would pass a tuple, instead of a list, to prevent the following warning:

<stdin>:1: UserWarning: Using a non-tuple sequence for multidimensional indexing is deprecated and will be changed in pytorch 2.9; use x[tuple(seq)] instead of x[seq]. In pytorch 2.9 this will be interpreted as tensor index, x[torch.tensor(seq)], which will result either in an error or a different result (Triggered internally at /pytorch/torch/csrc/autograd/python_variable_indexing.cpp:306.)

pablomlago · 2025-10-03T10:48:45Z

src/brevitas/graph/equalize.py

+class PermuteGraph():
+
+    def __init__(self):
+        super().__init__()


Is this needed? PermuteGraph does not seem to inherit from other class.

docsrc/source/user_guide/rotations.rst

nickfraser · 2025-10-07T16:51:24Z

src/brevitas/graph/equalize.py

One possibility would be to have a file with the base equalization functionality, and then other file with the code for scalar/rotation/permutation equalization, but other options could be explored.

Agreed, this would be better... Do we do that in this PR or in a future refactor of this section?

src/brevitas/graph/equalize.py

nickfraser · 2025-10-07T16:58:47Z

src/brevitas/graph/equalize.py

+    device = next(single_module.parameters()).device
+    dtype = next(single_module.parameters()).dtype
+
+    # If equalization criteria are not met, we return a scalar one to indicate that no equalization


Is this supposed to return a scalar 1?

nickfraser · 2025-10-07T17:03:44Z

src/brevitas/graph/equalize.py

+    if delay_rewriters:
+        return model
+
+    if not hasattr(model, '_hf_map'):


Let's move this to a method in accelerate_utils.

nickfraser · 2025-10-07T17:03:51Z

src/brevitas/graph/equalize.py

+    if offload_model is None or remove_hooks is None:
+        raise RuntimeError("Accelerate is not installed")
+    # if we use _hf_map to check and all the model is on a single GPU, then all rewriters are safe
+    if len(model._hf_map.values()) > 1:


Let's move this to a method in accelerate_utils.

nickfraser · 2025-10-07T17:04:37Z

tests/brevitas/graph/test_equalization.py

        equalized_layers.update(r[1])

    # Check that we found all the expected regions
+    print(len(regions))


Remove or use logger?

nickfraser · 2025-10-07T17:04:43Z

tests/brevitas/graph/test_equalization.py

        sinks = region.sinks_names
        sinks_check = set(sinks) == set(expected_region[1])
+        print(len(srcs), len(expected_region[0]))
+        print(srcs)


Remove or use logger?

src/brevitas/graph/equalize.py

pablomlago · 2025-10-08T09:15:41Z

src/brevitas/graph/equalize.py

        weight = weight.cpu().to(torch.float32)
        return scale_fn(weight.reshape(weight.size(0), -1))

+    def permute(self, permute_index):


This logic is repeated in SinkWrapper (excluding the bias permutation). I would factor out the common logic to ModuleWrapper (or even a standalone method) and then do the appropiate calls in Sink/SourceWrapper.

pablomlago · 2025-10-08T09:16:57Z

src/brevitas/graph/equalize.py

+        self.module.weight.data = self.module.weight.data[permutation_list]
+
+
+def new_axis(x, block_size=32):


I would add a docstring to this method, and potentially change its name.

pablomlago · 2025-10-08T09:18:49Z

src/brevitas/graph/equalize.py

+    # If act_val is enabled, use source or sink weights to determine the activation channel
+    # For example, if the source is BatchNorm, we need to use the information coming from the sinks
+    if not region.is_valid_activation_equalization:
+        return _no_permute()


Note that with the current implementation of _no_permute, this line is returning None.

src/brevitas/graph/equalize.py

pablomlago · 2025-10-08T09:39:06Z

src/brevitas/graph/equalize.py

+
+
+@torch.no_grad()
+def _permute(region, list_of_act_val):


I feel that at the moment there are too many functions related to permutations, which makes it difficult to follow: permute in both EqualizationSinkWrapper and EqualizationSourceWrapper, _permute as a standalone method and apply_permute for Region. From my point of view, we could either move the logic of _permute into apply_permute or viceversa. Also, I would consider refactoring the common logic of permute (EqualizationSink/SourceWrapper) into the base class.

pablomlago · 2025-10-08T09:42:06Z

src/brevitas/graph/equalize.py

+
+    # scale_fn = permute_op_type
+    single_module = region.get_module_from_name(next(iter(region.sinks_names)))
+    device = next(single_module.parameters()).device


Are device and dtype used in this method?

pablomlago · 2025-10-08T09:43:18Z

src/brevitas/graph/equalize.py

+    if not region.is_valid_activation_equalization:
+        return _no_permute()
+
+    list_of_act_val_shapes = [act_val.shape for act_val in list_of_act_val]


This piece of code does not seem self-explanatory. I would consider adding some comments to explain its functionality.

pablomlago · 2025-10-08T10:02:40Z

src/brevitas/graph/equalize.py


    # If act_val is enabled, use source or sink weights to determine the activation channel
    # For example, if the source is BatchNorm, we need to use the information coming from the sinks
    if list_of_act_val is not None:


There is a lot of logic around list_of_act_val scattered around this file. I feel that, we should stop using a dictionary for float_act_map and define a class for it to encapsulate its functionality. Probably, this should be done in a future PR.

pablomlago · 2025-10-08T10:08:42Z

src/brevitas/graph/equalize.py



+@torch.no_grad()
+def apply_rewriters(


If moving the accelerate-related functionality into equalize.py is not strictly needed to support permutations, I would leave this changes for a future PR, as we might need to discuss whether to introduce this dependency at src/brevitas level.

SpiritSeeker · 2025-10-22T07:00:09Z

src/brevitas/graph/equalize.py

+        self.model = model
+        self.rewriters = rewriters
+        self.rotation.permute_class.setup_permute()
+


Suggested change

return self

to use rewriters outside the context.

i-colbert reviewed Sep 25, 2025

View reviewed changes

src/brevitas/graph/equalize.py Outdated Show resolved Hide resolved

i-colbert reviewed Sep 25, 2025

View reviewed changes

src/brevitas/graph/equalize.py Outdated Show resolved Hide resolved

Giuseppe5 and others added 25 commits October 1, 2025 11:33

Tentative fix

a353cba

New change

dc8556b

Comment

7ffe6f7

Revert change

79932cf

Rename

e0526d4

Fix + comments

e47e91e

Typo

db6ca13

typo 2

1f46480

Fix condition

c375f3d

fix tests

655e51b

Update src/brevitas/graph/equalize.py

69b2354

Co-authored-by: Pablo Monteagudo Lago <44771380+pablomlago@users.noreply.github.com>

Docs

0d7136c

Merge branch 'refactor_rotation' of https://github.com/Giuseppe5/brev…

a5ad5af

…itas into refactor_rotation

Precommit

5ff9f07

Fix tests

2069773

Check accelerate installation

efc2873

Feat (graph/equalize): implement permute regions

2d1ea20

Fix

abf74c2

Indexes fix

d5085a8

More indexes fix

ed6e0d4

Correct return statement

fc9717b

Fix for region eq method and act_axis

d578792

Fix rotation test

3ebb858

Temp

15ad102

Cleanup

e0458a8

Giuseppe5 force-pushed the permutation branch from 8a84c01 to e0458a8 Compare October 2, 2025 09:41

Giuseppe5 added 2 commits October 2, 2025 11:28

New func

fa607c3

Fix tests

6b752db

Fix tests

86a3af8

pablomlago reviewed Oct 3, 2025

View reviewed changes

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

pablomlago reviewed Oct 3, 2025

View reviewed changes

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

pablomlago reviewed Oct 3, 2025

View reviewed changes

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

pablomlago reviewed Oct 3, 2025

View reviewed changes

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

pablomlago reviewed Oct 3, 2025

View reviewed changes

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

pablomlago reviewed Oct 3, 2025

View reviewed changes

nickfraser requested changes Oct 7, 2025

View reviewed changes

pablomlago reviewed Oct 8, 2025

View reviewed changes

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

pablomlago reviewed Oct 8, 2025

View reviewed changes

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

pablomlago reviewed Oct 8, 2025

View reviewed changes

SpiritSeeker reviewed Oct 22, 2025

View reviewed changes

Fix (ex/llm): Fix integration with Lighteval Python API (Xilinx#1379)

9df3e9c

nickfraser assigned Giuseppe5 Nov 3, 2025

	model, rewriters = self.rotation.apply(self.model)
	self.model, self.rewriters = self.rotation.apply(self.model)

		self.module.weight.data = self.module.weight.data[permutation_list]


		def new_axis(x, block_size=32):

Feat (graph/equalize): implement permute regions #1380

Are you sure you want to change the base?

Feat (graph/equalize): implement permute regions #1380

Uh oh!

Conversation

Giuseppe5 commented Sep 25, 2025

Reason for this PR

Changes Made in this PR

Testing Summary

Risk Highlight

Checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Giuseppe5 Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pablomlago Oct 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pablomlago Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pablomlago Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Giuseppe5 Oct 9, 2025 •

edited

Loading

pablomlago Oct 3, 2025 •

edited

Loading

pablomlago Oct 8, 2025 •

edited

Loading

pablomlago Oct 8, 2025 •

edited

Loading