Add LoRA support for FT API (#156)

* Add lora support

* bump up version

* fix docs

* fix style and typing

* Ignore spurious typing errors

---------

Co-authored-by: Max Ryabinin <[email protected]>

Author: artek0chumak

pyproject.toml

@@ -12,7 +12,7 @@ build-backend = "poetry.masonry.api"
 
 [tool.poetry]
 name = "together"
-version = "1.2.2"
+version = "1.2.3"
 authors = [
     "Together AI <[email protected]>"
 ]

src/together/cli/api/finetune.py

@@ -2,6 +2,7 @@
 from textwrap import wrap
 
 import click
+from click.core import ParameterSource  # type: ignore[attr-defined]
 from tabulate import tabulate
 
 from together import Together
@@ -26,7 +27,22 @@ def fine_tuning(ctx: click.Context) -> None:
     "--n-checkpoints", type=int, default=1, help="Number of checkpoints to save"
 )
 @click.option("--batch-size", type=int, default=32, help="Train batch size")
-@click.option("--learning-rate", type=float, default=3e-5, help="Learning rate")
+@click.option("--learning-rate", type=float, default=1e-5, help="Learning rate")
+@click.option(
+    "--lora/--no-lora",
+    type=bool,
+    default=False,
+    help="Whether to use LoRA adapters for fine-tuning",
+)
+@click.option("--lora-r", type=int, default=8, help="LoRA adapters' rank")
+@click.option("--lora-dropout", type=float, default=0, help="LoRA adapters' dropout")
+@click.option("--lora-alpha", type=float, default=8, help="LoRA adapters' alpha")
+@click.option(
+    "--lora-trainable-modules",
+    type=str,
+    default="all-linear",
+    help="Trainable modules for LoRA adapters. For example, 'all-linear', 'q_proj,v_proj'",
+)
 @click.option(
     "--suffix", type=str, default=None, help="Suffix for the fine-tuned model name"
 )
@@ -39,19 +55,44 @@ def create(
     n_checkpoints: int,
     batch_size: int,
     learning_rate: float,
+    lora: bool,
+    lora_r: int,
+    lora_dropout: float,
+    lora_alpha: float,
+    lora_trainable_modules: str,
     suffix: str,
     wandb_api_key: str,
 ) -> None:
     """Start fine-tuning"""
     client: Together = ctx.obj
 
+    if lora:
+        learning_rate_source = click.get_current_context().get_parameter_source(  # type: ignore[attr-defined]
+            "learning_rate"
+        )
+        if learning_rate_source == ParameterSource.DEFAULT:
+            learning_rate = 1e-3
+    else:
+        for param in ["lora_r", "lora_dropout", "lora_alpha", "lora_trainable_modules"]:
+            param_source = click.get_current_context().get_parameter_source(param)  # type: ignore[attr-defined]
+            if param_source != ParameterSource.DEFAULT:
+                raise click.BadParameter(
+                    f"You set LoRA parameter `{param}` for a full fine-tuning job. "
+                    f"Please change the job type with --lora or remove `{param}` from the arguments"
+                )
+
     response = client.fine_tuning.create(
         training_file=training_file,
         model=model,
         n_epochs=n_epochs,
         n_checkpoints=n_checkpoints,
         batch_size=batch_size,
         learning_rate=learning_rate,
+        lora=lora,
+        lora_r=lora_r,
+        lora_dropout=lora_dropout,
+        lora_alpha=lora_alpha,
+        lora_trainable_modules=lora_trainable_modules,
         suffix=suffix,
         wandb_api_key=wandb_api_key,
     )

src/together/resources/finetune.py

@@ -11,8 +11,11 @@
     FinetuneListEvents,
     FinetuneRequest,
     FinetuneResponse,
+    FullTrainingType,
+    LoRATrainingType,
     TogetherClient,
     TogetherRequest,
+    TrainingType,
 )
 from together.utils import normalize_key
 
@@ -30,6 +33,11 @@ def create(
         n_checkpoints: int | None = 1,
         batch_size: int | None = 32,
         learning_rate: float | None = 0.00001,
+        lora: bool = True,
+        lora_r: int | None = 8,
+        lora_dropout: float | None = 0,
+        lora_alpha: float | None = 8,
+        lora_trainable_modules: str | None = "all-linear",
         suffix: str | None = None,
         wandb_api_key: str | None = None,
     ) -> FinetuneResponse:
