Treinamento em pods da TPU

Visão geral

As TPUs foram projetadas para serem escalonadas 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 pode ter até 2.048 núcleos de TPU, o que permite distribuir a carga de processamento em várias TPUs. Cada placa de TPU é conectada a uma máquina host de alto desempenho baseada em CPU para itens como carregamento e pré-processamento de dados. Para aproveitar ao máximo um número maior de TPUs, ajuste vários parâmetros de tarefa de treinamento. Neste documento, explicamos alguns problemas comuns e as alterações que você precisa fazer nos seus modelos.

Como escalonar o tamanho do lote e as etapas de treinamento

Para conseguir 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 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.

Muitos modelos incluem uma sinalização train_steps em que uma etapa corresponde a um único lote de dados. Quando aumentar o tamanho do lote, reduza o número de etapas de treinamento para que o número total de exemplos de treinamento permaneça o mesmo.

Por exemplo, se a configuração original era de 1024 para 1.000 etapas, o tamanho do lote poderia ser aumentado para 4096 e o número de etapas poderia ser reduzido para 250 (1.000 / 4). Caso o modelo use uma sinalização epochs, não é necessário 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 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.