Feat (hadamard): support region expansion #1178

Giuseppe5 · 2025-02-11T10:59:18Z

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

Giuseppe5 · 2025-02-11T11:10:56Z

src/brevitas/graph/equalize.py

            # Only "weight" is rotated
-            tensor_names_axis = [("weight", _get_input_axis(module))]
+            if region.expand_region:
+                assert isinstance(module, nn.Linear), "Currently only Linear layers support expanded hadamard"


The code just below can support any layer type, the issue is more on the activation side since we need to define also the dimension along which the input needs padding. It is not complicated to do tbh, but not necessary atm

Giuseppe5 · 2025-02-11T11:12:35Z

src/brevitas/graph/equalize.py

+            model: nn.Module,
+            regions: List[Region],
+            prefix: str = '',
+            blacklist_layers: Optional[List[str]] = None):


The exclusion list needs to be passed since we want to exclude all the layers that might be expanded

src/brevitas/graph/hadamard.py

src/brevitas/nn/equalized_layer.py

Giuseppe5 · 2025-02-11T11:15:06Z

src/brevitas_examples/llm/main.py

@@ -350,20 +351,26 @@ def quantize_llm(args, extra_args=None):
            model(**calibration_loader[0])
        remove_hooks(model)

+    layers_to_expand = []


The idea is that if the user passes q_proj, this piece of code will catch all the layers layer.module...q_proj

src/brevitas_examples/llm/main.py

src/brevitas/graph/equalize.py

nickfraser · 2025-02-11T11:14:23Z

src/brevitas/graph/equalize.py

            # Only "weight" is rotated
-            tensor_names_axis = [("weight", _get_input_axis(module))]
+            if region.expand_region:
+                assert isinstance(module, nn.Linear), "Currently only Linear layers support expanded hadamard"


Save convs for a new PR?

I think so, yeah. I'll open an issue to keep track of that as good first issue

nickfraser · 2025-02-11T11:15:47Z

src/brevitas/graph/equalize.py

+            if region.expand_region:
+                assert isinstance(module, nn.Linear), "Currently only Linear layers support expanded hadamard"
+                hidden_dim = module.weight.shape[1]
+                new_hidden = find_closest_hadamard_number(hidden_dim, hidden_dim + 1)


+1 here? Maybe the previous comment should live here?

nickfraser · 2025-02-11T11:19:14Z

src/brevitas/nn/equalized_layer.py

+    def input_pad(self, inp):
+        # TODO: This only works for Linear layers. We have an assert in equalize.py to check for this
+        hidden_dim = inp.shape[-1]
+        new_hidden = find_closest_hadamard_number(hidden_dim, hidden_dim + 1)


Similar here - the +1 hardcoding...

src/brevitas/nn/equalized_layer.py

src/brevitas_examples/llm/main.py

nickfraser

Only one note about verbosity flags. I guess this can also be achieved by setting the log level higher globally.

nickfraser · 2025-02-11T19:14:41Z

src/brevitas/utils/logging.py

@@ -12,6 +13,25 @@
 from brevitas.utils.quant_utils import *


+def setup_logger(name):


nickfraser · 2025-02-11T19:16:36Z

src/brevitas/graph/equalize.py

+            parameter_number_pre = 0
+            for m in graph_model.parameters():
+                parameter_number_pre += m.numel()
+            logging.info(f"{len(expanded_regions)} layers will be expanded during rotation")


Personally I think all logging should be protected with a verbose flag. Sometimes we will want to turn these things off (for some specific experiments)

I will go the pytorch way and have an env variable to set the level.
Default will be "INFO" and user can select DEBUG with BREVITAS_LOGGING=DEBUG

nickfraser · 2025-02-11T19:20:27Z

src/brevitas/utils/logging.py

+def setup_logger(name):
+
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)


Maybe have the logging level as an optional parameter? You might want to turn off all 'INFO' messages at some point.

See the comment above. I think this piece of code will change a lot as we figure out what to do/how to do it, I want to have something reasonable to start and I believe your comment/my proposal above should accomplish just that

nickfraser

LGTM!

nickfraser · 2025-02-12T15:05:12Z

src/brevitas/config.py

@@ -20,6 +20,7 @@ def env_to_bool(name, default):
 JIT_ENABLED = env_to_bool('BREVITAS_JIT', False) and _enabled
 NATIVE_STE_BACKEND_ENABLED = env_to_bool('BREVITAS_NATIVE_STE_BACKEND', False)
 VERBOSE = env_to_bool('BREVITAS_VERBOSE', False)
+LOGGING_LEVEL = os.environ.get('BREVITAS_LOGGING', 'CRITICAL')


Love it! <3

Feat (hadamard): support region expansion

5dd0db9

Giuseppe5 force-pushed the exp_had branch from 2b33f6d to 5dd0db9 Compare February 11, 2025 11:02

Giuseppe5 commented Feb 11, 2025

View reviewed changes

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

Giuseppe5 commented Feb 11, 2025

View reviewed changes

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

Giuseppe5 commented Feb 11, 2025

View reviewed changes

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

Giuseppe5 commented Feb 11, 2025

View reviewed changes

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

Giuseppe5 commented Feb 11, 2025

View reviewed changes

src/brevitas/nn/equalized_layer.py Outdated Show resolved Hide resolved

remove indentation

1a7389f

Giuseppe5 commented Feb 11, 2025

View reviewed changes

src/brevitas_examples/llm/main.py Outdated Show resolved Hide resolved

remove rebase conflict

fb50b96

nickfraser requested changes Feb 11, 2025

View reviewed changes

Giuseppe5 added 5 commits February 11, 2025 13:29

fix

c29f94a

Review

447d4cd

Last fixes

54da318

Fix llm tests

e249c0f

Add logging

55580c7

nickfraser approved these changes Feb 11, 2025

View reviewed changes

nickfraser reviewed Feb 11, 2025

View reviewed changes

Optional logging

72bc6a0

nickfraser approved these changes Feb 12, 2025

View reviewed changes

Giuseppe5 requested a review from nickfraser February 12, 2025 15:34

Update equalize.py

3d4e600

Giuseppe5 merged commit 15135bc into Xilinx:dev Feb 12, 2025
23 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feat (hadamard): support region expansion #1178

Feat (hadamard): support region expansion #1178

Giuseppe5 commented Feb 11, 2025

Giuseppe5 Feb 11, 2025

Giuseppe5 Feb 11, 2025

Giuseppe5 Feb 11, 2025

nickfraser Feb 11, 2025

Giuseppe5 Feb 11, 2025

nickfraser Feb 11, 2025

nickfraser Feb 11, 2025

nickfraser left a comment

nickfraser Feb 11, 2025

nickfraser Feb 11, 2025

Giuseppe5 Feb 12, 2025

nickfraser Feb 11, 2025

Giuseppe5 Feb 12, 2025 •

edited

Loading

nickfraser left a comment

nickfraser Feb 12, 2025

		@@ -12,6 +13,25 @@
		from brevitas.utils.quant_utils import *


		def setup_logger(name):

Feat (hadamard): support region expansion #1178

Feat (hadamard): support region expansion #1178

Conversation

Giuseppe5 commented Feb 11, 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

nickfraser left a 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

Giuseppe5 Feb 12, 2025 • edited Loading

Choose a reason for hiding this comment

nickfraser left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Giuseppe5 Feb 12, 2025 •

edited

Loading