────────────────────────────────────────────────────

████████╗██████╗ ██╗ ███╗ ███╗██████╗ ███████╗ ╚══██╔══╝██╔══██╗██║ ████╗ ████║██╔══██╗██╔════╝ ██║ ██████╔╝██║ ██╔████╔██║██████╔╝███████╗ ██║ ██╔══██╗██║ ██║╚██╔╝██║██╔═══╝ ╚════██║ ██║ ██║ ██║███████╗██║ ╚═╝ ██║██║ ███████║ ╚═╝ ╚═╝ ╚═╝╚══════╝╚═╝ ╚═╝╚═╝ ╚══════╝

🚀 TRLmps: High-Performance fine-tuning framework for Apple Silicon 🚀

The first Mac-native Reinforcement Learning framework built on Hugging Face TRL.
Fine-tune LLMs with GRPO, SFT locally on Apple Unified Memory OOM panics.

Quick Start • Benchmarks • Features • Supported Models • Architecture • Citation

---

📌 Table of Contents

• News
• Why TRLmps?
• Benchmarks
• Key Features
• Supported Models
• Quick Start
• Minimal Working Example
• Architecture
• Citation
• License

---

TRLmps solves this by introducing synchronous memory boundaries inside generation loops, allowing stable fine-tuning within tight memory envelopes on Apple M series chips with UAW Architecture.

---

📊 Benchmarks

Tested on Apple M4 PRO with Qwen3-VL

Framework	Peak VRAM	Step Speed
Standard TRL (mps)	34.2 GB	142s / step
TRLmps	18.4 GB	12.1s / step
TRLmps + Metalliger	14.1 GB	9.8s / step

---

✨ Key Features

• 🛡️ Zero Swap Memory Bloat (_MPSCacheFlusher): Injects synchronous memory release boundaries (torch.mps.empty_cache()) into generate() loops to prevent Metal allocator hoard.
• ⚡ Mac-Optimized GRPOTrainer: Re-engineered GRPO trainer with memory-chunked log-probability calculations and advantage normalization tailored for MPS.
• 🧠 Native bfloat16 Precision: Full support for Apple Silicon native bfloat16 to prevent gradient underflow in RL policy updates.
• 📉 Non-Reentrant Gradient Checkpointing: Memory-saving activation checkpointing that avoids legacy PyTorch graph tracking bugs.

---

📋 Supported Models

Model Family	SFT	GRPO	Vision-Language (VL)
Qwen3.5	✅	✅	✅
Qwen3-VL	✅	✅	✅

---

🚀 Quick Start

Prerequisites

• macOS 14.0 (Sonoma) or newer
• Apple Silicon Mac (M3/M4 Series)
• Python 3.10+
• PyTorch 2.2+ with MPS backend enabled

Installation

git clone https://github.com/krrish-v/trlmps.git
git clone https://github.com/krrish-v/metalliger.git

cd trlmps
pip install -e .

cd metalliger
pip install -e .

Required Environment Variables

For stable MPS training, you must configure PyTorch's Metal memory allocator:

# Allow PyTorch to access full system memory (disables 70% soft-cap)
export PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0
export PYTORCH_MPS_LOW_WATERMARK_RATIO=0.0

# Optional: Faster math on M4 architectures
export PYTORCH_MPS_FAST_MATH=1

# Prevent OpenMP thread contention with accelerating dispatchers
export OMP_NUM_THREADS=1

---

💡 Usage Working Example

Save and run the following script to launch GRPO training on your Mac:

from trlmps.trl import GRPOConfig, GRPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer


model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-VL", torch_dtype="bfloat16", device_map="mps")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL")

dataset = load_dataset("dataset/dataset", split="train")

# 2. Configure Mac-Optimized GRPO Training Arguments
training_args = GRPOConfig(
    output_dir="./grpo_mac_output",
    
    # 🍎 MPS Memory Optimization Flags
    use_mps_optimization=True,
    use_metalliger=True,
    use_metalliger_compile=False,
    mps_cleanup_frequency=1,
    mps_eval_num_workers=0,
    mps_fused_loss_chunk_size=4096,
    mps_max_tokens_per_batch=4096,
    mps_memory_fraction=0.9,
    mps_prefetch_factor=2,
    
    # 📦 Batch & Rollout Settings
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_generations=4,            # Rollouts per prompt
    max_prompt_length=512,
    max_completion_length=768,
    
    # 🎯 Learning & Precision
    learning_rate=1e-5,
    bf16=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
)

# 3. Define Reward Function
def reward_len(completions, **kwargs):
    return [float(len(c)) for c in completions]

# 4. Initialize & Train
trainer = GRPOTrainer(
    model=model,
    reward_funcs=[reward_len],
    args=training_args,
    train_dataset=dataset,
)

trainer.train()

import torch
from trlmps.trl import SFTConfig, SFTTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-VL", torch_dtype="bfloat16", device_map="mps")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL")

dataset = load_dataset("dataset/dataset", split="train")

config = SFTConfig(
    output_dir="./output",
    bf16=True,                             # Use BFloat16 natively on Mac
    
    # --- TRL-MPS OPTIMIZATIONS ---
    mps_memory_fraction=0.98,             # 90% of system RAM for MPS
    mps_fused_loss_chunk_size=8192,      # Vocab chunk size (M4 Pro optimized)
    mps_cleanup_frequency=10,             # Predictable Memory > Optimized Queuing
    
    use_mps_optimization=True,
    use_metal_liger=True,           
    use_metal_liger_compile=True,

    dataloader_pin_memory=False,         # Must be False on Mac
    dataloader_num_workers=4,            # Maximize Performance Cores
    dataloader_prefetch_factor=4,        # Keep the GPU fed
)

trainer = SFTTrainer(
    model=model,
    args=config,
    train_dataset=dataset,
    processing_class=tokenizer,
)

trainer.train()

---

🏗️ Architecture

                                standard TRL
  Prompt ---> [ generate() ] ------------> [ Metal Cache Hoard ] ---> [ High memory consumtion ] ---> ❌ Crash
                    |
                 TRLmps
                    v
  Prompt ---> [ generate() ]
                    |
        [_MPSCacheFlusher Intercept]
                    |
                    v
        [ torch.mps.empty_cache() ] 
                    |
                    v
  [ Memory Yielded to macOS ] ---> [ Logits & Reward Calc ] ---> [ Policy Update ] ---> ✅ Success

---

📜 Proof of work

MadriMed Collection:

madrisight/MadriMed-VL-2B-enc

madrisight/MadriMed-VL-2B

---

📜 Citation

If you use TRLmps in your research or project, please cite:

@software{trlmps,
  title = {TRLmps:High-Performance fine-tuning framework for Apple Silicon},
  author = {Krrish},
  year = {2026},
  publisher = {GitHub},
  url = {https://github.com/krrish-v/trlmps}
}

---

📄 License

This project is licensed under the Apache 2.0 License, building upon Hugging Face's TRL framework.

---

Original TRL library by the Hugging Face team. TRL-MPS is a specialized Mac-Silicon optimization fork.