Adding grpo training

llms/mlx_lm/LORA.md

@@ -18,6 +18,7 @@ LoRA (QLoRA).[^qlora] LoRA fine-tuning works with the following model families:
 
 - [Run](#Run)
   - [Fine-tune](#Fine-tune)
+  - [GRPO](#GRPO)
   - [Evaluate](#Evaluate)
   - [Generate](#Generate)
 - [Fuse](#Fuse)
@@ -84,6 +85,33 @@ ignore the prompt and compute loss for just the completion by passing
 datasets. For `chat` datasets the final message in the message list is
 considered the completion. See the [dataset section](#Data) for more details. 
 
+### Group Relative Policy Optimization (GRPO)
+
+To fine-tune a model using GRPO, which optimizes policy using multiple responses per prompt, use:
+
+```shell
+mlx_lm.lora \
+    --model <path_to_model> \
+    --train \
+    --data <path_to_data> \
+    --fine-tune-type grpo \
+    --group-size 4
+```
+
+GRPO specific arguments:
+
+- `--group-size`: Number of responses generated per prompt (default: 4)
+- `--beta`: KL penalty coefficient for policy optimization (default: 0.1)
+- `--epsilon`: Small constant for numerical stability (default: 1e-4)
+- `--max-completion-length`: Maximum length of generated completions (default: 512)
+- `--reference-model-path`: Path to reference model weights. If not specified, uses the same model
+- `--temperature`: Sampling temperature for generations. Higher values increase randomness (default: 1.0)
+- `--reward-weights`: Optional list of weights for multiple reward functions. Must match number of reward functions. If not specified, all rewards weighted equally with 1.0
+
+The GRPO training method generates multiple responses for each prompt and optimizes the policy using relative rewards between responses. This approach helps improve response quality by learning from comparisons between different completions.
+
+Note that GRPO requires more compute resources than standard LoRA training since it generates multiple responses per prompt. Consider reducing batch size or using gradient checkpointing if running into memory issues.
+
 ### Evaluate
 
 To compute test set perplexity use:

llms/mlx_lm/lora.py

@@ -1,20 +1,21 @@
 # Copyright © 2024 Apple Inc.
 
+from pathlib import Path
 import argparse
+import types
 import math
 import os
 import re
-import types
-from pathlib import Path
 
-import mlx.nn as nn
 import mlx.optimizers as optim
+import mlx.nn as nn
 import numpy as np
 import yaml
 
+from .tuner.grpo_trainer import GRPOTrainingArgs, evaluate_grpo, train_grpo
+from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tokenizer_utils import TokenizerWrapper
 from .tuner.datasets import load_dataset
-from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
     build_schedule,
     linear_to_lora_layers,
@@ -42,6 +43,7 @@
 CONFIG_DEFAULTS = {
     "model": "mlx_model",
     "train": False,
+    "training_mode": "normal",
     "fine_tune_type": "lora",
     "data": "data/",
     "seed": 0,
@@ -63,6 +65,17 @@
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
     "mask_prompt": False,
+
+    # GRPO args
+    "reference_model_path": None,
+    "group_size": 4,
+    "beta": 0.1,
+    "epsilon": 1e-4,
+    "max_completion_length": 512,
+    "use_chat_template": False,
+    "use_prompt": False,
+    "temperature": 1.0,
+    "reward_weights": None
 }
 
 
@@ -103,6 +116,12 @@ def build_parser():
         default=None,
     )
 
+    parser.add_argument(
+        "--training-mode",
+        type=str,
+        choices=["normal", "grpo"],
+        help="Training mode: normal or GRPO",
+    )
     parser.add_argument(
         "--num-layers",
         type=int,
@@ -170,8 +189,93 @@ def build_parser():
         default=None,
     )
     parser.add_argument("--seed", type=int, help="The PRNG seed")
+
+    # GRPO args
+    parser.add_argument(
+        "--group-size",
+        type=int,
+        help="Number of generations.",
+        default=4,
+    )
+    parser.add_argument(
+        "--max-completion-length",
+        type=int,
+        help="Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left.",
+        default=512,
+    )
+    parser.add_argument(
+        "--beta",
+        type=float,
+        help="KL penalty coefficient.",
+        default=0.1,
+    )
+    parser.add_argument(
+        "--epsilon",
+        type=float,
+        help="The Epsilon for numerical stability.",
+        default=1e-4,
+    )
+    parser.add_argument(
+        "--use-chat-template",
+        action="store_true",
+        help="If the model is a Chat model, use the Chat template.",
+        default=None,
+    )
+    parser.add_argument(
+        "--use-prompt",
+        action="store_true",
+        help="Rather to use the prompt from the R1 paper.",
+        default=None,
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        help="Temperature for sampling. The higher the temperature, the more random the completions.",
+        default=1.0,
+    )
+    parser.add_argument(
+        "--reward-weights",
+        type=str,
+        help="Weights for each reward function. Must match the number of reward functions and be in this format [0.1, 0.2, 0.3, 0.4, 0.5]. If not given, all rewards are weighted equally with weight `1.0`.",
+        default=None,
+    )
     return parser
 
+def train_model_grpo(model, tokenizer, args, opt, train_set, valid_set, adapter_file, training_callback):
+    training_args = GRPOTrainingArgs(
+        batch_size=args.batch_size,
+        iters=args.iters,
+        val_batches=args.val_batches,
+        steps_per_report=args.steps_per_report,
+        steps_per_eval=args.steps_per_eval,
+        steps_per_save=args.save_every,
+        adapter_file=adapter_file,
+        max_seq_length=args.max_seq_length,
+        max_completion_length=args.max_completion_length,
+        grad_checkpoint=args.grad_checkpoint,
+        beta=args.beta,
+        group_size=args.group_size,
+        epsilon=args.epsilon,
+        reference_model_path=args.reference_model_path,
+        temperature=args.temperature,
+        reward_weights=[float(x) for x in args.reward_weights.strip('[]').split(',')] if args.reward_weights else None
+    )
+
+    if args.reference_model_path:
+        reference_model, _ = load(args.reference_model_path)
+    else:
+        reference_model, _ = load(args.model)
+
+    train_grpo(
+        model=model,
+        ref_model=reference_model.freeze(),
+        tokenizer=tokenizer,
+        optimizer=opt,
+        train_dataset=train_set,
+        val_dataset=valid_set,
+        args=training_args,
+        training_callback=training_callback,
+    )
 
 def train_model(
     args,
@@ -215,19 +319,6 @@ def train_model(
     adapter_file = adapter_path / "adapters.safetensors"
     save_config(vars(args), adapter_path / "adapter_config.json")
 
-    # init training args
-    training_args = TrainingArgs(
-        batch_size=args.batch_size,
-        iters=args.iters,
-        val_batches=args.val_batches,
-        steps_per_report=args.steps_per_report,
-        steps_per_eval=args.steps_per_eval,
-        steps_per_save=args.save_every,
-        adapter_file=adapter_file,
-        max_seq_length=args.max_seq_length,
-        grad_checkpoint=args.grad_checkpoint,
-    )
-
     model.train()
     opt = optim.Adam(
         learning_rate=(
@@ -236,32 +327,82 @@ def train_model(
     )
 
     # Train model
-    train(
-        model=model,
-        tokenizer=tokenizer,
-        args=training_args,
-        optimizer=opt,
-        train_dataset=train_set,
-        val_dataset=valid_set,
-        training_callback=training_callback,
-    )
+    if args.training_mode == "grpo":
+        train_model_grpo(
+            model,
+            tokenizer,
+            args,
+            opt,
+            train_set,
+            valid_set,
+            adapter_file,
+            training_callback
+        )
+    else:
+        training_args = TrainingArgs(
+            batch_size=args.batch_size,
+            iters=args.iters,
+            val_batches=args.val_batches,
+            steps_per_report=args.steps_per_report,
+            steps_per_eval=args.steps_per_eval,
+            steps_per_save=args.save_every,
+            adapter_file=adapter_file,
+            max_seq_length=args.max_seq_length,
+            grad_checkpoint=args.grad_checkpoint
+        )
+
+        train(
+            model=model,
+            tokenizer=tokenizer,
+            args=training_args,
+            optimizer=opt,
+            train_dataset=train_set,
+            val_dataset=valid_set,
+            training_callback=training_callback,
+        )
 
 
 def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
     model.eval()
 
-    test_loss = evaluate(
-        model=model,
-        dataset=test_set,
-        tokenizer=tokenizer,
-        batch_size=args.batch_size,
-        num_batches=args.test_batches,
-        max_seq_length=args.max_seq_length,
-    )
+    if args.training_mode == "grpo":
+        if args.reference_model_path:
+            reference_model, _ = load(args.reference_model_path)
+        else:
+            reference_model, _ = load(args.model)
+
+        test_loss, _, test_rewards = evaluate_grpo(
+            model=model,
+            ref_model=reference_model.freeze(),
+            dataset=test_set,
+            tokenizer=tokenizer,
+            batch_size=args.batch_size,
+            num_batches=args.test_batches,
+            max_seq_length=args.max_seq_length,
+            beta=args.beta,
+            group_size=args.group_size,
+            epsilon=args.epsilon,
+            temperature=args.temperature,
+            max_tokens=args.max_seq_length
+        )
+
+        test_ppl = math.exp(test_loss)
+
+        rewards_str = ", ".join([f"{k}: {v:.3f}" for k, v in test_rewards.items()])
+        print(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}, Rewards: {rewards_str}")
+    else:
+        test_loss = evaluate(
+            model=model,
+            dataset=test_set,
+            tokenizer=tokenizer,
+            batch_size=args.batch_size,
+            num_batches=args.test_batches,
+            max_seq_length=args.max_seq_length,
+        )
 
-    test_ppl = math.exp(test_loss)
+        test_ppl = math.exp(test_loss)
 
-    print(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
+        print(f"Test loss {test_loss:.3f}, Test ppl {test_ppl:.3f}.")
 
 
 def run(args, training_callback: TrainingCallback = None):
@@ -312,4 +453,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()