mlx-modernbert

AI Backend MLX Apple Silicon Python NLP

Project Details

Pure MLX training framework for ModernBERT on Apple Silicon, optimized for M-series neural accelerators.

mlx-modernbert

Table of Contents

List of all the sections in this post. Click to jump to that section.

Recruiter Scan (First 30 Seconds)

The Motivation: I wanted to train a ModernBERT model on my MacBook Pro M5 Pro (48GB RAM). I wanted to use Apple’s neural accelerators for AI tasks. However, I couldn’t find a good and efficient library. Existing options lacked solid metrics and basic features, like resuming training from checkpoints.

The Solution: I decided to build my own framework from scratch. It started as a small addition to existing tools, but it grew into a fully independent training library using Apple MLX. I rebuilt the ModernBERT architecture without using PyTorch or TensorFlow. It runs natively on Apple Metal GPUs.

The Impact: I successfully trained a 550M parameter model locally without any “out-of-memory” errors. The framework is now available as a PyPI package (mlx-modernbert). It helps create strong models for text classification (like prompt-injection guardrails) and token classification (like PII detection).

Overview

mlx-modernbert is an independent MLX training framework made specifically for ModernBERT on Apple Silicon. My main goal was to build a highly optimized tool to fine-tune large models directly on Apple Metal GPUs.

At first, the project used other MLX tools (mlx_raclate). Later, I rewrote the code to make it fully independent. This gave me complete control over memory management and the training loop. The main focus of this framework is sequence classification. Using it, I trained a model to detect prompt injections. It can tell the difference between safe inputs and malicious attacks. Later on, I added support for token classification tasks, like detecting PII (Personally Identifiable Information) in BIO format. This includes aligning subwords with their labels and using a weighted loss function to handle unbalanced datasets.

Architecture and Stack

The system is designed with a modular architecture. It separates the model, the configuration, the data processing, and the training engine into different parts. Data flows from HuggingFace datasets through custom collators into a natively compiled MLX computation graph.

  • Core Logic: Python 3.13, Apple MLX
  • Hardware Target: Apple Metal GPU (Optimized for M5 Pro)
  • Data & Ecosystem: HuggingFace Transformers, Datasets, seqeval, scikit-learn
  • Packaging: Published to PyPI using uv and hatchling

Engineering Challenges and Trade-offs

Challenge: RAM Optimization on Apple Silicon

  • The Issue: ModernBERT-large has about 550M parameters. Training and evaluating it in standard fp32 quickly uses up all the memory on Apple Silicon. Also, during long evaluation steps, keeping the data in the MLX backend created big memory leaks that caused OOM (Out Of Memory) crashes.
  • Trade-off: I had to accept slightly slower data transfers during evaluation to keep memory usage low. I also had to implement dynamic layer checkpointing during training.
  • Resolution: I solved the memory problem in three ways. First, I used fp16 training, which cuts memory usage in half. Second, I built a dynamic gradient checkpointing system. This system checks the model and changes the __call__ method of the layers, wrapping them in mx.checkpoint(). Finally, for the evaluation loop, I heavily optimized the memory: I manually deleted local variables (del loss, preds, batch), forced the system to clean memory (gc.collect()), cleared the MLX cache (mx.clear_cache()), and moved output data to NumPy arrays step by step. This stopped the evaluation process from filling up the Metal backend cache, making it faster and completely stable.

Challenge: The Custom “Compiled” Training Engine

  • The Issue: Python training loops can slow down the GPU if the process is rebuilt in every step. I needed a custom training loop that supported advanced features (like multiclass evaluation and early stopping) without losing performance.
  • Trade-off: Instead of using basic training templates, I designed a custom Trainer class. This class compiles its core operations into static graphs, which makes the code a bit more complex but much faster.
  • Resolution: I used mx.compile on the main functions (step_fn, update_fn, and eval_fn). This ensures that the whole training pass runs continuously on the GPU without interruptions. On top of this fast base, I added an Early Stopping system based on validation loss to prevent overfitting. I also included sklearn.metrics to easily handle multiclass evaluation.

Challenge: Loading Model Configurations and Checkpoints

  • The Issue: Because the framework is built entirely in MLX without PyTorch, I needed a way to load pretrained PyTorch weights from HuggingFace (answerdotai/ModernBERT-large) and adapt them to my custom MLX model.
  • Trade-off: I could have asked users to run a conversion script before training. Instead, I created a system that loads and converts the weights on the fly. This makes it easier for the user, even if the first load takes a little longer.
  • Resolution: I built an automatic loading module that downloads safetensors checkpoints, fixes the parameter names (mapping HuggingFace names to my MLX structure), and initializes the model. I also made sure the model head initializes correctly without relying on simple word matching. Finally, I built a reliable checkpoint system that saves the optimizer state, allowing the user to easily pause and resume training.

Key Optimizations and Results

  • Evaluation Stability: By moving data to NumPy and forcing Python to clean memory (gc.collect()), the evaluation loop runs perfectly on 550M parameter models without memory spikes.
  • Bare-Metal Speed via mx.compile: Compiling the training steps directly to the Metal backend removed Python delays, making the custom trainer extremely fast.
  • Production-Ready Training: I successfully built important MLOps features from scratch, including Early Stopping, Gradient Checkpointing, and checkpoint rotation.
  • Complete Independence: I removed all early third-party dependencies. I turned the project into a robust, independent mlx-modernbert PyPI package. I also included nice CLI inference scripts with emojis and speed metrics to easily showcase the models on social media.

Future Iterations

  • Multi-GPU distributed training support using mlx.core.distributed to run models across multiple Mac Studios.
  • Adding lower-precision quantization (like int8 or int4) to support even larger models or bigger batch sizes on the same computer.

Project Links