Улучшения WebAssembly и WebGPU для более быстрого веб-ИИ, часть 2

Этот документ является продолжением усовершенствований WebAssembly и WebGPU для более быстрого веб-ИИ, часть 1 . Прежде чем продолжить, мы рекомендуем вам прочитать этот пост или посмотреть выступление на IO 24 .

Остин Энг
Austin Eng
Дипти Гандлури
Deepti Gandluri
Франсуа Бофор
François Beaufort

Веб-графический процессор

WebGPU предоставляет веб-приложениям доступ к оборудованию графического процессора клиента для выполнения эффективных высокопараллельных вычислений. С момента запуска WebGPU в Chrome мы видели в сети невероятные демонстрации искусственного интеллекта (ИИ) и машинного обучения (МО).

Например, Web Stable Diffusion продемонстрировала, что можно использовать ИИ для создания изображений из текста прямо в браузере. Ранее в этом году собственная команда Google Mediapipe опубликовала экспериментальную поддержку вывода больших языковых моделей .

На следующей анимации показана Gemma , модель большого языка Google (LLM) с открытым исходным кодом, работающая полностью на устройстве в Chrome в режиме реального времени.

Следующая демонстрация Hugging Face модели Segment Anything Model компании Meta создает высококачественные маски объектов полностью на клиенте.

Это всего лишь пара замечательных проектов, демонстрирующих возможности WebGPU для искусственного интеллекта и машинного обучения. WebGPU позволяет этим и другим моделям работать значительно быстрее, чем на процессоре.

Тест Hugging Face WebGPU для встраивания текста демонстрирует огромное ускорение по сравнению с реализацией той же модели на ЦП. На ноутбуке Apple M1 Max WebGPU работал более чем в 30 раз быстрее. Другие сообщают, что WebGPU ускоряет тест более чем в 120 раз .

Улучшение функций WebGPU для искусственного интеллекта и машинного обучения

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

Среди многочисленных улучшений WebGPU за последний год мы продолжили добавлять новые возможности для повышения производительности машинного обучения и искусственного интеллекта в Интернете. Недавно мы запустили две новые функции: 16-битные числа с плавающей запятой и упакованные целочисленные скалярные произведения.

16-битная с плавающей запятой

Помните, что рабочие нагрузки машинного обучения не требуют точности . shader-f16 — это функция, которая позволяет использовать тип f16 в языке шейдеров WebGPU. Этот тип с плавающей запятой занимает 16 бит вместо обычных 32 бит. f16 имеет меньший диапазон и менее точен, но для многих моделей ML этого достаточно.

Эта функция повышает эффективность несколькими способами:

  • Уменьшение памяти : тензоры с элементами f16 занимают половину места, что сокращает использование памяти вдвое. Вычисления на графическом процессоре часто ограничивают пропускную способность памяти, поэтому использование половины памяти часто означает, что шейдеры работают в два раза быстрее. Технически вам не нужна клавиша f16 для экономии пропускной способности памяти. Можно хранить данные в формате низкой точности, а затем расширять их до полного f32 в шейдере для вычислений. Но графический процессор тратит дополнительную вычислительную мощность на упаковку и распаковку данных.

  • Сокращение преобразования данных : f16 использует меньше вычислений за счет минимизации преобразования данных. Данные низкой точности можно сохранять и затем использовать напрямую, без преобразования.

  • Повышенный параллелизм . Современные графические процессоры способны одновременно помещать больше значений в исполнительные блоки графического процессора, что позволяет ему выполнять большее количество параллельных вычислений. Например, графический процессор, поддерживающий до 5 триллионов операций с плавающей запятой f32 в секунду, может поддерживать 10 триллионов операций с плавающей запятой f16 в секунду.

Снимок экрана теста WebGPU для встраивания текста
Благодаря shader-f16 тест Hugging Face WebGPU для встраивания текста выполняет тест в 3 раза быстрее, чем f32 на ноутбуке Apple M1 Max.

WebLLM — это проект, который может запускать несколько больших языковых моделей. Он использует Apache TVM , среду компилятора машинного обучения с открытым исходным кодом.

Я попросил WebLLM спланировать поездку в Париж, используя модель Llama 3 с восемью миллиардами параметров. Результаты показывают, что на этапе предварительного заполнения модели f16 работает в 2,1 раза быстрее, чем f32. На этапе декодирования скорость более чем в 1,3 раза выше.

Приложения должны сначала подтвердить, что адаптер графического процессора поддерживает f16, и, если он доступен, явно включить его при запросе устройства графического процессора. Если f16 не поддерживается, вы не можете запросить его в массиве requiredFeatures .

// main.js

