Treinamento em pods da TPU
Visão geral
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 entre 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 para treinamento com pods de TPU é diferente para cada framework. Use os links abaixo para conferir informações detalhadas sobre o treinamento em pods com cada framework:
As seções a seguir explicam alguns problemas comuns, alterações que você precisa fazer nos seus 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 flag train_steps em que uma etapa corresponde a
processar 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 é importante principalmente 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 gradualmente e iterativamente o tamanho da TPU enquanto valida a escalonabilidade (funcionalidade e desempenho) até atingir o tamanho de TPU desejado.