Глубокое погружение в SparseCore для больших встраиваемых моделей (LEM)

SparseCore — это специализированный тайловый процессор, разработанный для высокопроизводительного ускорения рабочих нагрузок, связанных с нерегулярным доступом к разреженной памяти и вычислениями, особенно с большими наборами данных, хранящимися в памяти с высокой пропускной способностью (HBM). Он отлично справляется с такими задачами, как встраивание поиска, но его возможности распространяются и на ускорение множества других динамических и разреженных рабочих нагрузок.

1. Введение в SparseCore

Основные архитектурные особенности:

  • Плиточная архитектура : включает в себя несколько вычислительных плиток (каждая плитка представляет собой законченный блок потока данных со своей локальной памятью и блоком обработки), что позволяет выполнять параллельную обработку.
  • Динамическое выполнение : встроенная поддержка зависящего от данных потока управления и доступа к памяти, что крайне важно для разреженных данных.
  • Векторная обработка : использует небольшие векторные задачи (8-элементные или 16-элементные, в зависимости от версии оборудования) для эффективных вычислений.
  • Централизованное управление : единый секвенсор SparseCore координирует задачи на всех плитках, обеспечивая синхронизированные операции.
  • Поддержка суммирования данных : включает специализированные кросс-полосные операции, полезные для таких задач, как сортировка, фильтрация и суммирование префиксов.
  • Иерархия памяти : стратегически использует HBM для хранения больших наборов данных и локальную сверхоперативную память (SPMEM) для размещения часто используемых данных, что значительно сокращает задержку HBM.

Краткие характеристики:

Атрибут ТПУ v4 ТПУ v5p Триллиум
SparseCores/Chip 4 4 2
Плитка/Разреженное ядро 16 16 16
Ширина SIMD 8 8 8 (Ф32) 16 (БФ16)
Емкость HBM 32 ГиБ 96 ГиБ 32 ГиБ

2. Предварительная обработка хоста SparseCore

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

  • Преобразование данных :
    • Примените необходимые преобразования к необработанным входным данным.
    • Управление преобразованиями идентификаторов, что особенно важно при работе с функциями или стекированием таблиц.
    • Преобразуйте входные данные в разреженный формат координат (COO), подробно описанный в следующем разделе.
    • Разделите данные для эффективного распределения по различным разреженным ядрам, доступным на чипе.
  • Проверка предела :
    • Убедитесь, что характеристики входных данных (например, количество идентификаторов) соответствуют предопределенным эксплуатационным ограничениям SparseCore, таким как max_ids_per_partition и max_unique_ids_per_partition .
    • Если входные данные превышают эти ограничения, ваш уровень предварительной обработки может попытаться сегментировать данные на более мелкие мини-пакеты, которые соответствуют ограничениям.
  • Передача данных :
    • Эффективно копируйте обработанные и проверенные данные в высокоскоростную память TPU (HBM), подготавливая ее к выполнению SparseCore.

Понимание укладки таблиц:

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

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

Основным преимуществом стекирования таблиц является создание большего эффективного размера пакета для операций над этими стекированными таблицами. Это снижает вычислительные затраты и может эффективно скрывать задержки межкристального обмена данными (ICI). Для оптимальной производительности рекомендуется использовать умеренное количество стекированных таблиц (обычно от 5 до 100).

3. Преобразование в тензоры COO

Перед обработкой данных в SparseCore они обычно преобразуются в формат разреженного тензора координат (COO). Формат COO — это способ эффективного представления разреженных матриц, обычно использующий три массива:

  • row_ids : массив, содержащий индексы строк для каждого ненулевого элемента. В контексте пакетной обработки это часто соответствует размерности пакета.
  • col_ids : массив, содержащий индексы столбцов для каждого ненулевого элемента. Для вложений это часто значения признаков или идентификаторов.
  • values (необязательно): массив, содержащий фактические значения ненулевых элементов в соответствующих координатах ( row , col ). При вычислении пределов (обсуждается далее), связанных с количеством идентификаторов, эти значения (приросты) часто не учитываются.

Наглядный пример:

Рассмотрим входную разреженную матрицу, представляющую пакеты идентификаторов:

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

После преобразования в формат COO (и, возможно, после дедупликации идентификаторов в пределах одной выборки):

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

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

4. SparsecoreConfig: высокоуровневый API

API-интерфейсы для встраивания, специфичные для фреймворка:

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

  • disable_table_stacking: bool = False
    • Объяснение : Этот флаг управляет тем, будет ли автоматическое укладывание таблиц препятствовать фреймворку укладывать таблицы, что может привести к снижению производительности из-за увеличения накладных расходов и уменьшения возможности скрыть задержку межкристального соединения (ICI).
    • По умолчанию : False (подразумевается, что стекирование таблиц обычно включено по умолчанию, если фреймворк его поддерживает).
  • max_ids_per_chip_per_sample: int = 64
    • Пояснение : Этот параметр устанавливает глобальный верхний предел общего количества идентификаторов встраивания, которые один чип может обработать из одного образца во входном пакете, агрегированного по всем таблицам. Это механизм управления ресурсами на уровне чипа, до того как будут учтены более детальные ограничения по таблицам или разделам. Тонкая настройка этого значения обычно зависит от характеристик конкретной модели и общей производительности системы.
    • По умолчанию : 64 .
  • max_ids_per_table: Optional[Dict[str, int]] = None
    • Пояснение : Этот параметр определяет максимальное количество идентификаторов вложений (включая дубликаты), которые могут быть обработаны для каждой логической таблицы с учётом всех её разделов по всем SparseCore. Это ограничение шире, чем max_ids_per_partition . Если таблица T разделена на P разделов, это ограничение применяется к сумме идентификаторов, направленных во все P разделов. Оно часто связано с max_ids_per_partition_per_sample и общим размером пакета.
    • Настройка : обычно настраивается с помощью файла ограничений (например, с помощью флага xla_sparse_core_max_ids_file ), в котором определен параметр max_ids_per_partition . Эта концепция на уровне таблицы представляет собой метод установки ограничений на уровне разделов ( max_ids и max_uniques ).
    • По умолчанию : None (значение может быть выведено из ограничений для каждого раздела или других конфигураций, если не указано явно).
  • max_unique_ids_per_table: Optional[Dict[str, int]] = None
    • Пояснение : Аналогично max_ids_per_table , но этот параметр определяет максимальное количество уникальных идентификаторов для каждой логической таблицы. Это критически важный параметр для правильного определения размера буферов на устройстве, используемых при обработке уникальных идентификаторов и последующих векторных операциях.
    • Настройка : Также обычно определяется в файле ограничений или выводится из max_unique_ids_per_partition_per_sample .
    • По умолчанию : None .
  • allow_id_dropping: bool = False
    • Объяснение : Этот логический флаг управляет отбрасыванием идентификаторов, когда количество идентификаторов, встречающихся во входных данных (наблюдаемые пределы), превышает пределы, установленные во время компиляции (например, max_ids_per_partition ).
      • Если True : идентификаторы, которые могут привести к превышению лимитов, автоматически отбрасываются. Как правило, идентификаторы внутри раздела обрабатываются в порядке сортировки, и любой идентификатор, который может привести к превышению текущего счётчика лимита для назначенного мини-пакета, отбрасывается. Это позволяет программе продолжить выполнение, но может негативно повлиять на точность модели.
      • Если False : возникает ошибка, и процесс, вероятно, будет прерван, если наблюдаемые ограничения выйдут за пределы скомпилированных ограничений. Такой подход гарантирует обработку всех данных, но требует более консервативной настройки ограничений.
    • По умолчанию : False (вызывает ошибку при переполнении, а не молчаливое удаление данных).

  • initialize_tables_on_host: bool = True

    • Пояснение : Этот флаг определяет, инициализируются ли встроенные таблицы на центральном процессоре перед их последующей передачей в высокоскоростную память (HBM) TPU. Стандартная практика заключается в том, что таблицы инициализируются на хосте. Установка значения True соответствует этому соглашению. Если бы значение было установлено равным False , это означало бы механизм инициализации на устройстве, что может иметь различные последствия для производительности или особые требования к инициализации.

  • enable_fast_table_initialization: bool = False

    • Объяснение : Инициализирует таблицы непосредственно на TPU. Это может помочь сократить время запуска модели.

5. Конвейеризация для повышения производительности

Конвейеризация — это метод оптимизации производительности, позволяющий одновременно выполнять операции на тензорном ядре (TensorCore, TC) и разреженном ядре (SparseCore, SC). Перекрытие этих вычислений позволяет значительно повысить общую производительность.

  • Механизм : На стандартном этапе обучения, включающем разреженные поисковые запросы вложений (обрабатываемые SC) и плотные вычисления слоев (обрабатываемые TC), конвейеризация позволяет SC работать над своей частью шага i (например, прямой или обратный проход), в то время как TC параллельно обрабатывает другую часть того же шага i или даже части соседних шагов, например i-1 или i+1 .
  • Влияние на градиенты : SparseCore может работать с устаревшими градиентами. Например, градиенты, рассчитанные на этапе обратного распространения ошибки на шаге i могут быть не полностью обновлены и не видны SC до шага i+2 .
  • Компромисс между производительностью и численными данными : такое перекрывающееся выполнение может привести к существенному ускорению, потенциально до двукратного сокращения времени выполнения шага устройства. Однако незначительные изменения в числовых данных (embedding_weights), возникающие из-за использования устаревших градиентов, могут повлиять на сходимость модели или на конечную точность. Приемлемость этого компромисса сильно зависит от модели и часто требует эмпирической проверки.
  • Флаг управления : конвейеризацией можно управлять с помощью tf_xla_disable_full_embedding_pipelining . Установка этого флага в значение true отключает полную конвейеризацию (перекрывающиеся вычисления TensorCore и SparseCore), тогда как установка в false (или если семантика флага подразумевает включение при значении false) активирует её.

Концептуальный поток конвейеризации:

  • Без конвейеризации (упрощенный последовательный поток) :

    Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

  • С конвейеризацией (упрощенный перекрывающийся поток) :

    Time ->
Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1

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

6. Роль XLA

XLA (Accelerated Linear Algebra) — это предметно-ориентированный компилятор, который транслирует высокоуровневые вычислительные графы, как правило, из таких фреймворков, как TensorFlow, в высокооптимизированный машинный код, адаптированный для TPU. Это включает в себя генерацию инструкций для операций, предназначенных для SparseCore.

Ключевые функции в контексте SparseCore:

  • Компиляция разреженных операций : XLA отвечает за компиляцию встраиваемых операций поиска (таких как SparseDenseMatmulOp ) и других разреженных вычислений в низкоуровневые исполняемые программы SparseCore.
  • Интеграция ограничений : использует настроенные рабочие ограничения (например, max_ids_per_partition , max_unique_ids_per_partition , часто предоставляемые через файл ограничений, заданный такими флагами, как xla_sparse_core_max_ids_file ) для статического определения размеров и выделения буферов памяти на устройстве, особенно в SPMEM.
  • Целевые оптимизации : XLA выполняет набор оптимизаций, специально разработанных для архитектуры SparseCore. Они могут включать планирование инструкций, преобразование структуры памяти и объединение операций для максимального повышения эффективности.
  • Управление с помощью флагов : многие аспекты поведения SparseCore, параметры настройки и стратегии оптимизации раскрываются и контролируются с помощью флагов XLA (например, xla_sparse_core_estimate_max_ids для оценки пределов или xla_sc_detect_nan для отладки).

Статус открытого исходного кода:

В настоящее время реализация Sparsecore является внутренней и обслуживается с помощью libtpu.so .

Отчеты об ошибках и диагностика:

Ошибки компиляции, связанные с конфигурациями SparseCore или ограничениями ресурсов, часто проявляются в виде ошибок компиляции XLA:TPU . Эти сообщения об ошибках могут предоставить ценную информацию о таких проблемах, как слишком высокие ограничения для доступного SPMEM или использование неподдерживаемых конфигураций.

7. Как ограничения переносятся на таблицы в SparseCore

В SparseCore «лимиты» — это основополагающие параметры конфигурации, которые в первую очередь относятся к двум настройкам для каждого раздела каждой таблицы, которая сегментируется (распределяется) по доступным SparseCore:

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

Перевод в физическую табличную структуру и обработку:

  • Стратегия шардинга таблиц : Встраиваемые таблицы обычно шардируются по всем разреженным ядрам (SparseCore) в системе. Это означает, что каждое разреженное ядро (SparseCore) отвечает за отдельное подмножество словаря (строк) каждой таблицы. Идентификатор j обычно присваивается ядру SparseCore_k по формуле вида k = j % num_total_sparse_cores .
  • Определение «раздела» : В данном контексте «раздел» относится к определенному сегменту таблицы внедрения, для которого один SparseCore обрабатывает поиск.
  • Выделение буфера SPMEM : эти ограничения используются компилятором XLA для статического определения размера и выделения буферов в оперативной памяти устройства (SPMEM). Размеры буферов определяются таким образом, чтобы все необходимые данные, связанные с идентификаторами данного раздела (в пределах заданных значений max_ids и max_unique_ids ), могли быть загружены в SPMEM для обработки. Это особенно важно для неэлементных вычислений, таких как сокращение числа дублирующихся идентификаторов в разделе (например, при создании представления Compressed Sparse Row (CSR)), где весь релевантный набор данных для идентификаторов этого раздела должен быть легко доступен в быстрой памяти.

  • Компилированные пределы в сравнении с наблюдаемыми пределами :

    • Наблюдаемые ограничения : это фактическое количество идентификаторов, обнаруженных для каждого раздела во время выполнения, на основе обрабатываемых входных данных.
    • Если наблюдаемые пределы превышают скомпилированные пределы, это может привести к отбрасыванию идентификатора (если включено allow_id_dropping ) или ошибкам.

  • Вычисление пределов : процесс определения подходящих пределов включает в себя тщательный анализ распределения входных данных. Для любой заданной таблицы (назовём её T1 , которая сама может быть частью более крупной составной таблицы T ):

    1. Входной пакет (например, двумерный SparseTensor с размером [BatchSize, MaxSequenceLength] ) изначально распределяется между доступными SparseCore. Например, если TensorCore используется в паре с двумя SparseCore, каждый SparseCore может получить подпакет с размером [BatchSize/2, MaxSequenceLength] .
    2. Затем этот подпакет преобразуется в формат COO, в результате чего получаются row_ids и col_ids .
    3. Дублирующиеся идентификаторы в пределах одного образца (т. е. записи с одинаковыми row_id и col_id ) удаляются.
    4. Для каждого оставшегося уникального col_id (в пределах выборки) целевой SparseCore, ответственный за этот ID, определяется с использованием правила mod-sharding: target_sc_id = col_id % num_total_sparse_cores .
    5. Ведется подсчет общего количества идентификаторов ( ids_per_sparse_core[target_sc_id]++ ) и количества уникальных идентификаторов ( unique_ids_per_sparse_core[target_sc_id]++ ) после обеспечения уникальности для конкретного target_sc_id ), предназначенных для каждого target_sc_id .
    6. Затем max_ids_per_partition для таблицы T1 устанавливается равным max(ids_per_sparse_core_array) .
    7. Аналогично, max_unique_ids_per_partition для таблицы T1 устанавливается равным max(unique_ids_per_sparse_core_array) .
    8. Если таблица T1 является компонентом стековой таблицы, перед суммированием статистики всех составляющих таблиц к распределениям идентификаторов могут быть применены дополнительные преобразования, такие как повороты или сдвиги. Это помогает сбалансировать нагрузку между чипами.

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

8. Как взаимодействует каждое SparseCore

Взаимодействие SparseCore, особенно в контексте обработки списка идентификаторов для встраивания поисковых запросов, опирается на несколько скоординированных механизмов:

  • Модифицированное шардинг и неявная маршрутизация :
    • Таблицы встраивания распределены по всем SparseCores в системе.
    • Когда хост предоставляет пакет входных данных (который впоследствии предварительно обрабатывается в формате COO, включая col_ids ), значение col_id используется для определения того, какой SparseCore отвечает за этот конкретный идентификатор: target_sc_id = col_id % num_total_sparse_cores .
    • Каждый SparseCore эффективно получает и обрабатывает только подмножество идентификаторов, соответствующих назначенным ему разделам словаря. Этап предварительной обработки на хосте критически важен для подготовки данных таким образом, чтобы каждый SparseCore мог легко идентифицировать и обрабатывать соответствующие идентификаторы.
  • Распределение данных по хостам :
    • Логика предварительной обработки хоста разделяет общий входной пакет и распределяет соответствующие части row_ids и col_ids (вместе с любыми связанными характеристиками или весами, если применимо) либо в память (HBM), к которой каждый SparseCore имеет прямой доступ, либо в общую HBM, из которой SparseCore будут извлекать необходимые данные.
  • Обработка Intra-SparseCore :
    • Получив назначенный набор идентификаторов для заданного раздела таблицы, SparseCore выполняет такие операции, как дедупликация этих идентификаторов и сбор соответствующих векторов встраивания. Эти операции, в основном, представляют собой локальные вычисления, выполняемые в собственных тайлах SparseCore с использованием его локальной SPMEM.
  • Связь между разреженными ядрами (All-to-All) :
    • После этапа начальной обработки (например, встраивания поиска) можно использовать шаблон связи «все ко всем» для объединения или перераспределения результатов между разреженными ядрами (например, перед передачей активаций в слой TensorCore, ожидающий входные данные, соответствующие всем исходным позициям выборки). Это крайне важно для реконструкции полного набора активаций, если исходный пакет входных данных был распределён для параллельной обработки.
  • Связь с TensorCores :
    • Ядра SparseCores взаимодействуют с ядрами TensorCores для отправки активаций встраивания (во время прямого прохода) и получения градиентов (во время обратного прохода). Это взаимодействие организуется программой, скомпилированной в XLA, и часто использует HBM в качестве промежуточного буфера. Стратегия конвейеризации (обсуждавшаяся ранее) существенно влияет на время и синхронизацию этого взаимодействия SC-TC.

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

9. Управление памятью SparseCore

Каждое ядро SparseCore эффективно управляет несколькими различными типами памяти для выполнения своих вычислений:

  • Блокнотная память (SPMEM) :
    • Природа : Относительно небольшая, но очень быстрая локальная SRAM-память, доступная исключительно каждому SparseCore. Важно отметить, что SPMEM не является кэшем; её использование явно контролируется и регулируется компилятором XLA.
    • Назначение : SPMEM используется для «оптимального размещения данных». Это включает в себя входные, выходные и промежуточные результаты, необходимые для текущих вычислений SC. Размещение данных в SPMEM значительно снижает высокую задержку, обычно связанную с доступом к HBM.
    • Размер : Как обсуждалось в разделе «Ограничения», буферы SPMEM статически задаются во время компиляции. Этот размер основан на таких параметрах, как max_ids_per_partition и max_unique_ids_per_partition . Статическое распределение гарантирует, что для любой операции с разделом таблицы (например, сокращения CSR) все необходимые данные для идентификаторов этого раздела (в пределах заданных ограничений) поместятся в SPMEM.
    • Оптимизации компилятора : компилятор XLA использует сложные оптимизации для точного определения того, какой объем данных и какие конкретные элементы данных необходимо поместить в SPMEM, чтобы эффективно скрыть задержку HBM и максимально повысить производительность.
    • Ограничение динамического выделения памяти : компилятор SparseCore в настоящее время не поддерживает динамическое выделение памяти во временной области. Это подчёркивает критическую важность статического определения размера посредством тщательной настройки ограничений.
  • Память с высокой пропускной способностью (HBM) :
    • Природа : большой общий ресурс памяти, доступный всем разреженным и тензорным ядрам, а также хост-системе. Первичные таблицы встраивания хранятся в HBM.
    • Использование стека : операции SparseCore часто требуют временного хранения в HBM промежуточных результатов, которые либо не помещаются в ограниченный SPMEM, либо должны передаваться между более крупными этапами конвейера обработки. Использование стека HBM во время прямого и обратного проходов можно оценить следующим образом:
      • Стек HBM прямого прохода (одна таблица) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 байта
      • Стек HBM обратного прохода (одна таблица) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 байта
    • Использование кучи : HBM также использует кучу, которой управляет хост. В куче хранятся такие данные, как веса плотных слоёв, константы, используемые моделью, и предварительно выбранные входные данные. Использование кучи, как правило, увеличивается с увеличением количества шагов, на которые хост предварительно выбирает данные (контролируется флагом maximum_parallel_iterations ). Хотя более частая предварительная выборка может повысить производительность за счёт перекрытия передачи данных между хостом и устройством с вычислениями на устройстве, она также потребляет больше HBM.
    • Сериализация для оптимизации HBM : флаг xla_sc_num_serialized_tables_to_optimize_hbm предоставляет механизм управления количеством таблиц, хранящихся в памяти стека HBM в любой момент времени. Увеличение этого числа фактически сериализует обработку большего количества таблиц, что может снизить пиковую нагрузку на стек HBM, но может привести к снижению производительности из-за снижения параллелизма.
  • Векторная память (VMEM) :
    • VMEM — это локальная сверхоперативная память, используемая исключительно TC (TensorCore). Хотя VMEM не управляется напрямую SparseCore, она является неотъемлемой частью экосистемы памяти, с которой взаимодействует SC, в основном через TensorCore.

Общая стратегия управления памятью:

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

10. Узкие места производительности и памяти

Для достижения оптимальной производительности SparseCore необходимо чёткое понимание потенциальных узких мест и способов их устранения. Они могут возникнуть на хосте, внутри самого SparseCore или при его взаимодействии с TensorCore.

Распространенные узкие места производительности:

  • Узкое место хоста :
    • Проблема : центральный процессор может не справляться с предварительной обработкой данных и их передачей в TPU достаточно быстро, что приводит к недоиспользованию разреженных и тензорных ядер. Это часто ограничивает производительность.
    • Снижение риска : отслеживайте загрузку ЦП хоста и показатели входного конвейера. Оптимизируйте процедуры загрузки и предварительной обработки данных на стороне хоста (см. советы по преобразованию COO). Настройте флаг maximum_parallel_iterations для точной настройки предварительной выборки данных.
  • Неоптимальная синхронизация TC/SC (отсутствие конвейеризации) :
    • Проблема : Если конвейеризация между TensorCore и SparseCore отключена или работает неэффективно, один блок может тратить значительное время на ожидание другого, тем самым снижая общую пропускную способность системы.
    • Смягчение : убедитесь, что конвейеризация включена (например, tf_xla_disable_full_embedding_pipelining = false или эквивалент).
  • Узкие места, вызванные ограничениями :
    • Проблема :
      • Слишком низкие лимиты : могут привести к чрезмерному мини-пакетированию (разделению входных пакетов на множество более мелких подпакетов для соблюдения жёстких ограничений). Хотя это обеспечивает корректность, каждый мини-пакет вносит определённые накладные расходы на обработку, что потенциально замедляет общее выполнение. Если allow_id_dropping имеет значение true, слишком низкие лимиты также могут привести к отбрасыванию идентификаторов, что влияет на точность модели.
      • Слишком высокие пределы (но всё ещё подходят) : хотя очень высокие пределы могут помешать мини-пакетированию, они могут привести к неоправданному увеличению нагрузки на SPMEM, если фактические характеристики данных редко приближаются к этим пиковым значениям. Они также могут привести к более интенсивному использованию стека HBM, чем это строго необходимо.
      • Ошибки компиляции : если настроенные ограничения требуют большего объема стека SPMEM или HBM, чем доступная физическая память, компиляция завершится ошибкой.
    • Смягчение : убедитесь, что ограничения установлены правильно.
  • Перекос распределения данных :
    • Проблема : если некоторые разделы SparseCore постоянно получают непропорционально большее количество идентификаторов по сравнению с другими (что указывает на плохое распределение идентификаторов), эти перегруженные разделы SparseCore станут узкими местами производительности.
    • Смягчение : перетасовка идентификаторов в процессе мини-пакетирования может помочь смягчить эту проблему для таблиц, расположенных в стопке, особенно с «горячими» пользовательскими таблицами. Тщательно проанализируйте распределение идентификаторов, чтобы установить адекватные и сбалансированные ограничения для каждой таблицы.
  • Проблемы с укладкой таблиц :
    • Проблема :
      • Слишком мало таблиц, сложенных друг на друга : этого может быть недостаточно для эффективного сокрытия задержки ICI или адекватного снижения накладных расходов на обработку.
      • Слишком много таблиц размещено друг на друге : может привести к созданию очень больших логических таблиц, которыми будет трудно управлять, или которые могут превысить доступные ограничения ресурсов.
    • Смягчение :
      • Обеспечьте оптимальное количество столов для штабелирования. Согласно общим рекомендациям, оптимальное количество столов для штабелирования составляет от 5 до 100.
  • Неэффективная нумерация/квантование :
    • Проблема : использование полной точности FP32, когда достаточно форматов с меньшей точностью, таких как BF16 или квантованные целые числа (и обеспечивающих более быстрые вычисления), может стать узким местом производительности.
    • Снижение риска : рассмотрите варианты с меньшей точностью. Однако имейте в виду, что само квантование имеет определённые накладные расходы и может потребовать тщательной настройки параметров квантования для поддержания точности модели.
  • Насыщенность полосы пропускания HBM :
    • Проблема : Чрезмерное перемещение данных в HBM и из него, потенциально вызванное очень малой шириной объектов (что приводит к высоким накладным расходам на заполнение), неэффективными шаблонами доступа к памяти или чрезвычайно большим количеством поисков, может привести к переполнению доступной полосы пропускания HBM.
    • Смягчение : Масштабирование количества TPU может помочь справиться с насыщением полосы пропускания HBM.

Распространенные узкие места памяти:

  • Переполнение SPMEM (ошибка компиляции) :
    • Проблема : Если значения max_ids_per_partition и max_unique_ids_per_partition слишком велики, компилятор XLA может не выделить достаточное количество SPMEM, что приводит к ошибкам компиляции, например: "Fixed size allocations (...) do not fit in TileSpmem (...)" . Кроме того, если значение (sample_count * feature_width) / kNumTiles (где kNumTiles — количество тайлов на SC) слишком велико для размещения операндов сборки в SPMEM тайла, могут возникать ошибки, например: "Gather operand too large..." .
    • Смягчение : Уменьшите размер партии или увеличьте количество чипов, используемых для обработки.
  • Переполнение стека HBM (во время выполнения или компиляции) :
    • Проблема : Если сочетание feature_width , max_unique_nz_per_row и logical_replica_count приводит к тому, что требования к памяти стека HBM превышают доступный HBM, это может вызвать ошибки Out-Of-Memory (OOM) либо во время выполнения, либо во время компиляции.
    • Смягчение : настройте флаг xla_sc_num_serialized_tables_to_optimize_hbm , чтобы уменьшить использование стека HBM за счет сериализации обработки таблиц (это обычно приводит к снижению производительности).
  • Истощение кучи HBM :
    • Проблема : в первую очередь вызвана очень большими плотными весами слоев, большим количеством констант, хранящихся в памяти, или чрезмерно агрессивной предварительной выборкой входных данных (высоким maximum_parallel_iterations ).
    • Смягчение : отслеживайте использование кучи с помощью таких инструментов, как XProf Memory Viewer.
  • Накладные расходы на заполнение :
    • Проблема : Таблицы встраивания дополняются до 32 байт (что эквивалентно 8 плавающим элементам) в измерении признаков. Следовательно, небольшая ширина элементов (например, 1 плавающий элемент) влечет за собой значительные накладные расходы на дополнение (например, 7/8 выделенного буферного пространства используется для дополнения), что приводит к неэффективному использованию HBM. Словарный размер таблиц также дополняется до значения, кратного количеству разреженных ядер в системе; однако это влияние обычно незначительно для таблиц с достаточно большим размером словаря.

Общие факторы, влияющие на производительность и память:

  • Топология : количество доступных микросхем и архитектура их взаимосвязей.
  • Размер пакета : напрямую влияет на sample_count на SparseCore, что, в свою очередь, влияет на потребление памяти и вычислительную нагрузку.
  • Форматирование данных : обеспечение эффективной компоновки данных на устройстве имеет решающее значение для оптимальной производительности.

