Для получения дополнительных галантерейных товаров, пожалуйста, обратите внимание на публичный аккаунт WeChat «AI Frontline» (ID: ai-front)
TVM — это метод автоматической генерации кода с глубоким обучением, предложенный Ченом Тяньци, аспирантом Вашингтонского университета, и др. Эта технология может автоматически генерировать пригодный для развертывания оптимизированный код для большинства вычислительных аппаратных средств, и ее производительность можно сравнить с оптимизированная вычислительная библиотека, предоставляемая лучшими на данный момент поставщиками, и может поддерживать новые специализированные серверные части ускорителей.
Нейронный машинный перевод (NMT) — это сквозной метод автоматического перевода, который может преодолеть недостатки традиционных систем перевода фраз. В последнее время Alibaba Group пытается внедрить услуги NMT в глобальные бизнес-сценарии электронной коммерции.
В настоящее время мы пытаемся разработать Transformer[1] в качестве основного компонента нашей системы NMT, поскольку он более подходит для эффективного автономного обучения, чем классические модели на основе RNN/LSTM, с равной (или даже более высокой) точностью. В то время как Transformer разрушает зависимости между временными шагами и поэтому очень удобен для этапа автономного обучения, он не так эффективен для онлайн-вывода. В нашей производственной среде первоначальная версия вывода Transformer была примерно в 1,5–2 раза медленнее, чем версия LSTM. Чтобы улучшить производительность логического вывода, мы пробовали различные оптимизации, такие как объединение на уровне графа, движение узлов, не зависящее от циклов [3]. Мы обнаружили, что пакетный матмул является критической точкой производительности в Transformer, а текущая реализация cuBLAS плохо оптимизирована.
Наши практические результаты показывают, что ядра, сгенерированные TVM (с оптимизированным расписанием), могут выполнять пакетные вычисления matmul как минимум в 13 раз быстрее, а при включении слияния операторов еще больше ускоряются.
В Transformer пакетный матмул широко используется при вычислении многоголового внимания. Благодаря пакетному матмулированию несколько головок на уровне внимания могут работать параллельно, что помогает повысить вычислительную эффективность оборудования.
Наше углубленное профилирование модели Transformer на этапе логического вывода показывает, что на пакетные вычисления matmul приходится около 30% времени выполнения ядра графического процессора. Мы выполнили первопринципный анализ пакетного ядра matmul cuBLAS с помощью nvprof[2], и результаты показывают, что текущая реализация очень плохая, а также мы обнаружили некоторые интересные явления.
В общем, пакетное вычисление matmul относится к выполнению матричного умножения пакета матриц. где пакеты считаются «однородными», т. е. все экземпляры имеют одинаковые размеры (M, N, K), начальные размеры (lda, ldb, ldc) и скорости перехода для соответствующих матриц A, B и C.
Пакетное вычисление матмул можно описать более конкретно следующим образом:
В задаче языкового перевода форма пакетного вычисления matmul значительно меньше, чем обычное вычисление matmul при других рабочих нагрузках. Форма преобразователя связана с длиной входного предложения и длиной шагов декодера. Обычно меньше 30.
Если для вывода задан размер пакета, размер пакета является фиксированным значением. Например, если размер пакета установлен равным 16, а размер луча равен 4, размер пакета равен 16*4*#головки (количество головок в мультиголовке, обычно 8), а формы матриц M, K и N находятся в [1, максимальная длина декодирования] или [1, максимальная длина кодирования].
Во-первых, мы проводим теоретический анализ FLOP на пакетных ядрах Matmul. Результаты весьма интересны: все пакетные матмулы очень ограничены в вычислительном отношении (менее 1 TFLOPS).
Затем мы использовали программное обеспечение nvprof для анализа cuBLAS пакетного матмула различной формы. В таблице ниже показаны некоторые метрики, полученные на графическом процессоре NVIDIA M40 с использованием CUDA 8.0.
Даже при различных формах (изменение M, N, K) все вызовы Maxwell_sgembatched_128x128_raggedMn_tn выполняют одинаковое количество FLOP, а фактическое значение намного больше теоретического. Можно сделать вывод, что все эти различные формы могут быть заполнены до определенной формы. Во всех этих формах, даже в лучшем случае, теоретические FLOP составляют всего 2,74% от фактически выполненных FLOP, поэтому большая часть вычислений довольно избыточна. Точно так же такое же явление происходит с Maxwell_sgembatched_64x64_raggedMn_tn, вызывающим другое ядро cuBLAS.
Очевидно, что пакетная реализация cuBLAS неэффективна. Поэтому мы используем TVM для создания эффективных пакетных ядер matmul для рабочих нагрузок NMT.
В TVM обычно пакетное вычисление matmul может быть объявлено как:
После объявления пакетных вычислений matmul нам необходимо тщательно спроектировать расписание, чтобы полностью использовать его потенциал производительности.
Мы объединяем внешние измерения пакетного матмула, а именно BB и FF операционного измерения, обычно называемого «пакетным» измерением в пакетных вычислениях матмула. Затем мы разделяем внешнее и внутреннее измерения на коэффициент (number_thread*vthread).
Stridged режим не требуется в пакетном matmul, поэтому количество виртуальных потоков (vthread_y и vthread_x) установлено равным 1.
На рисунке ниже показан результат, полученный на графическом процессоре NVIDIA M40 с использованием CUDA8.0.
Исходя из прошлого опыта, способ найти наилучшую комбинацию num_thread_y и num_thread_x — это поиск методом грубой силы. Путем перебора можно найти наилучшую комбинацию текущей формы (shape), и наилучшая комбинация, которую мы получаем в текущем расчете, это num_thread_y=8, num_thread_x=32.
Текущие вызовы библиотеки cuBLAS «черного ящика» часто используются в качестве границ для общих стратегий оптимизации «op fusion». Однако при сгенерированном эффективном пакетном ядре matmul эта граница слияния может быть легко нарушена, и для дальнейшего повышения производительности можно объединять не только поэлементные операции.
Как видно из графика вычислений, за пакетным матмул всегда следует широковещательное добавление или транспонирование. Накладные расходы на запуск ядра и время доступа к избыточной памяти можно сократить, объединив операции «добавить» или «транспонировать» с пакетной программой matmul.
Пакетные матмульные и широковещательно-аддитивные вычисления слияния могут быть объявлены следующим образом:
Пакетные вычисления matmul и transpose fusion могут быть объявлены следующим образом:
Мы выбираем форму [партия = 64, головки = 8, M = 1, N = 17, K = 128], чтобы проверить производительность сгенерированного выше кода. Мы устанавливаем длину последовательности равной 17, поскольку это значение является средней длиной входных данных в нашем производственном сценарии.
Результаты теста следующие:
tf-r1.4 BatchMatmul: 513,9 мкс
tf-r1.4 BatchMatmul+Transpose (отдельно): 541,9 мкс
TVM BatchMatmul: 37,62 мкс
TVM BatchMatmul+Transpose (объединенный): 38,39 мкс
Оптимизация слияния ядра дополнительно увеличивает скорость в 1,7 раза.
В нашей рабочей нагрузке входная форма для пакета matmul ограничена, и ее легко предварительно перечислить. Используя эти предопределенные формы, мы можем заранее генерировать высокооптимизированные ядра CUDA (вычисления с фиксированной формой обеспечивают наилучший потенциал оптимизации). В то же время мы также будем генерировать общие пакетные ядра matmul для большинства форм, предоставляя резервный механизм для форм, которые не генерируют ядра заранее.
Мы интегрировали сгенерированные эффективные ядра для конкретных форм и резервные механизмы в структуру TensorFlow. Мы разработали операции слияния, такие как BatchMatMulTranspose или BatchMatMulAdd, чтобы использовать API среды выполнения TVM для запуска ядра определенного поколения для определенной входной формы или вызова резервного ядра. Выполняя проход оптимизации графа, исходный пакетный шаблон matmul+add/transpose может быть автоматически заменен объединенной операцией. Между тем, включив более агрессивный проход оптимизации графики, мы пытаемся использовать TVM для создания более эффективных ядер слияния для режимов работы с длинным хвостом, чтобы еще больше повысить сквозную производительность.
Внутри Alibaba мы обнаружили, что TVM — очень эффективный инструмент, с помощью которого мы можем разрабатывать высокопроизводительные ядра GPU для внутренних нужд.
В этой статье модель NMT Transformer используется в качестве примера для подробного описания стратегии оптимизации TVM. Во-первых, ключевые проблемы модели Transformer идентифицируются с помощью анализа первых принципов. Затем используйте TVM для создания высокооптимизированных ядер CUDA для замены версии cuBLAS (на данный момент достигнуто 13-кратное ускорение). Затем мы используем механизм слияния ядра TVM для слияния операций до/после пакетной обработки matmul, что приводит к дальнейшему повышению производительности (дополнительное повышение производительности в 1,7 раза). На этой основе мы разработали канал оптимизации графики, который автоматически заменяет исходный режим вычислений ядром слияния TVM, обеспечивая прозрачность процесса оптимизации для конечного пользователя. Как поставщик инфраструктуры ИИ, мы считаем, что прозрачность важна для продвижения применения стратегий оптимизации.
Наконец, все эти оптимизации слабо связаны в TensorFlow, демонстрируя один из возможных способов интеграции TVM с различными средами глубокого обучения. Кроме того, в настоящее время мы работаем над интеграцией TVM в качестве бэкенда для генерации кода TensorFlow и надеемся поделиться с сообществом другими результатами в будущем.
Реализация TVM объединенного пакетного кода matmul + код проекта транспонированных вычислений
https://github.com/Orion34C/tvm-batch-matmul-example/blob/master/tvm_batch_matmul_transpose_m1_kX.py
Представление команды PAI на китайском языке
https://zhuanlan.zhihu.com/p/33513874
[1] Attention is All You Need
https://arxiv.org/pdf/1706.03762.pdf
[2] nvprof is Your Handy Universal GPU Profiler
https://devblogs.nvidia.com/cuda-pro-tip-nvprof-your-handy-universal-gpu-profiler/
[3] Add Loop Invariant Node Motion Optimization in GraphOptimizer
https://github.com/tensorflow/tensorflow/pull/16306
Оригинальная ссылка:
http://www.tvmlang.org/2018/03/23/nmt-transformer-optimize.html
Для получения дополнительных галантерейных товаров, пожалуйста, обратите внимание на публичный аккаунт WeChat «AI Frontline» (ID: ai-front)