Код ошибки: E3000

Категория: Время компиляции: Сбой выделения разреженного ядра

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

Примеры сообщений об ошибках:

INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>

XLA Backends: TPU

Обзор

SparseCore (SC) — это специализированный процессор для работы с разреженными рабочими нагрузками. Он использует определённые иерархии памяти для эффективного управления перемещением данных. Компилятор XLA пытается статически задать размер и выделить буферы на основе аппаратных ограничений и заданных пользователем форм. Эта ошибка указывает на нехватку памяти (OOM) на данном этапе выделения памяти.

В сообщении об ошибке обычно указывается идентификатор области памяти. Ниже приведены распространенные области памяти и их целочисленные кодировки:

Идентификатор пространства памяти Имя Описание
0 Смем Локальная скалярная память. Используется для скалярных регистров и управления потоком выполнения.
201 TileSpmem Память временного хранилища, привязанная к конкретному тайлу. Быстрая локальная SRAM, доступная для конкретного тайла SC.
202 Спмем Общая буферная память. Используется для оперативного размещения данных (входных, выходных, промежуточных) с целью скрытия задержки HBM.
203 HBM Память с высокой пропускной способностью. Большая разделяемая память, используемая для встраивания таблиц, куч и стеков.
204 Флаги синхронизации Примитивы синхронизации, используемые для координации.

Для более подробного изучения SC и его иерархии памяти обратитесь к документации SparseCore .

Отладка

Решение зависит от того, какое именно пространство памяти не удалось выделить.

Сценарий 1. Сбои при распределении HBM.

Идентификатор пространства памяти: 203

Эта ошибка возникает, если запрошенное программой SparseCore временное выделение памяти слишком велико, чтобы поместиться в доступный HBM. В стандартных рабочих нагрузках встраивания и коллективных задачах с разгрузкой SC чрезвычайно большие разделы для каждого ядра или некорректные спецификации сегментирования могут заставить компилятор запрашивать огромные буферы.

Рекомендуемые действия:

  • Проверка сегментирования: Убедитесь, что ваши таблицы встраивания и входные/выходные тензоры SC правильно разделены/сегментированы. Если одно ядро ​​обрабатывает слишком большой объем данных, выделение памяти может завершиться неудачей.
  • Настройка ограничений: проверьте значения max_ids_per_partition и max_unique_ids_per_partition . Если они установлены слишком высоко, компилятор резервирует больше памяти, чем необходимо. См. раздел «Как ограничения преобразуются в таблицы» .

Сценарий 2. Сбои внутренней памяти.

Идентификаторы пространства памяти: 0, 201, 202, 204

Сбои при выделении памяти в Smem , TileSpmem , Spmem или Sync Flags обычно происходят из-за ошибок компилятора или ограничений в стратегии выделения памяти, когда компилятор не учитывает все требования к памяти.

Рекомендуемые действия:

  1. Выделите сбойную операцию XLA: чтобы определить конкретное ядро ​​SC HLO или Mosaic, вызывающее сбой, сгенерируйте промежуточные представления компилятора:
    • Для создания дампа MLIR для SparseCore установите флаг --xla_sc_dump_mlir_to=/path/to/dump . Это сгенерирует MLIR для программы SparseCore, позволяя увидеть, какой размер выделения соответствует сообщению об ошибке.
    • Дамп Mosaic LLO: Для пользовательских ядер используйте --xla_mosaic_dump_to=/path/to/dump , чтобы проверить все программы низкоуровневого оптимизатора (LLO), генерируемые Mosaic.
  2. Уменьшите размеры временных файлов (для пользователей Pallas): Если сбой происходит в ядре Mosaic, проверьте конфигурацию scratch_shapes . Убедитесь, что ваши запросы pltpu.SMEM соответствуют аппаратным характеристикам вашего конкретного поколения TPU.
  3. Отключение коллективной разгрузки: Если ошибка возникает из-за коллективных операций, выполняемых с использованием функции разгрузки SC, попробуйте отключить функции разгрузки SC:
    • --xla_tpu_enable_sparse_core_collective_offload_all_gather=false
    • --xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
  4. Сообщить об ошибке: Если описанные выше шаги не решат проблему, скорее всего, это ошибка компилятора. Пожалуйста, отправьте отчет об ошибке.