Categoria:tempo de execução: incompatibilidade do buffer de entrada do programa
Esse erro ocorre quando o tempo de execução do XLA detecta uma incompatibilidade entre o tamanho de um buffer de memória esperado por um programa compilado e o tamanho do buffer fornecido no momento da execução.
Exemplo de mensagem de erro:
XlaRuntimeError: INVALID_ARGUMENT: Executable(jit_embedding_pipeline_step_fn) expected parameter 2482 of size 5242880 (bf16[16,1280,40]{2,1,0:T(8,128)(2,1)}) but got buffer with incompatible size 1638400 (bf16[16,1280,40]{1,2,0:T(8,128)(2,1)}): while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
Backends do XLA:TPU
Visão geral
A mensagem de erro indica os tamanhos esperados e reais, bem como as formas e os layouts dos tensores. Esses erros podem ocorrer mesmo que dois tensores tenham o mesmo formato, mas o tamanho deles na memória pode ser diferente se o layout físico (como os dados são organizados e dispostos no hardware) for diferente.
Esses erros são causados principalmente por:
- Incompatibilidade de configuração de checkpoint e XLA: um modelo é treinado e um checkpoint é salvo. O layout físico dos pesos nesse ponto de verificação é determinado pela versão e configuração exatas do XLA (por exemplo, flags do XLA) naquele momento. Mais tarde, esse ponto de verificação é carregado em um ambiente diferente, em que a configuração mudou. Uma nova flag, um valor padrão diferente ou uma mudança no modelo/código XLA podem fazer com que o tempo de execução espere um layout físico diferente para os pesos. Quando o buffer antigo do checkpoint é transmitido para o novo programa XLA compilado, o tempo de execução gera um erro.
- Layouts específicos de hardware/topologia: o compilador XLA pode escolher layouts físicos diferentes para tensores e otimizar o desempenho em diferentes hardwares. Um layout ideal para TPU v4 pode ser diferente de uma TPU v5 ou até mesmo para diferentes frações de pod do mesmo chip (por exemplo, 4x4x4 vs 4x8). O erro ocorre quando um modelo é compilado com uma proposição sobre o layout de uma topologia, mas, no tempo de execução, ele é programado em uma topologia diferente ou há um bug na lógica de layout do compilador para uma peça específica de hardware.
Depuração
- Garanta a consistência da configuração entre a exportação do modelo e as novas execuções de
checkpoints:
- Evite usar checkpoints antigos com código novo, a menos que você tenha certeza de que não foram feitas mudanças que afetam o layout.
- Reexporte o modelo salvo: se você suspeitar de uma incompatibilidade de checkpoint/configuração, a solução mais confiável é reexportar o modelo salvo usando exatamente a mesma base de código e configuração (atuais) que você está usando para inferência ou ajuste refinado.
- Verifique se há mudanças de configuração (por exemplo, flags XLA) entre as duas execuções.
- Layouts específicos de hardware/topologia:
- Verifique se há incompatibilidades de versão e topologia de hardware ao trocar de hardware ou topologias.