Организация параллельных вычислений Multi-Head Self-Attention

Мы продолжаем наше уверенное движение вперед по пути познания и строительства библиотеки для создания моделей машинного обучения в среде MQL5. В данном разделе мы планируем завершить работу еще над одним классом нейронных слоев CNeuronMHAttention, в котором реализован алгоритм Multi-Head Self-Attention. В предыдущих разделах мы уже полностью реализовали алгоритм стандартными средствами MQL5. Сейчас дополним его функционал возможностью использования технологии OpenCL для организации процесса вычислений в многопоточном режиме с задействованием ресурсов GPU.

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

Как обычно, начинаем с создания программы OpenCL. И сразу в голове всплывает вопрос: а надо ли нам создавать новую программу? Для чего нам создавать новые кернелы? Ответ, конечно, очевиден: для реализации функционала. Но давайте вспомним, что наш класс мы наследовали от похожего класса реализации алгоритма Self-Attention. Мы не раз говорили о преемственности указанных алгоритмов. Можем ли мы использовать созданные ранее кернелы для реализации процессов в этом классе?

На самом деле, учитывая схожесть процессов, нам было бы выгоднее использовать одни кернелы для обоих реализаций. Во-первых, это сокращает количество OpenCL-объектов, число которых в системе тоже не безгранично. Во-вторых, всегда удобнее поддерживать и оптимизировать один объект, чем копировать общие блоки между несколькими схожими объектами. Будь-то кернелы или классы.

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

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

Первое, что мы уже сделали в самом начале при создании класса, так это полное использование конкатенированных буферов данных. То есть все наши буферы данных содержат данные сразу всех голов внимания. Передавая данные в память контекста, мы передаем данные всех голов внимания. А значит, на стороне OpenCL мы можем параллельно работать со всеми головами внимания. Нам лишь надо правильно определить смещение в буфере данных до нужных значений. И это те изменения, которые мы должны внести в кернел.

Чтобы определить это смещение, нам необходимо понимание общего количества используемых голов внимания и порядковый номер рабочей головы внимания. Если общее количество мы можем передать в параметрах, то порядковый номер текущей — нет. Для организации передачи таких данных нам потребовалось бы создавать цикл с последовательным вызовом кернела для каждой головы внимания, а мы пытаемся уйти от этого.

Давайте вспомним, как организована функция постановки кернела в очередь выполнения. У функции CLExecute есть параметр work_dim, отвечающий за размерность пространства задач. Также функция в параметрах принимает динамический массив global_work_size[], в котором указывается общее количество выполняемых задач в каждом измерении.

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

Что ж, решение найдено, и мы можем приступать к реализации. Но тут есть еще один вопрос: создавать новый кернел или нет. Все говорит за изменение ранее созданного. Однако в этом случае после завершения работы нам придется сделать шаг назад и скорректировать методы класса CNeuronAttention. Иначе получим критическую ошибку при попытке запуска кернела.

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

Приступаем к работе. Посмотрим на внесенные в кернел прямого прохода изменения.

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

Тут же мы определяем смещение до начала анализируемого вектора в тензоре запросов и матрице коэффициентов зависимости.

Как вы понимаете, от предыдущей версии кернел отличается наличием второго измерения, учитывающего головы внимания. Соответственно, расчет смещения также посчитан с учетом многоголового внимания.

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

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

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

Во вложенном цикле с числом итераций равным размеру вектора ключей мы выполним операцию умножения вектора запросов на вектор ключей.

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

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

Для выполнения этой операции создадим еще один цикл с количеством итераций равным количеству элементов в последовательности.

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

Вначале определим смещение в тензоре результатов до начала вектора анализируемого элемента.

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

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

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

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

Проведем аналогичную работу с кернелами обратного прохода. Как вы помните, в блоке Self-Attention, в отличие от реализации других нейронных слоев, распределение градиента ошибки через внутреннее пространство скрытого нейронного слоя мы организовали двумя последовательными кернелами. И нам необходимо «перевести на рельсы» многоголового внимания оба кернела. Но давайте смотреть по порядку.

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

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

Алгоритм кернела мы не меняем. Как и при реализации алгоритма Self-Attention, логически кернел можно разбить на два блока.

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

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

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