11. Анализ профиля SparseCore

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

  1. Получить след :
    • Используйте инструменты профилирования, такие как XProf , для получения подробной трассировки выполнения во время обучения модели или вывода. Эта трассировка предоставит временную шкалу операций, выполняемых на хосте, тензорных ядрах (TensorCores) и разреженных ядрах (SparseCores).
  2. Проверьте Trace Viewer (например, в XProf или TensorBoard) :
    • Активность хоста : внимательно изучите активность хоста. Есть ли существенные пробелы в активности TPU? Такие пробелы могут указывать на то, что хост является узким местом, не справляясь с подачей данных достаточно быстро. Проанализируйте производительность входного конвейера.
    • Активность TensorCore (TC) и SparseCore (SC) :
      • Посмотрите на временные шкалы выполнения для TC и SC. Работают ли они параллельно, что свидетельствует об эффективной конвейеризации? Или наблюдаются длительные периоды, когда один блок простаивает, ожидая другой?
      • Определите операции, которые занимают больше всего времени (самые длительные операции) как на SC, так и на TC.
      • Визуальные результаты трассировки (часто отображающие цветные блоки, представляющие различные операции с течением времени, например, TPU:0 SparseCore 1 (pid 1005) ) бесценны для визуального определения доминирующих операций и периодов простоя.
    • Анализ времени шага : проследите за общим временем шага и поймите, как оно распределяется или разбивается между обработкой хоста, вычислениями SC и вычислениями TC.
  3. Анализ памяти (XProf Memory Viewer) :
    • Использование кучи . Используйте такие инструменты, как вкладка Xprof «Просмотр памяти», чтобы осмотреть использование кучи HBM. Это может помочь определить, потребляют ли большие веса модели, константы или чрезмерно агрессивные входные предварительные выбранные, чрезмерное количество HBM. Включение флагов, таких как --vmodule=best_fit_allocator=1 может обеспечить журналы пикового использования кучи.
    • Использование стека (косвенное) : В то время как прямое профилирование стека HBM может быть сложным, если вы столкнетесь с ошибками вне памяти, а использование кучи представляется разумным, истощение стека HBM (часто из-за чрезмерно больших пределов или ширины элементов) является сильным подозреваемым. Формулы, предусмотренные для использования стека HBM, могут помочь в оценке этого.
  4. Ищите определенные шаблоны :
    • Мини-оценка : если ограничения часто превышаются, вы можете наблюдать доказательства мини-обозначения в трассировке (например, большее количество меньших операций SC, чем ожидалось для глобального размера партии). Это часто можно вывести из журналов или наблюдения за количеством вызовов определенных операций.
    • Отброшение идентификатора : если сброс идентификатора включен и происходит, системные журналы могут дать указания об этом. Это также было бы явным признаком того, что настроенные пределы слишком ограничивают для входных данных.
    • Время компиляции : расширенное время перекомпиляции, особенно если включена в обратная связь (FDO) и часто регулируют ограничения, может добавить значительные накладные расходы на общее время обучения.
  5. Коррелируйте с флагами и конфигурацией :
    • Свяжите наблюдаемое поведение в профиле обратно с вашими конфигурациями Sparsecore (настройки в файлах ограничений, флаги XLA). Например, если xla_sc_num_serialized_tables_to_optimize_hbm установлен на высокое значение, вы можете ожидать более медленную производительность SC, но снизит потребление стека HBM.
  6. Итеративный процесс :
    • Профилирование часто является итеративным процессом уточнения. Сделайте определенное изменение (отрегулируйте предел, включите или отключите функцию), захватите новый профиль, а затем сравните его с предыдущим профилем, чтобы увидеть влияние вашей модификации.

