Feat (ex/common): learned round refactor #1323

pablomlago · 2025-06-23T09:51:02Z

TBA

Giuseppe5 · 2025-10-14T08:33:38Z

Is there a chance to split this into more manageable chunks?

Giuseppe5

First review

Giuseppe5 · 2026-01-13T10:03:39Z

src/brevitas_examples/common/learned_round/learned_round_args.py

+from brevitas_examples.common.learned_round.learned_round_method import MSELoss
+from brevitas_examples.common.learned_round.learned_round_method import RegularisedMSELoss
+
+OPTIMIZER_MAP = {


Maybe we should implement a way to register new and/or custom optimizer to this map?

I opted for removing this dict, and having the optimizers retrieved from either torch.optim or brevitas.optim.

Giuseppe5 · 2026-01-13T10:07:15Z

src/brevitas_examples/common/learned_round/learned_round_args.py

+
+OPTIMIZER_MAP = {
+    "sign_sgd": SignSGD,}
+LR_SCHEDULER_MAP = {}


Same as above

I opted for removing this dict since, at the moment, there is no plan to introduce a custom one. Moreover, this could be registered in brevitas.optim.lr_scheduler and use a logic similar as the one for the optimizers.

Giuseppe5 · 2026-01-13T10:15:26Z

src/brevitas_examples/common/learned_round/learned_round_args.py

+
+    @property
+    def param_fn(self) -> Callable[[nn.Module, OrderedDict, str], bool]:
+        return {


Bit ugly, we can have a separate attribute/property with the dict, and this just return self.pre_dict[self.value]

I created a Registry class in python_utils.py, and I use that for tracking/retrieving the different implementations.

Giuseppe5 · 2026-01-13T10:15:33Z

src/brevitas_examples/common/learned_round/learned_round_args.py

+
+    @property
+    def loss_class(self) -> Type[BlockLoss]:
+        return {


Giuseppe5 · 2026-01-13T10:19:08Z

src/brevitas_examples/common/learned_round/learned_round_args.py

+
+# Both `get_round_parameters` and `get_scale_parameters` are meant to be passed as the argument `get_target`
+# of `_get_target_parameters`, which iterates over the modules of a model in a recursive function.
+# In the case of `get_round_parameters` the return value indicates whether the submodules of a given module


You say:

In the case of get_round_parameters

as if it is different in get_scale_parameters but it is not really the case (plus/minus some very small differences), but overall the comment can be adjusted (and shortened)

I polished it a bit to make it clearer that it is an example.

Giuseppe5 · 2026-01-13T10:38:53Z

src/brevitas_examples/common/learned_round/learned_round_optimizer.py

        raise StopFwdException


+def get_blocks(model: nn.Module, block_check_fn: Callable[[nn.Module, str],


I believe we might have something similar to this for quantization

Giuseppe5 · 2026-01-13T10:58:10Z

src/brevitas_examples/common/learned_round/learned_round_optimizer.py

-                optimizer.zero_grad()
-        for lr_scheduler in lr_schedulers:
-            if lr_scheduler:
+class LearnedRoundOptimizer:


This should be a Trainer class of some sorts

Giuseppe5 · 2026-01-13T11:02:12Z

src/brevitas_examples/common/learned_round/learned_round_args.py

+
+
+@dataclass
+class OptimizerArgs:


Split optimizer/scheduler

I'm not sure if the case is strong enough for that, a lr scheduler instance is semantically tied to an optimizer, so it is sensible for the dataclass hierarchy to reflect this.

Giuseppe5 · 2026-01-13T11:02:50Z

src/brevitas_examples/common/learned_round/learned_round_optimizer.py

+            state_dict = self._get_target_parameters(model, optimizer_target.param_fn, state_dict)
+        return state_dict
+
+    def _create_optimizer_and_scheduler(


split optimizer/scheduler

Same as above.

Giuseppe5 · 2026-01-13T11:03:04Z

src/brevitas_examples/common/learned_round/learned_round_optimizer.py

+            if lr_scheduler_args is not None else None)
+        return optimizer, lr_scheduler
+
+    def _create_optimizers_lr_schedulers(


name of the function is confusingly similar with the one above

I opted for removing it, after the changes to the optimizer instantiation.

pablomlago force-pushed the feat-round-refactor branch from 24a1c36 to 224d626 Compare June 24, 2025 08:39

pablomlago force-pushed the feat-round-refactor branch from 224d626 to 577907e Compare September 10, 2025 08:53

nickfraser mentioned this pull request Sep 29, 2025

v0.13.0 Minor Release #1290

Open

28 tasks

pablomlago marked this pull request as ready for review September 30, 2025 13:51

nickfraser added the next release PRs which should be merged for the next release label Oct 13, 2025

nickfraser assigned pablomlago Oct 13, 2025

nickfraser requested a review from Giuseppe5 December 1, 2025 14:37

Giuseppe5 requested changes Jan 13, 2026

View reviewed changes

pablomlago added 21 commits January 19, 2026 09:48

Args and cache

3cd9936

Working

6879e53

Refactor training loop

647d596

Refactor optimizers

4c92234

Refactor block preparation

1753c04

Cache refactoring

3685304

Remove unused methods

f1b9982

Cleanup

ccf5de2

Remove unused imports

2118f12

Minor refactoring

9230fb5

amp refactor

53e80af

Improved type hints

5bb7372

Refactor argument logic

827190b

Remove notebook

a414055

Remove maps

c032988

Remove deprecated grad scaler

b335acf

Minor fixs

50dd00f

Clean-up

994543c

Remove dacite dependency

800c06a

Address review comments

9b0247d

Tie parameters to optimizer

f580604

pablomlago force-pushed the feat-round-refactor branch from b700947 to f580604 Compare January 20, 2026 11:08

pablomlago added 2 commits January 20, 2026 11:33

Fix incompatiblity with older PT versions

116ecaa

Use literal type hint for GradScaler

513c0b9

		raise StopFwdException


		def get_blocks(model: nn.Module, block_check_fn: Callable[[nn.Module, str],

Feat (ex/common): learned round refactor #1323

Are you sure you want to change the base?

Feat (ex/common): learned round refactor #1323

Uh oh!

Conversation

pablomlago commented Jun 23, 2025

Uh oh!

Giuseppe5 commented Oct 14, 2025

Uh oh!

Giuseppe5 left a comment

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!

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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants