Categoría: Tiempo de compilación: No coincide la descarga del host
Este error ocurre cuando se devuelve un tensor descargado explícitamente en la memoria del host como resultado del programa, pero la firma de salida del programa no está configurada para esperar la memoria del host.
Ejemplos de mensajes de error:
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.
Back-ends de XLA: TPU, GPU
Descripción general
Cuando el compilador encuentra una anotación para descargar un tensor en el host (CPU), hace un seguimiento de la ubicación de ese tensor a través del grafo de procesamiento hasta que ocurre uno de los siguientes tres eventos:
- Move to Device: Una anotación coincidente mueve el tensor de vuelta al acelerador.
- Cálculo del host: La operación del host consume el tensor.
- Fin del programa: El tensor llega al final del programa y se convierte en un resultado.
Este error se activa en la situación #3. El tensor se encuentra físicamente en la memoria del host al final de la ejecución, pero la firma de cálculo de entrada del programa XLA define que ese resultado específico reside en la memoria del dispositivo. Como el compilador no puede cambiar de forma implícita la interfaz del cálculo de entrada, genera un error.
Depuración
Para resolver este error, determina si querías que este tensor fuera un resultado en el host o si debería haberse movido de nuevo al dispositivo antes de devolverse.
Devolución prevista en el host: Si deseas explícitamente que este tensor se devuelva en la memoria del host (para evitar una transferencia de vuelta al dispositivo), debes establecer de forma explícita el espacio de memoria de salida del cálculo de entrada en Memoria del host para esta salida específica.
Intended to return on Device: Si el tensor debía permanecer en el dispositivo o volver a él antes de que finalizara el programa, es probable que hayas omitido una anotación. Inserta una anotación coincidente para volver a mover el tensor al dispositivo.
Si no está claro el origen del tensor descargado o no puedes encontrar dónde falta la anotación "move to device", usa el registro de XLA para hacer un seguimiento de las instrucciones.
- Habilita el registro: Si estás en Google Cloud TPU, vuelve a ejecutar tu programa con la siguiente marca:
--vmodule=host_offloader=1. - Analiza los registros: Busca el resultado "trace" en los registros. Se mostrará la ruta del tensor a partir de la instrucción de descarga. Úsalo para identificar con precisión dónde el tensor alcanza el límite del programa sin volver al dispositivo.