Большинство инженеров, создающих модели ИИ сегодня, работают с графическими процессорами (GPU). Они пишут модели на PyTorch, запускают их и измеряют задержку и загрузку. То, что происходит между фреймворком и кремнием, часто остается скрытым. Это было приемлемо, когда доминировала одна архитектура. Сейчас это гораздо сложнее, так как аппаратное обеспечение для ИИ диверсифицировалось и теперь включает в себя GPU, Google TPU, чип Cerebras Wafer-Scale Engine (WSE), процессор Groq Language Processing Unit (LPU) и ускоритель Spyre от IBM.

Новая специализация на Coursera под названием «Архитектура потока данных для ИИ: от принципов к практике» создана для инженеров, которым необходимо работать во всем этом диапазоне. Программу из шести курсов ведет Каутар Эль Маграуи (Kaoutar El Maghraoui), ведущий научный сотрудник IBM Research и адъюнкт-профессор компьютерных наук Колумбийского университета. Она охватывает модели выполнения, компиляцию, системы памяти, отладку и моделирование производительности посредством практических лабораторных работ.

«Иногда сама модель не является сложной частью, — говорит Эль Маграуи. — Сложность заключается в понимании того, как эта модель отображается на конкретную архитектуру и почему то, что отлично работает на одной системе, может вести себя совершенно иначе на другой».

Одна модель PyTorch и пять моделей выполнения

Специализация начинается с одной модели PyTorch. На уровне фреймворка это тот же самый nn.Module на любой целевой платформе. Как только компилятор отображает его на аппаратное обеспечение, выполнение начинает выглядеть совершенно иначе, и каждый стек поддерживает собственный набор операторов, макетов и графовых преобразований. GPU запускает модель в виде ядер (kernels) во множестве параллельных потоков с кэшами, управляемыми аппаратно, и специализированными матричными блоками. TPU распределяет работу по систолическим массивам, а компилятор XLA управляет тайлингом (разбиением на плитки) и перемещением данных. Cerebras распределяет вычисления и более 40 ГБ памяти SRAM на кристалле по всей пластине. Groq полагается на свой компилятор для планирования запуска каждой операции и перемещения данных. Spyre компилирует модель в граф потока данных, причем его компилятор определяет макет и размещение в scratchpad-памяти, в то время как среда выполнения организует передачи и запуск. Детали варьируются в зависимости от поколения продукта.

Эти особенности определяют, какие оптимизации имеют значение. Слияние (fusion) устраняет промежуточный трафик памяти, но степень экономии зависит от того, сохраняет ли компилятор промежуточные данные на кристалле, как тензоры разбиваются на плитки и создает ли слитое ядро дополнительную нагрузку на память. Размерность тензора, которая идеально отображается на одной целевой платформе, может потребовать дополнения (padding) или перестановки данных на другой. В специализации выполнение потока данных рассматривается как архитектурный класс с общими принципами. Каждая концепция вводится в общих чертах до того, как обсуждается какая-либо конкретная система. Spyre служит основным практическим примером, поскольку команда Эль Маграуи работает с ним напрямую, но это лишь одна из пяти систем, сравниваемых на протяжении всего курса.

Две реализации одной и той же операции также могут выполняться с сильно различающейся скоростью по причинам, которые имеют мало общего с арифметикой. Для матричного умножения и внимания (attention) тайлинг, размещение данных, синхронизация и перемещение между памятью на кристалле и вне его часто определяют достижимую производительность.

Рисунок 1: Карта ландшафта ускорителей. Студенты используют эту таблицу во всех шести курсах и заполняют обоснование для каждой ячейки по мере прохождения.

«Не существует единой архитектуры, которая побеждает во всех категориях, — отметила Эль Маграуи. — GPU отдают приоритет гибкости и широкой программируемости. Другие системы жертвуют частью этой гибкости ради эффективности, детерминизма, масштабируемости или более жесткого контроля над выполнением».

Компилятору уделяется особое внимание. На специализированных ускорителях он отвечает за тайлинг, слияние, планирование, размещение и управление памятью. Именно здесь выигрывается или проигрывается большая часть производительности при адаптации модели. «У вас может быть ускоритель, который чрезвычайно эффективен для поддерживаемых им операций, — сказала она, — но один неподдерживаемый оператор, одна форма тензора с неудачным отображением или одна проблема с размещением памяти могут оказать серьезное влияние на всю модель».

Рисунок 2: Упрощенный вид компиляции PyTorch, рассматриваемый в курсе 3. Выбор запускаемых преобразований зависит от режима компиляции и бэкенда. С помощью torch.compile модуль Inductor выполняет трансляцию по умолчанию, а torch-spyre расширяет его для Spyre. TPU транслируют код через XLA. Бэкенды потока данных принимают большинство решений о слиянии, тайлинге и размещении во время трансляции.

Чему посвящена специализация

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

Курс 1 — Ландшафт ускорителей: предметно-ориентированные архитектуры, управление потоком и выполнение потока данных, иерархии памяти и модели программирования для всех пяти систем.
Курс 2 — Модели выполнения потока данных: узлы, ребра, правила срабатывания, статическое и динамическое планирование. Блок внимания трансформера, прослеженный через Spyre и TPU.
Курс 3 — Компиляция для потока данных: TorchDynamo, AOTAutograd и генерация кода бэкенда. Слияние, тайлинг и планирование в XLA и компиляторе Spyre.
Курс 4 — Системы памяти: кэши, scratchpad-память и архитектуры только на базе SRAM. Как каждая иерархия памяти меняет управление кэшем KV для инференса LLM.
Курс 5 — Отладка в стеке потока данных: восьмиуровневая таксономия сбоев, включая ошибки захвата графа, тупики планирования, тайм-ауты DMA и числовые проблемы.
Курс 6 — Моделирование производительности и оптимизация: анализ Roofline, адаптированный для памяти, управляемой программным обеспечением, где планирование измеряется как третье «бутылочное горлышко» наряду с вычислениями и памятью.