const adapter = await navigator.gpu.requestAdapter();
const supportsF16 = adapter.features.has('shader-f16');
if (supportsF16) {
  // Use f16.
  const device = await adapter.requestDevice({
    requiredFeatures: ['shader-f16'],
  });
  initApp(device);
}

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

// my-shader.wgsl

enable f16;

struct Data {
  values : array<vec4<f16>>
}
@group(0) @binding(0) var<storage, read> data : Data;
@compute @workgroup_size(64) fn main(@builtin(global_invocation_id) gid : vec3u) {
  let value : vec4<f16> = data.values[gid.x];
  ...
}

Упакованные целочисленные скалярные произведения

Многие модели по-прежнему хорошо работают с точностью всего 8 бит (половина f16). Это популярно среди LLM и моделей изображений для сегментации и распознавания объектов. Тем не менее, качество вывода моделей ухудшается с меньшей точностью, поэтому 8-битное квантование подходит не для каждого приложения.

Относительно немногие графические процессоры изначально поддерживают 8-битные значения. Именно здесь на помощь приходят упакованные целочисленные скалярные произведения. Мы включили DP4a в Chrome 123 .

Современные графические процессоры имеют специальные инструкции для того, чтобы взять два 32-битных целых числа, интерпретировать каждое из них как 4 последовательно упакованных 8-битных целых числа и вычислить скалярное произведение между их компонентами.

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

Например, давайте умножим матрицу 4 x 8 на вектор 8 x 1. Для этого необходимо взять 4 скалярных произведения для вычисления каждого значения выходного вектора; А, Б, С и Д.

Схема примера умножения матрицы на вектор

Процесс вычисления каждого из этих выходных данных одинаков; мы рассмотрим шаги, необходимые для вычисления одного из них. Перед любыми вычислениями нам сначала нужно преобразовать 8-битные целочисленные данные в тип, с которым мы можем выполнять арифметические действия, например f16. Затем мы выполняем поэлементное умножение и, наконец, складываем все произведения вместе. В общей сложности для всего умножения матрицы на вектор мы выполняем 40 преобразований целых чисел в числа с плавающей запятой для распаковки данных, 32 умножения с плавающей запятой и 28 операций сложения с плавающей запятой.

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

Для каждого из выходных данных вектора результатов мы выполняем две операции упакованного скалярного произведения, используя встроенный язык шейдинга WebGPU dot4U8Packed , а затем суммируем результаты. Итого, для всего умножения матрицы на вектор мы не производим никакого преобразования данных. Мы выполняем 8 упакованных скалярных произведений и 4 сложения целых чисел.

Схема примера умножения упакованной целочисленной матрицы на вектор

Мы протестировали упакованные целочисленные скалярные произведения с 8-битными данными на различных потребительских графических процессорах. По сравнению с 16-битной плавающей запятой мы видим, что 8-битная операция работает в 1,6–2,8 раза быстрее. Когда мы дополнительно используем упакованные целочисленные скалярные произведения, производительность становится еще лучше. Это в 1,7–2,9 раза быстрее.

Скриншот ускорения умножения матрицы на вектор: f16 против u8
Диаграмма 1. Ускорение матричного вектора при сравнении f16 с U8 и U8 с dot4U8Packed.

Проверьте поддержку браузера с помощью свойства wgslLanguageFeatures . Если графический процессор изначально не поддерживает упакованные скалярные произведения, браузер полифилизирует свою собственную реализацию.

// main.js

if (navigator.gpu.wgslLanguageFeatures.has('packed_4x8_integer_dot_product')) {
  // Use dot4U8Packed, dot4I8Packed builtin
  // functions in the shaders.
}

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

Раньше — шейдер WebGPU, который накапливает частичные скалярные произведения в переменную sum. В конце цикла сумма сохраняет полное скалярное произведение между вектором и одной строкой входной матрицы.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f16;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : vec4<f16> = vector.values[i];
    let v2 : vec4<f16> = matrix.values[start + i];
    sum += dot(v1, v2);
  }
}

After — шейдер WebGPU, написанный для использования скалярных произведений упакованных целых чисел. Основное отличие состоит в том, что вместо загрузки 4 значений с плавающей запятой из вектора и матрицы этот шейдер загружает одно 32-битное целое число. Это 32-битное целое число содержит данные четырех 8-битных целочисленных значений. Затем мы вызываем dot4U8Packed для вычисления скалярного произведения двух значений.

// my-dot-product.wgsl

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid : vec3u) {
  var sum : f32;
  let start = gid.x * uniforms.dim;
  for (var i = 0u; i < uniforms.dim; i++) {
    let v1 : u32 = vector.values[i];
    let v2 : u32 = matrix.values[start + i];
    sum += dot4U8Packed(v1, v2);
  }
}

