Use subclasses instead of validation

Author: Arsh Zahed

src/together/resources/finetune.py

@@ -22,6 +22,8 @@
     TogetherRequest,
     TrainingType,
     FinetuneLRScheduler,
+    FinetuneLinearLRScheduler,
+    FinetuneCosineLRScheduler,
     FinetuneLinearLRSchedulerArgs,
     FinetuneCosineLRSchedulerArgs,
     TrainingMethodDPO,
@@ -132,19 +134,20 @@ def createFinetuneRequest(
             f"training_method must be one of {', '.join(AVAILABLE_TRAINING_METHODS)}"
         )
 
+    # Default to generic lr scheduler
+    lrScheduler: FinetuneLRScheduler = FinetuneLRScheduler(lr_scheduler_type="linear")
+
     if lr_scheduler_type == "cosine":
         if num_cycles <= 0.0:
             raise ValueError("Number of cycles should be greater than 0")
 
-        lrScheduler = FinetuneLRScheduler(
-            lr_scheduler_type="cosine",
+        lrScheduler = FinetuneCosineLRScheduler(
             lr_scheduler_args=FinetuneCosineLRSchedulerArgs(
                 min_lr_ratio=min_lr_ratio, num_cycles=num_cycles
             ),
         )
     else:
-        lrScheduler = FinetuneLRScheduler(
-            lr_scheduler_type="linear",
+        lrScheduler = FinetuneLinearLRScheduler(
             lr_scheduler_args=FinetuneLinearLRSchedulerArgs(min_lr_ratio=min_lr_ratio),
         )
 

src/together/types/__init__.py

@@ -34,12 +34,14 @@
     TrainingMethodDPO,
     TrainingMethodSFT,
     FinetuneCheckpoint,
+    FinetuneCosineLRScheduler,
     FinetuneCosineLRSchedulerArgs,
     FinetuneDownloadResult,
+    FinetuneLinearLRScheduler,
     FinetuneLinearLRSchedulerArgs,
+    FinetuneLRScheduler,
     FinetuneList,
     FinetuneListEvents,
-    FinetuneLRScheduler,
     FinetuneRequest,
     FinetuneResponse,
     FinetuneTrainingLimits,
@@ -70,7 +72,9 @@
     "FinetuneListEvents",
     "FinetuneDownloadResult",
     "FinetuneLRScheduler",
+    "FinetuneLinearLRScheduler",
     "FinetuneLinearLRSchedulerArgs",
+    "FinetuneCosineLRScheduler",
     "FinetuneCosineLRSchedulerArgs",
     "FileRequest",
     "FileResponse",

src/together/types/finetune.py

@@ -176,7 +176,7 @@ class FinetuneRequest(BaseModel):
     # training learning rate
     learning_rate: float
     # learning rate scheduler type and args
-    lr_scheduler: FinetuneLRScheduler | None = None
+    lr_scheduler: FinetuneLinearLRScheduler | FinetuneCosineLRScheduler | None = None
     # learning rate warmup ratio
     warmup_ratio: float
     # max gradient norm
@@ -239,7 +239,7 @@ class FinetuneResponse(BaseModel):
     # training learning rate
     learning_rate: float | None = None
     # learning rate scheduler type and args
-    lr_scheduler: FinetuneLRScheduler | None = None
+    lr_scheduler: FinetuneLinearLRScheduler | FinetuneCosineLRScheduler | None = None
     # learning rate warmup ratio
     warmup_ratio: float | None = None
     # max gradient norm
@@ -354,44 +354,18 @@ class FinetuneCosineLRSchedulerArgs(BaseModel):
     num_cycles: float | None = 0.5
 
 
-LRSchedulerTypeToArgs = {
-    "linear": FinetuneLinearLRSchedulerArgs,
-    "cosine": FinetuneCosineLRSchedulerArgs,
-}
-
-FinetuneLRSchedulerArgs = Union[
-    FinetuneLinearLRSchedulerArgs, FinetuneCosineLRSchedulerArgs, None
-]
-
-
 class FinetuneLRScheduler(BaseModel):
     lr_scheduler_type: str
-    lr_scheduler_args: FinetuneLRSchedulerArgs | None = None
 
-    @field_validator("lr_scheduler_type")
-    @classmethod
-    def validate_scheduler_type(cls, scheduler_type: str) -> str:
-        if scheduler_type not in LRSchedulerTypeToArgs:
-            raise ValueError(
-                f"Scheduler type must be one of: {LRSchedulerTypeToArgs.keys()}"
-            )
-        return scheduler_type
-
-    @field_validator("lr_scheduler_args")
-    @classmethod
-    def validate_scheduler_args(
-        cls, args: FinetuneLRSchedulerArgs, info: ValidationInfo
-    ) -> FinetuneLRSchedulerArgs:
-        scheduler_type = str(info.data.get("lr_scheduler_type"))
 
-        if args is None:
-            return args
+class FinetuneLinearLRScheduler(FinetuneLRScheduler):
+    lr_scheduler_type: Literal["linear"] = "linear"
+    lr_scheduler: FinetuneLinearLRSchedulerArgs | None = None
 
-        expected_type = LRSchedulerTypeToArgs[scheduler_type]
-        if not isinstance(args, expected_type):
-            raise TypeError(f"Expected {expected_type}, got {type(args)}")
 
-        return args
+class FinetuneCosineLRScheduler(FinetuneLRScheduler):
+    lr_scheduler_type: Literal["cosine"] = "cosine"
+    lr_scheduler: FinetuneCosineLRSchedulerArgs | None = None
 
 
 class FinetuneCheckpoint(BaseModel):