Support prompt-completion data for LoRA fine-tuning

lora/README.md

@@ -186,6 +186,15 @@ file should look like:
 
 Note other keys will be ignored by the loader.
 
+For instruction tuning, each line can also provide a prompt and response:
+
+```
+{"prompt": "[INST] Your input prompt here [/INST]", "text": "The expected output result here"}
+```
+
+When `prompt` is present, the loader trains on the tokens from `text` only. The
+prompt is included as context, but prompt tokens are masked out of the loss.
+
 ## Memory Issues
 
 Fine-tuning a large model with LoRA requires a machine with a decent amount

lora/lora.py

@@ -16,9 +16,6 @@
 from mlx.utils import tree_flatten
 from models import LoRALinear
 
-# Disable output buffering to see print statements in real-time
-sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
-
 
 def build_parser():
     parser = argparse.ArgumentParser(description="LoRA or QLoRA finetuning.")
@@ -142,7 +139,10 @@ def __init__(self, path: Path, key: str = "text"):
         self._key = key
 
     def __getitem__(self, idx: int):
-        return self._data[idx][self._key]
+        item = self._data[idx]
+        if "prompt" in item:
+            return item
+        return item[self._key]
 
     def __len__(self):
         return len(self._data)
@@ -175,21 +175,54 @@ def load_and_check(name):
     return train, valid, test
 
 
-def loss(model, inputs, targets, lengths):
+def loss(model, inputs, targets, loss_masks):
     # Run model on inputs
     logits, _ = model(inputs)
     logits = logits.astype(mx.float32)
 
-    # Mask padding tokens
-    length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]
-
     # Calculate the loss
-    ce = nn.losses.cross_entropy(logits, targets) * length_mask
-    ntoks = length_mask.sum()
+    ce = nn.losses.cross_entropy(logits, targets) * loss_masks
+    ntoks = loss_masks.sum()
     ce = ce.sum() / ntoks
     return ce, ntoks
 
 
+def _encode(tokenizer, text, **kwargs):
+    try:
+        return tokenizer.encode(text, **kwargs)
+    except TypeError:
+        return tokenizer.encode(text)
+
+
+def _remove_trailing_eos(tokens, tokenizer):
+    eos_token_id = getattr(tokenizer, "eos_token_id", None)
+    if eos_token_id is not None and tokens and tokens[-1] == eos_token_id:
+        return tokens[:-1]
+    return tokens
+
+
+def encode_dataset_item(item, tokenizer):
+    if not isinstance(item, dict):
+        tokens = _encode(tokenizer, item)
+        return tokens, np.ones(len(tokens) - 1, dtype=np.float32)
+
+    if "prompt" not in item or "text" not in item:
+        raise ValueError("Prompt-completion examples must contain 'prompt' and 'text'.")
+
+    prompt_tokens = _remove_trailing_eos(_encode(tokenizer, item["prompt"]), tokenizer)
+    completion_tokens = _encode(tokenizer, item["text"], add_special_tokens=False)
+    tokens = prompt_tokens + completion_tokens
+
+    if getattr(tokenizer, "add_eos_token", False):
+        eos_token_id = getattr(tokenizer, "eos_token_id", None)
+        if eos_token_id is not None and (not tokens or tokens[-1] != eos_token_id):
+            tokens.append(eos_token_id)
+
+    loss_mask = np.zeros(len(tokens) - 1, dtype=np.float32)
+    loss_mask[max(len(prompt_tokens) - 1, 0) :] = 1
+    return tokens, loss_mask
+
+
 def iterate_batches(dset, tokenizer, batch_size, train=False):
     # Shuffle indices
     while True:
@@ -200,7 +233,12 @@ def iterate_batches(dset, tokenizer, batch_size, train=False):
         # Collect batches from dataset
         for i in range(0, len(indices) - batch_size + 1, batch_size):
             # Encode batch
-            batch = [tokenizer.encode(dset[indices[i + j]]) for j in range(batch_size)]
+            batch, loss_masks = zip(
+                *[
+                    encode_dataset_item(dset[indices[i + j]], tokenizer)
+                    for j in range(batch_size)
+                ]
+            )
             lengths = [len(x) for x in batch]
 
             # Check if any sequence is longer than 2048 tokens
@@ -212,11 +250,13 @@ def iterate_batches(dset, tokenizer, batch_size, train=False):
 
             # Pad to the max length
             batch_arr = np.zeros((batch_size, max(lengths)), np.int32)
+            loss_mask_arr = np.zeros((batch_size, max(lengths) - 1), np.float32)
 
             for j in range(batch_size):
                 batch_arr[j, : lengths[j]] = batch[j]
+                loss_mask_arr[j, : lengths[j] - 1] = loss_masks[j]
             batch = mx.array(batch_arr)
-            yield batch[:, :-1], batch[:, 1:], mx.array(lengths)
+            yield batch[:, :-1], batch[:, 1:], mx.array(loss_mask_arr)
 
         if not train:
             break
@@ -327,6 +367,9 @@ def generate(model, prompt, tokenizer, args):
 
 
 if __name__ == "__main__":
+    # Disable output buffering to see print statements in real-time
+    sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1)
+
     parser = build_parser()
     args = parser.parse_args()
 

lora/test_lora.py

@@ -0,0 +1,68 @@
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+
+
+def _module(name, **attrs):
+    module = types.ModuleType(name)
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    return module
+
+
+sys.modules["mlx"] = _module("mlx")
+sys.modules["mlx.core"] = _module("mlx.core")
+sys.modules["mlx.nn"] = _module(
+    "mlx.nn",
+    losses=_module("mlx.nn.losses", cross_entropy=lambda logits, targets: logits),
+    value_and_grad=lambda model, loss: None,
+)
+sys.modules["mlx.optimizers"] = _module("mlx.optimizers", Adam=object)
+sys.modules["mlx.utils"] = _module("mlx.utils", tree_flatten=lambda tree: [])
+sys.modules["models"] = _module("models", LoRALinear=object)
+sys.modules["utils"] = _module("utils")
+
+sys.path.insert(0, str(Path(__file__).parent))
+spec = importlib.util.spec_from_file_location(
+    "lora_example", Path(__file__).with_name("lora.py")
+)
+lora = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(lora)
+
+
+class DummyTokenizer:
+    eos_token_id = 0
+    add_eos_token = True
+
+    def encode(self, text, add_special_tokens=True):
+        tokens = [ord(c) for c in text]
+        if add_special_tokens:
+            tokens = [1] + tokens
+            if self.add_eos_token:
+                tokens.append(self.eos_token_id)
+        return tokens
+
+
+def test_encode_dataset_item_with_text_trains_on_all_tokens():
+    tokens, loss_mask = lora.encode_dataset_item("abc", DummyTokenizer())
+
+    assert tokens == [1, 97, 98, 99, 0]
+    np.testing.assert_array_equal(loss_mask, np.ones(4, dtype=np.float32))
+
+
+def test_encode_dataset_item_with_prompt_masks_prompt_tokens():
+    tokens, loss_mask = lora.encode_dataset_item(
+        {"prompt": "[INST] Say hi [/INST]", "text": "Hi."},
+        DummyTokenizer(),
+    )
+
+    prompt_tokens = DummyTokenizer().encode("[INST] Say hi [/INST]")[:-1]
+    assert tokens[: len(prompt_tokens)] == prompt_tokens
+    assert tokens[-4:] == [72, 105, 46, 0]
+    np.testing.assert_array_equal(
+        loss_mask,
+        np.array([0] * (len(prompt_tokens) - 1) + [1, 1, 1, 1], dtype=np.float32),
+    )