Пошаговое обучение LLM с нуля на MacBook: полный гайд с PyTorch и MLX

Почему все говорят, что на Mac нельзя обучать LLM?

Открою секрет: могут. Просто не хотят. Классический путь — арендовать кластер в облаке за $200 в час, три недели отлаживать скрипты, а потом понять, что забыл нормализовать данные. Знакомо? А ведь на вашем MacBook с M3 или M4 Pro сидит железо, которое отлично справляется с матричными операциями. Просто его нужно правильно использовать.

Проблема не в железе. Проблема в том, что экосистема машинного обучения заточена под CUDA и NVIDIA. Apple давно это поняла и выпустила MLX — фреймворк, который превращает Neural Engine и GPU в единый вычислительный блок. Сегодня мы заставим его работать.

MLX: секретное оружие Apple Silicon, о котором все молчат

MLX — это не просто еще один фреймворк. Это способ говорить с чипом Apple на его языке. Вместо того чтобы гонять данные между CPU и GPU, как в старых добрых PyTorch с бэкендом MPS, MLX создает единое адресное пространство. Данные лежат в одной памяти, а операции выполняются там, где это эффективнее.

Мы будем использовать гибридный подход: PyTorch для построения графа вычислений и знакомого API, а MLX — как движок для выполнения. Это дает баланс между удобством и производительностью.

Что нам понадобится перед стартом

• MacBook с Apple Silicon (M1, M2, M3, M4). На Intel будет больно и медленно.
• Не менее 16 ГБ оперативной памяти. 32 ГБ — комфортно. 8 ГБ — забудьте.
• macOS Sonoma 14.4 или новее.
• Python 3.11 или 3.12. Старые версии не дружат с MLX.
• Терпение и готовность к тому, что что-то сломается. Это нормально.

1 Установка Python и менеджера пакетов

Забудьте про системный Python. Мы используем pyenv или conda. Я предпочитаю conda для ML-проектов — он лучше управляет native-зависимостями.

# Устанавливаем Miniconda (если нет)
curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
sh Miniconda3-latest-MacOSX-arm64.sh

# Создаем отдельное окружение
conda create -n llm-from-scratch python=3.12
conda activate llm-from-scratch

2 Настройка виртуального окружения и установка PyTorch

Теперь ставим PyTorch с поддержкой MLX. Официальный pip-пакет torch-mlx появился в конце 2025.

# Устанавливаем PyTorch с бэкендом MLX
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
pip install mlx

# Проверяем, что все видит MLX
python -c "import torch; import mlx.core as mx; print(f'PyTorch: {torch.__version__}'); print(f'MLX available: {mx.metal.is_available()}')"

3 Подготовка датасета: где взять текст и как его обработать

Для претренинга нужны гигабайты текста. Брать первый попавшийся дамп Wikipedia — плохая идея. Качество данных определяет 80% успеха. Я использую очищенную версию The Pile, но для начала хватит и части.

import datasets
from transformers import AutoTokenizer

# Загружаем небольшую часть датасета для теста
dataset = datasets.load_dataset("EleutherAI/the_pile", split="train", streaming=True, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")

def tokenize_function(examples):
    # Простейшая токенизация, на практике нужно добавить чанкинг
    return tokenizer(examples["text"], truncation=True, max_length=512)

# Токенизируем потоково
tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
# Сохраняем для ускорения следующих запусков
tokenized_dataset.save_to_disk("./tokenized_pile_part")

Токенизатор берем подходящий под архитектуру. Для GPT-like моделей — GPT-2 токенизатор. Не используйте BERT-токенизатор для генеративных задач — это частая ошибка новичков.

4 Создание архитектуры модели: не повторяйте моих ошибок

Мы не будем писать трансформер с нуля — возьмем готовую реализацию из библиотеки и адаптируем под MLX. Но важно понимать, что происходит внутри.

import torch
import torch.nn as nn
import mlx.core as mx
import mlx.nn as mlnn

# Простейший блок трансформера для демонстрации
class SimpleTransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x):
        # Внимание
        attn_output, _ = self.self_attn(x, x, x, need_weights=False)
        x = x + self.dropout(attn_output)
        x = self.norm1(x)
        
        # Feed-forward
        ff_output = self.linear2(torch.relu(self.linear1(x)))
        x = x + self.dropout(ff_output)
        x = self.norm2(x)
        return x

