From c3d819cc8bfcb97f8909c36196ec01721d460d19 Mon Sep 17 00:00:00 2001 From: Tanisha Chawada Date: Mon, 24 Nov 2025 17:34:41 +0000 Subject: [PATCH 1/3] Adding optimizer registry and its test cases Signed-off-by: Tanisha Chawada --- .../finetune/experimental/core/optimizer.py | 22 +++++++ .../experimental/tests/test_optimizer.py | 62 +++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 QEfficient/finetune/experimental/tests/test_optimizer.py diff --git a/QEfficient/finetune/experimental/core/optimizer.py b/QEfficient/finetune/experimental/core/optimizer.py index d647b73a6..2304a1525 100644 --- a/QEfficient/finetune/experimental/core/optimizer.py +++ b/QEfficient/finetune/experimental/core/optimizer.py @@ -4,3 +4,25 @@ # SPDX-License-Identifier: BSD-3-Clause # # ----------------------------------------------------------------------------- + +""" +Optimizer components for the training system. +""" + +from typing import Type + +import torch.optim as optim +from torch.optim import Optimizer + +from QEfficient.finetune.experimental.core.component_registry import registry + +registry.optimizer("adam")(optim.Adam) +registry.optimizer("adamw")(optim.AdamW) +registry.optimizer("sgd")(optim.SGD) + + +def get_optimizer_cls(optimizer_name: str) -> Type[Optimizer]: + optimizer_cls = registry.get_optimizer(optimizer_name) + if optimizer_cls is None: + raise ValueError(f"Unknown optimizer: {optimizer_name}") + return optimizer_cls diff --git a/QEfficient/finetune/experimental/tests/test_optimizer.py b/QEfficient/finetune/experimental/tests/test_optimizer.py new file mode 100644 index 000000000..b1f6f82be --- /dev/null +++ b/QEfficient/finetune/experimental/tests/test_optimizer.py @@ -0,0 +1,62 @@ +# ----------------------------------------------------------------------------- +# +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +# SPDX-License-Identifier: BSD-3-Clause +# +# ----------------------------------------------------------------------------- + +import sys +from pathlib import Path + +import pytest +import torch.optim as optim + +from QEfficient import QEFFAutoModelForCausalLM +from QEfficient.finetune.experimental.core.component_registry import ComponentFactory + +sys.path.insert(0, str(Path(__file__).parent.parent)) +OPTIMIZER_CONFIGS = { + "adam": { + "name": "adam", + "opt_cls": optim.Adam, + "lr": 1e-4, + "weight_decay": 0.01, + "betas": (0.9, 0.999), + "eps": 1e-8, + "amsgrad": False, + }, + "adamw": { + "name": "adamw", + "opt_cls": optim.AdamW, + "lr": 1e-4, + "weight_decay": 0.01, + "betas": (0.9, 0.999), + "eps": 1e-8, + "amsgrad": False, + }, + "sgd": { + "name": "sgd", + "opt_cls": optim.SGD, + "lr": 1e-4, + "momentum": 0.9, + "weight_decay": 0.01, + "dampening": 0.0, + "nesterov": False, + }, +} + + +@pytest.fixture +def ref_model(): + return QEFFAutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B") + + +@pytest.mark.parametrize("opt_name", OPTIMIZER_CONFIGS.keys()) +def test_optimizers(opt_name, ref_model): + """Test that all optimizers can be created with their configs.""" + # Create optimizer using the factory + config = OPTIMIZER_CONFIGS[opt_name] + opt_inst = ComponentFactory.create_optimizer(**config, model_params=ref_model.model.parameters()) + assert opt_inst is not None + assert isinstance(opt_inst, optim.Optimizer) + assert len(list(opt_inst.param_groups)) == 1 From 9e402bcfbaa69280d598eb95aa129dc09209a232 Mon Sep 17 00:00:00 2001 From: Tanisha Chawada Date: Thu, 27 Nov 2025 07:00:25 +0000 Subject: [PATCH 2/3] Adding optimizer registry and its test cases Signed-off-by: Tanisha Chawada --- .../finetune/experimental/core/optimizer.py | 12 +++-- .../experimental/tests/test_optimizer.py | 47 ++++++++++++++----- 2 files changed, 45 insertions(+), 14 deletions(-) diff --git a/QEfficient/finetune/experimental/core/optimizer.py b/QEfficient/finetune/experimental/core/optimizer.py index 2304a1525..64c110ec5 100644 --- a/QEfficient/finetune/experimental/core/optimizer.py +++ b/QEfficient/finetune/experimental/core/optimizer.py @@ -16,9 +16,15 @@ from QEfficient.finetune.experimental.core.component_registry import registry -registry.optimizer("adam")(optim.Adam) -registry.optimizer("adamw")(optim.AdamW) -registry.optimizer("sgd")(optim.SGD) + +def register_optimizer(optimizer_name: str, cls: Type[Optimizer]) -> None: + """Register a new optimizer class.""" + registry.optimizer(optimizer_name)(cls) + + +register_optimizer("adam", optim.Adam) +register_optimizer("adamw", optim.AdamW) +register_optimizer("sgd", optim.SGD) def get_optimizer_cls(optimizer_name: str) -> Type[Optimizer]: diff --git a/QEfficient/finetune/experimental/tests/test_optimizer.py b/QEfficient/finetune/experimental/tests/test_optimizer.py index b1f6f82be..d84a2a524 100644 --- a/QEfficient/finetune/experimental/tests/test_optimizer.py +++ b/QEfficient/finetune/experimental/tests/test_optimizer.py @@ -5,19 +5,20 @@ # # ----------------------------------------------------------------------------- +import inspect import sys from pathlib import Path import pytest +import torch.nn as nn import torch.optim as optim -from QEfficient import QEFFAutoModelForCausalLM -from QEfficient.finetune.experimental.core.component_registry import ComponentFactory +from QEfficient.finetune.experimental.core.optimizer import get_optimizer_cls, register_optimizer sys.path.insert(0, str(Path(__file__).parent.parent)) OPTIMIZER_CONFIGS = { "adam": { - "name": "adam", + "optimizer_name": "adam", "opt_cls": optim.Adam, "lr": 1e-4, "weight_decay": 0.01, @@ -26,7 +27,7 @@ "amsgrad": False, }, "adamw": { - "name": "adamw", + "optimizer_name": "adamw", "opt_cls": optim.AdamW, "lr": 1e-4, "weight_decay": 0.01, @@ -35,7 +36,7 @@ "amsgrad": False, }, "sgd": { - "name": "sgd", + "optimizer_name": "sgd", "opt_cls": optim.SGD, "lr": 1e-4, "momentum": 0.9, @@ -47,16 +48,40 @@ @pytest.fixture -def ref_model(): - return QEFFAutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B") +def dummy_model(): + return nn.Sequential( + nn.Linear(10, 5), + nn.ReLU(), + nn.Linear(5, 1), + ) @pytest.mark.parametrize("opt_name", OPTIMIZER_CONFIGS.keys()) -def test_optimizers(opt_name, ref_model): +def test_optimizers(opt_name, dummy_model): """Test that all optimizers can be created with their configs.""" - # Create optimizer using the factory + # Register optimizer class config = OPTIMIZER_CONFIGS[opt_name] - opt_inst = ComponentFactory.create_optimizer(**config, model_params=ref_model.model.parameters()) - assert opt_inst is not None + register_optimizer(config["optimizer_name"], config["opt_cls"]) + optimizer_class = get_optimizer_cls(config["optimizer_name"]) + assert optimizer_class is not None + assert optimizer_class == config["opt_cls"] + valid_params = inspect.signature(optimizer_class).parameters + filtered_config = {k: v for k, v in config.items() if k in valid_params} + opt_inst = optimizer_class(dummy_model.parameters(), **filtered_config) assert isinstance(opt_inst, optim.Optimizer) assert len(list(opt_inst.param_groups)) == 1 + assert opt_inst.param_groups[0]["lr"] == config["lr"] + if "weight_decay" in config: + assert opt_inst.param_groups[0]["weight_decay"] == config["weight_decay"] + if "betas" in config: + assert opt_inst.param_groups[0]["betas"] == config["betas"] + if "eps" in config: + assert opt_inst.param_groups[0]["eps"] == config["eps"] + if "momentum" in config: + assert opt_inst.param_groups[0]["momentum"] == config["momentum"] + if "dampening" in config: + assert opt_inst.param_groups[0]["dampening"] == config["dampening"] + if "nesterov" in config: + assert opt_inst.param_groups[0]["nesterov"] == config["nesterov"] + if "amsgrad" in config: + assert opt_inst.param_groups[0]["amsgrad"] == config["amsgrad"] From 34f15c4f96aa97eaa392e6c6bbcc9e275bdcc881 Mon Sep 17 00:00:00 2001 From: Tanisha Chawada Date: Fri, 28 Nov 2025 06:00:44 +0000 Subject: [PATCH 3/3] Adding optimizer registry and its test cases Signed-off-by: Tanisha Chawada --- QEfficient/finetune/experimental/core/optimizer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/QEfficient/finetune/experimental/core/optimizer.py b/QEfficient/finetune/experimental/core/optimizer.py index 64c110ec5..de28848ae 100644 --- a/QEfficient/finetune/experimental/core/optimizer.py +++ b/QEfficient/finetune/experimental/core/optimizer.py @@ -32,3 +32,11 @@ def get_optimizer_cls(optimizer_name: str) -> Type[Optimizer]: if optimizer_cls is None: raise ValueError(f"Unknown optimizer: {optimizer_name}") return optimizer_cls + + +def get_optimizer(opt_config): + opt_name = opt_config.pop("optimizer_name") + opt_cls = get_optimizer_cls(opt_name) + opt_config["lr"] = float(opt_config["lr"]) + optimizer_cls_and_kwargs = (opt_cls, opt_config) + return optimizer_cls_and_kwargs