Портирование llm.c Андрея Карпаты на Mojo с ручной реализацией Metal-ядер дало неоднозначные результаты. bf16-конфигурация на M4 Max достигает 503.3 мс/шаг, что эквивалентно 8138 токенов/с. Это в 1.71 раза быстрее PyTorch MPS bf16, но на 24% медленнее нативного MLX, который показывает 406.5 мс/шаг. Причина разрыва - неэффективная работа bf16 matmul на Metal: ускорение относительно fp32 составляет лишь 1.1x против 2x у MLX, использующего тензорные ядра. Дополнительным препятствием стал троттлинг GPU M4 Max после ~8 секунд нагрузки, вынуждающий делать 30-секундные паузы между замерами.
Репозиторий включает 235 тестов корректности, а обученный чекпоинт с 29.53% на HellaSwag статистически неотличим от эталонного llm.c с 29.9%. Mojo-компилятор не обеспечил автоматической переносимости ядер между вендорами - потребовалась ручная ветвящаяся логика под Metal и CUDA, что контрастирует с компактностью исходных CUDA-ядер Карпаты.
Если вы работаете с Apple Silicon и ищете способы выжать максимум из железа, этот разбор даст полную картину: цифры, узкие места, практические ограничения и рекомендации по выбору инструмента. Для более широкого контекста по оптимизации инференса на нестандартных конфигурациях рекомендуем материал о запуске MoE-моделей на одной видеокарте - там разбираются схожие проблемы работы с памятью и пропускной способностью.
Ключевые результаты: Mojo с Metal-ядрами против PyTorch MPS и MLX
Прямое сравнение трёх конфигураций на M4 Max даёт чёткую иерархию производительности. За базу взят bf16-режим обучения GPT-2 124M - стандартный бенчмарк для низкоуровневых оптимизаций.
Mojo с ручными Metal-ядрами показывает 503.3 мс/шаг (8138 tok/s). PyTorch MPS bf16 на той же задаче тратит ~861 мс/шаг. Ускорение в 1.71x - значимое, но не революционное. MLX, нативный фреймворк Apple, выполняет шаг за 406.5 мс, оказываясь быстрее Mojo-решения на 24%. Разница в 96.8 мс на каждом шаге обучения быстро накапливается: за тысячу итераций MLX экономит более полутора минут чистого вычислительного времени.
Эти цифры ставят практический вопрос: оправдана ли сложность ручного написания Metal-ядер на Mojo, если готовый MLX даёт лучший результат без низкоуровневых усилий? Ответ зависит от контекста. Mojo выигрывает у PyTorch MPS, но проигрывает специализированному решению Apple. Для разработчиков, уже использующих Mojo в кроссплатформенных проектах, прирост относительно PyTorch может быть достаточным аргументом. Для тех, кто работает исключительно на Apple Silicon, MLX остаётся более производительным выбором.
Схожие дилеммы выбора инструмента возникают и при работе с большими MoE-моделями - в статье про стриминг-инференс для MoE мы разбирали, когда оправдано использовать TensorRT-LLM вместо стандартных решений.
Почему matmul стал узким местом: bf16 на Metal не даёт ожидаемого прироста
Умножение матриц - вычислительное ядро любого трансформера. В GPT-2 124M на matmul приходится подавляющая часть операций, поэтому эффективность этой операции напрямую определяет общую скорость обучения. Результаты порта вскрыли фундаментальную проблему: Metal-реализация bf16 matmul работает лишь в 1.1x быстрее fp32-версии. Это аномально низкий показатель для формата с половинной точностью.
Сравнение операций matmul: Metal bf16 vs fp32 vs MLX
Теоретически bf16 должен давать двукратный прирост производительности относительно fp32 за счёт уменьшенного вдвое объёма данных и соответствующего снижения нагрузки на память и вычислительные блоки. На практике Metal-ядра этого не обеспечивают. Причина - отсутствие эффективного использования тензорных ядер GPU Apple.
MLX решает эту задачу иначе. Фреймворк напрямую задействует тензорные ядра через специализированные низкоуровневые примитивы, получая честное 2x ускорение bf16 matmul относительно fp32. Это стандартный подход для современных GPU: тензорные ядра аппаратно оптимизированы под операции с пониженной точностью. Ручные Metal-ядра, написанные в рамках порта llm.c, не смогли воспроизвести эту оптимизацию.
Разрыв в производительности matmul объясняет почти всё отставание Mojo-решения от MLX. Остальные компоненты пайплайна обучения - нормализация, активации, attention-механизмы - вносят меньший вклад в общее время. Фокусировка на оптимизации именно matmul является приоритетной задачей для будущих итераций порта.
Этот паттерн знаком и по другим задачам инференса: в разборе DeepSeek-V4-Flash на B300 мы видели, как отказ MoE-ядра без expert parallel режет производительность в разы. Узкое место всегда одно - вычислительное ядро, и его оптимизация даёт максимальный эффект.
Троттлинг M4 Max: почему замеры требуют 30-секундных пауз
M4 Max - производительный чип, но его GPU-подсистема подвержена тепловому троттлингу при длительных нагрузках. В ходе экспериментов выяснилось, что стабильная производительность сохраняется лишь первые ~8 секунд непрерывной работы. После этого температура достигает порогового значения, и чип снижает частоты для защиты от перегрева.
Для получения воспроизводимых замеров между тестовыми прогонами требуются 30-секундные паузы. Это не баг реализации, а физическое ограничение системы охлаждения. Пассивное охлаждение MacBook Pro с M4 Max справляется с пиковыми нагрузками, но не с продолжительным обучением нейросетей на полной мощности GPU.
Практические следствия:
- Время экспериментов увеличивается кратно. Каждый замер требует почти минуты с учётом паузы, что замедляет итеративную разработку и отладку.
- Длительное обучение без пауз приводит к прогрессирующему падению производительности. Через 30-60 секунд непрерывной работы скорость может упасть на 20-30% от пиковой.
- Сравнение разных конфигураций требует строгого контроля температуры. Замеры, сделанные на горячем GPU, несопоставимы с холодным стартом.
Для продакшен-обучения на M4 Max это означает необходимость либо вводить принудительные паузы в тренировочный цикл, либо использовать внешнее охлаждение. Ни один из этих вариантов неудобен: паузы снижают общую пропускную способность, а внешнее охлаждение противоречит мобильной природе устройства.
Корректность порта: 235 тестов и качество модели на HellaSwag
Производительность важна, но она ничего не стоит без корректности вычислений. Порт llm.c на Mojo включает 235 тестов, покрывающих все ключевые операции: прямой и обратный проход, градиенты, функции потерь, инициализацию весов. Все тесты проходят, подтверждая численную эквивалентность реализации эталонному коду Карпаты.
Финальная проверка - качество обученной модели. Чекпоинт, полученный на Mojo с Metal-ядрами, показывает 29.53% на HellaSwag. Эталонный llm.c даёт 29.9%. Разница в 0.37 процентных пункта находится в пределах статистической погрешности, обусловленной недетерминизмом операций с плавающей точкой на GPU. Две модели, обученные с разными сидами на одном коде, могут показать схожий разброс.
Этот результат принципиален: он доказывает, что ручная реализация Metal-ядер не внесла скрытых ошибок в вычислительный граф. Пользователи порта могут быть уверены, что получают ту же модель, что и при обучении на CUDA, без компромиссов в качестве.
Переносимость Mojo между Metal и CUDA: миф или реальность?
Mojo позиционируется как язык, способный заменить Python и C++ в высокопроизводительных вычислениях, с потенциалом кроссплатформенной переносимости. Практика портирования llm.c показала, что это ожидание не оправдывается на уровне GPU-ядер.
Исходные CUDA-ядра Карпаты компактны и выразительны - несколько десятков строк кода реализуют matmul, attention и другие операции. При портировании на Mojo потребовалась ручная ветвящаяся логика для поддержки Metal и CUDA. Компилятор не предоставил абстракции, автоматически транслирующей один диалект GPU-программирования в другой. Разработчику пришлось писать два набора ядер и переключаться между ними в зависимости от целевой платформы.
Это контрастирует с обещаниями унифицированного программирования GPU. Metal и CUDA имеют разную модель памяти, разные примитивы синхронизации, разные соглашения об индексации потоков. Mojo-компилятор не скрывает эти различия - он транслирует код в соответствующий бэкенд, но логику адаптации оставляет программисту.
Вывод прагматичный: Mojo полезен как замена Python для высокоуровневой логики и как замена C++ для CPU-вычислений, но для GPU-ядер он не устраняет vendor lock-in. Если вы пишете под Metal, вы пишете под Metal. Если под CUDA - под CUDA. Кроссплатформенность достигается дублированием кода, а не автоматической трансляцией.
Практические выводы: когда стоит использовать Mojo с Metal-ядрами
Суммируем плюсы и минусы подхода на основе полученных данных.
Плюсы:
- Ускорение в 1.71x относительно PyTorch MPS bf16 - значимый выигрыш для тех, кто по каким-то причинам привязан к PyTorch и не хочет переходить на MLX.
- Полный контроль над вычислительным графом и ядрами. Это ценно для исследовательских задач, где требуется нестандартная логика обучения.
- Подтверждённая корректность: 235 тестов и идентичное качество модели на HellaSwag.
Минусы:
- Отставание на 24% от MLX. Нативный фреймворк Apple быстрее и не требует ручного написания ядер.
- Высокая сложность разработки. Ручные Metal-ядра требуют экспертизы в GPU-программировании и отладки низкоуровневого кода.
- Троттлинг M4 Max, вынуждающий делать 30-секундные паузы. Это снижает практическую применимость для длительного обучения.
- Отсутствие автоматической переносимости между Metal и CUDA. Кроссплатформенная разработка требует двойной работы.
Рекомендации в зависимости от сценария:
- Обучение небольших моделей на Mac. Используйте MLX. Он быстрее, проще в освоении и не требует пауз из-за троттлинга при правильной настройке.
- Кроссплатформенная разработка с GPU. Mojo оправдан, если вы готовы писать отдельные ядра под каждый бэкенд. Выигрыш относительно PyTorch есть, но он не перекрывает затраты на разработку для большинства проектов.
- Исследовательские проекты с нестандартными операциями. Ручные ядра на Mojo дают максимальную гибкость. Если MLX не поддерживает нужную вам операцию, Mojo с Metal - рабочий путь.
- Продакшен-обучение на Apple Silicon. Троттлинг M4 Max делает этот сценарий проблематичным. Рассмотрите Mac Studio с более мощным охлаждением или облачные GPU.
Порт llm.c на Mojo с Metal-ядрами - ценный эксперимент, который дал чёткие цифры и развеял несколько иллюзий о кроссплатформенности и производительности. Практическая польза для сообщества - в детальном понимании узких мест и ограничений, а не в готовом инструменте для повседневной работы. Если вас интересуют смежные темы оптимизации, обратите внимание на бенчмарк Multi-Token Prediction на RTX 5090 - там разбирается, когда продвинутые техники дают прирост, а когда их лучше отключить.