Add a trainer by ToluClassics · Pull Request #29 · ToluClassics/mlx-transformers

ToluClassics · 2026-03-23T02:53:23Z

No description provided.

Copilot

Pull request overview

This PR introduces a lightweight MLX-native training stack (generic Trainer + SFTTrainer) with LoRA adapter support, adds unit tests and example scripts, and updates several model heads to return scalar (mean) losses.

Changes:

Add Trainer/TrainingArguments with batching, training/eval loops, logging, and checkpoint saving.
Add LoRA utilities (apply_lora, save/load adapters) and SFTTrainer with dataset preparation and LM data collator.
Add tests + docs/examples; adjust several model loss computations to use reshape and mx.mean.

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
`src/mlx_transformers/trainer.py`	New generic trainer (train/eval/checkpointing, collator, metrics)
`src/mlx_transformers/sft_trainer.py`	New SFT trainer with dataset preprocessing + LM collator + LoRA-aware saving
`src/mlx_transformers/lora.py`	New LoRA layer wrapper and adapter save/load utilities
`src/mlx_transformers/__init__.py`	Export trainer/SFT/LoRA public API
`src/mlx_transformers/models/bert.py`	Adjust loss reductions / reshaping for MLX
`src/mlx_transformers/models/roberta.py`	Adjust loss reductions / reshaping; QA loss path updated
`src/mlx_transformers/models/xlm_roberta.py`	Adjust loss reductions / reshaping; QA loss path updated
`tests/test_trainer.py`	Coverage for trainer loop, collator, checkpoints, custom loss
`tests/test_sft_trainer.py`	Coverage for SFT preprocessing, collator masking, LoRA saving
`tests/test_lora.py`	Coverage for LoRA application and adapter save/load
`examples/bert/imdb_training.py`	Example: BERT fine-tuning via new `Trainer`
`examples/sft/qwen3_wildchat_sft.py`	Example: LoRA SFT for Qwen3 on WildChat subset
`examples/text_generation/benchmark_generation.py`	Change benchmark CLI defaults
`README.md`	Document minimal training loop usage
`benchmark_results-phi3-4k.md`	Remove benchmark results file

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-25T16:48:27Z

+    def _num_train_batches(self) -> int:
+        if self.train_dataset is None:
+            return 0
+        if not self._is_indexable_dataset(self.train_dataset):
+            if self.args.max_steps <= 0:
+                raise ValueError("Iterable train datasets require max_steps to be set.")
+            return self.args.max_steps
+


There’s no test exercising Trainer.train() with a non-indexable iterable dataset (the code even has a dedicated _num_train_batches() branch for it). Adding a small unit test with a generator-based dataset would help prevent regressions once iterable training is supported/fixed.

Copilot · 2026-03-25T16:48:27Z

+    def _peek_example(self, dataset):
+        if self._is_indexable_dataset(dataset) and len(dataset) > 0:
+            return dataset[0]
+
+        iterator = iter(dataset)
+        try:
+            return next(iterator)
+        except StopIteration:
+            return None
+


The dataset auto-preparation path relies on _peek_example(), but there’s no unit test covering iterable/generator datasets where peeking can consume data. A focused test with an iterator dataset would guard against silently dropping the first example.

Copilot · 2026-03-25T16:48:28Z

