31,18 ГиБ весов модели генерации изображений Krea 2 Turbo удалось уместить в 32 ГБ unified memory MacBook Pro. Без Python и без облака. Порт на Swift и MLX загружает компоненты поэтапно: сначала Qwen кодирует промпт, затем DiT выполняет диффузию, и только потом VAE декодирует латенты в картинку. Пик потребления памяти не превышает доступный лимит.
Этот разбор полезен тем, кто изучает локальный AI-инференс на MacBook или другом Apple Silicon, хочет запускать diffusion-модели без Python и пытается понять, где проходит практическая граница 32 ГБ unified memory. Здесь показано, как устроен такой порт, почему поэтапная загрузка важнее номинального размера весов и какие проверки нужны, чтобы не принять почти совпадающие эмбеддинги за корректный результат.
В процессе порта обнаружился критический баг с неправильными RoPE-позициями для пяти токенов. Старые тесты на cosine similarity показывали ~1.0 - почти полное совпадение. Визуально разница была незаметна. Но исправление ошибки изменило 99,32% пикселей итогового изображения. Этот случай - жёсткое напоминание: полагаться только на метрики близости эмбеддингов при валидации порта нельзя.
Одинаковый seed не гарантирует идентичность персонажа. Причина - недетерминированность в порядке загрузки и выгрузки компонентов между запусками. Реализованные Regional Prompts и Character Sheet работают в однопроходном режиме, давая контроль над композицией и внешностью персонажа без файнтюнинга. Разберём архитектуру, баг и практические уроки порта.
Архитектура порта: как 31,18 ГиБ весов уместились в 32 ГБ unified memory
Прямой запуск Krea 2 Turbo на потребительском Apple Silicon выглядит невозможным. 31,18 ГиБ весов при доступных 32 ГБ unified memory - запаса всего 0,82 ГБ на системные нужды, буферы и промежуточные тензоры. Python-стек с PyTorch и MPS бэкендом в эту память не влезает: накладные расходы на аллокаторы, фрагментация и копирования между CPU и GPU съедят доступный запас. Решение - нативный Swift-стек с MLX и поэтапная загрузка компонентов.
Почему Swift и MLX: MLX vs PyTorch/MPS для unified memory
PyTorch на Apple Silicon работает через MPS-бэкенд, который организует вычисления с разделением ролей CPU и GPU, хотя физически память у платформы едина. Дополнительные буферы, кэширование, фрагментация и перемещения тензоров увеличивают пиковое потребление. Для модели на 31 ГиБ этого запаса недостаточно, поэтому типичный результат - out-of-memory ещё до завершения полного пайплайна.
MLX рассчитан на вычисления с unified memory и позволяет работать с массивами без обязательного повторного размещения между отдельными пулами памяти. Ленивые вычисления откладывают создание промежуточных тензоров до момента реальной необходимости, что помогает снизить пики потребления. Это не означает, что MLX отменяет стоимость латентов, кэшей и временных буферов: тяжёлый компонент всё равно должен помещаться вместе с рабочими данными.
Swift в этом стеке нужен не только для интерфейса. Он даёт явный контроль над жизненным циклом объектов: компонент загрузился, отработал, веса выгружены - память освобождена детерминированно, без ожидания сборщика мусора. Поэтому преимущество Swift/MLX в данном кейсе связано прежде всего с управлением unified memory и структурой пайплайна, а не с универсальным превосходством над PyTorch по скорости вычислений.
SwiftUI в этом стеке решает единственную задачу - минимальный UI для ввода промпта и отображения результата. Никаких Electron-подобных прослоек, которые тоже жрут память. Всё нативное, всё в unified memory.
Поэтапная загрузка: Qwen → DiT → VAE
Схема загрузки трёх компонентов последовательно - ключевой архитектурный паттерн порта. Каждый компонент живёт в памяти ровно столько, сколько нужно для его задачи:
- Qwen (~8 ГБ) загружается первым. Кодирует текстовый промпт в эмбеддинги. Результат - компактный тензор, который остаётся в памяти.
- Qwen выгружается. Загружается DiT (~20 ГБ) - Diffusion Transformer. Выполняет процесс диффузии: от шума к латентному представлению изображения. Пик памяти здесь максимальный: веса DiT + латенты + эмбеддинги от Qwen.
- DiT выгружается. Загружается VAE (~3 ГБ). Декодирует латенты в итоговое изображение.
Суммарный пик памяти не превышает ~28-29 ГБ. Оставшиеся 3-4 ГБ уходят на системные нужды и фреймворк MLX. Веса могут подгружаться с SSD напрямую через memory-mapped файлы, но это замедляет инференс - для продакшен-сценариев лучше держать модель в RAM.
Memory-mapped загрузка полезна, когда нужно уменьшить разовый объём копирования или не держать весь файл в активной памяти, но она не превращает 32 ГБ unified memory в память большего объёма. Во время вычислений всё равно нужны рабочие буферы, латенты и активный набор весов. Поэтому при диагностике OOM важно смотреть не только на размер файла модели, но и на пик самого тяжёлого этапа.
Признаки похожего OOM-сценария на Mac:
- модель проходит загрузку отдельных файлов, но падает при создании первого большого рабочего тензора;
- система начинает активно использовать SSD и приложение резко замедляется;
- ошибка возникает только на этапе DiT, хотя Qwen и VAE загружаются отдельно;
- после повторных запусков доступная память уменьшается из-за неосвобождённых буферов или кэшей;
- уменьшение разрешения помогает, но не устраняет проблему полностью, потому что основной объём занимают веса.
Этот подход применим к diffusion-модели с разделяемыми компонентами. Если текстовый энкодер, дифа и декодер - отдельные модули, их можно загружать последовательно. Похожая схема подходит для пайплайнов, где тяжёлый Transformer можно освободить до декодирования результата. Она хуже подходит для монолитных архитектур, которым требуется одновременный доступ к нескольким крупным блокам, а также для случаев, когда один компонент вместе с промежуточными тензорами уже превышает доступную память. Ограничение - сумма пика весов самого тяжёлого компонента плюс промежуточные тензоры должна влезать в доступную unified memory. Для 32 ГБ машин потолок - модели с DiT до ~22 ГБ.
Критический баг RoPE: когда cosine similarity ~1.0 недостаточно
Баг обнаружился случайно. Порт выдавал визуально корректные изображения. Промежуточные эмбеддинги при сравнении с эталонным Python-раннером давали cosine similarity ~1.0 - расхождение в пятом-шестом знаке после запятой, которое обычно списывают на разницу в реализации операций. Но при попиксельном сравнении итоговых изображений выяснилось: различается 99,32% пикселей.
Как пять токенов сломали воспроизводимость
Проблема сидела в позиционном кодировании RoPE (Rotary Position Embedding). При токенизации промпта пять специальных токенов - <|startoftext|>, <|endoftext|> и три паддинг-токена - получали индексы позиций, смещённые на единицу относительно эталонной реализации. Ошибка в индексации цикла: сдвиг начинался не с нулевой позиции, а с первой.
Влияние на attention оказалось коварным. Модель «видела» неправильные расстояния между токенами. Для семантически значимых токенов ошибка в одну позицию на коротких последовательностях почти не меняла attention-веса - отсюда cosine similarity ~1.0. Но в процессе диффузии, где позиционная информация влияет на пространственное распределение признаков, накопленная ошибка давала совершенно другую картину латентов.
Связь с проблемой идентичности персонажа прямая. Даже при фиксированном seed, разные запуски порта могли давать разные результаты из-за недетерминированности в порядке загрузки компонентов. Пять токенов с неправильными позициями добавляли дополнительный источник вариативности, который не ловился стандартными тестами.
RoPE bug: симптом, причина, поиск и валидация
- Симптом: эмбеддинги почти совпадают по cosine similarity, но латенты или итоговые изображения заметно расходятся.
- Причина: смещение индексов позиций, различия в обработке специальных токенов или паддинга, а также несовпадение 0-based и 1-based индексации.
- Как искать: сравнивать токены, position ids и результаты RoPE отдельно, затем проверять латенты после каждого шага диффузии.
- Как валидировать: фиксировать seed, параметры генерации, типы данных и порядок загрузки, а итог проверять попиксельно, а не только по эмбеддингам.
Уроки для тестирования генеративных моделей
Cosine similarity на эмбеддингах не работает как единственный критерий корректности порта. Промежуточные представления могут совпадать с точностью до пятого знака, а итоговое изображение - различаться почти полностью. Необходимо попиксельное сравнение эталонных и тестовых изображений.
Фиксация всех случайных факторов критична: seed, порядок загрузки модулей, типы данных (float32 vs bfloat16), даже версия MLX. Разница в реализации exp2 на GPU разных поколений Apple Silicon может давать расхождение в младших битах, которое каскадно расходится через 50 шагов диффузии.
Практический чек-лист для валидации порта diffusion-модели:
- Зафиксировать seed и все параметры генерации.
- Прогнать один и тот же промпт на эталонной и тестовой реализации.
- Сравнить попиксельно итоговые изображения (допустимое расхождение - 0 пикселей при одинаковых типах данных).
- Если расхождение есть - сравнивать латенты после каждого шага диффузии, чтобы локализовать шаг расхождения.
- Проверить position ids для специальных токенов и паддинга.
- Не полагаться на cosine similarity эмбеддингов как на достаточный тест.
Этот баг - не экзотика. При портировании моделей между фреймворками ошибки позиционного кодирования возникают регулярно. Разница в индексации (0-based vs 1-based), обработка специальных токенов, паддинг - типичные места слома.
Regional Prompts и Character Sheet в однопроходном режиме
Порт Krea 2 Turbo реализует две продвинутые функции контролируемой генерации без многократных прогонов модели. Обе работают за один проход диффузии.
Regional Prompts: контроль композиции без дообучения
Regional Prompts позволяют задать разные текстовые описания для разных областей изображения. Пример: «слева - рыжий кот на диване, справа - чёрная собака у окна». Области задаются координатами или масками. Механизм работает через модификацию attention-масок в DiT: токены левого промпта «видят» только левую часть латентного пространства, токены правого - только правую.
Взаимодействие с основным промптом настраивается весами. Основной промпт задаёт общую сцену и стиль, региональные - уточняют детали в конкретных зонах. Артефакты на стыках возможны, если маски имеют резкие границы. Размытие границ маски на 10-15% ширины региона сглаживает переходы.
Character Sheet: стабильность персонажа между генерациями
Character Sheet фиксирует внешность персонажа через эмбеддинги референсного изображения. Референс прогоняется через VAE-энкодер, полученные латенты конкатенируются с текстовыми эмбеддингами промпта. DiT получает сигнал «этот персонаж должен выглядеть вот так».
Ограничения метода: поза и освещение на референсе сильно влияют на результат. Если референс снят анфас при дневном свете, а промпт описывает персонажа в профиль при неоновом освещении, консистентность падает. Решение - несколько референсов с разных ракурсов, но это увеличивает объём эмбеддингов и пик памяти.
Обе функции реализованы в однопроходном режиме: не требуется генерировать изображение, потом маскировать, потом генерировать снова. Это критично для скорости - один проход диффузии на M2 Max занимает 8-12 секунд, каждый дополнительный проход удваивает время.
Практические выводы: когда стоит портировать модели на Swift/MLX
Порт на Swift/MLX оправдан в трёх случаях:
- Память - узкое место. Python-стек не влезает в доступную unified memory, а модель физически помещается при поэтапной загрузке. Выигрыш - возможность запуска на потребительском железе без облаков.
- Продакшен в Apple-экосистеме. Нативное приложение на SwiftUI с инференсом на MLX не требует серверной части, работает офлайн и использует все оптимизации чипа (ANE, AMX-блоки).
- Контроль над пайплайном. Поэтапная загрузка, детерминированное управление памятью, отсутствие фреймворк-специфичных сюрпризов вроде неявных копирований тензоров.
Минусы: меньше готовых инструментов. Приходится руками реализовывать токенизатор, attention-маски, позиционное кодирование. Отладка ошибок вроде бага с RoPE-позициями занимает дни. Привязка к Apple Silicon означает, что порт не запустится на NVIDIA-железе или в облаке.
Сравнение по скорости: на M2 Max с 32 ГБ unified memory порт Krea 2 Turbo генерирует изображение 1024x1024 за 8-12 секунд. Python-аналог с PyTorch и MPS на той же машине либо не запускается вовсе (out-of-memory), либо требует агрессивного оффлоада на SSD и выдаёт результат за 40-60 секунд. Выигрыш - в эффективном использовании unified memory, а не в сырой производительности вычислений.
Частые ограничения локального инференса на Mac
Можно ли портировать таким способом любую diffusion-модель? Нет. Проще всего работать с архитектурой, где Qwen или другой текстовый энкодер, DiT и VAE являются отдельными компонентами. Монолитный пайплайн или модель, которой нужно одновременно держать несколько крупных блоков, может не получить преимущества от поэтапной загрузки.
Заменяет ли MLX PyTorch/MPS во всех сценариях? Нет. MLX уместен, когда важны unified memory, нативный Apple Silicon-стек и контроль над памятью. PyTorch/MPS остаётся практичнее там, где нужны готовые реализации, совместимость с существующим кодом или переносимость на другие GPU.
Достаточно ли 32 ГБ unified memory? Только при подходящем пике потребления. Номинальный размер весов не равен объёму памяти, необходимому для инференса: добавляются латенты, эмбеддинги, временные буферы и системные процессы. Поэтапная загрузка снижает пик, но не устраняет это ограничение.
Гарантирует ли одинаковый seed одинаковый результат? Нет. Нужно контролировать порядок загрузки и выгрузки компонентов, типы данных, версию MLX и другие источники недетерминированности. Для проверки порта важнее воспроизводимость всего пайплайна, чем совпадение одного промежуточного вектора.
Для разработчиков, которые рассматривают локальный инференс diffusion-моделей на Mac, порт Krea 2 Turbo - прецедент. Он показывает: модели, которые «не влезают» в 32 ГБ по спецификации, могут работать на потребительском Apple Silicon при подходящей декомпозиции пайплайна. Цена - ручная реализация компонентов, ограничения совместимости и тщательное тестирование каждого слоя.