Как совместно разработать архитектуру программного и аппаратного обеспечения для искусственного интеллекта и машинного обучения в новой эре?

Совместная разработка архитектуры программного и аппаратного обеспечения для искусственного интеллекта и машинного обучения в новой эре секреты и практические рекомендации

Гармонический подход к созданию эффективной архитектуры для ИИ/МО

Современные инновационные технологии генеративного ИИ в компьютерном зрении, обработке естественного языка и т.д. недавно начали экспоненциально развиваться с научными прорывами в области инновационных модельных архитектур, включая стабильную диффузию, нейронную визуализацию (NeRF), текст в 3D, большие языковые модели (LLM) и т.д. Для этих передовых технологий требуются более сложные нейронные сети ИИ и гораздо больше вычислительных ресурсов и энергоэффективных архитектур.

Для выполнения вышеуказанных архитектурных требований важно сокращение разрыва между аппаратным (HW) и программным (SW)/алгоритмическим проектированием. Совместное проектирование включает в себя итеративный процесс проектирования, тестирования и совершенствования как аппаратных, так и программных компонентов до достижения требуемых производственных характеристик системы. Однако SW и HW обычно разрабатываются независимо друг от друга, поскольку программисты SW редко задумываются о том, на каком HW выполнить задачу, а HW, как правило, проектируется для поддержки широкого спектра SW. Требуется очень мало исследований совместного проектирования SW/HW, что необходимо для эффективной поддержки рабочих нагрузок ИИ.

Источник: Изображение автора.

Нам нужно разработать эффективный HW, осведомлённый о SW/алгоритмах, и SW/алгоритмы, осведомлённые о HW, чтобы они тесно взаимодействовали, чтобы выжать максимум из наших ограниченных вычислительных ресурсов. Как нам это сделать? Вот развивающаяся методология, которую можно использовать в качестве справочного материала, если вы создаёте новые возможности HW/SW для ИИ.

1. Идентифицировать прокси-рабочие нагрузки ИИ.

Для начала анализа нам понадобятся прокси-модели ИИ и список приоритетов для дальнейшего исследования. Вы можете использовать различные ресурсы, включая последние научные работы (CVPR, Siggraph, исследовательские лаборатории крупных технологических компаний) с открытым исходным кодом, отзывы клиентов или запросы, тенденции в отрасли, и т.д. Отфильтруйте несколько представительных моделей на основе вашего экспертного мнения. Этот шаг крайне важен, так как вы будете использовать их для проектирования «архитектуры будущего».

2. Тщательный анализ архитектуры модели.

Обязательно проведите всеобъемлющий анализ архитектуры модели, чтобы понять ее функциональность, инновации и разобрать ее на детальные элементы настолько, насколько возможно. Есть ли новые операторы, которые не поддерживаются в текущем технологическом стеке? Где находятся вычислительно интенсивные слои? Является ли модель интенсивной по передаче данных (памяти)? Какого типа данные требуются, и какие методы квантования можно применить, не потеряв точность? В какой части модели можно ускорить HW, и где есть потенциальные оптимизации производительности?

Например, в нейронной визуализации модели требуется работающая параллельно визуализация и вычисление (умножение матрицы). Вам нужно проверить, поддерживается ли это текущим SW стеком одновременно визуализацию или вычисление. В больших языковых моделях (LLMs) размер кеша ключ-значение (KV) растет в зависимости от длины последовательности ввода, и важно понять требования к памяти и потенциальные оптимизации иерархии передачи данных/памяти для обработки больших кешей ключ-значение.

Efficient Memory Management for Large Language Model Serving with Paged Attention [1]” src=”https://ai.miximages.com/miro.medium.com/v2/resize:fit:640/format:webp/1*hwEzEru80a5ZKJJXHI3wfA.png”/>

3. Включение и прототипирование SW

Загрузите открытый код для модели, определенной на Шаге 2, и запустите его на целевой рабочей платформе SW/HW. Этот шаг не является простым, особенно для новых/революционных моделей. Целью является создание работоспособного решения для анализа производительности, поэтому на этом этапе не обязательно предоставлять код продуктового качества. Нерабочий SW фикс без настройки производительности допустим для перехода к Шагу 4. Одним из важных шагов является преобразование предварительно обученной модели в среде разработки (Pytorch) в новый формат, требуемый целевым новым фреймворком.

torch.onnx.export(model, dummy_input, "resnet50.onnx", verbose=False, input_names=input_names,outputnames=output_names, export_params=True)

Однако часто возникают случаи, когда требуется значительная поддержка. Например, для выполнения моделей дифференцируемого рендеринга требуется поддержка авторегра. Очень вероятно, что эта функция не готова в новой среде и требует месяцев усилий от команды разработчиков. Другим примером является квантование GPTQ для LLMs, которое изначально может не поддерживаться в системе вывода. Вместо ожидания команды инженерии архитекторы могут запускать рабочую нагрузку на системе Nvidia для анализа производительности, так как Nvidia является выбором для академической разработки. Это позволяет создать список требований к ПО на основе наблюдаемых разрывов при обеспечении ПО.

4. Анализ производительности и архитектурные инновации.

Существует множество метрик для оценки производительности модели ИИ. Ниже рассмотрены основные, которые следует учитывать.

4.1 FLOPs (операции с плавающей запятой)и MACs (умножение-накопление операций).

Эти метрики обычно используются для расчета вычислительной сложности моделей глубокого обучения. Они предоставляют быстрый и простой способ понять количество арифметических операций, необходимых для работы модели. FLOPs можно рассчитать с помощью таких методов, как анализ статьи, отчеты Vtune или инструменты, например flops-counter.pytorch и pytorch-OpCounter.

4.2 Потребление памяти и пропускная способность (BW)

В основном потребление памяти состоит из весов (параметров сети) и входных данных. Например, модель Llama с 13 млрд параметров в FP16 потребляет около 13*2 (FP16=2 байта) = 26 ГБ памяти (входные данные можно пренебречь, так как вес занимает гораздо больше места). Еще одним ключевым фактором для LLMs является размер кэша KV. Кэш KV занимает до 30% от общей памяти, и он динамический (см. рисунок в Шаге 2). Большие модели обычно ограничены памятью, поскольку скорость зависит от того, насколько быстро данные перемещаются из системной памяти в локальную память или из локальной памяти в локальные кэши/регистры. Доступная пропускная способность памяти намного лучше предсказывает задержку вывода (время генерации токена для LLMs) по сравнению с пиковыми вычислительными TOPS. Один из показателей производительности – использование пропускной способности памяти (MBU), которое определяется как фактическая пропускная способность/пиковая пропускная способность. Идеально, если MBU близко к 100%, это указывает на полное использование пропускной способности памяти.

Достаточно памяти недостаточно!

Nvidia увеличивает FLOPs в разы, но не пропускную способность памяти. Источник: Substack/SemiAnalysis

Поскольку память является узким местом, требуются исследования в области сжатия продвинутых моделей и технологий кэширования/хранения памяти. Ниже перечислены несколько пионерских работ:

  • MemGPT: он использует ресурсы разных уровней иерархии памяти, таких как комбинация быстрой и маленькой ОЗУ и большой и медленной памяти хранения. Информация должна быть явно передана между ними. [2]
  • Квантование с низкой точностью (GPTQ, AWQ, GGML) для уменьшения потребления памяти моделями
  • Вычисления в памяти (PIM): снижение энергопотребления и улучшение производительности путем устранения необходимости перемещения данных.

4.3 Задержка/пропускная способность.

