VLM Finetuning support (#411)

* Support Multimodal datasets
* Support VLM finetuning

Author: nikita-smetanin

pyproject.toml

@@ -12,7 +12,7 @@ build-backend = "poetry.masonry.api"
 
 [tool.poetry]
 name = "together"
-version = "1.5.33"
+version = "1.5.34"
 authors = ["Together AI <[email protected]>"]
 description = "Python client for Together's Cloud Platform! Note: SDK 2.0 is now available at https://github.com/togethercomputer/together-py"
 readme = "README.md"

src/together/cli/api/finetune.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-import re
 from datetime import datetime, timezone
 from textwrap import wrap
 from typing import Any, Literal
@@ -14,18 +13,11 @@
 
 from together import Together
 from together.cli.api.utils import BOOL_WITH_AUTO, INT_WITH_MAX, generate_progress_bar
-from together.types.finetune import (
-    DownloadCheckpointType,
-    FinetuneEventType,
-    FinetuneTrainingLimits,
-    FullTrainingType,
-    LoRATrainingType,
-)
+from together.types.finetune import DownloadCheckpointType, FinetuneTrainingLimits
 from together.utils import (
     finetune_price_to_dollars,
     format_timestamp,
     log_warn,
-    log_warn_once,
     parse_timestamp,
 )
 
@@ -203,6 +195,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,
@@ -258,6 +256,7 @@ def create(
     lora_dropout: float,
     lora_alpha: float,
     lora_trainable_modules: str,
+    train_vision: bool,
     suffix: str,
     wandb_api_key: str,
     wandb_base_url: str,
@@ -299,6 +298,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,
@@ -368,6 +368,10 @@ def create(
             "You have specified a number of evaluation loops but no validation file."
         )
 
+    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/constants.py

@@ -1,5 +1,6 @@
 import enum
 
+
 # Session constants
 TIMEOUT_SECS = 600
 MAX_SESSION_LIFETIME_SECS = 180
@@ -40,6 +41,11 @@
 # the number of bytes in a gigabyte, used to convert bytes to GB for readable comparison
 NUM_BYTES_IN_GB = 2**30
 
+# 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/resources/finetune.py

@@ -2,7 +2,7 @@
 
 import re
 from pathlib import Path
-from typing import List, Dict, Literal
+from typing import Dict, List, Literal
 
 from rich import print as rprint
 
@@ -18,10 +18,11 @@
     FinetuneList,
     FinetuneListEvents,
     FinetuneLRScheduler,
-    FinetuneRequest,
-    FinetuneResponse,
+    FinetuneMultimodalParams,
     FinetunePriceEstimationRequest,
     FinetunePriceEstimationResponse,
+    FinetuneRequest,
+    FinetuneResponse,
     FinetuneTrainingLimits,
     FullTrainingType,
     LinearLRScheduler,
@@ -73,6 +74,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,
@@ -252,6 +254,15 @@ 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,
@@ -272,6 +283,7 @@ def create_finetune_request(
         wandb_project_name=wandb_project_name,
         wandb_name=wandb_name,
         training_method=training_method_cls,
+        multimodal_params=multimodal_params,
         from_checkpoint=from_checkpoint,
         from_hf_model=from_hf_model,
         hf_model_revision=hf_model_revision,
@@ -342,6 +354,7 @@ def create(
         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,
@@ -387,6 +400,7 @@ def create(
             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".
+            train_vision (bool, optional): Whether to train vision encoder in multimodal models. Defaults to False.
             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.
@@ -464,6 +478,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,
@@ -906,6 +921,7 @@ async def create(
         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,
@@ -951,6 +967,7 @@ async def create(
             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".
+            train_vision (bool, optional): Whether to train vision encoder in multimodal models. Defaults to False.
             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.
@@ -1028,6 +1045,7 @@ async 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,
@@ -1046,7 +1064,11 @@ async def create(
             hf_output_repo_name=hf_output_repo_name,
         )
 
-        if from_checkpoint is None and from_hf_model is None:
+        if (
+            from_checkpoint is None
+            and from_hf_model is None
+            and not model_limits.supports_vision
+        ):
             price_estimation_result = await self.estimate_price(
                 training_file=training_file,
                 validation_file=validation_file,

src/together/types/__init__.py

@@ -7,17 +7,18 @@
     AudioSpeechStreamChunk,
     AudioSpeechStreamEvent,
     AudioSpeechStreamResponse,
+    AudioTimestampGranularities,
     AudioTranscriptionRequest,
-    AudioTranslationRequest,
     AudioTranscriptionResponse,
+    AudioTranscriptionResponseFormat,
     AudioTranscriptionVerboseResponse,
+    AudioTranslationRequest,
     AudioTranslationResponse,
     AudioTranslationVerboseResponse,
-    AudioTranscriptionResponseFormat,
-    AudioTimestampGranularities,
     ModelVoices,
     VoiceListResponse,
 )
+from together.types.batch import BatchEndpoint, BatchJob, BatchJobStatus
 from together.types.chat_completions import (
     ChatCompletionChunk,
     ChatCompletionRequest,
@@ -31,6 +32,19 @@
 )
 from together.types.embeddings import EmbeddingRequest, EmbeddingResponse
 from together.types.endpoints import Autoscaling, DedicatedEndpoint, ListEndpoint
+from together.types.evaluation import (
+    ClassifyParameters,
+    CompareParameters,
+    EvaluationCreateResponse,
+    EvaluationJob,
+    EvaluationRequest,
+    EvaluationStatus,
+    EvaluationStatusResponse,
+    EvaluationType,
+    JudgeModelConfig,
+    ModelRequest,
+    ScoreParameters,
+)
 from together.types.files import (
     FileDeleteResponse,
     FileList,
@@ -41,49 +55,32 @@
     FileType,
 )
 from together.types.finetune import (
-    TrainingMethodDPO,
-    TrainingMethodSFT,
-    FinetuneCheckpoint,
     CosineLRScheduler,
     CosineLRSchedulerArgs,
+    FinetuneCheckpoint,
+    FinetuneDeleteResponse,
     FinetuneDownloadResult,
-    LinearLRScheduler,
-    LinearLRSchedulerArgs,
-    FinetuneLRScheduler,
     FinetuneList,
     FinetuneListEvents,
-    FinetuneRequest,
-    FinetuneResponse,
+    FinetuneLRScheduler,
+    FinetuneMultimodalParams,
     FinetunePriceEstimationRequest,
     FinetunePriceEstimationResponse,
-    FinetuneDeleteResponse,
+    FinetuneRequest,
+    FinetuneResponse,
     FinetuneTrainingLimits,
     FullTrainingType,
+    LinearLRScheduler,
+    LinearLRSchedulerArgs,
     LoRATrainingType,
+    TrainingMethodDPO,
+    TrainingMethodSFT,
     TrainingType,
 )
 from together.types.images import ImageRequest, ImageResponse
 from together.types.models import ModelObject, ModelUploadRequest, ModelUploadResponse
 from together.types.rerank import RerankRequest, RerankResponse
-from together.types.batch import BatchJob, BatchJobStatus, BatchEndpoint
-from together.types.evaluation import (
-    EvaluationType,
-    EvaluationStatus,
-    JudgeModelConfig,
-    ModelRequest,
-    ClassifyParameters,
-    ScoreParameters,
-    CompareParameters,
-    EvaluationRequest,
-    EvaluationCreateResponse,
-    EvaluationJob,
-    EvaluationStatusResponse,
-)
-from together.types.videos import (
-    CreateVideoBody,
-    CreateVideoResponse,
-    VideoJob,
-)
+from together.types.videos import CreateVideoBody, CreateVideoResponse, VideoJob
 
 
 __all__ = [
@@ -131,6 +128,7 @@
     "RerankRequest",
     "RerankResponse",
     "FinetuneTrainingLimits",
+    "FinetuneMultimodalParams",
     "AudioSpeechRequest",
     "AudioResponseFormat",
     "AudioLanguage",

src/together/types/finetune.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import List, Literal, Any
+from typing import Any, List, Literal
 
 from pydantic import Field, StrictBool, field_validator
 
 from together.types.abstract import BaseModel
-from together.types.common import (
-    ObjectType,
-)
+from together.types.common import ObjectType
 
 
 class FinetuneJobStatus(str, Enum):
@@ -175,6 +173,14 @@ class TrainingMethodDPO(TrainingMethod):
     simpo_gamma: float | None = None
 
 
+class FinetuneMultimodalParams(BaseModel):
+    """
+    Multimodal parameters
+    """
+
+    train_vision: bool = False
+
+
 class FinetuneProgress(BaseModel):
     """
     Fine-tune job progress
@@ -231,6 +237,8 @@ class FinetuneRequest(BaseModel):
     )
     # from step
     from_checkpoint: str | None = None
+    # multimodal parameters
+    multimodal_params: FinetuneMultimodalParams | None = None
     # hf related fields
     hf_api_token: str | None = None
     hf_output_repo_name: str | None = None
@@ -313,6 +321,8 @@ class FinetuneResponse(BaseModel):
     training_file_size: int | None = Field(None, alias="TrainingFileSize")
     train_on_inputs: StrictBool | Literal["auto"] | None = "auto"
     from_checkpoint: str | None = None
+    # multimodal parameters
+    multimodal_params: FinetuneMultimodalParams | None = None
 
     progress: FinetuneProgress | None = None
 
@@ -409,6 +419,7 @@ class FinetuneTrainingLimits(BaseModel):
     min_learning_rate: float
     full_training: FinetuneFullTrainingLimits | None = None
     lora_training: FinetuneLoraTrainingLimits | None = None
+    supports_vision: bool = False
 
 
 class LinearLRSchedulerArgs(BaseModel):