+        return self._iter_dataset(
+            self.train_dataset,
+            batch_size=self.args.per_device_train_batch_size,
+            shuffle=True,


get_train_dataloader() always requests shuffle=True, but _iter_dataset() raises for non-indexable (iterable) datasets when shuffle is enabled. This makes iterable train datasets unusable even though the type hints and _num_train_batches() suggest they’re supported. Consider disabling shuffle automatically when the dataset isn’t indexable (or implement buffered shuffling), so Trainer.train() can run on iterable datasets when max_steps is set.

Suggested change

return self._iter_dataset(

self.train_dataset,

batch_size=self.args.per_device_train_batch_size,

shuffle=True,

dataset = self.train_dataset

shuffle = False

if dataset is not None and self._is_indexable_dataset(dataset):

shuffle = True

return self._iter_dataset(

dataset,

batch_size=self.args.per_device_train_batch_size,

shuffle=shuffle,

Copilot · 2026-03-25T16:48:28Z

+                predictions.append(self._to_numpy(batch_predictions))
+                labels.append(self._to_numpy(self._extract_labels(batch)))


When compute_metrics is set, evaluation extracts labels via _extract_labels(batch) but doesn’t validate that any label keys are present. If a user passes a label-less eval dataset, label_ids becomes None and compute_metrics will likely fail with a less clear error later. Consider raising a ValueError when labels are missing (similar to the predictions check) to make failures actionable.

Suggested change

predictions.append(self._to_numpy(batch_predictions))

labels.append(self._to_numpy(self._extract_labels(batch)))

batch_labels = self._extract_labels(batch)

if batch_labels is None:

raise ValueError(

"Evaluation batch does not expose labels for compute_metrics."

)

predictions.append(self._to_numpy(batch_predictions))

labels.append(self._to_numpy(batch_labels))

Copilot · 2026-03-25T16:48:28Z

+    def _peek_example(self, dataset):
+        if self._is_indexable_dataset(dataset) and len(dataset) > 0:
+            return dataset[0]
+
+        iterator = iter(dataset)
+        try:
+            return next(iterator)
+        except StopIteration:
+            return None


_peek_example() consumes one element from iterable datasets by calling next(iter(dataset)), but it returns only the sample and not an iterator that re-includes it. For generator/streaming datasets this will silently drop the first training example. Consider using itertools.tee (or a small wrapper) to peek without consumption, or require indexable datasets for auto-preprocessing.

Copilot · 2026-03-25T16:48:28Z

    parser.add_argument(
        "--samples-per-bucket",
        type=int,
-        default=3,
+        default=10,
        help="How many prompts to sample from each dataset bucket.",
    )
    parser.add_argument(
        "--dataset-max-samples",
        type=int,
-        default=5000,
+        default=100,
        help="Maximum dataset rows to scan while filling prompt buckets.",
    )


This PR changes the default benchmarking behavior substantially (--samples-per-bucket from 3→10 and --dataset-max-samples from 5000→100). Since these defaults impact result quality and runtime, it would help to either keep prior defaults or add a short note in the help text / README explaining why the new defaults are preferred.

Copilot · 2026-03-25T16:48:29Z

+            ignored_index = start_logits.shape[1]
+            start_positions = mx.clip(start_positions, 0, ignored_index)
+            end_positions = mx.clip(end_positions, 0, ignored_index)

            loss_fct = nn.losses.cross_entropy
-            start_loss = loss_fct(start_logits, start_positions)
-            end_loss = loss_fct(end_logits, end_positions)
+            start_loss = mx.mean(loss_fct(start_logits, start_positions))
+            end_loss = mx.mean(loss_fct(end_logits, end_positions))
            total_loss = (start_loss + end_loss) / 2


The QA loss path clips positions to ignored_index = start_logits.shape[1] (inclusive) and then calls nn.losses.cross_entropy. If any position was out of range, clipping can produce a target equal to ignored_index, which is outside [0, ignored_index-1] and can cause an index error / invalid loss. If you want to ignore out-of-range positions (as the original HF logic does), either clip to ignored_index - 1 or implement an explicit mask/ignore-index behavior before averaging.

Copilot · 2026-03-25T16:48:29Z

+            ignored_index = start_logits.shape[1]
+            start_positions = mx.clip(start_positions, 0, ignored_index)
+            end_positions = mx.clip(end_positions, 0, ignored_index)

            loss_fct = nn.losses.cross_entropy
-            start_loss = loss_fct(start_logits, start_positions)
-            end_loss = loss_fct(end_logits, end_positions)
+            start_loss = mx.mean(loss_fct(start_logits, start_positions))
+            end_loss = mx.mean(loss_fct(end_logits, end_positions))
            total_loss = (start_loss + end_loss) / 2


The QA loss path clips positions to ignored_index = start_logits.shape[1] (inclusive) and then calls nn.losses.cross_entropy. If any position was out of range, clipping can produce a target equal to ignored_index, which is outside [0, ignored_index-1] and can cause an index error / invalid loss. If you want to ignore out-of-range positions, either clip to ignored_index - 1 or implement an explicit mask/ignore-index behavior before averaging.

Copilot · 2026-03-25T16:48:29Z

            ignored_index = start_logits.shape[1]
            start_positions = mx.clip(start_positions, 0, ignored_index)
            end_positions = mx.clip(end_positions, 0, ignored_index)

            loss_fct = nn.losses.cross_entropy
-            start_loss = loss_fct(start_logits, start_positions)
-            end_loss = loss_fct(end_logits, end_positions)
+            start_loss = mx.mean(loss_fct(start_logits, start_positions))
+            end_loss = mx.mean(loss_fct(end_logits, end_positions))
            total_loss = (start_loss + end_loss) / 2


The QA loss computation averages cross_entropy(start_logits, start_positions) / cross_entropy(end_logits, end_positions) after clipping positions up to ignored_index = start_logits.shape[1] (inclusive). For out-of-range labels this can still produce a target == ignored_index which is outside the valid class range and can crash or yield incorrect loss. Consider changing the clipping/ignore logic so out-of-range positions are truly ignored (e.g., mask them out) or clip to ignored_index - 1.

working on a trainer

843f973

ToluClassics requested a review from Copilot March 25, 2026 16:39

Copilot started reviewing on behalf of ToluClassics March 25, 2026 16:40 View session

Copilot AI reviewed Mar 25, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add a trainer#29

Add a trainer#29
ToluClassics wants to merge 1 commit into
mainfrom
trainer

ToluClassics commented Mar 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Copilot AI Mar 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		predictions.append(self._to_numpy(batch_predictions))
		labels.append(self._to_numpy(self._extract_labels(batch)))

Conversation

ToluClassics commented Mar 23, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants