diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 3c44c5ecb..b9686e733 100644 --- a/ACKNOWLEDGMENTS.md +++ b/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)`. diff --git a/mlx_lm/LORA.md b/mlx_lm/LORA.md index 7bfb42539..bfaf3d7b5 100644 --- a/mlx_lm/LORA.md +++ b/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 `. +#### 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 diff --git a/mlx_lm/examples/lora_config.yaml b/mlx_lm/examples/lora_config.yaml index 158d15173..a4fe069ba 100644 --- a/mlx_lm/examples/lora_config.yaml +++ b/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 diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 990bbd9a7..32a09a1c0 100644 --- a/mlx_lm/lora.py +++ b/mlx_lm/lora.py @@ -1,5 +1,3 @@ -# Copyright © 2024 Apple Inc. - import argparse import math import os @@ -13,6 +11,7 @@ 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 @@ "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) diff --git a/mlx_lm/tuner/callbacks.py b/mlx_lm/tuner/callbacks.py new file mode 100644 index 000000000..8e3a15788 --- /dev/null +++ b/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) diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py index 0c134c201..6c60ece99 100644 --- a/mlx_lm/tuner/trainer.py +++ b/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,