Treinamento em Pods da TPU

Informações gerais

As TPUs foram projetadas para escalonar horizontalmente em um Pod de TPU. Um Pod de TPU é um conjunto de dispositivos de TPU conectados por interfaces de rede dedicadas de alta velocidade. Um pod de TPU permite distribuir a carga de processamento em várias TPUs. Cada placa de TPU é conectada a uma máquina host com alto desempenho baseada em CPU para itens como carregamento e pré-processamento de dados. Para aproveitar ao máximo os números maiores de TPUs, é necessário ajustar vários parâmetros de tarefas de treinamento.

A configuração do treinamento com pods de TPU é diferente para cada framework. Use os links a seguir para ver informações detalhadas sobre o treinamento em pods com cada framework:

As seções a seguir explicam alguns problemas comuns, mudanças que você precisa fazer nos modelos e práticas recomendadas para reduzir ou evitar falhas no pod.

Como escalonar o tamanho do lote e as etapas de treinamento

Para alcançar o dimensionamento linear em tipos de TPU maiores, mantenha o tamanho do lote por núcleo igual.

Por exemplo, se você usar um tamanho de lote de 1.024 em uma v2-8, use um tamanho de lote de 4.096 (4 * 1024) em uma v2-32. Isso utiliza totalmente o hardware da TPU. É possível usar tamanhos de lote menores, mas o treinamento não será dimensionado linearmente se você fizer isso.

Alguns modelos incluem uma sinalização train_steps em que uma etapa corresponde ao processamento de um único lote de dados. Quando você aumentar o tamanho do lote, reduza verticalmente o número de etapas de treinamento. Assim, o número total de exemplos de treinamento permanece o mesmo.

Por exemplo, se você tem um tamanho de lote de 1.000 para 100 etapas, 100.000 exemplos são processados durante o treinamento. Se agora você tem quatro workers e um tamanho de lote efetivo de 4.000, é necessário ajustar o número de etapas para 25 para processar os mesmos 100.000 exemplos. Se o modelo usar uma sinalização epochs, não será preciso escalonar o número de etapas.

Tamanhos de lote maiores podem alterar o comportamento de convergência do modelo. Portanto, talvez você também precise ajustar alguns hiperparâmetros, como taxa de aprendizado.

Como usar buckets regionais do Google Cloud Storage na mesma região do pod de TPU

Em geral, a prática recomendada para treinamento de TPU é sempre usar recursos na mesma região. A região de recursos é particularmente importante ao usar pods de TPU porque as taxas de transferência de dados são mais altas quando o bucket do Google Cloud Storage e a TPU estão na mesma região.

Verifique se você está usando um bucket regional do Google Cloud Storage na mesma região da TPU para treinar conjuntos de dados e checkpoints.

Práticas recomendadas para fluxo de trabalho de desenvolvimento em Pods de TPU

Ao desenvolver uma nova carga de trabalho de TPUs, geralmente é ideal iniciar o desenvolvimento nas menores TPUs e iterar progressivamente até TPU maiores. Comece usando uma versão pequena de TPU (por exemplo, v2-8 ou v3-8).

  • Testar a funcionalidade da carga de trabalho
  • Testar e validar o desempenho usando as ferramentas de desempenho

Quando a carga de trabalho estiver funcionando e atingir suas metas de desempenho, escalonar verticalmente para um tipo de TPU maior, como v2-32 ou v3-32. Aumente de maneira gradual e iterativa o tamanho da TPU e, ao mesmo tempo, valida a escalonabilidade (funcionalidade e desempenho) até atingir o tamanho de TPU desejado.