From 25b8dd7cc4ac566cfcbc092ee664ffbb7ad0271c Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 14 Mar 2025 09:44:44 +0100 Subject: [PATCH 01/13] update lora_config.yaml + LORA.md + lora.py --- mlx_lm/LORA.md | 4 ++++ mlx_lm/examples/lora_config.yaml | 3 +++ mlx_lm/lora.py | 26 ++++++++++++++++++++++++++ 3 files changed, 33 insertions(+) diff --git a/mlx_lm/LORA.md b/mlx_lm/LORA.md index e863abc46..e6f8d3a5d 100644 --- a/mlx_lm/LORA.md +++ b/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) diff --git a/mlx_lm/examples/lora_config.yaml b/mlx_lm/examples/lora_config.yaml index 36bc1dff8..b1db8a81b 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 +report_to_wand: true + # Number of training steps between loss reporting. steps_per_report: 10 diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 042b40e2b..72160ef4a 100644 --- a/mlx_lm/lora.py +++ b/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,25 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set def run(args, training_callback: TrainingCallback = None): np.random.seed(args.seed) + 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) From abbbdcabb18ddfbe731f5f50c11259fdae90b411 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Tue, 18 Mar 2025 23:04:41 +0100 Subject: [PATCH 02/13] code formatting --- mlx_lm/lora.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 72160ef4a..63f5440b4 100644 --- a/mlx_lm/lora.py +++ b/mlx_lm/lora.py @@ -296,21 +296,23 @@ def run(args, training_callback: TrainingCallback = None): 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") From c39532f757f387900a262fba9bdd0aa0cd308626 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Mon, 31 Mar 2025 21:21:27 +0200 Subject: [PATCH 03/13] udpaet Acknowledgements.md --- ACKNOWLEDGMENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 54725af3d..ac2404f53 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`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`. +- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`: Added the following features: `reporting training metrics to WandB (Weights & Biases)`. From a0bbe1171dab81754510a05448e8c35c05e6a953 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Mon, 31 Mar 2025 21:23:49 +0200 Subject: [PATCH 04/13] nits --- mlx_lm/lora.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 63f5440b4..b1a8e3dd6 100644 --- a/mlx_lm/lora.py +++ b/mlx_lm/lora.py @@ -299,19 +299,18 @@ def run(args, training_callback: TrainingCallback = None): 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) + training_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.on_val_loss_report(val_info) training_callback = WandBCallback() From 0304b2836cf27f5986128c771aed9a3078ab3705 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Sat, 5 Apr 2025 14:17:06 +0200 Subject: [PATCH 05/13] 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. --- mlx_lm/lora.py | 36 +++++++++++------------------------- mlx_lm/tuner/callbacks.py | 33 +++++++++++++++++++++++++++++++++ mlx_lm/tuner/trainer.py | 17 ++--------------- 3 files changed, 46 insertions(+), 40 deletions(-) create mode 100644 mlx_lm/tuner/callbacks.py diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 115e0dcb0..269df2cca 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 @@ -23,6 +21,7 @@ print_trainable_parameters, ) from .utils import load, save_config +from .tuner.callbacks import WandBCallback yaml_loader = yaml.SafeLoader yaml_loader.add_implicit_resolver( @@ -69,7 +68,7 @@ "lr_schedule": None, "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}, "mask_prompt": False, - "report_to_wandb": False, + "wandb": None, } @@ -182,10 +181,10 @@ def build_parser(): default=None, ) parser.add_argument( - "--report-to-wandb", - action="store_true", - help="Report the training args to WandB.", + "--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 @@ -292,25 +291,12 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set def run(args, training_callback: TrainingCallback = None): np.random.seed(args.seed) - if args.report_to_wandb: - import wandb - - wandb.init(project="mlx-finetuning", config=vars(args)) - - original_callback = training_callback - - class WandBCallback(TrainingCallback): - def on_train_loss_report(self, train_info: dict): - wandb.log(train_info) - if original_callback: - training_callback.on_train_loss_report(train_info) - - def on_val_loss_report(self, val_info: dict): - wandb.log(val_info) - if original_callback: - training_callback.on_val_loss_report(val_info) - - training_callback = WandBCallback() + if args.wandb is not None: + training_callback = WandBCallback( + project_name=args.wandb, + 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..18fe0b23a --- /dev/null +++ b/mlx_lm/tuner/callbacks.py @@ -0,0 +1,33 @@ +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 + + +try: + import wandb +except ImportError: + wandb = None + + +class WandBCallback(TrainingCallback): + def __init__(self, project_name: 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, 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 af0c24507..41c6239c6 100644 --- a/mlx_lm/tuner/trainer.py +++ b/mlx_lm/tuner/trainer.py @@ -1,21 +1,19 @@ # 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 .datasets import CacheDataset +from .callbacks import TrainingCallback def grad_checkpoint(layer): @@ -187,17 +185,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, tokenizer, From a3a6d495756ff9ad94278212dc999b1655690a66 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Sat, 5 Apr 2025 14:21:17 +0200 Subject: [PATCH 06/13] 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. --- mlx_lm/lora.py | 1 + mlx_lm/tuner/callbacks.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index 269df2cca..bc3001765 100644 --- a/mlx_lm/lora.py +++ b/mlx_lm/lora.py @@ -294,6 +294,7 @@ def run(args, training_callback: TrainingCallback = None): 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 ) diff --git a/mlx_lm/tuner/callbacks.py b/mlx_lm/tuner/callbacks.py index 18fe0b23a..3a77414e6 100644 --- a/mlx_lm/tuner/callbacks.py +++ b/mlx_lm/tuner/callbacks.py @@ -16,11 +16,11 @@ def on_val_loss_report(self, val_info: dict): class WandBCallback(TrainingCallback): - def __init__(self, project_name: str, config: dict, wrapped_callback: TrainingCallback = None): + 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, config=config) + wandb.init(project=project_name, dir=log_dir, config=config) def on_train_loss_report(self, train_info: dict): wandb.log(train_info) From e7102055af79d1e7ff9d1b8ef31589804ed230cd Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Sat, 5 Apr 2025 14:36:06 +0200 Subject: [PATCH 07/13] nits --- mlx_lm/tuner/callbacks.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mlx_lm/tuner/callbacks.py b/mlx_lm/tuner/callbacks.py index 3a77414e6..7ab12eb22 100644 --- a/mlx_lm/tuner/callbacks.py +++ b/mlx_lm/tuner/callbacks.py @@ -1,3 +1,9 @@ +try: + import wandb +except ImportError: + wandb = None + + class TrainingCallback: def on_train_loss_report(self, train_info: dict): @@ -9,12 +15,6 @@ def on_val_loss_report(self, val_info: dict): pass -try: - import wandb -except ImportError: - wandb = None - - class WandBCallback(TrainingCallback): def __init__(self, project_name: str, log_dir: str, config: dict, wrapped_callback: TrainingCallback = None): if wandb is None: From 1db8c2867c508325e0577f1cf70ac7bb7df1e320 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Wed, 9 Apr 2025 19:22:49 +0200 Subject: [PATCH 08/13] formating --- mlx_lm/lora.py | 4 ++-- mlx_lm/tuner/callbacks.py | 12 ++++++++++-- mlx_lm/tuner/trainer.py | 2 +- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/mlx_lm/lora.py b/mlx_lm/lora.py index bc3001765..a5e72ff44 100644 --- a/mlx_lm/lora.py +++ b/mlx_lm/lora.py @@ -12,6 +12,7 @@ import yaml from .tokenizer_utils import TokenizerWrapper +from .tuner.callbacks import WandBCallback from .tuner.datasets import load_dataset from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train from .tuner.utils import ( @@ -21,7 +22,6 @@ print_trainable_parameters, ) from .utils import load, save_config -from .tuner.callbacks import WandBCallback yaml_loader = yaml.SafeLoader yaml_loader.add_implicit_resolver( @@ -296,7 +296,7 @@ def run(args, training_callback: TrainingCallback = None): project_name=args.wandb, log_dir=args.adapter_path, config=vars(args), - wrapped_callback=training_callback + wrapped_callback=training_callback, ) print("Loading pretrained model") diff --git a/mlx_lm/tuner/callbacks.py b/mlx_lm/tuner/callbacks.py index 7ab12eb22..8e3a15788 100644 --- a/mlx_lm/tuner/callbacks.py +++ b/mlx_lm/tuner/callbacks.py @@ -16,9 +16,17 @@ def on_val_loss_report(self, val_info: dict): class WandBCallback(TrainingCallback): - def __init__(self, project_name: str, log_dir: str, config: dict, wrapped_callback: TrainingCallback = None): + 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.") + 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) diff --git a/mlx_lm/tuner/trainer.py b/mlx_lm/tuner/trainer.py index 41c6239c6..98a3da192 100644 --- a/mlx_lm/tuner/trainer.py +++ b/mlx_lm/tuner/trainer.py @@ -12,8 +12,8 @@ from mlx.nn.utils import average_gradients from mlx.utils import tree_flatten -from .datasets import CacheDataset from .callbacks import TrainingCallback +from .datasets import CacheDataset def grad_checkpoint(layer): From 7fd78fda2eb4589d869a19b79b6fb500488c0ca0 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 9 May 2025 22:55:56 +0200 Subject: [PATCH 09/13] README.md --- mlx_lm/LORA.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/LORA.md b/mlx_lm/LORA.md index 92335bd9c..ab4bc0051 100644 --- a/mlx_lm/LORA.md +++ b/mlx_lm/LORA.md @@ -389,7 +389,7 @@ 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. +You can log training metrics to Weights & Biases by adding the `--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. From 35d9d6812c799b9f1754e5272cfc84ec113061a5 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 9 May 2025 22:57:18 +0200 Subject: [PATCH 10/13] update example yaml --- mlx_lm/examples/lora_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/examples/lora_config.yaml b/mlx_lm/examples/lora_config.yaml index 826f3e31d..ea2d93d6b 100644 --- a/mlx_lm/examples/lora_config.yaml +++ b/mlx_lm/examples/lora_config.yaml @@ -38,7 +38,7 @@ val_batches: 25 learning_rate: 1e-5 # Whether to report the logs to WandB -report_to_wand: true +wand: null # Number of training steps between loss reporting. steps_per_report: 10 From e5ed455cc1764163133f72eefa2d28eef89a7c22 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 9 May 2025 22:58:05 +0200 Subject: [PATCH 11/13] nits --- mlx_lm/examples/lora_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/examples/lora_config.yaml b/mlx_lm/examples/lora_config.yaml index ea2d93d6b..2e9f58d82 100644 --- a/mlx_lm/examples/lora_config.yaml +++ b/mlx_lm/examples/lora_config.yaml @@ -38,7 +38,7 @@ val_batches: 25 learning_rate: 1e-5 # Whether to report the logs to WandB -wand: null +# wand: 'wandb-project' # Number of training steps between loss reporting. steps_per_report: 10 From 325df0e4f5af8b09e8fbe7eb79d0a5b6cc9c7171 Mon Sep 17 00:00:00 2001 From: Goekdeniz-Guelmez Date: Fri, 9 May 2025 22:58:23 +0200 Subject: [PATCH 12/13] nits --- mlx_lm/examples/lora_config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx_lm/examples/lora_config.yaml b/mlx_lm/examples/lora_config.yaml index 2e9f58d82..a4fe069ba 100644 --- a/mlx_lm/examples/lora_config.yaml +++ b/mlx_lm/examples/lora_config.yaml @@ -38,7 +38,7 @@ val_batches: 25 learning_rate: 1e-5 # Whether to report the logs to WandB -# wand: 'wandb-project' +# wand: 'wandb-project" # Number of training steps between loss reporting. steps_per_report: 10 From 50f5381afc529577756bfdc1bc146bd60df61c28 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Fri, 9 May 2025 14:02:21 -0700 Subject: [PATCH 13/13] nits in readme --- mlx_lm/LORA.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mlx_lm/LORA.md b/mlx_lm/LORA.md index ab4bc0051..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 @@ -387,10 +392,6 @@ 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 `--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)