Categoria:tempo de compilação: Vmem OOM com escopo
Esse erro indica que o programa exige mais memória de vetor com escopo (Vmem) do que o 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 de XLA:TPU
Visão geral
As TPUs têm memória de vetor (VMEM), que é uma memória de rascunho local usada exclusivamente pelo TensorCore (TC). O compilador gerencia a Vmem para diferentes tipos de alocações:
- Alocações com escopo de instrução:armazenamento temporário na Vmem durante a execução de uma única instrução HLO. Isso inclui o buffer de extensão de operandos (por exemplo, para buffer duplo) e derramamentos de registro.
- Alocações com escopo de programa:alocações que vão além do escopo de uma única instrução HLO. Geralmente, são temporários de HLO e resultados intermediários que são entradas e/ou saídas de instruções HLO.
Um Vmem OOM com escopo de tempo de compilação ocorre quando as alocações com escopo de instrução excedem o limite de alocação para essa 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 cuidadosamente a mensagem de erro para identificar se ela é causada por um kernel personalizado ou 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" ...
- Vmem OOM com escopo de kernel personalizado: se o erro apontar para um kernel personalizado, acesse Ajustar o kernel.
- Problemas de Vmem não relacionados ao kernel: se o Vmem OOM ocorrer devido a uma operação não relacionada ao kernel personalizado, é provável que seja um bug interno do compilador. Preencha um relatório de bug no XLA com um despejo de HLO.
Ajustar 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 Vmem com escopo por kernel:solicite explicitamente a quantidade de memória necessária para esse kernel específico usando o parâmetro vmem_limit_bytes
- Modificar a coloração da memória: colore/restrinja explicitamente as entradas/saídas do kernel para Vmem usando pallas.tpu.with_memory_space_constraint. Tenha cuidado para não colorir muitas saídas de entradas para Vmem, porque isso pode causar um Vmem OOM geral.
- Ajustar o limite de Vmem: se o ajuste específico do kernel for difícil ou se o problema afetar muitos kernels, ajuste o limite global de Vmem usando a flag --xla_tpu_scoped_vmem_limit_kib.