enable memory efficient fine tuning for very long sequences

mlx_lm/lora.py

@@ -168,6 +168,12 @@ def build_parser():
         type=int,
         help="Maximum sequence length.",
     )
+    parser.add_argument(
+        "--seq-step-size",
+        type=int,
+        default=None,
+        help="",
+    )
     parser.add_argument(
         "-c",
         "--config",
@@ -238,6 +244,7 @@ def train_model(
         adapter_file=adapter_file,
         max_seq_length=args.max_seq_length,
         grad_checkpoint=args.grad_checkpoint,
+        seq_step_size=args.seq_step_size,
     )
 
     # Initialize the selected optimizer

mlx_lm/tuner/trainer.py

@@ -16,6 +16,15 @@ from mlx.utils import tree_flatten
 from transformers import PreTrainedTokenizer
 
 from .datasets import CacheDataset
+from ..models.cache import make_prompt_cache, KVCache
+
+
+def reset_prompt_cache(cache):
+    for e, c in enumerate(cache):
+        if isinstance(c, KVCache):
+            cache[e] = KVCache()
+        else:
+            raise ValueError("Unsupported cache")
 
 
 def grad_checkpoint(layer):
@@ -65,16 +74,21 @@ class TrainingArgs:
         default=False,
         metadata={"help": "Use gradient checkpointing to reduce memory use."},
     )
+    seq_step_size : Optional[int] = field(
+        default=None,
+        metadata={"help": "The examples are processsed in seq_step_size chunks."},
+    )
 
 
-def default_loss(model, batch, lengths):
+def default_loss(model, batch, lengths, cache):
     inputs = batch[:, :-1]
     targets = batch[:, 1:]
 
-    logits = model(inputs)
+    offset = cache[0].offset
+    logits = model(inputs, cache=cache)
     logits = logits.astype(mx.float32)
 
-    steps = mx.arange(1, targets.shape[1] + 1)
+    steps = mx.arange(1, targets.shape[1] + 1) + offset
     mask = mx.logical_and(steps >= lengths[:, 0:1], steps <= lengths[:, 1:])
 
     ce = nn.losses.cross_entropy(logits, targets) * mask
@@ -160,6 +174,7 @@ def evaluate(
     max_seq_length=2048,
     loss: callable = default_loss,
     iterate_batches: callable = iterate_batches,
+    seq_step_size: Optional[int] = None,
 ):
     model.eval()
     all_losses = mx.array(0.0)
@@ -167,6 +182,8 @@ def evaluate(
 
     index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)
 
+    seq_step_size = seq_step_size or max_seq_length
+
     for _, batch in zip(
         index_iterator,
         iterate_batches(
@@ -176,10 +193,14 @@ def evaluate(
             max_seq_length=max_seq_length,
         ),
     ):
-        losses, toks = loss(model, *batch)
-        all_losses += losses * toks
-        ntokens += toks
-        mx.eval(all_losses, ntokens)
+        cache = make_prompt_cache(model)
+        seq_length = batch[0].shape[1]
+        for s in range(0, seq_length, seq_step_size):
+            local_batch = (batch[0][:, s:s+seq_step_size], batch[1])
+            losses, toks = loss(model, *local_batch, cache)
+            all_losses += losses * toks
+            ntokens += toks
+            mx.eval(all_losses, ntokens)
 
     all_losses = mx.distributed.all_sum(all_losses, stream=mx.cpu)
     ntokens = mx.distributed.all_sum(ntokens, stream=mx.cpu)
@@ -220,12 +241,13 @@ def train(
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
 
+    cache = make_prompt_cache(model)
     state = [model.state, optimizer.state, mx.random.state]
 
     @partial(mx.compile, inputs=state, outputs=state)
     def step(batch):
         # Forward and backward pass
-        (lvalue, toks), grad = loss_value_and_grad(model, *batch)
+        (lvalue, toks), grad = loss_value_and_grad(model, *batch, cache)
 
         # All reduce the gradients if running in distributed mode
         grad = average_gradients(grad)
@@ -241,6 +263,7 @@ def train(
     loss_value_and_grad = nn.value_and_grad(model, loss)
 
     model.train()
+    seq_step_size = args.seq_step_size or args.max_seq_length
     losses = 0
     n_tokens = 0
     steps = 0
@@ -262,6 +285,7 @@ def train(
         # is always measured before any training.
         if it == 1 or it % args.steps_per_eval == 0 or it == args.iters:
             tic = time.perf_counter()
+            val_loss = 0.0
             val_loss = evaluate(
                 model=model,
                 dataset=val_dataset,
@@ -271,6 +295,7 @@ def train(
                 num_batches=args.val_batches,
                 max_seq_length=args.max_seq_length,
                 iterate_batches=iterate_batches,
+                seq_step_size=seq_step_size,
             )
             model.train()
             val_time = time.perf_counter() - tic
@@ -292,11 +317,15 @@ def train(
 
             tic = time.perf_counter()
 
-        lvalue, toks = step(batch)
-        losses += lvalue
-        n_tokens += toks
-        steps += 1
-        mx.eval(state, losses, n_tokens)
+        seq_length = batch[0].shape[1]
+        for s in range(0, seq_length, seq_step_size):
+            local_batch = (batch[0][:, s:s+seq_step_size], batch[1])
+            lvalue, toks = step(local_batch)
+            losses += lvalue
+            n_tokens += toks
+            steps += 1
+            mx.eval(state, losses, n_tokens)
+        reset_prompt_cache(cache)
         train_time += time.perf_counter() - tic
 
         # Report training loss if needed