Treinamento em TPU

A TPU (Tensor Processing Unit) é um tipo de hardware projetado especificamente para treinar modelos de forma eficiente. O Accelerate suporta treinamento em TPU, mas há alguns pontos que você deve ter em mente, nomeadamente a compilação de grafos. Este tutorial discute brevemente a compilação e, para mais detalhes, consulte o guia Training on TPUs with Accelerate.

Compilação

A TPU cria um grafo de todas as operações na etapa de treinamento, como a passagem direta, a passagem reversa e a etapa do otimizador. É por isso que a primeira etapa de treinamento sempre leva um tempo, pois a construção e a compilação desse grafo consomem tempo. Mas, uma vez que a compilação é concluída, ela é armazenada em cache e todas as etapas subsequentes são muito mais rápidas.

A chave é evitar recompilar seu código, caso contrário o treinamento fica extremamente lento. Isso significa que todas as suas operações devem ser exatamente as mesmas:

• todos os tensores em seus lotes devem ter o mesmo comprimento (por exemplo, sem preenchimento dinâmico para tarefas de PLN)
• seu código deve ser estático (por exemplo, sem camadas com loops for que tenham comprimentos diferentes dependendo da entrada, como uma LSTM)

Vinculação de pesos

Um design comum para modelos de linguagem é atrelar os pesos das camadas de incorporação (embedding) e softmax. No entanto, mover o modelo para um TPU (seja por sua conta ou passando-o para o método prepare()) quebra o atrelamento dos pesos e você precisará atrelá-los novamente.

Para adicionar um comportamento especial (como a amarração de pesos) no seu script para TPUs, defina distributed_type como DistributedType.TPU primeiro. Em seguida, você pode usar o método tie_weights para amarrar os pesos.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()