adding report-to-wandb (#9)

* update lora_config.yaml + LORA.md + lora.py

* code formatting

* udpaet Acknowledgements.md

* nits

* Refactor WandB integration in lora.py and trainer.py

- Updated WandB reporting mechanism to use a project name argument instead of a boolean flag.
- Removed the old TrainingCallback class definition from trainer.py and imported it from callbacks.
- Adjusted argument parsing to accommodate the new WandB configuration.

* Enhance WandBCallback to include log directory in initialization

- Added log_dir parameter to WandBCallback constructor for specifying the logging directory.
- Updated lora.py to pass adapter_path as log_dir when initializing WandBCallback.

* nits

* formating

* README.md

* update example yaml

* nits

* nits

* nits in readme

---------

Co-authored-by: Awni Hannun <[email protected]>

ACKNOWLEDGMENTS.md

@@ -9,4 +9,4 @@ MLX LM was developed with contributions from the following individuals:
 
 - Shunta Saito: Added support for PLaMo models.
 - Prince Canuma: Helped add support for `Starcoder2` models.
-- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`, and `reporting training metrics to WandB (Weights & Biases)`.

mlx_lm/LORA.md

@@ -76,6 +76,11 @@ You can specify the output location with `--adapter-path`.
 You can resume fine-tuning with an existing adapter with
 `--resume-adapter-file <path_to_adapters.safetensors>`.
 
+#### Logging
+
+You can log training metrics to Weights & Biases by passing a project name with
+the `--wandb` flag. Make sure to install wandb with `pip install wandb`.
+
 #### Prompt Masking
 
 The default training computes a loss for every token in the sample. You can

mlx_lm/examples/lora_config.yaml

@@ -37,6 +37,9 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5
 
+# Whether to report the logs to WandB
+# wand: 'wandb-project"
+
 # Number of training steps between loss reporting.
 steps_per_report: 10
 

mlx_lm/lora.py

@@ -1,5 +1,3 @@
-# Copyright © 2024 Apple Inc.
-
 import argparse
 import math
 import os
@@ -13,6 +11,7 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml
 
+from .tuner.callbacks import WandBCallback
 from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
@@ -68,6 +67,7 @@ CONFIG_DEFAULTS = {
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 10.0},
     "mask_prompt": False,
+    "wandb": None,
 }
 
 
@@ -179,6 +179,12 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
+    parser.add_argument(
+        "--wandb",
+        type=str,
+        default=None,
+        help="WandB project name to report training metrics. Disabled if None.",
+    )
     parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
@@ -281,6 +287,14 @@ def evaluate_model(args, model: nn.Module, test_set):
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
 
+    if args.wandb is not None:
+        training_callback = WandBCallback(
+            project_name=args.wandb,
+            log_dir=args.adapter_path,
+            config=vars(args),
+            wrapped_callback=training_callback,
+        )
+
     print("Loading pretrained model")
     model, tokenizer = load(args.model)
 

mlx_lm/tuner/callbacks.py

@@ -0,0 +1,41 @@
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+class TrainingCallback:
+
+    def on_train_loss_report(self, train_info: dict):
+        """Called to report training loss at specified intervals."""
+        pass
+
+    def on_val_loss_report(self, val_info: dict):
+        """Called to report validation loss at specified intervals or the beginning."""
+        pass
+
+
+class WandBCallback(TrainingCallback):
+    def __init__(
+        self,
+        project_name: str,
+        log_dir: str,
+        config: dict,
+        wrapped_callback: TrainingCallback = None,
+    ):
+        if wandb is None:
+            raise ImportError(
+                "wandb is not installed. Please install it to use WandBCallback."
+            )
+        self.wrapped_callback = wrapped_callback
+        wandb.init(project=project_name, dir=log_dir, config=config)
+
+    def on_train_loss_report(self, train_info: dict):
+        wandb.log(train_info)
+        if self.wrapped_callback:
+            self.wrapped_callback.on_train_loss_report(train_info)
+
+    def on_val_loss_report(self, val_info: dict):
+        wandb.log(val_info)
+        if self.wrapped_callback:
+            self.wrapped_callback.on_val_loss_report(val_info)

mlx_lm/tuner/trainer.py

@@ -1,20 +1,18 @@
 # Copyright © 2024 Apple Inc.
 
-import glob
-import shutil
+
 import time
 from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
-from typing import List, Optional, Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten
-from transformers import PreTrainedTokenizer
 
+from .callbacks import TrainingCallback
 from .datasets import CacheDataset
 
 
@@ -183,17 +181,6 @@ def evaluate(
     return (all_losses / ntokens).item()
 
 
-class TrainingCallback:
-
-    def on_train_loss_report(self, train_info: dict):
-        """Called to report training loss at specified intervals."""
-        pass
-
-    def on_val_loss_report(self, val_info: dict):
-        """Called to report validation loss at specified intervals or the beginning."""
-        pass
-
-
 def train(
     model,
     optimizer,