# Обертка для конвертации весов в MLX-формат
def convert_to_mlx(state_dict):
    mlx_weights = {}
    for k, v in state_dict.items():
        mlx_weights[k] = mx.array(v.numpy())  # Конвертируем torch.Tensor в mx.array
    return mlx_weights

5 Написание цикла обучения: где экономить память и время

Самый критичный этап. Ошибки здесь приводят к нулевым лоссам или переполнению памяти. Я покажу упрощенный цикл, а потом расскажу, как его улучшить.

import torch.optim as optim
from torch.utils.data import DataLoader
import mlx.optimizers as mlx_opt

# Подготовка данных
dataloader = DataLoader(tokenized_dataset, batch_size=4, shuffle=True)

# Модель, оптимизатор
model = SimpleTransformerBlock(d_model=512, nhead=8)
optimizer = mlx_opt.Adam(learning_rate=3e-4)

# Цикл обучения
for epoch in range(3):  # 3 эпохи для старта
    model.train()
    total_loss = 0
    
    for batch_idx, batch in enumerate(dataloader):
        # Переносим данные в MLX
        inputs = mx.array(batch['input_ids'].numpy())
        targets = mx.array(batch['input_ids'].numpy())  # Для language modeling цель — те же токены
        
        # Forward pass
        def loss_fn(model):
            outputs = model(inputs)
            loss = mlnn.losses.cross_entropy(outputs, targets)
            return mx.mean(loss)
        
        # Вычисляем loss и градиенты
        loss, grads = mlnn.value_and_grad(model, loss_fn)
        
        # Обновляем веса
        optimizer.update(model, grads)
        
        total_loss += loss.item()
        
        if batch_idx % 100 == 0:
            print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")
    
    print(f"Epoch {epoch} finished. Average loss: {total_loss / len(dataloader):.4f}")

Это наивная реализация. В реальности нужно добавить gradient accumulation, mixed precision, и сохранение чекпоинтов. И да, batch_size=4 — это для MacBook Air. На MacBook Pro с 38 GPU-ядрами можно выжать 8-12.

6 Запуск обучения и мониторинг: как понять, что все не зря

Запускаем скрипт и смотрим на лосс. Если он не падает — вы где-то накосячили. Первые 1000 шагов лосс должен уверенно уменьшаться. Для мониторинга я использую простой вывод в консоль, но для долгих запусков подключаю mlflow или wandb.

# Запуск обучения
python train.py --batch_size 4 --epochs 3 --save_dir ./checkpoints

# Мониторинг использования памяти (в отдельном терминале)
sudo powermetrics --samplers smc | grep -i "gpu|cpu"

Где собака зарыта: нюансы, которые съедят ваше время

• Память. MLX не умеет в gradient checkpointing из коробки. Придется реализовывать вручную или уменьшать размер модели.
• Данные. Токенизированный датасет в оперативке — путь к краху. Используйте потоковую загрузку и кэширование на диск.
• Отладка. Ошибки в MLX часто малопонятные. Начинайте с крошечной модели и датасета, чтобы убедиться, что цикл работает.
• Валидация. Без validation loss вы не поймете, переобучилась ли модель. Выделите 5% данных для валидации.

Если уперлись в потолок производительности, посмотрите статью про кастомные CUDA ядра. Принципы те же, только вместо CUDA — Metal Shading Language.

Частые вопросы (FAQ)

Что дальше? (Спойлер: самое интересное)

Вы обучили модель с нуля. Она генерирует текст, пусть и корявый. Теперь начинается магия — alignment. Той самой, которая превращает хаотичный генератор текста в полезного ассистента. RLHF, DPO, и прочие страшные аббревиатуры.

Мой совет: не пытайтесь сразу сделать идеальную модель. Запустите этот пайплайн на маленьком датасете, поймите каждую строчку кода, а потом масштабируйте. И да, сохраняйте чекпоинты каждую эпоху. Однажды это спасет вам неделю работы.

В следующей части разберем, как заставить модель следовать инструкциям и не галлюцинировать. А пока — удачного обучения. И помните: если что-то пошло не так, это не баг, это фича для следующего поста.