Обучаем LLM на Swift: от гигафлопсов к терафлопсам с помощью матричного умножения

Кто сказал, что Swift живёт только в мире мобильных приложений и не умеет в high performance computing? Один отважный разработчик решил бросить вызов устоям и попробовать обучить языковую модель (LLM) на чистом Swift. И, как это обычно бывает, упёрся в матричное умножение — ту самую операцию, на которой держится каждый уважающий себя нейросетевой фреймворк.

В первой части своего цикла автор показывает, как можно выжать из матричного умножения на Swift производительность от единиц гигафлопсов до единиц терафлопсов. Это как если бы ваша бабушка вдруг начала бегать марафоны — разница колоссальная. Секрет в том, чтобы задействовать все ядра процессора и не давать кэшу скучать.

Конечно, без боли не обошлось. Автор честно признаётся, что стандартная реализация на Swift даёт жалкие 0.1 Gflop/s — примерно как печатать код на клавиатуре с западающими клавишами. Но после ряда оптимизаций (ручная векторизация, работа с памятью, параллелизм) ему удалось достичь 1.2 Tflop/s на M1 Max. Это уже похоже на настоящую битву с бубном, только вместо бубна — LLVM и ассемблерные вставки.

Особенно забавно читать про то, как автор боролся с компилятором, который упорно не хотел генерировать оптимальный код. Знакомая боль, да? Кажется, каждый разработчик хоть раз ловил себя на мысли: "Я лучше знаю, как оптимизировать, чем этот умник из LLVM". В итоге пришлось использовать встроенный ассемблер и явно указывать инструкции NEON.

Впрочем, автор не обещает, что его подход заменит PyTorch или TensorFlow. Это скорее исследовательский проект, показывающий, что Swift может быть не только элегантным, но и быстрым. Если вам надоело ждать, пока Python прогреет нейронку, и вы готовы променять удобство на скорость — добро пожаловать в мир Swift для ML.

Комментарий студии METABYTE: Мы тоже любим оптимизировать код до последнего такта, но обычно предпочитаем не писать ассемблер вручную — для этого есть мы, а вы занимайтесь бизнесом. Хотя, если ваш проект на Swift и требует скорости, мы с радостью поможем разогнать его до терафлопсов.