12. Общие флаги отладки

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

  • Идентификационные проверки (вне расстояния) :
    • Флаг : xla_sparse_core_enable_id_bound_check = true
    • Цель : позволяет проверять систему хоста обнаружить, если какие -либо идентификаторы встраивания в входные данные выходят за пределы достоверного локального диапазона, определенного для данной таблицы внедрения. Это помогает поймать проблемы, связанные с неправильными или поврежденными входными данными.
  • НАНКА ПЕРЕДАЧА :
    • Флаг : xla_sc_detect_nan = true
    • Цель : позволяет обнаружить значения NAN (не число) в рамках данных с плавающей точкой, обрабатываемыми на sparsecore. Если NAN обнаруживается на входах или выходах различных проходов компилятора, этот флаг вызовет ошибку. Такие ошибки обычно предоставляют информацию о том, где встречалась NAN.
  • Проверка границ (доступ к памяти) :
    • Флаг : xla_sc_assert_level=bounds
    • Цель : Этот флаг включает инструмент ASAN (AddressSanitizer) -Sele, который переписывает инструкции по доведению памяти (например, нагрузки/хранилища VMEM и операции DMA), чтобы включить динамические проверки. Эти проверки проверяют, находится ли доступ к памяти в пределах выделенных границ целевой области памяти.
    • Поведение : если обнаружен доступ за пределами памяти, выполнение не удастся.
    • Осторожно : этот шахтек может произвести ложные срабатывания, например, из -за сложных моделей укрепления доступа, которые не полностью поняты контролером. Это преобразование применяется на поздней стадии в процессе компиляции бэкэнд.
  • Буферный шахер (коррупция памяти) :
    • Флаги :
      • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
      • xla_tpu_verify_launch_id_across_cores=true
    • Цель : Эти флаги помогают гарантировать, что буферы памяти не были непреднамеренно повреждены или перезаписываются не связанными операциями. Дезинфицирующее средство буфера проверяет содержимое буферов, чтобы убедиться, что они не изменяются неожиданно.

13. Поддержка квантования

Sparsecore SparseDenseMatmulOp предназначен для поддержки операций в таблицах встраивания, используя как 32-разрядную точку с плавающей запятой (FP32), так и типы целочисленных данных. В то время как модельное обучение обычно выполняется с использованием точности FP32 для таблиц встраивания, может применяться квантование после тренировки (PTQ). PTQ позволяет использовать данные дата с более низкой конкретной позицией (например, 8-битные целые числа) для вывода, что потенциально может привести к улучшению производительности и уменьшенной следов памяти.

Моделируемое квантование:

SparseDenseMatmulOp может быть настроен для выполнения «смоделированного квантования». В этом эксплуатационном режиме векторы встраивания сначала определяются до более низкой точности, а затем отбрасывают обратно на более высокую точность (например, FP32), прежде чем они будут использованы в последующих вычислениях. Этот метод позволяет обучать модели при учете эффектов шума квантования. Обучение с имитированным квантованием может повысить точность окончательной модели, когда она полностью квантована для вывода.

Атрибуты конфигурации для SparseDenseMatmulOp (для квантования):

  • quantization_config_num_buckets = 256
    • В этом атрибуте указывается количество дискретных ведер или уровней, на которых будет квантовано 32-разрядное число с плавающей запятой. Например, при квантовании до 8-битных целых чисел можно было бы указать 2^8 = 256 ведра.
  • quantization_config_low = -XX
    • Этот атрибут определяет минимальное значение с плавающей точкой в диапазоне квантования. Любые входные значения ниже этого указанного минимума будут обрезаны до этого минимального значения во время квантования.
  • quantization_config_high = YY
    • Этот атрибут определяет максимальное значение с плавающей точкой в диапазоне квантования. Любые входные значения выше этого указанного максимума будут обрезаны до этого максимального значения во время квантования.

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

Числовое поведение модели может измениться в зависимости от того, включено ли трубопровод между Tensorcore и Sparsecore. Если конвейер активен, градиенты, обработанные Sparsecore, могут быть «несвежими» (с предыдущей итерации). Это может взаимодействовать с процессом квантования и потенциально повлиять на динамику обучения модели или окончательную точность.

14. Предстоящие функции и недавние улучшения

Экосистема Sparsecore подвержена постоянному развитию и улучшению.

Дорожная карта:

  • Мини-оценка измерения образца :
    • Это запланировано в качестве дополнительной функции существующих возможностей мини-размерного словаря.
    • Это позволило бы сделать дальнейшее разделение ввода входных данных вдоль размерности образца. Это будет достигнуто путем введения петлей на устройстве, которые могут фильтровать и обрабатывать поиск из подмножества образцов одновременно. Такая функция может быть полезна для управления очень большим количеством идентификаторов для для выборки или для улучшения балансировки нагрузки между единицами обработки.
  • Улучшенная поддержка встраивания с менее чем 8 целыми числами в строке (ширина небольших элементов) :
    • Текущий дизайн часто использует значительную прокладку для ширины функций встраивания, которая составляет менее 8 поплавок (что соответствует 32 байтам). Эта заполнение может привести к потраченному впустую HBM и потенциально недостаточно используемым вычислительным ресурсам. Будущие улучшения направлены на то, чтобы смягчить эту неэффективность для таблиц с небольшими размерами объекта.

Последние улучшения:

  • Постановка операндов сборов в HBM :
    • Эта оптимизация помогает снизить давление на общую память царапин (SPMEM), позволяя некоторым операциям сбора входов или выходов в большую HBM.
  • Снижение использования памяти стека :
    • Улучшения были реализованы для снижения потребления памяти стека HBM во время операций Sparsecore, в идеале без отрицательного влияния на общую производительность или пропускную способность.

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