Категория: Несоответствие выходных данных при разгрузке хоста
Эта ошибка возникает, когда тензор, явно выгруженный в память хоста, возвращается в качестве выходных данных программы, но сигнатура выходных данных программы не настроена на ожидание памяти хоста.
Примеры сообщений об ошибках:
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
Обзор
Когда компилятор встречает аннотацию, указывающую на необходимость переноса тензора на хост (процессор), он отслеживает местоположение этого тензора в графе вычислений до тех пор, пока не произойдет одно из трех событий:
- Перемещение на устройство: соответствующая аннотация перемещает тензор обратно на ускоритель.
- Вычисления на стороне хоста: Тензор обрабатывается операцией на стороне хоста.
- Завершение программы: Тензор достигает конца программы и становится выходным значением.
This error is triggered in scenario #3 . The tensor is physically located in host memory at the end of the execution, but the XLA program's entry computation signature defines that specific output as residing in Device Memory . Because the compiler cannot implicitly change the entry computation's interface, it raises an error.
Отладка
Для устранения этой ошибки определите, предполагали ли вы, что этот тензор будет выводиться на хост-компьютер, или же его следовало переместить обратно на устройство перед возвратом.
Проверка намерения вывода и трассировка пути {#verify_intent}
Intended to return on Host: If you explicitly want this tensor to be returned in host memory (avoiding a transfer back to device), you should explicitly set the output memory space of the entry computation to Host Memory for this specific output.
Intended to return on Device: If the tensor was meant to stay on the device or return to it before the program ends, you likely missed an annotation. Insert a matching annotation to move the tensor back to the device.
Если источник выгруженного тензора неясен или вы не можете найти, где отсутствует аннотация «переместить на устройство», используйте логирование XLA для отслеживания инструкций.
- Включить логирование: Если вы используете Google Cloud TPU, перезапустите программу со следующим флагом:
--vmodule=host_offloader=1 - Analyze Logs: Look for the "trace" output in the logs. This will show the path of the tensor starting from the offload instruction. Use this to pinpoint exactly where the tensor reaches the program boundary without being moved back to the device.