Отладка и производительность в стеке

На ускорителе потока данных неправильный результат или замедление часто проявляются далеко от своей первопричины. Значение NaN на выходе модели может быть вызвано декомпозицией оператора, этапом слияния, несоответствием макета тензора или ошибкой DMA. В Курсе 5 эти сбои упорядочены в виде восьмиуровневой таксономии: от ошибок захвата графа во фронтенде до конфликтов планирования, взаимоблокировок, тайм-аутов DMA и числовых проблем в бэкенде.

Для проверки корректности студенты сводят сбой к небольшому воспроизводимому примеру, сравнивают вывод каждого слоя с эталоном на CPU или GPU и находят первую точку расхождения результатов. Для оценки производительности в Курсе 6 ограничивают область замеров, отделяют время ядра от времени передачи, сравнивают результат с идеалом компилятора и относят остаток на счет перемещенных байтов на каждом уровне памяти. Затем каждая операция заносится на график Roofline.

В одном из багов зарегистрированных в проекте с открытым исходным кодом torch-spyre, клонирование 65-элементного fp16 тензора возвращало неверное последнее значение. Журнал передачи был корректным, но скомпилированная копия охватывала только первые 64 элемента, а сравнение выводов каждого слоя с эталоном позволило обнаружить ошибку на этом шаге.

Рисунок 3: Методы отладки и анализа производительности, изучаемые в курсах 5 и 6, взятые из руководств по отладке и профилированию с открытым исходным кодом torch-spyre. Оба проходят через один и тот же стек компилятора. Работа с правильностью заключается в поиске первого слоя, выходные данные которого отличаются от эталона. Работа с производительностью соотносит измеренное время со слоем и уровнем памяти, которые его вызвали.

Рисунок 4: Анализ по методу Roofline, изучаемый в Курсе 6. График Roofline показывает, насколько операция находится ниже своего предела. Он не показывает причину. На целевой платформе потока данных причиной может быть планирование: перестановка данных, обход графа или простои конвейера. Профилировщики локализуют время до уровня ядер, передачи данных и связи. Насколько детально они распределяют время внутри ядра, зависит от инструмента.

Лабораторные работы на Google Colab и Spyre

Каждый курс завершается основной лабораторной работой, которая выполняется в Google Colab со стандартным PyTorch. Каждая лабораторная работа также имеет дополнительную практическую часть Spyre, выполняемую на аппаратном обеспечении Spyre. Лабораторные работы разработаны Рашедом Бхатти (Rashed Bhatti), кандидатом наук (PhD), старшим техническим сотрудником IBM Research и приглашенным преподавателем лабораторий в рамках специализации. Каждая работа занимает от 90 минут до двух часов, при этом студенты пишут и запускают собственный код. В первой лабораторной работе они создают среду тестирования производительности с нуля с помощью torch.utils.benchmark и объясняют разницу между CPU и GPU в терминах параллелизма SIMT и пропускной способности HBM. В последующих лабораторных работах они трассируют блок трансформера с помощью torch.fx, сравнивают режим eager с torch.compile, читают временную шкалу памяти PyTorch Profiler во время инференса LLM, диагностируют три намеренно заложенных бага и строят график Roofline для рабочей нагрузки LLM.

Для основных лабораторных работ не требуются проприетарные инструменты. В дополнительной практической части Spyre, называемой Dataflow Deep Dive, студенты в рамках партнерских программ выполняют те же лабораторные работы на оборудовании Spyre. Там они напрямую работают с концепциями потока данных из каждого курса: они читают граф потока данных, создаваемый компилятором Spyre, и классифицируют слияния и ребра перемещения данных, которые он добавляет, прослеживают, как тензоры разбиваются на блоки (sticks) и размещаются в кэш-памяти (scratchpad), измеряют трафик между HBM и кэш-памятью с помощью аппаратных счетчиков и строят график Roofline по этим счетчикам для программно управляемой иерархии памяти. Углубленное изучение (deep dive) не является обязательным для завершения специализации или получения сертификата.

Кому следует пройти этот курс

Специализация предназначена для учащихся среднего и продвинутого уровня, включая аспирантов в области компьютерных наук или электротехники, инженеров-разработчиков компиляторов, специалистов по аппаратной архитектуре и практикующих специалистов по машинному обучению, которые оптимизируют рабочие нагрузки для различных аппаратных платформ. Участники должны уверенно владеть Python и C/C++ на среднем уровне, быть знакомы с архитектурами CNN и трансформеров, а также иметь базовые знания по архитектуре компьютеров.

Слушатели, успешно завершившие программу, получают профессиональный сертификат в области инженерии производительности систем искусственного интеллекта (AI Systems Performance Engineering).

«К концу курса, когда вы увидите новый ускоритель ИИ, я хочу, чтобы вы знали, какие вопросы о нем задать», — сказал Эль Маграуи (El Maghraoui). — «Какова модель выполнения? Где находится память? Кто планирует операции? Что контролирует компилятор?»

Специализация открыта для записи на Coursera.