improve user experience with additional lr scheduler configuration in…

…spection (using an allowlist approach) and enhanced documentation. Expand use of allow_untested to allow use of unsupported/untested schedulers

docs/source/advanced/lr_scheduler_reinitialization.rst

@@ -30,12 +30,8 @@ annotated yaml schedule below and can be explored using the
 :ref:`advanced usage example<advanced-fine-tuning-lr-example>`.
 
 When specifying an LR scheduler configuration for a given phase, the ``new_lr_scheduler`` dictionary requires at minimum
-an ``lr_scheduler_init`` dictionary containing a ``class_path`` key indicating the class of the lr scheduler to be
-instantiated and wrapped around your optimizer. Currently,
-all :class:`~pytorch_lightning.utilities.types.LRSchedulerType` s are supported with the exception of
-:external+torch:class:`~torch.optim.lr_scheduler.ChainedScheduler` and
-:external+torch:class:`~torch.optim.lr_scheduler.SequentialLR` (due to the configuration complexity and semantic
-conflicts supporting them would introduce).
+an ``lr_scheduler_init`` dictionary containing a ``class_path`` key indicating the class of the lr scheduler
+(:ref:`list of supported schedulers<supported_lr_schedulers>`) to be instantiated and wrapped around your optimizer.
 
 Any arguments you would like to pass to initialize the specified lr scheduler with should be specified in the
 ``init_args`` key of the ``lr_scheduler_init`` dictionary.
@@ -115,8 +111,9 @@ sanity-checked prior to training initiation.
     introspected/simulated in the current :class:`~finetuning_scheduler.fts.FinetuningScheduler` version.
 
 Note that specifying LR scheduler reinitialization configurations is only supported for phases >= ``1``. This is because
-for fine-tuning phase ``0``, the LR scheduler configuration will be the scheduler that you initiate your training session
-with, usually via the ``configure_optimizer`` method of :external+pl:class:`~pytorch_lightning.core.module.LightningModule`.
+for fine-tuning phase ``0``, the LR scheduler configuration will be the scheduler that you initiate your training
+session with, usually via the ``configure_optimizer`` method of
+:external+pl:class:`~pytorch_lightning.core.module.LightningModule`.
 
 .. tip::
 

docs/source/index.rst

@@ -266,6 +266,7 @@ There are plenty of options for customizing
 :ref:`scheduled fine-tuning for SuperGLUE<scheduled-fine-tuning-superglue>` below for examples of composing different
 configurations.
 
+.. _supported_strategies:
 
 .. note::
    Currently, :class:`~finetuning_scheduler.fts.FinetuningScheduler` supports the following
@@ -280,17 +281,44 @@ configurations.
       * :external+pl:class:`~pytorch_lightning.strategies.sharded_spawn.DDPSpawnShardedStrategy`
       * :external+pl:class:`~pytorch_lightning.strategies.dp.DataParallelStrategy`
 
+.. _supported_lr_schedulers:
+
+.. note::
+   Currently, :class:`~finetuning_scheduler.fts.FinetuningScheduler` officially supports the following torch lr
+   schedulers:
+
+   .. hlist::
+      :columns: 2
+
+      * :external+torch:class:`~torch.optim.lr_scheduler.StepLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.MultiStepLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.CosineAnnealingWarmRestarts`
+      * :external+torch:class:`~torch.optim.lr_scheduler.ReduceLROnPlateau`
+      * :external+torch:class:`~torch.optim.lr_scheduler.LambdaLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.ConstantLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.LinearLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.ExponentialLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.CosineAnnealingLR`
+      * :external+torch:class:`~torch.optim.lr_scheduler.MultiplicativeLR`
+
 .. tip::
-    Custom or officially unsupported strategies can be used by setting
+    Custom or officially unsupported strategies and lr schedulers can be used by setting
     :paramref:`~finetuning_scheduler.fts.FinetuningScheduler.allow_untested` to ``True``.
 
     Some officially unsupported strategies may work unaltered and are only unsupported due to
-    the ``Fine-Tuning Scheduler`` project's lack of CI/testing resources for that strategy (e.g.
-    ``single_tpu``).
-
-    Most unsupported strategies, however, are currently unsupported because they require varying degrees of modification
-    to be compatible (e.g. ``deepspeed`` requires an ``add_param_group`` method, ``tpu_spawn`` an override of the
-    current broadcast method to include python objects).
+    the Fine-Tuning Scheduler project's lack of CI/testing resources for that strategy (e.g. ``single_tpu``). Most
+    unsupported strategies and schedulers, however, are currently unsupported because they require varying degrees of
+    modification to be compatible.
+
+    For instance, with respect to strategies, ``deepspeed`` requires an
+    ``add_param_group`` method, ``tpu_spawn`` an override of the current broadcast method to include python objects.
+
+    Regarding lr schedulers, :external+torch:class:`~torch.optim.lr_scheduler.ChainedScheduler` and
+    :external+torch:class:`~torch.optim.lr_scheduler.SequentialLR` are examples of schedulers not currently supported
+    due to the configuration complexity and semantic conflicts supporting them would introduce. If a supported torch lr
+    scheduler does not meet your requirements, one can always subclass a supported lr scheduler and modify it as
+    required (e.g. :external+torch:class:`~torch.optim.lr_scheduler.LambdaLR` is especially useful for this). PRs are
+    also always welcome!
 
 ----------
 

