Integrate finetune with wandb

clients/python/llmengine/data_types.py

@@ -383,6 +383,16 @@ class CreateFineTuneRequest(BaseModel):
     )
     """Hyperparameters to pass in to training job."""
 
+    wandb_config: Optional[Dict[str, Any]] = Field(
+        default=None, description="Configuration for Weights and Biases."
+    )
+    """
+    A dict of configuration parameters for Weights & Biases. See [Weights & Biases](https://docs.wandb.ai/ref/python/init) for more information.
+    Set `hyperparameter["report_to"]` to `wandb` to enable automatic finetune metrics logging.
+    Must include `api_key` field which is the wandb API key.
+    Also supports setting `base_url` to use a custom Weights & Biases server.
+    """
+
     suffix: Optional[str] = Field(
         default=None,
         description="Optional user-provided identifier suffix for the fine-tuned model.",

clients/python/llmengine/fine_tuning.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 from llmengine.api_engine import DEFAULT_TIMEOUT, APIEngine
 from llmengine.data_types import (
@@ -29,6 +29,7 @@ def create(
         training_file: str,
         validation_file: Optional[str] = None,
         hyperparameters: Optional[Dict[str, Union[str, int, float]]] = None,
+        wandb_config: Optional[Dict[str, Any]] = None,
         suffix: Optional[str] = None,
     ) -> CreateFineTuneResponse:
         """
@@ -66,6 +67,12 @@ def create(
                 * `epochs`: Number of fine-tuning epochs. This should be less than 20. (Default: 5)
                 * `weight_decay`: Regularization penalty applied to learned weights. (Default: 0.001)
 
+            wandb_config (`Optional[Dict[str, Any]]`):
+                A dict of configuration parameters for Weights & Biases. See [Weights & Biases](https://docs.wandb.ai/ref/python/init) for more information.
+                Set `hyperparameter["report_to"]` to `wandb` to enable automatic finetune metrics logging.
+                Must include `api_key` field which is the wandb API key.
+                Also supports setting `base_url` to use a custom Weights & Biases server.
+
             suffix (`Optional[str]`):
                 A string that will be added to your fine-tuned model name. If present, the entire fine-tuned model name
                 will be formatted like `"[model].[suffix].[YYYY-MM-DD-HH-MM-SS]"`. If absent, the
@@ -134,6 +141,7 @@ def create(
             training_file=training_file,
             validation_file=validation_file,
             hyperparameters=hyperparameters,
+            wandb_config=wandb_config,
             suffix=suffix,
         )
         response = cls.post_sync(

docs/model_zoo.md

@@ -10,6 +10,7 @@ Scale hosts the following models in the LLM Engine Model Zoo:
 | `llama-2-13b`         | ✅                       |                            |
 | `llama-2-13b-chat`    | ✅                       |                            |
 | `llama-2-70b`         | ✅                       |                            |
+| `llama-2-70b-chat`    | ✅                       |                            |
 | `falcon-7b`           | ✅                       |                            |
 | `falcon-7b-instruct`  | ✅                       |                            |
 | `falcon-40b`          | ✅                       |                            |