Categoria:tempo de compilação: incompatibilidade de saída de descarregamento do host
Esse erro ocorre quando um tensor explicitamente descarregado para a memória do host é retornado como uma saída do programa, mas a assinatura de saída do programa não está configurada para esperar a memória do host.
Exemplos de mensagens de erro :
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 do XLA:TPU, GPU
Visão geral
Quando o compilador encontra uma anotação para descarregar um tensor no host (CPU), ele rastreia o local desse tensor no gráfico de computação até que um dos três eventos ocorra:
- Mover para o dispositivo:uma anotação correspondente move o tensor de volta para o acelerador.
- Computação do host:o tensor é consumido por uma operação do lado do host.
- Fim do programa:o tensor chega ao fim do programa e se torna uma saída.
Esse erro é acionado no cenário nº 3. O tensor está fisicamente localizado na memória do host no final da execução, mas a assinatura de computação de entrada do programa XLA define essa saída específica como residente na memória do dispositivo. Como o compilador não pode mudar implicitamente a interface da computação de entrada, ele gera um erro.
Depuração
Para resolver esse erro, determine se você pretendia que esse tensor fosse uma saída no host ou se ele deveria ter sido movido de volta para o dispositivo antes de retornar.
Pretende retornar no host:se você quiser explicitamente que esse tensor seja retornado na memória do host (evitando uma transferência de volta para o dispositivo), defina explicitamente o espaço de memória de saída da computação de entrada como memória do host para essa saída específica.
Pretende retornar no dispositivo:se o tensor deveria permanecer no dispositivo ou retornar a ele antes do término do programa, provavelmente você perdeu uma anotação. Insira uma anotação correspondente para mover o tensor de volta para o dispositivo.
Se a origem do tensor descarregado não estiver clara ou você não conseguir encontrar onde a anotação "mover para o dispositivo" está ausente, use o registro do XLA para rastrear as instruções.
- Ativar o registro:se você estiver no Google Cloud TPU, execute o programa novamente com a seguinte flag:
--vmodule=host_offloader=1. - Analisar registros:procure a saída "trace" nos registros. Isso vai mostrar o caminho do tensor a partir da instrução de descarregamento. Use isso para identificar exatamente onde o tensor atinge o limite do programa sem ser movido de volta para o dispositivo.