rpeft 0.1.0 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 74527e43eb7c5d06b33fa44b6bc80fc57d7af6b1039248b0dd889f64cd6a7a5b
-  data.tar.gz: 497bff43821c0b8b25059de10321ac93d61ca3eacf3428d40e2de8db37b86b37
+  metadata.gz: d6162af3628a1e82e72c85f6f1ebf42c1f4bd58298992d4bc7036766628595c4
+  data.tar.gz: 52b2ac5be8272a31978e751e6d388e010cc97625b8cabeea4f6f6690bad3cb52
 SHA512:
-  metadata.gz: 5b1d932dc735db4961c1c6532139c750f423e1ca83404bca9f7a7e041a680e42037be253ec75198bcd5cef501da1174d6673ef1c8ec4a3c0f28de8d8b85be78e
-  data.tar.gz: 6d0bca19e47a92cbc975ee0fb73c06f5e4d337b3d6164371f43fcbd2ab7b1065cdd0b1a4730068e895312435c6f2d6bc91a64b62ff45f87a54fa16ca2d3fe32a
+  metadata.gz: ac75ae251f4f228ee5e3fca1b104f7819d509c10644754fb15754ade5bc82d03af67e3619c1492c9fbf715ca26f82e3e673caaf493f46114fd100271e28760bc
+  data.tar.gz: a3a06a0a30dd6419f92180c4ce3fa12b20f1a0a0ae45e503f0cd9b25afd23611f3bb1ed580f7e60899ab42c4d05ec79f9d317919e01b12aa00cc5aaa8f5c3e5f

data/README.md

@@ -1,5 +1,8 @@
 # RPEFT
 
+[![Gem Version](https://badge.fury.io/rb/rpeft.svg)](https://badge.fury.io/rb/rpeft)
+[![CI](https://github.com/koic/rpeft/actions/workflows/main.yml/badge.svg)](https://github.com/koic/rpeft/actions/workflows/main.yml)
+
 RPEFT brings LoRA fine-tuning of large language models to Ruby. It provides an API and CLI that model the LoRA lifecycle,
 while the heavy lifting is done by the Python ecosystem (PyTorch, Hugging Face Transformers, PEFT, Datasets, Accelerate) through PyCall.
 It is an independent project, not affiliated with Hugging Face.
@@ -15,16 +18,16 @@ It is an independent project, not affiliated with Hugging Face.
 
 ## The model behind the API
 
-LoRA keeps the base model weights $W_0$ frozen and learns a low-rank update of rank $r$:
+[LoRA](https://arxiv.org/abs/2106.09685) keeps the base model weights $W_0$ frozen and learns a low-rank update of rank $r$, scaled by $\alpha / r$:
 
 ```math
 \begin{aligned}
 W &= W_0 + \Delta W \\
-\Delta W &\approx BA
+\Delta W &\approx \frac{\alpha}{r} BA
 \end{aligned}
 ```
 
-$A$ and $B$ are small matrices that together typically hold well under 1% of the model's parameters. Training produces them; merging adds $BA$ back into $W_0$.
+$A$ and $B$ are small matrices that together typically hold well under 1% of the model's parameters. Training produces them; merging adds $\frac{\alpha}{r} BA$ back into $W_0$.
 The API is the same shape:
 
 ```mermaid
@@ -34,7 +37,7 @@ flowchart LR
 
 Merging is optional: an adapter can also generate on top of its base model as is.
 
-The [`rank:` option](#training-options) of `RPEFT::LoRA.train` is exactly the $r$ above.
+The [`rank:` and `alpha:` options](#training-options) of `RPEFT::LoRA.train` are exactly the $r$ and $\alpha$ above.
 
 ## Requirements
 
@@ -96,7 +99,7 @@ adapter = RPEFT::LoRA.train(
 )
 adapter.save('./ruby-lora')
 
-model = adapter.merge # W = W0 + B A
+model = adapter.merge # W = W0 + (alpha / r) B A
 model.save('./ruby-model')
 ```
 
@@ -176,21 +179,22 @@ The rule behind this also applies to choosing a dataset format: the prompt at in
 
 ### Training options
 
-| Option            | Default | Meaning                                                                                 |
-| ----------------- | ------- | --------------------------------------------------------------------------------------- |
-| `rank:`           | `8`     | LoRA rank $r$, the rank of the low-rank update $BA$                                     |
-| `alpha:`          | `16`    | LoRA scaling factor (`lora_alpha` in PEFT)                                              |
-| `dropout:`        | `0.05`  | LoRA dropout (`lora_dropout` in PEFT)                                                   |
-| `target_modules:` | `nil`   | Modules to adapt, e.g. `%w(q_proj v_proj)`. `nil` uses PEFT's per-architecture defaults |
-| `epochs:`         | `3`     | Training epochs                                                                         |
-| `batch_size:`     | `4`     | Per-device batch size                                                                   |
-| `learning_rate:`  | `2e-4`  | Learning rate                                                                           |
-| `max_length:`     | `512`   | Max token length per example                                                            |
-| `eval_dataset:`   | `nil`   | Held-out dataset, in the same format as `dataset:`. Enables per-epoch evaluation        |
-| `checkpoints:`    | `false` | Checkpoint every epoch so an interrupted run can resume; removed once training ends     |
-| `resume:`         | `false` | Resume from the last checkpoint in `output_dir`, which must be given explicitly         |
-| `output_dir:`     | `nil`   | Adapter output directory: new, empty, or an earlier output. `nil` uses a temporary one  |
-| `device:`         | `:auto` | `:auto`, `:cuda`, `:mps`, or `:cpu`                                                     |
+| Option             | Default | Meaning                                                                                 |
+| ------------------ | ------- | --------------------------------------------------------------------------------------- |
+| `rank:`            | `8`     | LoRA rank $r$, the rank of the low-rank update $BA$                                     |
+| `alpha:`           | `16`    | LoRA scaling factor $\alpha$ (`lora_alpha` in PEFT)                                     |
+| `dropout:`         | `0.05`  | LoRA dropout (`lora_dropout` in PEFT)                                                   |
+| `target_modules:`  | `nil`   | Modules to adapt, e.g. `%w(q_proj v_proj)`. `nil` uses PEFT's per-architecture defaults |
+| `epochs:`          | `3`     | Training epochs                                                                         |
+| `batch_size:`      | `4`     | Per-device batch size                                                                   |
+| `learning_rate:`   | `2e-4`  | Learning rate                                                                           |
+| `max_length:`      | `512`   | Max token length per example                                                            |
+| `train_on_prompt:` | `false` | Compute the loss on prompts too, not only on answers (instruction and chat formats)     |
+| `eval_dataset:`    | `nil`   | Held-out dataset, in the same format as `dataset:`. Enables per-epoch evaluation        |
+| `checkpoints:`     | `false` | Checkpoint every epoch so an interrupted run can resume; removed once training ends     |
+| `resume:`          | `false` | Resume from the last checkpoint in `output_dir`, which must be given explicitly         |
+| `output_dir:`      | `nil`   | Adapter output directory: new, empty, or an earlier output. `nil` uses a temporary one  |
+| `device:`          | `:auto` | `:auto`, `:cuda`, `:mps`, or `:cpu`                                                     |
 
 ### Dataset
 
@@ -215,6 +219,10 @@ the ones every chat template understands:
 {"messages":[{"role":"user","content":"How do I reverse an array?"},{"role":"assistant","content":"[1, 2, 3].reverse"}]}
 ```
 
+In instruction and chat data, the prompt is read as context, but the loss is computed on the answers alone: the `output`, or every assistant turn after the first message with its end-of-turn marker.
+A `text` line has no such boundary, so it is trained on as a whole. Pass `train_on_prompt: true` to compute the loss on every token instead.
+Examples whose answer does not fit within `max_length:` are skipped with a warning, and the number skipped is recorded in the adapter metadata.
+
 `dataset:` also accepts any Enumerable of record Hashes in the same three shapes, so data built in Ruby needs no file of its own:
 
 ```ruby
@@ -248,7 +256,7 @@ RPEFT.check # => a Hash describing Python, packages, and devices
 
 ```console
 $ rpeft doctor
-Ruby:         4.0.6 (rpeft 0.1.0)
+Ruby:         4.0.6 (rpeft 0.2.0)
 Python:       3.12.13 (/path/to/.venv/bin/python)
 torch:        OK 2.13.0
 transformers: OK 5.15.0

data/lib/rpeft/cli/lora_train.rb

@@ -4,8 +4,9 @@ module RPEFT
   module CLI
     # @api private
     class LoRATrain < Command
-      PASSTHROUGH_OPTIONS = %i(rank alpha dropout target_modules epochs batch_size
-                               learning_rate max_length device).freeze
+      PASSTHROUGH_OPTIONS = %i(
+        rank alpha dropout target_modules epochs batch_size learning_rate max_length train_on_prompt device
+      ).freeze
       REQUIRED_OPTIONS = %i(model dataset output).freeze
       DEVICES = %w(auto cuda mps cpu).freeze
 
@@ -26,7 +27,7 @@ module RPEFT
           o.on('--rank N', Integer, 'LoRA rank r (default: 8)') do |v|
             options[:rank] = v
           end
-          o.on('--alpha N', Float, 'LoRA alpha (default: 16)') do |v|
+          o.on('--alpha N', Float, 'LoRA scaling factor alpha (default: 16)') do |v|
             options[:alpha] = v
           end
           o.on('--dropout F', Float, 'LoRA dropout (default: 0.05)') do |v|
@@ -48,6 +49,9 @@ module RPEFT
           o.on('--max-length N', Integer, 'Max token length (default: 512)') do |v|
             options[:max_length] = v
           end
+          o.on('--train-on-prompt', 'Also train on the prompt of instruction and chat examples') do
+            options[:train_on_prompt] = true
+          end
           o.on('--device DEVICE', DEVICES,
                'auto, cuda, mps, or cpu (default: auto)') { |v| options[:device] = v.to_sym }
         end

data/lib/rpeft/dataset.rb

@@ -10,9 +10,11 @@ module RPEFT
   #
   # Three JSONL line formats are accepted and detected from the first line: `{"text": ...}` is
   # passed through as is, `{"instruction": ..., "input": ..., "output": ...}` is rendered into
-  # a text prompt, and `{"messages": [{"role": ..., "content": ...}, ...]}` is passed through
+  # a prompt and its completion, and `{"messages": [{"role": ..., "content": ...}, ...]}` is passed through
   # for the backend to render with the tokenizer's chat template.
   class Dataset
+    BACKEND_FORMATS = {text: 'text', instruction: 'completion', messages: 'messages'}.freeze
+
     def self.coerce(value)
       return value if value.is_a?(Dataset)
       return from_records(value) if value.is_a?(Enumerable)
@@ -63,10 +65,14 @@ module RPEFT
       @temporary_dirs.clear
     end
 
-    # The format the backend sees. Instruction datasets are normalized to text here,
-    # so only text and messages ever cross the boundary.
+    # The format the backend sees. Instruction datasets are normalized to prompt and completion pairs here,
+    # which keeps the boundary the backend needs to train on the completion alone.
     def backend_format
-      format == :messages ? 'messages' : 'text'
+      BACKEND_FORMATS.fetch(format)
+    end
+
+    def messages?
+      format == :messages
     end
 
     private
@@ -112,8 +118,8 @@ module RPEFT
       raise DatasetError, "dataset is not JSONL (#{e.message}): #{path}"
     end
 
-    # Rewrites a non-text dataset into text-format JSONL in
-    # a temporary directory, so the backend only ever sees one format.
+    # Rewrites an instruction dataset into prompt and completion JSONL in a temporary directory,
+    # so the backend needs no knowledge of the instruction template.
     def normalize(source_path)
       dir = Dir.mktmpdir('rpeft-dataset-')
       @temporary_dirs << dir
@@ -124,7 +130,9 @@ module RPEFT
           next if line.strip.empty?
 
           record = parse_line(line, source_path)
-          out.puts JSON.generate('text' => render_instruction(record, lineno, source_path))
+          prompt, completion = render_instruction(record, lineno, source_path)
+
+          out.puts JSON.generate('prompt' => prompt, 'completion' => completion)
         end
       end
 
@@ -142,10 +150,11 @@ module RPEFT
       end
 
       input = record['input']
-      text = "### Instruction:\n#{instruction}\n"
-      text << "\n### Input:\n#{input}\n" if input && !input.to_s.empty?
-      text << "\n### Response:\n#{output}"
-      text
+      prompt = "### Instruction:\n#{instruction}\n"
+      prompt << "\n### Input:\n#{input}\n" if input && !input.to_s.empty?
+      prompt << "\n### Response:\n"
+
+      [prompt, output]
     end
   end
 end

data/lib/rpeft/lora/adapter.rb

@@ -84,8 +84,8 @@ module RPEFT
         self
       end
 
-      # Merges the learned B A update into the base model weights
-      # (W = W0 + B A) and returns the resulting standalone Model.
+      # Merges the learned update into the base model weights
+      # (W = W0 + (alpha / r) B A) and returns the resulting standalone Model.
       def merge(output_dir: nil)
         unless base_model
           raise MergeError, 'base model is unknown; pass base_model: to Adapter.load ' \

data/lib/rpeft/lora.rb

@@ -7,17 +7,18 @@ require 'tmpdir'
 
 module RPEFT
   # LoRA fine-tuning. Keeps the base model weights W0 frozen and learns
-  # a low-rank update delta W ~= B A of rank r:
+  # a low-rank update of rank r, scaled by alpha / r:
   #
-  #   W = W0 + B A
+  #   W = W0 + delta W
+  #   delta W ~= (alpha / r) B A
   #
-  # The `rank:` option is exactly that r.
+  # The `rank:` and `alpha:` options are exactly that r and alpha.
   module LoRA
     POSITIVE_INTEGER_OPTIONS = %i(rank epochs batch_size max_length).freeze
     POSITIVE_NUMBER_OPTIONS = %i(alpha learning_rate).freeze
 
     class << self
-      def train(model:, dataset:, eval_dataset: nil, rank: 8, alpha: 16, dropout: 0.05, target_modules: nil, epochs: 3, batch_size: 4, learning_rate: 2e-4, max_length: 512, output_dir: nil, device: :auto, checkpoints: false, resume: false)
+      def train(model:, dataset:, eval_dataset: nil, rank: 8, alpha: 16, dropout: 0.05, target_modules: nil, epochs: 3, batch_size: 4, learning_rate: 2e-4, max_length: 512, train_on_prompt: false, output_dir: nil, device: :auto, checkpoints: false, resume: false)
         options = {
           rank: rank,
           alpha: alpha,
@@ -27,6 +28,7 @@ module RPEFT
           batch_size: batch_size,
           learning_rate: learning_rate,
           max_length: max_length,
+          train_on_prompt: train_on_prompt,
           checkpoints: checkpoints,
           resume: resume
         }
@@ -80,6 +82,13 @@ module RPEFT
         end
         validate_dropout(options[:dropout])
         validate_target_modules(options[:target_modules])
+        validate_train_on_prompt(options[:train_on_prompt])
+      end
+
+      def validate_train_on_prompt(train_on_prompt)
+        return if [true, false].include?(train_on_prompt)
+
+        raise ConfigurationError, "train_on_prompt must be true or false, got #{train_on_prompt.inspect}"
       end
 
       def validate_dropout(dropout)
@@ -102,12 +111,13 @@ module RPEFT
         raise ConfigurationError, 'resume: true requires an explicit output_dir: to find the previous checkpoints'
       end
 
-      # The backend tokenizes both datasets the same way, so mixing formats would silently mis-render the evaluation set.
+      # A chat transcript is rendered with the chat template and plain text is not, so a loss measured on one
+      # says little about a model trained on the other.
       def coerce_eval_dataset(eval_dataset, training_dataset)
         return nil unless eval_dataset
 
         dataset = Dataset.coerce(eval_dataset)
-        return dataset if dataset.backend_format == training_dataset.backend_format
+        return dataset if dataset.messages? == training_dataset.messages?
 
         # Raised before train assigns the dataset, so its ensure clause cannot discard it.
         dataset.discard unless dataset.equal?(eval_dataset)
@@ -133,6 +143,7 @@ module RPEFT
           'epochs' => options[:epochs],
           'batch_size' => options[:batch_size],
           'learning_rate' => options[:learning_rate],
+          'train_on_prompt' => options[:train_on_prompt],
           'logging_steps' => 5,
           'seed' => 42,
           'trust_remote_code' => false,
@@ -145,7 +156,10 @@ module RPEFT
             'target_modules' => options[:target_modules]
           }
         }
-        payload['eval_dataset_path'] = options[:eval_dataset].path if options[:eval_dataset]
+        if options[:eval_dataset]
+          payload['eval_dataset_path'] = options[:eval_dataset].path
+          payload['eval_dataset_format'] = options[:eval_dataset].backend_format
+        end
 
         payload
       end
@@ -192,6 +206,9 @@ module RPEFT
             'batch_size' => options[:batch_size],
             'learning_rate' => options[:learning_rate],
             'max_length' => options[:max_length],
+            'train_on_prompt' => options[:train_on_prompt],
+            'skipped_examples' => result['skipped_examples'],
+            'eval_skipped_examples' => result['eval_skipped_examples'],
             'device' => result['device'] || device.to_s,
             'train_loss' => result['final_loss'],
             'eval_loss' => result['eval_loss'],

data/lib/rpeft/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RPEFT
-  VERSION = '0.1.0'
+  VERSION = '0.2.0'
 end

data/python/rpeft_backend/__init__.py

@@ -14,7 +14,7 @@ import math
 import os
 import sys
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 
 def _dumps(value):

data/python/rpeft_backend/_chat.py

@@ -22,6 +22,40 @@ def render(tokenizer, messages, add_generation_prompt=False):
     return rendered
 
 
+def segments(tokenizer, messages):
+    """Splits a rendered transcript into (text, trainable) pairs, trainable for the assistant's turns.
+
+    Each assistant turn is located by rendering the transcript up to its generation prompt and up to the
+    end of the turn: the difference is what the model writes at inference, the turn's end marker included.
+    """
+    rendered = render(tokenizer, messages)
+    pairs = []
+    position = 0
+    for index, message in enumerate(messages):
+        # A turn opening the transcript answers nothing, such as a greeting, so it is read as context.
+        # Locating it would also mean rendering an empty transcript, which chat templates refuse.
+        if index == 0 or message.get("role") != "assistant":
+            continue
+
+        start = render(tokenizer, messages[:index], add_generation_prompt=True)
+        through = render(tokenizer, messages[: index + 1])
+        # Some templates render earlier turns differently once the conversation goes on, such as dropping
+        # reasoning from past replies. The turn then cannot be located, and guessing would train on the wrong text.
+        if not (rendered.startswith(through) and through.startswith(start) and len(start) >= position):
+            raise ValueError(
+                "the chat template renders earlier turns differently as the conversation goes on, so the "
+                "assistant's replies cannot be located; pass train_on_prompt: true (--train-on-prompt in the CLI) "
+                "to train on the whole transcript"
+            )
+
+        pairs.append((rendered[position : len(start)], False))
+        pairs.append((rendered[len(start) : len(through)], True))
+        position = len(through)
+
+    pairs.append((rendered[position:], False))
+    return [(text, trainable) for text, trainable in pairs if text]
+
+
 def writes_special_tokens(tokenizer):
     """Whether render's output already holds the special tokens, BOS included, that the model expects.
 

data/python/rpeft_backend/_data.py

@@ -1,23 +1,31 @@
 import json
+import sys
 
 from . import _chat
 from ._errors import BackendError, stage
 
+IGNORED = -100
 
-def load_tokenized_dataset(path, tokenizer, text_field, max_length, dataset_format="text"):
+# The record fields each backend format reads. text_field stands in for the configurable text column.
+FIELDS = {"text": ("text_field",), "completion": ("prompt", "completion"), "messages": ("messages",)}
+
+
+def load_tokenized_dataset(
+    path, tokenizer, text_field, max_length, dataset_format="text", train_on_prompt=False, name="dataset"
+):
+    """Returns the tokenized dataset and the number of examples left out of it.
+
+    Every example carries its own labels: the prompt positions hold IGNORED unless train_on_prompt is set,
+    so the loss covers only what the model is meant to write. name tells the datasets apart in messages.
+    """
     with stage("dataset"):
         from datasets import Dataset
 
         # Built in memory rather than with load_dataset("json", ...), which writes an Arrow copy of every
         # file into the Hugging Face cache. Each run reads a fresh file, so that copy would never be reused.
-        dataset = Dataset.from_list(_read_jsonl(path))
-        field = "messages" if dataset_format == "messages" else text_field
-        if field not in dataset.column_names:
-            raise BackendError(
-                "dataset",
-                "field '%s' not found in dataset (available: %s)"
-                % (field, ", ".join(dataset.column_names)),
-            )
+        records = _read_jsonl(path)
+        fields = [text_field if field == "text_field" else field for field in FIELDS[dataset_format]]
+        _check_fields(records, fields)
 
         add_special_tokens = not (dataset_format == "messages" and _chat.writes_special_tokens(tokenizer))
         # A chat template closes every turn with its own end marker. Plain text has none, so without an EOS
@@ -25,25 +33,95 @@ def load_tokenized_dataset(path, tokenizer, text_field, max_length, dataset_form
         eos = tokenizer.eos_token_id if dataset_format != "messages" else None
         content_length = max_length - 1 if eos is not None else max_length
 
-        def tokenize(batch):
-            if dataset_format == "messages":
-                texts = [_chat.render(tokenizer, messages) for messages in batch[field]]
-            else:
-                texts = batch[field]
+        examples = []
+        for record in records:
+            segments = _segments(tokenizer, record, dataset_format, fields, train_on_prompt)
+            example = _encode(tokenizer, segments, add_special_tokens, content_length, train_on_prompt, eos)
+            if example is not None:
+                examples.append(example)
 
-            # No padding here: the collator pads each batch dynamically.
-            encoded = tokenizer(
-                texts, truncation=True, max_length=content_length, add_special_tokens=add_special_tokens
-            )
-            if eos is not None:
-                for input_ids, attention_mask in zip(encoded["input_ids"], encoded["attention_mask"]):
-                    input_ids.append(eos)
-                    attention_mask.append(1)
-            return encoded
-
-        return dataset.map(
-            tokenize, batched=True, remove_columns=dataset.column_names, keep_in_memory=True
+        skipped = len(records) - len(examples)
+        reason = (
+            "no answer left to train on: a transcript without an assistant reply after its first message, "
+            "or an answer cut off by max_length %d" % max_length
         )
+        if not examples:
+            raise BackendError("dataset", "every example in %s has %s" % (name, reason))
+        if skipped:
+            print(
+                "rpeft: skipped %d of %d examples in %s with %s" % (skipped, len(records), name, reason),
+                file=sys.stderr,
+            )
+
+        return Dataset.from_list(examples), skipped
+
+
+def _check_fields(records, fields):
+    available = set()
+    for record in records:
+        available.update(record)
+    for field in fields:
+        if field not in available:
+            raise BackendError(
+                "dataset",
+                "field '%s' not found in dataset (available: %s)" % (field, ", ".join(sorted(available))),
+            )
+
+
+def _segments(tokenizer, record, dataset_format, fields, train_on_prompt):
+    # (text, trainable) pairs whose concatenation is the whole example.
+    if dataset_format == "messages":
+        if train_on_prompt:
+            return [(_chat.render(tokenizer, record[fields[0]]), True)]
+        return _chat.segments(tokenizer, record[fields[0]])
+    if dataset_format == "completion":
+        return [(record[fields[0]], False), (record[fields[1]], True)]
+    return [(record[fields[0]], True)]
+
+
+def _encode(tokenizer, segments, add_special_tokens, content_length, train_on_prompt, eos):
+    def tokenize(text):
+        return tokenizer(
+            text, truncation=True, max_length=content_length, add_special_tokens=add_special_tokens
+        )["input_ids"]
+
+    # The example is tokenized as a whole, so every token is the one the model sees in the full text.
+    # Tokenizing segments one by one would not do: a SentencePiece tokenizer marks the start of each piece
+    # with a space, turning an answer "hello" into " hello". A segment ends where the text up to its end,
+    # tokenized on its own as the prompt is at inference, stops matching the whole.
+    text = "".join(segment for segment, _ in segments)
+    input_ids = tokenize(text)
+    labels = []
+    answers = []
+    start = 0
+    offset = 0
+    for segment, trainable in segments:
+        offset += len(segment)
+        end = len(input_ids) if offset == len(text) else max(start, _common_prefix(tokenize(text[:offset]), input_ids))
+        ids = input_ids[start:end]
+        labels.extend(ids if trainable or train_on_prompt else [IGNORED] * len(ids))
+        answers.extend([trainable] * len(ids))
+        start = end
+
+    if eos is not None:
+        input_ids.append(eos)
+        labels.append(eos)
+
+    # Without prompt training, an answer cut away entirely leaves an example that teaches nothing but to stop.
+    answer_cut_away = any(trainable and segment for segment, trainable in segments) and not any(answers)
+    if (not train_on_prompt and answer_cut_away) or all(label == IGNORED for label in labels):
+        return None
+
+    return {"input_ids": input_ids, "attention_mask": [1] * len(input_ids), "labels": labels}
+
+
+def _common_prefix(ids, other):
+    length = 0
+    for left, right in zip(ids, other):
+        if left != right:
+            break
+        length += 1
+    return length
 
 
 def _read_jsonl(path):

data/python/rpeft_backend/_training.py

@@ -21,22 +21,11 @@ def train_lora(config):
             device=device,
             trust_remote_code=config.get("trust_remote_code", False),
         )
-        dataset = _data.load_tokenized_dataset(
-            config["dataset_path"],
-            tokenizer,
-            text_field=config.get("text_field", "text"),
-            max_length=config.get("max_length", 512),
-            dataset_format=config.get("dataset_format", "text"),
-        )
+        dataset, skipped = _load_dataset(config, tokenizer, "dataset")
         eval_dataset = None
+        eval_skipped = None
         if config.get("eval_dataset_path"):
-            eval_dataset = _data.load_tokenized_dataset(
-                config["eval_dataset_path"],
-                tokenizer,
-                text_field=config.get("text_field", "text"),
-                max_length=config.get("max_length", 512),
-                dataset_format=config.get("dataset_format", "text"),
-            )
+            eval_dataset, eval_skipped = _load_dataset(config, tokenizer, "eval_dataset")
 
         with stage("training"):
             model, target_modules, param_counts = _apply_lora(model, lora)
@@ -54,7 +43,7 @@ def train_lora(config):
             shutil.rmtree(_trainer_dir(output_dir), ignore_errors=True)
 
         return _result(
-            config, device, dtype, target_modules, param_counts, metrics, steps, len(dataset)
+            config, device, dtype, target_modules, param_counts, metrics, steps, len(dataset), skipped, eval_skipped
         )
     finally:
         model = None
@@ -62,6 +51,18 @@ def train_lora(config):
         _device.empty_cache(device)
 
 
+def _load_dataset(config, tokenizer, name):
+    return _data.load_tokenized_dataset(
+        config[name + "_path"],
+        tokenizer,
+        text_field=config.get("text_field", "text"),
+        max_length=config.get("max_length", 512),
+        dataset_format=config.get(name + "_format") or config.get("dataset_format", "text"),
+        train_on_prompt=bool(config.get("train_on_prompt", False)),
+        name=name,
+    )
+
+
 def _apply_lora(model, lora):
     from peft import LoraConfig, TaskType, get_peft_model
 
@@ -117,7 +118,7 @@ def _build_trainer(model, tokenizer, dataset, eval_dataset, config, device, dtyp
     )
 
 
-def _result(config, device, dtype, target_modules, param_counts, metrics, steps, num_examples):
+def _result(config, device, dtype, target_modules, param_counts, metrics, steps, num_examples, skipped, eval_skipped):
     trainable_params, total_params = param_counts
     effective = dict(config)
     effective["device"] = device
@@ -133,6 +134,8 @@ def _result(config, device, dtype, target_modules, param_counts, metrics, steps,
         "final_loss": float(metrics.get("train_loss", 0.0)),
         "eval_loss": float(metrics["eval_loss"]) if "eval_loss" in metrics else None,
         "num_examples": int(num_examples),
+        "skipped_examples": int(skipped),
+        "eval_skipped_examples": int(eval_skipped) if eval_skipped is not None else None,
         "steps": int(steps),
         "trainable_params": int(trainable_params),
         "total_params": int(total_params),
@@ -156,15 +159,23 @@ def _resume_point(config):
 
 
 def _collator(tokenizer):
-    # DataCollatorForLanguageModeling ignores every position holding the pad token. Tokenizers without
-    # a pad token of their own pad with EOS (see _models.load_tokenizer), and many that have one reuse
-    # EOS too, so the EOS that ends each example would never be learned. Only the padding the batch
-    # added is ignored here, found through the attention mask instead of the token id.
+    # Labels come with each example rather than from the pad token id, as DataCollatorForLanguageModeling
+    # derives them. Tokenizers without a pad token of their own pad with EOS (see _models.load_tokenizer),
+    # and many that have one reuse EOS too, so the EOS that ends each example would never be learned.
     def collate(features):
-        batch = tokenizer.pad(features, return_tensors="pt")
-        labels = batch["input_ids"].clone()
-        labels[batch["attention_mask"] == 0] = -100
-        batch["labels"] = labels
+        import torch
+
+        labels = [list(feature["labels"]) for feature in features]
+        batch = tokenizer.pad(
+            [{key: value for key, value in feature.items() if key != "labels"} for feature in features],
+            return_tensors="pt",
+        )
+        width = batch["input_ids"].shape[1]
+        padded = []
+        for example in labels:
+            padding = [_data.IGNORED] * (width - len(example))
+            padded.append(padding + example if tokenizer.padding_side == "left" else example + padding)
+        batch["labels"] = torch.tensor(padded, dtype=torch.long)
         return batch
 
     return collate

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: rpeft
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - Koichi ITO