Categoria:tempo de compilação: incompatibilidade de saída de descarga 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.
Backends do XLA:TPU, GPU
Visão geral
Quando o compilador encontra uma anotação para descarregar um tensor no host (CPU), ele rastreia a localização desse tensor no gráfico de computação até que um destes 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 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 do cálculo da entrada, ele gera um erro.
Depuração
Para resolver esse erro, determine se você queria que esse tensor fosse uma saída no host ou se ele deveria ter sido movido de volta para o dispositivo antes de retornar.
Verificar a intenção de saída e o caminho do rastreamento
Retorno pretendido 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 do cálculo de entrada como Memória do host para essa saída específica.
Retorno pretendido no dispositivo:se o tensor deveria permanecer no dispositivo ou retornar a ele antes do fim 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 se você não conseguir encontrar onde a anotação "mover para o dispositivo" está faltando, use o registro em log do XLA para rastrear as instruções.
- Ativar o registro em log:se você estiver usando o 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 começando pela instrução de descarga. Use isso para identificar exatamente onde o tensor atinge o limite do programa sem ser movido de volta para o dispositivo.