Руководство по совместимости XLA:GPU AOT

В этом руководстве определяется контракт совместимости между компилятором XLA:GPU, создающим сериализованный исполняемый файл, и средой выполнения, загружающей его, а также рассказывается, как внести изменения, не нарушая этот контракт. О том, что такое AOT-компиляция и как структурируется и загружается артефакт, см. раздел «Предварительная компиляция XLA:GPU» .

Данный контракт распространяется исключительно на бэкенд XLA:GPU ( GpuExecutableProto ). Исполняемые файлы AOT для CPU и TPU используют другие механизмы сериализации и не входят в сферу действия контракта.

Совместимость окон {#compatibility}

XLA:GPU предоставляет два окна совместимости:

  • Обратная совместимость (6 месяцев): среда выполнения, содержащая ваши изменения, должна загружать и выполнять с идентичной семантикой все артефакты, созданные компиляторами за предыдущие 6 месяцев. Это позволяет загружать модели, скомпилированные некоторое время назад в обновленной среде выполнения.
  • Обратная совместимость (2 недели): Компилятор, содержащий ваши изменения, не должен генерировать никаких конструкций сериализации, которые среда выполнения, выпущенная не более 2 недель назад, не сможет десериализовать или неправильно интерпретирует. Это позволяет загружать модели, скомпилированные с помощью новейшего компилятора, на более старых, долго работающих серверах вывода во время поэтапного развертывания.

Каждое протокольное сообщение и поле, транзитивно достижимое из GpuExecutableProto , является частью этого контракта.

Оценка сдачи: безопасна ли моя сдача? {#change_triage}

Чтобы определить, требуется ли поэтапное внедрение ваших изменений, следуйте этой последовательности действий:

  1. Затрагивает ли ваше изменение что-либо из перечисленного ниже?

    • прототип, доступный из GpuExecutableProto
    • ToProto / FromProto
    • Опкоды HLO или прототипы HLO, которые могут оказаться в исполняемом файле графического процессора.
    • имена или семантика зарегистрированных символов ядра, пользовательских вызовов или обработчиков FFI
    • Функция загрузки thunk передает или помечает данные, считываемые средой выполнения.
    • семантика выполнения существующих thunk-функций

    Нет: изменение является внутренним для компилятора (например, оптимизации, логика объединения, переписывание генерации кода) или представляет собой чисто внутреннюю рефакторизацию во время выполнения с неизменным поведением для существующих thunk-функций. Его можно безопасно внедрить напрямую.

    Да: переходите к шагу 2.

  2. Изменяет ли это изменение способ интерпретации существующих артефактов средой выполнения (выполнение существующих thunk-функций, проходы загрузки, значения по умолчанию для флагов чтения во время выполнения)?

    Да: Разрешено только в том случае, если это корректно для каждого артефакта, выпущенного за последние 6 месяцев; если семантическое значение меняется, вместо этого следует ввести новое поле или флаг.

    Нет: переходите к шагу 3.

  3. Удаляет ли это изменение существующий тип thunk, поле proto или резервную ветвь десериализации?

    Да: Нарушает обратную совместимость, если выпуск компилятором не прекратился как минимум 6 месяцев назад. Сохраняйте резервный вариант десериализации до истечения 6-месячного периода и помечайте удаленные теги и имена протоколов как reserved .

    Нет: переходите к шагу 4.

  4. В результате изменения в сериализованный артефакт добавляется новый тип thunk-кода, новое поле proto, необходимое для корректности, или новый HLO-опкод?

    Да: Немедленный выпуск нарушает обратную совместимость . Следуйте поэтапному развертыванию .

    Нет: переходите к шагу 5.

  5. В результате изменения происходит переименование или удаление зарегистрированного символа ядра, пользовательского вызова или имени обработчика FFI?

    Да: Нарушает обратную совместимость . Сохраняйте обратно совместимые регистрации под устаревшими именами как минимум 6 месяцев.

    Нет: переходите к шагу 6.

  6. Добавляет ли это изменение необязательное поле протокола, которое более старые среды выполнения могут безопасно игнорировать без ухудшения корректности?

    Да: Безопасно для посадки, при условии, что незаданные поля прототипа сохраняют устаревшую семантику среды выполнения.

    Нет: Необходимо переработать изменения, чтобы сохранить обратную и прямую совместимость.

Безопасные изменения против критических изменений

Изменять Вердикт Влияние на совместимость Безопасная процедура
Оптимизация на уровне компилятора, слияние или генерация кода (одна и та же схема) Безопасный Никто Приземлитесь прямо.
Рефакторинг среды выполнения в оперативной памяти с сохранением неизменного поведения существующих промежуточных элементов (кэшей, управления буферами). Безопасный Никто Внедрите модульные тесты напрямую.
Рефакторинг ToProto / FromProto с использованием идентичных сериализованных байтов Безопасный Никто Приземление без пересадок; включить тесты на перелет туда и обратно.
Добавить новое поле протокола, незаданное значение которого сохраняет устаревшее поведение. Безопасный Никто Убедитесь, что FromProto корректно обрабатывает незаданное состояние.
Добавлено новое поле протокола, необходимое для корректной работы во время выполнения. Срочный Обратная совместимость Поэтапное внедрение (в первую очередь поддержка во время выполнения).
Вывести новый тип thunk, вариант oneof или значение перечисления. Срочный Обратная совместимость Поэтапное развертывание (сначала добавить thunk-функцию во время выполнения).
Вывести новый код операции HLO в сериализованный модуль Срочный Обратная совместимость Поэтапное внедрение; эмиссия через шлюзы продолжается до истечения временного окна выполнения.
Удалите тип thunk, поле proto или ветку FromProto Срочный Обратная совместимость Подождите 6 месяцев после прекращения выбросов; сохраните номерные знаки.
Изменить существующий номер поля или тип данных протокола. Срочный Прямая и обратная совместимость Никогда не изменяйте существующие теги полей; присваивайте новый тег.
Изменить семантическое значение существующего поля. Срочный Прямая и обратная совместимость Добавьте новое поле, отражающее новое поведение.
Переименовать или удалить зарегистрированный символ ядра или обработчик FFI. Срочный Обратная совместимость Сохраните регистрацию псевдонима под прежним именем на 6 месяцев.
Изменить параметр thunk во время загрузки или переключить флаг чтения во время выполнения по умолчанию. Зависит от Обратная совместимость Разрешается только при условии корректности всех артефактов, выпущенных за последние 6 месяцев.

Поэтапное внедрение критических изменений {#staged_rollout}

При внедрении новой конструкции сериализации, её следует реализовывать в три этапа, чтобы обеспечить соблюдение обоих окон совместимости:

  1. Поддержка среды выполнения. Добавьте логику десериализации и выполнения в среду выполнения. Если в том же запросе на слияние включена генерация компилятором, оставьте генерацию отключенной по умолчанию за экспериментальным флагом. Затем подождите не менее 2 недель после слияния этого изменения, чтобы обновленная среда выполнения распространилась по всем средам развертывания.
  2. Генерация кода компилятором. Включите генерацию кода компилятором по умолчанию (или измените значение по умолчанию флага функции). Тогда сохраните резервный вариант десериализации в среде выполнения как минимум на 6 месяцев , считая с этого этапа (когда компилятор перестал генерировать устаревшую форму), а не с этапа 1.
  3. Очистка. Удалите устаревший путь десериализации и пометьте устаревшие номера тегов протоколов и имена полей как reserved .

Фазы 1 и 2 могут представлять собой два запроса на слияние (второй отправляется через 2 недели после первого) или один запрос на слияние, в котором генерация по умолчанию отключена с помощью флага и изменена в последующем запросе.

Комментарии и примеры удаления

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

// Backward-compatibility fallback for legacy AOT-compiled kernels without
// an explicit CollectiveKernelSpec.
// Can be removed in <month year> (6 months backward compatibility window).

Примеры публичного развертывания OpenXLA:

  • PR #49046 : Введено выделение выделенного потока данных в среде выполнения, при этом передача данных по умолчанию остается отключенной на протяжении всего периода внедрения.
  • PR #46865 : Добавлена ​​поддержка thunk-функций во время выполнения для CollectiveReduce , при этом генерация компилятором защищена экспериментальным флагом.

Разработка кода времени выполнения для десериализованных исполняемых файлов

При реализации функций отладки во время выполнения и процедур десериализации:

  • Не допускаются типы, зависящие только от компилятора: функции обработки данных во время выполнения не должны зависеть от структур данных компилятора, таких как HloInstruction* или BufferAssignment .
  • Проверяйте границы десериализованных данных: всегда проверяйте индексы выделения буфера, смещения срезов и идентификаторы потоков на соответствие границам выделения перед индексированием.
  • Не дублируйте полезную нагрузку: не храните одну и ту же большую полезную нагрузку (например, константы) более одного раза в артефакте; это многократно увеличивает объем памяти, используемой для загрузки.
  • Защититесь от выделения памяти при инициализации: среды выполнения, предоставляющие сервер, работают в условиях строгих ограничений по объему памяти. Избегайте безусловного выделения памяти в куче во время инициализации с помощью FromProto или thunk.

Тестирование

Все изменения в сериализации должны быть охвачены автоматизированными тестами:

  • Модульные тесты с обратной связью: проверка точности сериализации и десериализации с помощью тестов ToProto и FromProto (см. ThunkProtoDeserializationTest ).
  • «Золотые» AOT-тесты: сквозные тесты совместимости, доступные в xla/tests/aot_compatibility/gpu (например, collective_ops_aot_test.cc ), выполняются с использованием сериализованных «золотых» исполняемых файлов в executables/<target>/v<N>/ .
  • Тестирование всех исторических версий: по умолчанию эталонные тесты проверяются на соответствие граничным версиям (самой старой и второй по новизне). Установите параметр XLA_AOT_TEST_ALL_VERSIONS=1 в вашей тестовой среде, чтобы выполнить проверку на соответствие всем историческим эталонным версиям.