Здесь мы также внесли изменения только в части определения смещения до анализируемых элементов в буферах данных.

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

Для выполнения операции умножения организуем вложенный цикл.

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

Результат выполнения операций записываем в соответствующие элементы тензора градиентов ошибки.

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

Переходим ко второму кернелу алгоритма обратного распространения ошибки AttentionCalcHiddenGradient. В нем нам предстоит распределить градиент ошибки от матрицы коэффициентов зависимости до буферов внутренних нейронных слоев m_cQuerys и m_cKeys.

С точки зрения математики, операция не сложная. Градиент ошибки на уровне матрицы коэффициентов зависимости мы уже определили в предыдущем кернеле. Теперь нам нужно умножить матрицу коэффициентов зависимости на противоположный тензор.

Как и в предыдущем кернеле, шапка и параметры кернела абсолютно не изменились. Здесь мы видим тот же набор буферов и параметров.

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

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

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

Завершаем работу с кернелами программы OpenCL. Итак, мы лишь сделали, так сказать, «косметические» правки в кернелах алгоритма Self-Attention для их перевода в область многоголового внимания.

Теперь нам предстоит дополнить основную программу функционалом вызова данных кернелов из методов как класса CNeuronAttention, так и класса CNeuronMHAttention. Обычно мы начинаем эту работу с создания констант для работы с кернелами. Но в данном случае константы уже созданы.

Далее мы создавали кернелы в контексте OpenCL. Но сейчас мы не создавали новые кернелы. А те, которые мы немного скорректировали, уже объявлены в теле основной программы. Поэтому мы пропускаем и этот шаг.

Переходим к внесению изменений непосредственно в методы классов. Чтобы новые кернелы работали в классе CNeuronAttention, мы добавляем второй элемент в массивы смещения и пространства задач. Для смещения укажем 0 в обоих измерениях. Для пространства задач первое значение оставляем без изменения, во второй элемент массива внесем 1 (используется одна голова внимания). Также при постановке кернела в очередь выполнения укажем двухмерность пространства задач.

После этого мы можем полноценно использовать обновленный кернел прямого прохода.

Такие несложные манипуляции делаем для вызова всех трех кернелов в методах класса CNeuronAttention.

На этом мы восстановили работоспособность методов класса CNeuronAttention, в котором реализован алгоритм Self-Attention. На стороне основной программы изменений тоже немного.

Переходим к работе над нашим классом CNeuronMHAttention с реализацией алгоритма Multi-Head Self-Attention. Как обычно, начнем работу с метода прямого прохода. Прежде чем поставить кернел в очередь выполнения операций, необходимо провести подготовительную работу. В первую очередь проверяем наличие необходимых буферов в памяти контекста OpenCL.

После проверки всех необходимых буферов мы передаем указатели на буферы в параметры кернела. Туда же мы передаем константы, необходимые для работы кернела.

Обратите внимание, что предавая параметры кернелу, мы дважды указали переменную m_iKeysSize, в которой содержится размер вектора ключей одного элемента последовательности. Мы указали его как для параметра размера вектора ключей, так и для параметра размера вектора значений. Два параметра в кернеле — вынужденная мера. При использовании одной головы внимания для размера вектора значений мы должны были бы указать размер вектора исходных данных. Это требование алгоритма Self-Attention. Но при использовании технологии многоголового внимания матрица W0 позволяет нам использовать различные варианты размера вектора значений.

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

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

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

После проверки буферов мы передаем указатели на них и необходимые константы в параметры кернела.

Осуществляем постановку кернела в очередь выполнения операций. Как и в случае прямого прохода, мы создаем двухмерное пространство задач. В первом измерении мы указываем количество анализируемых элементов в последовательности, а во втором — количество голов внимания.

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

Передадим указатели на буферы данных в параметры второго кернела. Туда же добавим необходимые константы.

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

На этом завершается наша работа над реализацией алгоритма Multi-Head Self-Attention как в части организации многопоточных вычислений, так и в целом. Мы с вами реализовали весь функционал класса CNeuronMHAttention. Теперь можем перейти к всестороннему тестированию его работы на обучающей и тестовой выборках.