Аннотация: Краткое введение в то, как работает XLA и его использование в Pytorch.
Эта статья опубликована в сообществе HUAWEI CLOUD.«Введение в принципы оптимизации XLA», Автор: Пионер 01.
Знакомство с XLA
Полное название XLA — Accelerated Linear Algebra, то есть ускоренная линейная алгебра. Как компилятор глубокого обучения, он долгое время разрабатывался как экспериментальная функция фреймворка Tensorflow.Прошло более двух-трех лет.С выпуском Tensorflow 2.X XLA окончательно превратился из экспериментальной функции. быть включенным по умолчанию. Кроме того, сообщество Pytorch также активно продвигает разработку XLA в рамках Pytorch.Теперь существует версия PyTorch/XLA TPU, которая поддерживается только на платформе Google TPU.
LLVM
Говоря о компиляторах, нельзя не упомянуть знаменитый LLVM. LLVM — это структура компилятора, написанная на C++, которая включает в себя ряд подмодулей, многоразовые инструменты компиляции.
Основными компонентами фреймворка LLVM являются:
1. Внешний интерфейс: отвечает за преобразование исходного кода в промежуточное представление.
2. Оптимизатор: отвечает за оптимизацию промежуточного кода.
3. Backend: модули, генерирующие исполняемый машинный код
LLVM обеспечивает одинаковое промежуточное представление LLVM IR для разных языков, поэтому, если нам нужно разработать новый язык, нам нужно только реализовать соответствующий интерфейсный модуль, если мы хотим поддерживать новое оборудование, нам нужно только реализовать соответствующий внутренний модуль, другие части можно использовать повторно.
XLA компиляция
XLA также разработан на основе LLVM framework.Фронтальный ввод — это Graph, а внешний интерфейс не конвертирует напрямую Graph в LLVM IR. Прежде всего, функции XLA в основном отражаются в двух аспектах:
1. Своевременная компиляция
2. Опережающая компиляция (Aheda-of-time)
Независимо от того, какую функцию он выполняет, он служит следующим целям:
1. Улучшить скорость выполнения кода
2. Оптимизируйте использование хранилища
Кроме того, у XLA есть мечта, которая есть у большинства компиляторов глубокого обучения: избавиться от ограничений вычислительных библиотек, автоматически генерировать код оператора и обеспечить хорошую переносимость на несколько аппаратных средств.
В качестве компилятора XLA отвечает за оптимизацию графа вычислений, определяемого внешним интерфейсом. Как показано на рисунке выше, процесс оптимизации XLA можно разделить на два аспекта: оптимизация, независимая от цели, и оптимизация, зависящая от цели. Между шагами оптимизации передается промежуточное представление вычислительного графа,HLO, а именно High Level Optimizer, XLA использует это промежуточное представление для представления оптимизируемого вычислительного графа, имеющего собственную грамматику и семантику, которые здесь не будут подробно описываться.
Преимущество XLA
-
Компилируйте графы подвычислений, чтобы сократить время выполнения кратковременных операций, тем самым устраняя накладные расходы во время выполнения; объединяйте конвейерные операции, чтобы уменьшить накладные расходы на память; и выполняйте специализированную оптимизацию для известных форм тензора для поддержки более агрессивного постоянного распространения.
-
Улучшите использование памяти: анализируйте и планируйте использование памяти, устраняя множество промежуточных буферов хранения.
-
Снижение зависимости от настраиваемых операций. Устраните необходимость в нескольких настраиваемых операциях, повысив производительность автоматически объединенных низкоуровневых операций до уровня производительности пользовательских операций, объединенных вручную.
-
Улучшенная переносимость: позволяет относительно легко писать новые серверные части для нового оборудования, и большинство программ будут работать без изменений при запуске на новом оборудовании. В отличие от разработки отдельных монолитных операций специально для нового оборудования, этот режим позволяет эффективно использовать эти операции без необходимости перепрограммирования.
Как работает XLA
Давайте сначала посмотрим, как XLA действует на вычислительном графе.Ниже приведен простой вычислительный граф
Здесь мы предполагаем, что XLA поддерживает только matmul и add. XLA использует методы оптимизации графа для поиска областей вычислительного графа, подходящих для JIT-компиляции.
XLA определяет эту область как кластер, как независимую единицу JIT-компиляции, которая отмечена атрибутом узла на графике расчета.
Затем другой метод оптимизации графа преобразует кластер в функциональный подграф TensorFlow. Используйте узел вызывающего абонента на исходном изображении, чтобы представить расположение функции на исходном изображении.
Наконец, вызывается метод оптимизации графа TensorFlow (BuildXlaOps) для преобразования узла Function в специальный узел Xla.
Когда TensorFlow запущен, при запуске в XlaCompile скомпилируйте подграф кластера Xla, а затем передайте скомпилированный исполняемый файл в XlaRun для запуска через XlaExecutableClosure.
Затем выделяем поток GPU и видеопамять в соответствии с виртуальными инструкциями, а затем IrEmitter преобразует HLO Graph в ядро GPU, представленное промежуточным выражением компилятора LLVM IR. Наконец, nvPTX (виртуальная базовая форма выражения инструкции, определенная Nvidia) генерируется LLVM, а затем NVCC генерирует исполняемый код CuBin.
АОТ и JIT
JIT, динамическая (точно в срок) компиляция, компилируется во время работы; AOT относится к предварительной компиляции. Основное различие между этими двумя методами компиляции заключается в том, следует ли компилировать во время выполнения. Для моделей обучения ИИ режим AOT имеет больше преимуществ в производительности. Конкретный процесс выглядит следующим образом:
Для большинства моделей ИИ график не сильно меняется в процессе обучения, поэтому использование режима AOT, когда обучение пропускается в процессе выполнения, может значительно повысить скорость обучения.
Pytorch/XLA
Создание тензоров XLA: PyTorch/XLA добавляет в PyTorch новый тип устройства xla. Этот тип устройства работает как обычный тип устройства PyTorch. Например, вот как создать и распечатать тензор XLA:
Этот код должен выглядеть знакомо. PyTorch/XLA использует тот же интерфейс, что и обычный PyTorch, но с некоторыми дополнениями. Импортируйте torch_xla для инициализации PyTorch/XLA, xm.xla_device() возвращает текущее устройство XLA. Это может быть ЦП или ГП, в зависимости от вашей среды.
Тензоры XLA — это тензоры PyTorch: операции PyTorch могут выполняться на тензорах XLA, точно так же, как тензоры CPU или CUDA. Например, можно добавить тензоры XLA:
Запуск моделей на устройствах XLA: для создания новой сети PyTorch или преобразования существующей сети для работы на устройстве XLA требуется всего несколько строк кода, специфичного для XLA, и в настоящее время поддерживается только режим JIT. На картинке фрагмент кода при запуске на одном устройстве XLA в официальной версии.
Этот код показывает, как легко переключать модели для работы на XLA. Определения моделей, загрузчики данных, оптимизаторы и циклы обучения работают на любом устройстве. Единственный специальный код — это несколько строк кода для получения устройства XLA и отметки шага. Потому что прогоны тензора XLA ленивы. Таким образом, в графе записываются только операции, и пока результат не понадобится, вызов xm.mark_step() будет выполнять свой текущий граф, чтобы получить результат выполнения и обновить параметры модели.
Нажмите «Подписаться», чтобы впервые узнать о новых технологиях HUAWEI CLOUD~