Categoria:tempo de compilação: Vmem OOM com escopo
Esse erro indica que o programa exige mais memória vetorial (Vmem) com escopo do que foi alocado.
Exemplos de mensagens de erro:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
Back-ends do XLA:TPU
Visão geral
As TPUs têm memória de vetor (VMEM), que é uma memória de bloco de notas local usada exclusivamente pelo TensorCore (TC). O compilador gerencia a Vmem para diferentes tipos de alocações:
- Alocações no escopo da instrução:armazenamento temporário em Vmem durante a execução de uma única instrução HLO. Isso inclui o buffer de extensão do operando (por exemplo, para buffer duplo) e o transbordamento de registros.
- Alocações no escopo do programa:alocações que vão além do escopo de uma única instrução HLO. Geralmente, são resultados temporários e intermediários de HLO que são entradas e/ou saídas de instruções de HLO.
Uma OOM de Vmem com escopo de tempo de compilação ocorre quando as alocações com escopo de instrução excedem o limite de alocação dessa instrução. Esse limite é controlado
- globalmente para todo o programa usando a flag --xla_tpu_scoped_vmem_limit_kib e
- por kernel personalizado usando o parâmetro vmem_limit_bytes.
Esses erros geralmente são causados por um bug interno do compilador ou por um kernel personalizado que excede o limite de alocação.
Depuração
Analise com cuidado a mensagem de erro para identificar se ela vem de um kernel personalizado ou de um HLO padrão. Um erro devido a um kernel personalizado precisa ter a seguinte assinatura:
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- OOM de Vmem no escopo do kernel personalizado: se o erro apontar para um kernel personalizado, acesse Ajustar o kernel.
- Problemas de Vmem não relacionados ao kernel: se o OOM de Vmem ocorrer devido a uma operação não relacionada ao kernel personalizado, provavelmente é um bug interno do compilador. Registre um bug no XLA com um despejo de HLO.
Reajustar o kernel
Se o erro for originado de um kernel personalizado, use as técnicas a seguir para reduzir o requisito de memória do kernel:
- Ajustar tamanhos de bloco:reduza os tamanhos de bloco (tamanhos de bloco) na configuração do kernel para diminuir o uso de Vmem com escopo.
- Definir limites de memória virtual com escopo por kernel:solicite explicitamente a quantidade necessária de memória para esse kernel específico usando o parâmetro vmem_limit_bytes.
- Modificar a coloração da memória:pinte/restrinja explicitamente as entradas/saídas do kernel para VMEM usando pallas.tpu.with_memory_space_constraint. Mas tenha cuidado para não colorir muitas entradas/saídas para Vmem, porque isso pode causar um OOM geral da VMEM.
- Se for difícil fazer um ajuste específico do kernel ou se o problema afetar muitos kernels, ajuste o limite global de Vmem usando a flag --xla_tpu_scoped_vmem_limit_kib.