Категория: Время компиляции: Несоответствие выходных данных при разгрузке хоста
Эта ошибка возникает, когда тензор, явно выгруженный в память хоста, возвращается в качестве выходных данных программы, но сигнатура выходных данных программы не настроена на ожидание памяти хоста.
Примеры сообщений об ошибках:
INVALID_ARGUMENT: Tensor which is moved to host (starting from tuple.64) is returned from the entry computation but the layout for this output is not set to host memory.
XLA бэкэнды: TPU, GPU
Обзор
Когда компилятор встречает аннотацию, указывающую на необходимость переноса тензора на хост (процессор), он отслеживает местоположение этого тензора в графе вычислений до тех пор, пока не произойдет одно из трех событий:
- Перемещение на устройство: соответствующая аннотация перемещает тензор обратно на ускоритель.
- Вычисления на стороне хоста: Тензор обрабатывается операцией на стороне хоста.
- Завершение программы: Тензор достигает конца программы и становится выходным значением.
Эта ошибка возникает в сценарии №3 . В конце выполнения тензор физически находится в памяти хоста, но сигнатура вычисления точки входа программы XLA определяет этот конкретный выходной параметр как находящийся в памяти устройства . Поскольку компилятор не может неявно изменить интерфейс вычисления точки входа, он выдает ошибку.
Отладка
Для устранения этой ошибки определите, предполагали ли вы, что этот тензор будет выводиться на хост-компьютер, или же его следовало переместить обратно на устройство перед возвратом.
Проверьте намерение вывода и отследите путь.
Предполагается, что результат будет возвращен на хост: Если вы хотите, чтобы этот тензор был явно возвращен в память хоста (во избежание передачи обратно на устройство), вам следует явно установить объем памяти для выходных данных при вычислении входного параметра в память хоста для этого конкретного выходного параметра.
Предполагается, что тензор должен остаться на устройстве или вернуться на него до завершения программы: Если предполагалось, что тензор останется на устройстве или вернется на него до завершения программы, вы, вероятно, пропустили аннотацию. Вставьте соответствующую аннотацию, чтобы переместить тензор обратно на устройство.
Если источник выгруженного тензора неясен или вы не можете найти, где отсутствует аннотация «переместить на устройство», используйте логирование XLA для отслеживания инструкций.
- Включить логирование: Если вы используете Google Cloud TPU, перезапустите программу со следующим флагом:
--vmodule=host_offloader=1. - Анализ логов: Найдите в логах вывод "trace". Он покажет путь тензора, начиная с инструкции выгрузки. Используйте это, чтобы точно определить, где тензор достигает границы программы, не будучи перемещенным обратно на устройство.