Added LR schedulers (#16)

Author: spietras

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/__init__.py

@@ -1,46 +1,65 @@
-from kilroy_module_pytorch_py_sdk.resources import (
-    resource,
-    resource_bytes,
-    resource_text,
-)
 from kilroy_module_server_py_sdk import *
-from kilroy_module_pytorch_py_sdk.generator import Generator, GenerationResult
+from kilroy_module_pytorch_py_sdk.codec import Codec
+from kilroy_module_pytorch_py_sdk.generator import GenerationResult, Generator
+from kilroy_module_pytorch_py_sdk.models import LanguageModel, RewardModel
 from kilroy_module_pytorch_py_sdk.modules.basic import (
     BasicModule,
+    MetricsState as BasicModuleMetricsState,
+    ReportsState as BasicModuleReportsState,
     State as BasicModuleState,
 )
 from kilroy_module_pytorch_py_sdk.modules.reward import (
+    LanguageModelState as RewardModelModuleLanguageModelState,
+    MetricsState as RewardModelModuleMetricsState,
+    ReportsState as RewardModelModuleReportsState,
     RewardModelModule,
+    RewardModelState as RewardModelModuleRewardModelState,
     State as RewardModelModuleState,
 )
 from kilroy_module_pytorch_py_sdk.optimizers import (
-    Optimizer,
     AdamOptimizer,
-    SGDOptimizer,
+    Optimizer,
     RMSPropOptimizer,
+    SGDOptimizer,
+)
+from kilroy_module_pytorch_py_sdk.resources import (
+    resource,
+    resource_bytes,
+    resource_text,
 )
 from kilroy_module_pytorch_py_sdk.samplers import (
-    Sampler,
-    ProportionalSampler,
+    EpsilonNucleusSampler,
     EpsilonProportionalSampler,
-    TopKSampler,
     EpsilonTopKSampler,
     NucleusSampler,
-    EpsilonNucleusSampler,
+    ProportionalSampler,
+    Sampler,
+    TopKSampler,
+)
+from kilroy_module_pytorch_py_sdk.schedulers import (
+    ConstantScheduler,
+    CosineAnnealingScheduler,
+    CyclicScheduler,
+    ExponentialScheduler,
+    LinearScheduler,
+    MultiStepScheduler,
+    OneCycleScheduler,
+    ReduceOnPlateauScheduler,
+    Scheduler,
+    StepScheduler,
+    WarmRestartsScheduler,
 )
-from kilroy_module_pytorch_py_sdk.codec import Codec
-from kilroy_module_pytorch_py_sdk.models import LanguageModel, RewardModel
 from kilroy_module_pytorch_py_sdk.tokenizer import Tokenizer
 from kilroy_module_pytorch_py_sdk.utils import (
+    freeze,
+    pack_list,
+    pack_padded,
+    pad,
     slice_sequences,
+    squash_packed,
     truncate_first_element,
     truncate_last_element,
-    pad,
-    unpad,
-    pack_padded,
-    pack_list,
-    unpack_to_padded,
     unpack_to_list,
-    squash_packed,
-    freeze,
+    unpack_to_padded,
+    unpad,
 )

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/basic.py

@@ -16,6 +16,7 @@
 import torch
 from aiostream import stream
 from kilroy_module_server_py_sdk import (
+    CategorizableBasedOptionalParameter,
     CategorizableBasedParameter,
     JSONSchema,
     Metric,
@@ -33,6 +34,7 @@
 from kilroy_module_pytorch_py_sdk.generator import Generator
 from kilroy_module_pytorch_py_sdk.models import LanguageModel
 from kilroy_module_pytorch_py_sdk.optimizers import Optimizer
+from kilroy_module_pytorch_py_sdk.schedulers.base import Scheduler
 from kilroy_module_pytorch_py_sdk.tokenizer import Tokenizer
 from kilroy_module_pytorch_py_sdk.utils import (
     pack_list,
@@ -96,6 +98,8 @@ class State:
     tokenizer: Tokenizer
     optimizer: Optimizer
     optimizers_params: Dict[str, Dict[str, Any]]
+    scheduler: Optional[Scheduler]
+    schedulers_params: Dict[str, Dict[str, Any]]
     generator: Generator
     codec: Codec
     results_cache: Dict[UUID, Tuple[Tensor, Tensor]]
@@ -112,6 +116,22 @@ async def _get_params(self, state: State, category: str) -> Dict[str, Any]:
             **state.optimizers_params.get(category, {}),
         }
 
+    async def _set_categorizable(self, state: State, value: Optimizer) -> None:
+        await super()._set_categorizable(state, value)
+        if state.scheduler is not None:
+            optimizer = await value.get()
+            await state.scheduler.change_optimizer(optimizer)
+
+
+class SchedulerParameter(
+    CategorizableBasedOptionalParameter[State, Scheduler]
+):
+    async def _get_params(self, state: State, category: str) -> Dict[str, Any]:
+        return {
+            "optimizer": await state.optimizer.get(),
+            **state.schedulers_params.get(category, {}),
+        }
+
 
 class GeneratorParameter(NestedParameter[State, Generator]):
     pass
@@ -136,6 +156,7 @@ def post_schema(cls) -> JSONSchema:
     def parameters(cls) -> Set[Parameter]:
         return {
             OptimizerParameter(),
+            SchedulerParameter(),
             GeneratorParameter(),
             CodecParameter(),
             BatchSizeParameter(),
@@ -240,6 +261,8 @@ async def _reset_reports(state: State) -> None:
     async def step(self) -> None:
         async with self.state.write_lock() as state:
             await state.optimizer.step()
+            if state.scheduler is not None:
+                await state.scheduler.step()
             await self._report_mean_from_epoch(
                 state.metrics.supervised_loss_metric,
                 state.epoch,

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/reward.py

@@ -19,6 +19,7 @@
 from aiostream import stream
 from aiostream.aiter_utils import aiter, anext
 from kilroy_module_server_py_sdk import (
+    CategorizableBasedOptionalParameter,
     CategorizableBasedParameter,
     JSONSchema,
     Metric,
@@ -37,6 +38,7 @@
 from kilroy_module_pytorch_py_sdk.generator import Generator
 from kilroy_module_pytorch_py_sdk.models import LanguageModel, RewardModel
 from kilroy_module_pytorch_py_sdk.optimizers import Optimizer
+from kilroy_module_pytorch_py_sdk.schedulers.base import Scheduler
 from kilroy_module_pytorch_py_sdk.tokenizer import Tokenizer
 from kilroy_module_pytorch_py_sdk.utils import (
     freeze,
@@ -125,6 +127,8 @@ class LanguageModelState:
     tokenizer: Tokenizer
     optimizer: Optimizer
     optimizers_params: Dict[str, Dict[str, Any]]
+    scheduler: Optional[Scheduler]
+    schedulers_params: Dict[str, Dict[str, Any]]
 
 
 @dataclass
@@ -133,6 +137,8 @@ class RewardModelState:
     tokenizer: Tokenizer
     optimizer: Optimizer
     optimizers_params: Dict[str, Dict[str, Any]]
+    scheduler: Optional[Scheduler]
+    schedulers_params: Dict[str, Dict[str, Any]]
 
 
 @dataclass
@@ -177,6 +183,22 @@ async def _get_params(self, state: State, category: str) -> Dict[str, Any]:
             **state.language_model.optimizers_params.get(category, {}),
         }
 
+    async def _set_categorizable(self, state: State, value: Optimizer) -> None:
+        await super()._set_categorizable(state, value)
+        if state.language_model.scheduler is not None:
+            optimizer = await value.get()
+            await state.language_model.scheduler.change_optimizer(optimizer)
+
+
+class LanguageModelSchedulerParameter(
+    CategorizableBasedOptionalParameter[State, Scheduler]
+):
+    async def _get_params(self, state: State, category: str) -> Dict[str, Any]:
+        return {
+            "optimizer": await state.language_model.optimizer.get(),
+            **state.language_model.schedulers_params.get(category, {}),
+        }
+
 
 class RewardModelOptimizerParameter(
     CategorizableBasedParameter[State, Optimizer]
@@ -187,6 +209,22 @@ async def _get_params(self, state: State, category: str) -> Dict[str, Any]:
             **state.reward_model.optimizers_params.get(category, {}),
         }
 
+    async def _set_categorizable(self, state: State, value: Optimizer) -> None:
+        await super()._set_categorizable(state, value)
+        if state.reward_model.scheduler is not None:
+            optimizer = await value.get()
+            await state.reward_model.scheduler.change_optimizer(optimizer)
+
+
+class RewardModelSchedulerParameter(
+    CategorizableBasedOptionalParameter[State, Scheduler]
+):
+    async def _get_params(self, state: State, category: str) -> Dict[str, Any]:
+        return {
+            "optimizer": await state.reward_model.optimizer.get(),
+            **state.reward_model.schedulers_params.get(category, {}),
+        }
+
 
 class FrontendGeneratorParameter(NestedParameter[State, Generator]):
     pass
@@ -221,7 +259,9 @@ def post_schema(cls) -> JSONSchema:
     def parameters(cls) -> Set[Parameter]:
         return {
             LanguageModelOptimizerParameter(),
+            LanguageModelSchedulerParameter(),
             RewardModelOptimizerParameter(),
+            RewardModelSchedulerParameter(),
             FrontendGeneratorParameter(),
             BackendGeneratorParameter(),
             CodecParameter(),
@@ -425,7 +465,11 @@ async def _reset_reports(state: State) -> None:
     async def step(self) -> None:
         async with self.state.write_lock() as state:
             await state.language_model.optimizer.step()
+            if state.language_model.scheduler is not None:
+                await state.language_model.scheduler.step()
             await state.reward_model.optimizer.step()
+            if state.reward_model.scheduler is not None:
+                await state.reward_model.scheduler.step()
             await self._report_mean_from_epoch(
                 state.metrics.supervised_loss_metric,
                 state.epoch,

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/schedulers/__init__.py

@@ -0,0 +1,21 @@
+from kilroy_module_pytorch_py_sdk.schedulers.base import Scheduler
+from kilroy_module_pytorch_py_sdk.schedulers.constant import ConstantScheduler
+from kilroy_module_pytorch_py_sdk.schedulers.cosine import (
+    CosineAnnealingScheduler,
+)
+from kilroy_module_pytorch_py_sdk.schedulers.cyclic import CyclicScheduler
+from kilroy_module_pytorch_py_sdk.schedulers.exponential import (
+    ExponentialScheduler,
+)
+from kilroy_module_pytorch_py_sdk.schedulers.linear import LinearScheduler
+from kilroy_module_pytorch_py_sdk.schedulers.multistep import (
+    MultiStepScheduler,
+)
+from kilroy_module_pytorch_py_sdk.schedulers.onecycle import OneCycleScheduler
+from kilroy_module_pytorch_py_sdk.schedulers.plateau import (
+    ReduceOnPlateauScheduler,
+)
+from kilroy_module_pytorch_py_sdk.schedulers.step import StepScheduler
+from kilroy_module_pytorch_py_sdk.schedulers.warmrestarts import (
+    WarmRestartsScheduler,
+)