Adding reporting to wandb

llms/mlx_lm/LORA.md

@@ -387,6 +387,10 @@ tokens-per-second, using the MLX Example
 [`wikisql`](https://github.com/ml-explore/mlx-examples/tree/main/lora/data)
 data set.
 
+## Logging
+
+You can log training metrics to Weights & Biases by adding the `--report-to-wandb` flag. This requires installing wandb manually with `pip install wandb`. When enabled, all training and validation metrics will be logged to your wandb account.
+
 [^lora]: Refer to the [arXiv paper](https://arxiv.org/abs/2106.09685) for more details on LoRA.
 
 [^qlora]: Refer to the paper [QLoRA: Efficient Finetuning of Quantized LLMs](https://arxiv.org/abs/2305.14314)

llms/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
+report_to_wand: true
+
 # Number of training steps between loss reporting.
 steps_per_report: 10
 

llms/mlx_lm/lora.py

@@ -68,6 +68,7 @@
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
     "mask_prompt": False,
+    "report_to_wandb": False
 }
 
 
@@ -179,6 +180,12 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
+    parser.add_argument(
+        "--report-to-wandb",
+        action="store_true",
+        help="Report the training args to WandB.",
+        default=None,
+    )
     parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
@@ -287,6 +294,26 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
 
+    # Initialize WandB if requested
+    if args.report_to_wandb:
+        import wandb
+        wandb.init(project="mlx-finetuning", config=vars(args))
+
+        # Create a simple wandb callback that wraps the existing one
+        original_callback = training_callback
+        class WandBCallback(TrainingCallback):
+            def on_train_loss_report(self, train_info: dict):
+                wandb.log(train_info)
+                if original_callback:
+                    original_callback.on_train_loss_report(train_info)
+
+            def on_val_loss_report(self, val_info: dict):
+                wandb.log(val_info)
+                if original_callback:
+                    original_callback.on_val_loss_report(val_info)
+
+        training_callback = WandBCallback()
+
     print("Loading pretrained model")
     model, tokenizer = load(args.model)
 

llms/mlx_lm/tuner/trainer.py

@@ -1,20 +1,14 @@
 # Copyright © 2024 Apple Inc.
 
-import glob
-import shutil
 import time
 from dataclasses import dataclass, field
 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 .datasets import CompletionsDataset
 
 
 def grad_checkpoint(layer):