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 pode ter até 2.048 núcleos, permitindo 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. Neste documento, explicamos 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.

Muitos modelos incluem uma sinalização train_steps em que uma etapa corresponde a 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.

Como usar um NFS para armazenar dados

Quando você cria um Pod de TPU, uma VM separada é criada para cada nó de TPU. Por padrão, cada VM de TPU recebe um ID de usuário (UID, na sigla em inglês) diferente. Isso causa problemas quando você tenta acessar o mesmo diretório NFS de vários nós. O mesmo diretório teria proprietários diferentes em nós diferentes e as permissões padrão do Linux não se aplicam aos nós. Por exemplo, um processo em um nó não pode gravar no diretório de registros criado por outro nó.

Para contornar esse problema, ative o Login do SO. O Login do SO pode ser configurado para uma instância de VM específica ou para um projeto. Para mais informações, consulte Como configurar o Login do SO.

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 a menor 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, escalone para 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.