@@ -45,6 +53,11 @@ def create(
             batch_size (int, optional): Batch size for fine-tuning. Defaults to 32.
             learning_rate (float, optional): Learning rate multiplier to use for training
                 Defaults to 0.00001.
+            lora (bool, optional): Whether to use LoRA adapters. Defaults to True.
+            lora_r (int, optional): Rank of LoRA adapters. Defaults to 8.
+            lora_dropout (float, optional): Dropout rate for LoRA adapters. Defaults to 0.
+            lora_alpha (float, optional): Alpha for LoRA adapters. Defaults to 8.
+            lora_trainable_modules (str, optional): Trainable modules for LoRA adapters. Defaults to "all-linear".
             suffix (str, optional): Up to 40 character suffix that will be added to your fine-tuned model name.
                 Defaults to None.
             wandb_api_key (str, optional): API key for Weights & Biases integration.
@@ -58,13 +71,23 @@ def create(
             client=self._client,
         )
 
+        training_type: TrainingType = FullTrainingType()
+        if lora:
+            training_type = LoRATrainingType(
+                lora_r=lora_r,
+                lora_alpha=lora_alpha,
+                lora_dropout=lora_dropout,
+                lora_trainable_modules=lora_trainable_modules,
+            )
+
         parameter_payload = FinetuneRequest(
             model=model,
             training_file=training_file,
             n_epochs=n_epochs,
             n_checkpoints=n_checkpoints,
             batch_size=batch_size,
             learning_rate=learning_rate,
+            training_type=training_type,
             suffix=suffix,
             wandb_key=wandb_api_key,
         ).model_dump()

src/together/types/__init__.py

@@ -26,6 +26,9 @@
     FinetuneListEvents,
     FinetuneRequest,
     FinetuneResponse,
+    FullTrainingType,
+    LoRATrainingType,
+    TrainingType,
 )
 from together.types.images import (
     ImageRequest,
@@ -60,4 +63,7 @@
     "ImageRequest",
     "ImageResponse",
     "ModelObject",
+    "TrainingType",
+    "FullTrainingType",
+    "LoRATrainingType",
 ]

src/together/types/finetune.py

@@ -100,6 +100,34 @@ class FinetuneEvent(BaseModel):
     hash: str | None = None
 
 
+class TrainingType(BaseModel):
+    """
+    Abstract training type
+    """
+
+    type: str
+
+
+class FullTrainingType(TrainingType):
+    """
+    Training type for full fine-tuning
+    """
+
+    type: str = "Full"
+
+
+class LoRATrainingType(TrainingType):
+    """
+    Training type for LoRA adapters training
+    """
+
+    lora_r: int
+    lora_alpha: int
+    lora_dropout: float
+    lora_trainable_modules: str
+    type: str = "Lora"
+
+
 class FinetuneRequest(BaseModel):
     """
     Fine-tune request type
@@ -121,6 +149,7 @@ class FinetuneRequest(BaseModel):
     suffix: str | None = None
     # weights & biases api key
     wandb_key: str | None = None
+    training_type: FullTrainingType | LoRATrainingType | None = None
 
 
 class FinetuneResponse(BaseModel):
@@ -138,6 +167,8 @@ class FinetuneResponse(BaseModel):
     model: str | None = None
     # output model name
     output_name: str | None = Field(None, alias="model_output_name")
+    # adapter output name
+    adapter_output_name: str | None = None
     # number of epochs
     n_epochs: int | None = None
     # number of checkpoints to save
@@ -148,11 +179,8 @@ class FinetuneResponse(BaseModel):
     learning_rate: float | None = None
     # number of steps between evals
     eval_steps: int | None = None
-    # is LoRA finetune boolean
-    lora: bool | None = None
-    lora_r: int | None = None
-    lora_alpha: int | None = None
-    lora_dropout: int | None = None
+    # training type
+    training_type: FullTrainingType | LoRATrainingType | None = None
     # created/updated datetime stamps
     created_at: str | None = None
     updated_at: str | None = None