В компьютерном зрении задержка – это время генерации одного кадра. В контексте LLMs это время между первым токеном и генерацией следующего токена. Пропускная способность – это количество токенов/кадров в секунду. Задержка является важной метрикой для измерения производительности системы ИИ и является комплексным фактором как производительности ПО, так и аппаратного обеспечения. Существует различные стратегии оптимизации, о которых следует помнить, ниже перечислены некоторые из них:

  • Оптимизация операций с ограниченной пропускной способностью, таких как нормализации, операции поэлементно, SoftMax и ReLU. Предполагается, что нормализация и операции поэлементно потребляют примерно на 40% больше времени выполнения, чем умножение матриц, при этом достигают только 250x и 700x менее FLOPS, чем умножение матриц соответственно. Для решения этой проблемы можно использовать слияние ядер для слияния нескольких операторов с целью снижения затрат на передачу данных или заменить дорогостоящие операторы (softmax) на более легкие (ReLU).
Пропорции классов операторов. Источник: Все, что вам нужно - это перемещение данных.
  • Специализированная аппаратная архитектура. Интеграция специализированного оборудования (AVX, GPU, TPU, NPU) может привести к значительному ускорению работы и снижению энергопотребления, что особенно важно для приложений, требующих обработки в режиме реального времени на ресурсоемких устройствах. Например, инструкции Intel AVX могут ускорить работу на порядок выше, чем нативный Python код.
Ускорение за счет оптимизации производительности программы, умножающей две матрицы 4096 на 4096. Источник: Топлинда оптимальный шаг: Что будет побуждать вычислительную производительность после закона Мура?

Тензорные ядра на графических ускорителях Nvidia (V100, A100, H100 и др.) могут умножать и складывать две матрицы FP16 и/или FP32 за один такт, в отличие от ядер Cuda, которые могут выполнять только 1 операцию за такт. Однако использование тензорных ядер очень низкое (3% — 9% для полноценного обучения), что приводит к высоким затратам энергии и низкой производительности. В настоящее время идут активные исследования по улучшению использования систолического массива (FlexSA, многонаправленный SA и др.), о которых я напишу в следующих статьях.

Систолический массив. Источник: Telesens

Кроме того, так как память и передача данных всегда являются узким местом для больших моделей искусственного интеллекта, важно исследовать передовую архитектуру с учетом большей и более эффективной памяти на чипе. Примером является дизайн памяти Cerebras, где память адресуется независимо для каждого ядра.

Архитектура памяти Cerebras
  • Существует множество других оптимизаций: параллелизм, кэширование KV-кэша для LLM, разреженная активация и оптимизация от начала и до конца – я расскажу более подробно в следующих статьях

4.4 Мощность и энергоэффективность

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

Затраты энергии на вычисления и память. Источник: Song H из Стэнфорда

Вывод

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

Это совместное усилие по проведению тщательного архитектурного анализа новых моделей искусственного интеллекта. Цикл исследования может быть довольно долгим, и всегда существуют компромиссы между сроками выпуска продукта и глубиной анализа. В то же время требуется глубокая экспертиза в области аппаратных и программных технологий и искусственного интеллекта для создания эффективных решений в области ИИ.

Ссылки

[1] Ли Вусук, Ли Чжохан, Чжуан Сиюан, Шэн Йин, Ляньмин Чжэн, Ю Коди Хао, Джозеф Э. Гонсалес, Хао Чжан, Айон Стоица, Управление памятью для предоставления больших языковых моделей с пагинированным вниманием, 2023, arxiv

[2] Чарльз Пэкер, Вивиан Фэнг, Шишир Г. Патил, Кевин Лин, Сара Вудерс, Джозеф Э. Гонсалес, MemGPT: к LLM в качестве операционных систем, 2023, arxiv

[3] Чарльз Лейисерсон, Нил Томпсон, Джоэл Эмер, Брэдли Кузьмол, Батлер Лампсон, Дэниел Санчес и Тао Шардл, «Места на вершине хватит всем: что будет стимулировать производительность компьютеров после закона Мура?» 2020, Science

[4] Сон Хан, Джефф Пул, Джон Тран, Уильям Дж. Дэлли, Обучение весам и связям для эффективных нейронных сетей, 2015, arxiv