import functools from tml.projects.twhin.models.config import TwhinModelConfig from tml.projects.twhin.models.models import TwhinModel from tml.optimizers.optimizer import get_optimizer_class, LRShim from tml.optimizers.config import get_optimizer_algorithm_config, LearningRate from tml.ml_logging.torch_logging import logging from torchrec.optim.optimizers import in_backward_optimizer_filter from torchrec.optim import keyed FUSED_OPT_KEY = "fused_opt" TRANSLATION_OPT_KEY = "operator_opt" def _lr_from_config(optimizer_config): if optimizer_config.learning_rate is not None: return optimizer_config.learning_rate else: # treat None as constant lr lr_value = get_optimizer_algorithm_config(optimizer_config).lr return LearningRate(constant=lr_value) def build_optimizer(model: TwhinModel, config: TwhinModelConfig): """Builds an optimizer for a Twhin model combining the embeddings optimizer with an optimizer for per-relation translations. Args: model: TwhinModel to build optimizer for. config: TwhinConfig for model. Returns: Optimizer for model. """ translation_optimizer_fn = functools.partial( get_optimizer_class(config.translation_optimizer), **get_optimizer_algorithm_config(config.translation_optimizer).dict(), ) translation_optimizer = keyed.KeyedOptimizerWrapper( dict(in_backward_optimizer_filter(model.named_parameters())), optim_factory=translation_optimizer_fn, ) lr_dict = {} for table in config.embeddings.tables: lr_dict[table.name] = _lr_from_config(table.optimizer) lr_dict[TRANSLATION_OPT_KEY] = _lr_from_config(config.translation_optimizer) logging.info(f"***** LR dict: {lr_dict} *****") logging.info( f"***** Combining fused optimizer {model.fused_optimizer} with operator optimizer: {translation_optimizer} *****" ) optimizer = keyed.CombinedOptimizer( [ (FUSED_OPT_KEY, model.fused_optimizer), (TRANSLATION_OPT_KEY, translation_optimizer), ] ) # scheduler = LRShim(optimizer, lr_dict) scheduler = None logging.info(f"***** Combined optimizer after init: {optimizer} *****") return optimizer, scheduler