Как 16-битные скалярные произведения с плавающей запятой, так и упакованные целочисленные скалярные произведения — это встроенные функции Chrome, которые ускоряют работу искусственного интеллекта и машинного обучения. 16-битные числа с плавающей запятой доступны, если аппаратное обеспечение их поддерживает, а Chrome реализует упакованные целочисленные скалярные произведения на всех устройствах.

Вы можете использовать эти функции в Chrome Stable уже сегодня, чтобы повысить производительность.

Предлагаемые функции

Забегая вперед, мы исследуем еще две функции: подгруппы и кооперативное матричное умножение.

Функция подгрупп обеспечивает параллелизм на уровне SIMD для взаимодействия или выполнения коллективных математических операций, таких как суммирование более 16 чисел. Это обеспечивает эффективный обмен данными между потоками. Подгруппы поддерживаются современными API-интерфейсами графических процессоров с разными именами и в несколько разных формах.

Мы превратили общий набор в предложение, которое подали в группу стандартизации WebGPU. Мы создали прототипы подгрупп в Chrome под экспериментальным флагом и представили на обсуждение наши первоначальные результаты. Основной вопрос заключается в том, как обеспечить переносимость.

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

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

Чтобы оценить производительность операций с подгруппами, в реальном приложении мы интегрировали экспериментальную поддержку подгрупп в MediaPipe и протестировали ее с прототипом Chrome для операций с подгруппами.

Мы использовали подгруппы в ядрах графического процессора на этапе предварительного заполнения большой языковой модели, поэтому я сообщаю об ускорении только на этапе предварительного заполнения. На графическом процессоре Intel мы видим, что подгруппы работают в два с половиной раза быстрее, чем базовый уровень. Однако эти улучшения не одинаковы для разных графических процессоров.

Скриншот ускорения подгрупп в выводе MediaPipe LLM
Диаграмма 2. Подгруппы позволяют выполнять предварительное заполнение в 2,5 раза быстрее на графическом процессоре Intel Tiger Lake GT2 с экспериментальной поддержкой в ​​Chrome и Mediapipe.

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

Снимок экрана: ускорение подгруппы при умножении матрицы
Диаграмма 3. Применение подгрупп для умножения матриц может еще больше повысить производительность.

Оптимизация графического процессора сложна

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

Вы хотите минимизировать пропускную способность памяти, полностью используя вычислительные потоки графического процессора.

Шаблоны доступа к памяти также могут быть очень важны. Графические процессоры, как правило, работают намного лучше, когда вычислительные потоки обращаются к памяти по схеме, оптимальной для оборудования. Важно: на разном оборудовании графического процессора следует ожидать разных характеристик производительности. Возможно, вам придется запустить различные оптимизации в зависимости от графического процессора.

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

Вы можете видеть, что изменение памяти имеет значительное влияние; иногда это даже более эффективно, чем подгруппы. На графическом процессоре 6 swizzle обеспечивает ускорение в 12 раз, а подгруппы — в 13 раз. В совокупности они дают невероятное ускорение в 26 раз. Для других графических процессоров иногда объединение swizzle и подгрупп работает лучше, чем каждый из них по отдельности. А на других графических процессорах лучше всего работает исключительно использование swizzle.

Скриншот ускорения стратегий умножения матриц
График 4.

Настройка и оптимизация алгоритмов графического процессора для хорошей работы на любом оборудовании может потребовать большого опыта. Но, к счастью, ведется огромная талантливая работа над инфраструктурами библиотек более высокого уровня, такими как Mediapipe , Transformers.js , Apache TVM , ONNX Runtime Web и другие.

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

Вынос

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

Наша цель — максимально расширить возможности платформы, сохранив при этом лучшее от Интернета: охват, удобство использования и мобильность. И мы делаем это не в одиночку. Мы работаем в сотрудничестве с другими поставщиками браузеров в W3C и многими партнерами по разработке.

Мы надеемся, что при работе с WebAssembly и WebGPU вы помните следующее:

  • Выводы ИИ теперь доступны в Интернете на всех устройствах. Это дает преимущества работы на клиентских устройствах, такие как снижение стоимости сервера, низкая задержка и повышение конфиденциальности.
  • Хотя многие обсуждаемые функции актуальны в первую очередь для авторов фреймворка, ваши приложения могут получить от этого пользу без особых накладных расходов.
  • Веб-стандарты изменчивы и развиваются, и мы всегда ищем обратную связь. Поделитесь своими для WebAssembly и WebGPU .

Благодарности

Мы хотели бы поблагодарить команду Intel по веб-графике, которая сыграла важную роль в разработке WebGPU f16 и функций целочисленного скалярного произведения. Мы хотели бы поблагодарить других членов рабочих групп WebAssembly и WebGPU в W3C, включая других поставщиков браузеров.

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