Improve fine-tuning checkpoint support (#259)

Author: artek0chumak

README.md

@@ -67,7 +67,7 @@ print(response.choices[0].message.content)
 response = client.chat.completions.create(
     model="meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo",
     messages=[{
-        "role": "user", 
+        "role": "user",
         "content": [
             {
                 "type": "text",
@@ -91,7 +91,7 @@ response = client.chat.completions.create(
         "role": "user",
         "content": [
             {
-                "type": "text", 
+                "type": "text",
                 "text": "Compare these two images."
             },
             {

pyproject.toml

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

src/together/cli/api/finetune.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import json
-from datetime import datetime
+from datetime import datetime, timezone
 from textwrap import wrap
 from typing import Any, Literal
+import re
 
 import click
 from click.core import ParameterSource  # type: ignore[attr-defined]
@@ -17,8 +18,13 @@
     log_warn,
     log_warn_once,
     parse_timestamp,
+    format_timestamp,
+)
+from together.types.finetune import (
+    DownloadCheckpointType,
+    FinetuneTrainingLimits,
+    FinetuneEventType,
 )
-from together.types.finetune import DownloadCheckpointType, FinetuneTrainingLimits
 
 
 _CONFIRMATION_MESSAGE = (
@@ -126,6 +132,14 @@ 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(
+    "--from-checkpoint",
+    type=str,
+    default=None,
+    help="The checkpoint identifier to continue training from a previous fine-tuning job. "
+    "The format: {$JOB_ID/$OUTPUT_MODEL_NAME}:{$STEP}. "
+    "The step value is optional, without it the final checkpoint will be used.",
+)
 def create(
     ctx: click.Context,
     training_file: str,
@@ -152,6 +166,7 @@ def create(
     wandb_name: str,
     confirm: bool,
     train_on_inputs: bool | Literal["auto"],
+    from_checkpoint: str,
 ) -> None:
     """Start fine-tuning"""
     client: Together = ctx.obj
@@ -180,6 +195,7 @@ def create(
         wandb_project_name=wandb_project_name,
         wandb_name=wandb_name,
         train_on_inputs=train_on_inputs,
+        from_checkpoint=from_checkpoint,
     )
 
     model_limits: FinetuneTrainingLimits = client.fine_tuning.get_model_limits(
@@ -261,7 +277,9 @@ def list(ctx: click.Context) -> None:
 
     response.data = response.data or []
 
-    response.data.sort(key=lambda x: parse_timestamp(x.created_at or ""))
+    # Use a default datetime for None values to make sure the key function always returns a comparable value
+    epoch_start = datetime.fromtimestamp(0, tz=timezone.utc)
+    response.data.sort(key=lambda x: parse_timestamp(x.created_at or "") or epoch_start)
 
     display_list = []
     for i in response.data:
@@ -344,6 +362,34 @@ def list_events(ctx: click.Context, fine_tune_id: str) -> None:
     click.echo(table)
 
 
+@fine_tuning.command()
+@click.pass_context
+@click.argument("fine_tune_id", type=str, required=True)
+def list_checkpoints(ctx: click.Context, fine_tune_id: str) -> None:
+    """List available checkpoints for a fine-tuning job"""
+    client: Together = ctx.obj
+
+    checkpoints = client.fine_tuning.list_checkpoints(fine_tune_id)
+
+    display_list = []
+    for checkpoint in checkpoints:
+        display_list.append(
+            {
+                "Type": checkpoint.type,
+                "Timestamp": format_timestamp(checkpoint.timestamp),
+                "Name": checkpoint.name,
+            }
+        )
+
+    if display_list:
+        click.echo(f"Job {fine_tune_id} contains the following checkpoints:")
+        table = tabulate(display_list, headers="keys", tablefmt="grid")
+        click.echo(table)
+        click.echo("\nTo download a checkpoint, use `together fine-tuning download`")
+    else:
+        click.echo(f"No checkpoints found for job {fine_tune_id}")
+
+
 @fine_tuning.command()
 @click.pass_context
 @click.argument("fine_tune_id", type=str, required=True)
@@ -358,7 +404,7 @@ def list_events(ctx: click.Context, fine_tune_id: str) -> None:
     "--checkpoint-step",
     type=int,
     required=False,
-    default=-1,
+    default=None,
     help="Download fine-tuning checkpoint. Defaults to latest.",
 )
 @click.option(
@@ -372,7 +418,7 @@ def download(
     ctx: click.Context,
     fine_tune_id: str,
     output_dir: str,
-    checkpoint_step: int,
+    checkpoint_step: int | None,
     checkpoint_type: DownloadCheckpointType,
 ) -> None:
     """Download fine-tuning checkpoint"""

src/together/legacy/finetune.py

@@ -161,7 +161,7 @@ def download(
         cls,
         fine_tune_id: str,
         output: str | None = None,
-        step: int = -1,
+        step: int | None = None,
     ) -> Dict[str, Any]:
         """Legacy finetuning download function."""