VLM Fine-tuning support (#201)

* Support VLM finetuning

Author: nikita-smetanin

src/together/lib/cli/api/fine_tuning.py

@@ -176,6 +176,12 @@ def fine_tuning(ctx: click.Context) -> None:
     help="Whether to mask the user messages in conversational data or prompts in instruction data. "
     "`auto` will automatically determine whether to mask the inputs based on the data format.",
 )
+@click.option(
+    "--train-vision",
+    type=bool,
+    default=False,
+    help="Whether to train the vision encoder. Only supported for multimodal models.",
+)
 @click.option(
     "--from-checkpoint",
     type=str,
@@ -231,6 +237,7 @@ def create(
     lora_dropout: float | None,
     lora_alpha: float | None,
     lora_trainable_modules: str | None,
+    train_vision: bool,
     suffix: str | None,
     wandb_api_key: str | None,
     wandb_base_url: str | None,
@@ -272,6 +279,7 @@ def create(
         lora_dropout=lora_dropout,
         lora_alpha=lora_alpha,
         lora_trainable_modules=lora_trainable_modules,
+        train_vision=train_vision,
         suffix=suffix,
         wandb_api_key=wandb_api_key,
         wandb_base_url=wandb_base_url,
@@ -363,6 +371,10 @@ def create(
             simpo_gamma=simpo_gamma or 0,
         )
 
+    if model_limits.supports_vision:
+        # Don't show price estimation for multimodal models yet
+        confirm = True
+
     finetune_price_estimation_result = client.fine_tuning.estimate_price(
         training_file=training_file,
         validation_file=validation_file,

src/together/lib/constants.py

@@ -37,6 +37,12 @@
 # maximum number of GB sized files we support finetuning for
 MAX_FILE_SIZE_GB = 50.1
 
+# Multimodal limits
+MAX_IMAGES_PER_EXAMPLE = 10
+MAX_IMAGE_BYTES = 10 * 1024 * 1024  # 10MB
+# Max length = Header length + base64 factor (4/3) * image bytes
+MAX_BASE64_IMAGE_LENGTH = len("data:image/jpeg;base64,") + 4 * MAX_IMAGE_BYTES // 3
+
 # expected columns for Parquet files
 PARQUET_EXPECTED_COLUMNS = ["input_ids", "attention_mask", "labels"]
 

src/together/lib/resources/fine_tuning.py

@@ -22,6 +22,7 @@
     CosineLRSchedulerArgs,
     LinearLRSchedulerArgs,
     FinetuneTrainingLimits,
+    FinetuneMultimodalParams,
 )
 
 AVAILABLE_TRAINING_METHODS = {
@@ -51,6 +52,7 @@ def create_finetune_request(
     lora_dropout: float | None = 0,
     lora_alpha: float | None = None,
     lora_trainable_modules: str | None = "all-linear",
+    train_vision: bool = False,
     suffix: str | None = None,
     wandb_api_key: str | None = None,
     wandb_base_url: str | None = None,
@@ -207,6 +209,13 @@ def create_finetune_request(
             simpo_gamma=simpo_gamma,
         )
 
+    if model_limits.supports_vision:
+        multimodal_params = FinetuneMultimodalParams(train_vision=train_vision)
+    elif not model_limits.supports_vision and train_vision:
+        raise ValueError(f"Vision encoder training is not supported for the non-multimodal model `{model}`")
+    else:
+        multimodal_params = None
+
     finetune_request = FinetuneRequest(
         model=model,
         training_file=training_file,
@@ -227,6 +236,7 @@ def create_finetune_request(
         wandb_project_name=wandb_project_name,
         wandb_name=wandb_name,
         training_method=training_method_cls,  # pyright: ignore[reportPossiblyUnboundVariable]
+        multimodal_params=multimodal_params,
         from_checkpoint=from_checkpoint,
         from_hf_model=from_hf_model,
         hf_model_revision=hf_model_revision,

src/together/lib/types/fine_tuning.py

@@ -203,6 +203,7 @@ class FinetuneTrainingLimits(BaseModel):
     min_learning_rate: float
     full_training: Optional[FinetuneFullTrainingLimits] = None
     lora_training: Optional[FinetuneLoraTrainingLimits] = None
+    supports_vision: bool = False
 
 
 class LinearLRSchedulerArgs(BaseModel):
@@ -270,6 +271,14 @@ class UnknownLRScheduler(BaseModel):
 ]
 
 
+class FinetuneMultimodalParams(BaseModel):
+    """
+    Multimodal parameters
+    """
+
+    train_vision: bool = False
+
+
 class FinetuneProgress(BaseModel):
     """
     Fine-tune job progress
@@ -305,6 +314,9 @@ class FinetuneResponse(BaseModel):
     from_checkpoint: Optional[str] = None
     """Checkpoint used to continue training"""
 
+    multimodal_params: Optional[FinetuneMultimodalParams] = None
+    """Multimodal parameters"""
+
     from_hf_model: Optional[str] = None
     """Hugging Face Hub repo to start training from"""
 
@@ -469,6 +481,9 @@ class FinetuneRequest(BaseModel):
     training_method: TrainingMethod = Field(default_factory=TrainingMethodSFT)
     # from step
     from_checkpoint: Union[str, None] = None
+    # multimodal parameters
+    multimodal_params: Union[FinetuneMultimodalParams, None] = None
+    # hugging face related fields
     from_hf_model: Union[str, None] = None
     hf_model_revision: Union[str, None] = None
     # hf related fields