Treinamento em frações de TPU

As TPUs foram projetadas para escalonamento horizontal em um Pod de TPU. Um Pod de TPU é uma coleção 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 baseada em CPU e com alto desempenho para ações 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 tarefa 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:

• JAX

• PyTorch

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

Como escalonar o tamanho do lote e as etapas de treinamento

Para um escalonamento 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 v6e-8, use um tamanho de lote de 4.096 (4 * 1024) em uma v6e-32. Isso utiliza totalmente o hardware da TPU. É possível usar tamanhos de lote menores, mas o treinamento não será escalonado linearmente.

Muitos modelos incluem uma flag train_steps, em que uma etapa corresponde a um único lote de dados. Se 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. No entanto, se você tem quatro workers e um tamanho de lote efetivo de 4.000, precisa ajustar o número de etapas para 25 para processar os mesmos 100.000 exemplos. Se o modelo usar uma flag epochs, não será preciso escalonar o número de etapas.

Como tamanhos de lote maiores podem alterar o comportamento de convergência do modelo, também pode ser preciso ajustar alguns hiperparâmetros, como a taxa de aprendizado.

Como usar buckets regionais do 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 dos recursos é importante principalmente ao usar Pods de TPU porque as taxas de transferência de dados são mais altas quando o bucket do Cloud Storage e a TPU estão na mesma região.

Verifique se você está usando um bucket regional do 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 TPU, uma boa prática é iniciar o desenvolvimento nas menores TPUs e iterar progressivamente até as maiores. Comece usando uma versão de TPU pequena, como a v6e-8.

• Teste a carga de trabalho para conferir se ela tem o comportamento esperado.
• Teste e valide o desempenho usando as ferramentas de desempenho.

Quando a carga de trabalho estiver funcionando e atingir as metas de desempenho, escalone para um tipo de TPU maior, como a v6e-32. Aumente gradualmente e iterativamente o tamanho da TPU enquanto valida a escalonabilidade (funcionalidade e desempenho) até atingir o tamanho de TPU pretendido.