src/finetuning_scheduler/fts.py

@@ -118,12 +118,11 @@ def __init__(
                 :class:`~finetuning_scheduler.fts_supporters.FTSEarlyStopping` criteria only.
                 epoch_transitions_only defaults to ``False``.
             reinit_lr_cfg: A lr scheduler reinitialization configuration dictionary consisting of at minimum a nested
-                ``lr_scheduler_init`` dictionary with a ``class_path`` key specifying the class of the
-                :class:`~pytorch_lightning.utilities.types.LRSchedulerType` to be instantiated. Optionally, an
-                ``init_args`` dictionary of arguments to initialize the lr scheduler with may be included. Additionally,
-                one may optionally include arguments to pass to PyTorch Lightning's lr scheduler configuration
-                :class:`~pytorch_lightning.utilities.types.LRSchedulerConfig` in the ``pl_lrs_cfg`` dictionary. By way
-                of example, one could configure this dictionary via the
+                ``lr_scheduler_init`` dictionary with a ``class_path`` key specifying the class of the lr scheduler
+                to be instantiated. Optionally, an ``init_args`` dictionary of arguments to initialize the lr scheduler
+                with may be included. Additionally, one may optionally include arguments to pass to PyTorch Lightning's
+                lr scheduler configuration :class:`~pytorch_lightning.utilities.types.LRSchedulerConfig` in the
+                ``pl_lrs_cfg`` dictionary. By way of example, one could configure this dictionary via the
                 :external+pl:class:`~pytorch_lightning.utilities.cli.LightningCLI` with the following:
 
                 .. code-block:: yaml
@@ -139,19 +138,28 @@ def __init__(
                                 frequency: 1
                                 name: Implicit_Reinit_LR_Scheduler
 
-            allow_untested: If ``True``, allows the use of custom or unsupported training strategies (e.g.
-                ``single_tpu``, ``MyCustomStrategy``). Defaults to ``False``.
+            allow_untested: If ``True``, allows the use of custom or unsupported training strategies and lr schedulers
+                (e.g. ``single_tpu``, ``MyCustomStrategy``, ``MyCustomLRScheduler``) . Defaults to ``False``.
 
-                .. note:: Custom or officially unsupported strategies can be used by setting
+                .. note:: Custom or officially unsupported strategies and lr schedulers can be used by setting
                     :paramref:`~finetuning_scheduler.fts.FinetuningScheduler.allow_untested` to ``True``.
 
                     Some officially unsupported strategies may work unaltered and are only unsupported due to
                     the ``Fine-Tuning Scheduler`` project's lack of CI/testing resources for that strategy (e.g.
                     ``single_tpu``).
 
-                    Most unsupported strategies, however, are currently unsupported because they require varying degrees
-                    of modification to be compatible (e.g. ``deepspeed`` requires an ``add_param_group`` method,
-                    ``tpu_spawn`` an override of the current broadcast method to include python objects).
+                    Most unsupported strategies and schedulers, however, are currently unsupported because they require
+                    varying degrees of modification to be compatible.
+
+                    For instance, with respect to strategies, ``deepspeed`` requires an ``add_param_group`` method,
+                    ``tpu_spawn`` an override of the current broadcast method to include python objects.
+
+                    Regarding lr schedulers, :external+torch:class:`~torch.optim.lr_scheduler.ChainedScheduler` and
+                    :external+torch:class:`~torch.optim.lr_scheduler.SequentialLR` are examples of schedulers not
+                    currently supported due to the configuration complexity and semantic conflicts supporting them would
+                    introduce. If a supported torch lr scheduler does not meet your requirements, one can always
+                    subclass a supported lr scheduler and modify it as required
+                    (e.g. :external+torch:class:`~torch.optim.lr_scheduler.LambdaLR` is especially useful for this).
             apply_lambdas_new_pgs: If ``True``, applies most recent lambda in ``lr_lambdas`` list to newly added
                 optimizer groups for lr schedulers that have a ``lr_lambdas`` attribute. Note this option only applies
                 to phases without reinitialized lr schedulers. Phases with defined lr scheduler reinitialization configs
@@ -427,6 +435,7 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
         """
         if len(trainer.optimizers) > 1:
             raise MisconfigurationException("fts currently only supports a single-optimizer configuration")
+        self._is_supported_lr(type(trainer.lr_scheduler_configs[0].scheduler))
         super().on_fit_start(trainer, pl_module)
 
     def state_dict(self) -> Dict[str, Any]: