PyTorch MPS против Apple MLX: что выбрать для дообучения LLM на Mac с 16 ГБ

Исследование провела команда badlogicgames, авторы датасета pi-mono и утилиты Trackio для логирования экспериментов. Они обучали модели Qwen 2.5 трёх размеров (0,5 B, 1,5 B и 3 B параметров) на реальных многоходовых диалогах автономных ИИ-агентов (агент, программа, которая сама решает, какой инструмент вызвать для выполнения задачи). До этого большинство гайдов по дообучению (fine-tuning, обучение модели на ваших примерах под узкую задачу) использовали короткие одноходовые примеры на 256-512 токенов (токен, минимальная единица текста для модели, примерно три четверти слова). Реальные сессии агентов гораздо длиннее и сложнее, поэтому именно на них видна разница между фреймворками.

Что понадобится

• Mac на Apple Silicon (M1, M2, M3 или M4) с 16 ГБ объединённой памяти
• Python 3.10 или новее
• PyTorch с поддержкой MPS (устанавливается стандартно через pip)
• Библиотеки Hugging Face Transformers и TRL (для варианта с PyTorch MPS)
• Фреймворк MLX и пакет mlx-lm (для варианта с Apple MLX)
• Модели Qwen 2.5 нужного размера: с Hugging Face Hub для MPS, из коллекции mlx-community для MLX с 4-битным квантованием (QLoRA, дообучение квантованной модели через маленькие адаптеры, когда основные веса сжаты до 4 бит)
• Обучающие данные в формате многоходовых диалогов (пример: badlogicgames/pi-mono)
• Время: один эксперимент на модели 1,5 B занимает от нескольких минут до получаса в зависимости от длины контекста

Пошаговая инструкция

1. Определитесь с целью. Если вам нужна скорость и модель не крупнее 1,5 B параметров, берите PyTorch MPS. Если нужна модель от 3 B или длинный контекст от 4 096 токенов, единственный рабочий вариант на 16 ГБ это Apple MLX.

2. Установите окружение для PyTorch MPS.

pip install torch torchvision torchaudio
pip install transformers trl peft accelerate

Убедитесь, что MPS доступен:

import torch
print(torch.backends.mps.is_available())  # должно вернуть True

1. Или установите окружение для Apple MLX.

pip install mlx mlx-lm

Для MLX загрузите предварительно квантованные веса из коллекции mlx-community на Hugging Face Hub. Это модели, уже сжатые до 4 бит, что критично для экономии памяти.

1. Подготовьте данные. Многоходовые диалоги агентов содержат три типа сообщений: пользовательские запросы, рассуждения ассистента с вызовами инструментов и ответы окружения (вывод терминала, содержимое файлов). Для PyTorch MPS в SFTTrainer из TRL включите параметр completion_only_loss=True, чтобы модель училась только на ответах ассистента, а не на логах сборки.

# Пример конфигурации SFTTrainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    completion_only_loss=True,  # считаем потери только по токенам ассистента
)

1. Для MLX подготовьте маски вручную. В mlx-lm нет встроенного флага, аналогичного completion_only_loss. Перед обучением преобразуйте данные так, чтобы пользовательские запросы и вывод инструментов были замаскированы значением -100 в целевых токенах.

2. Настройте LoRA-адаптеры (LoRA, метод дообучения через маленькие матрицы-вставки, не трогая основные веса модели). Для PyTorch MPS: ранг r=16, alpha=32, адаптация всех семи проекций внимания и MLP (q, k, v, o, gate, up, down). Для MLX: стандартные параметры mlx_lm.lora, ранг 8, scale=20.0, dropout=0, адаптация последних 24 слоёв трансформера.

3. Запустите обучение и следите за памятью. На macOS CPU, GPU и сама система используют общий пул памяти. Если во время обучения система начинает «задыхаться» (резкое замедление, swap), значит, вы вышли за пределы доступной памяти. Для MPS это случается при переходе к моделям крупнее 1,5 B или длинному контексту. Для MLX порог значительно выше.

4. Сравните результаты. Функция потерь на валидации (validation loss, насколько хорошо модель предсказывает правильные токены на отложенных примерах) у обоих фреймворков оказалась близкой: 0,462 для MPS против 0,478 для MLX на модели 1,5 B с контекстом 2 048 токенов. Квантование до 4 бит не испортило сходимость.

Что делать с этим прямо сейчас?

Авторам Дзена и копирайтерам. Если вы хотите дообучить небольшую модель (0,5-1,5 B) под свой стиль или тематику, берите PyTorch MPS. Скрипты стандартные, документация обширная, и на коротких текстах 16 ГБ хватит с запасом. Это проще всего для первого опыта.

Маркетологам и тем, кто работает с агентами. Если задача сложнее, модель должна уметь вызывать инструменты, обрабатывать длинные контексты, работать в многоходовых сценариях, выбирайте MLX. Только он позволяет запустить модель от 3 B на Mac с 16 ГБ.

Предпринимателям в РФ и СНГ. Оба фреймворка доступны без ограничений, модели Qwen 2.5 имеют открытые веса (open weights, веса модели опубликованы и их можно скачать свободно). Российские аналоги для инференса (инференс, использование обученной модели для генерации ответов), YandexGPT и GigaChat, пока не дают возможности дообучения на своём железе. Локальное дообучение на Mac остаётся одним из немногих способов получить кастомную модель без облачного GPU.

Выбор между PyTorch MPS и Apple MLX на Mac с 16 ГБ сводится к одному вопросу: какого размера модель вам нужна. До 1,5 B параметров MPS быстрее и проще. От 3 B и выше, только MLX. Validation loss практически одинаковый, значит, качество обучения от выбора фреймворка не страдает, страдает только ваше терпение или ваша память.