Treinamento elástico: TPU cai e treino volta em 2 minutos
O Google matou de propósito uma máquina TPU durante o treinamento de um modelo de linguagem e o job não reiniciou. A falha de hardware chegou ao código como exceção Python, o mesmo processo se recuperou e o passo de treino seguinte saiu em menos de dois minutos. A maior parte desse tempo foi o Kubernetes agendando um pod substituto. O relato saiu no blog de desenvolvedores do Google, como tutorial de treinamento elástico com MaxText, Pathways e JAX em Cloud TPU. É fonte única, e a fonte é a dona da pilha.
O que é treinamento elástico
Treinamento elástico é a técnica que transforma a queda de uma máquina, no meio de um treino distribuído, em exceção que o código captura e trata. Sem ela, a queda encerra todos os processos e obriga a relançar o job a partir do último checkpoint.
O termo cobre coisas diferentes. No caso descrito pelo Google, significa que o processo de treino segue vivo, com config e imports já carregados, e as slices de TPU sobreviventes continuam de pé esperando instrução.
Uma máquina cai e leva o job inteiro junto
Treino distribuído reparte os pesos do modelo entre várias máquinas. A cada passo, cada uma calcula o gradiente da sua fatia e entra numa operação de all-reduce, em que todas trocam gradientes para manter o modelo em sincronia.
O all-reduce exige todos os participantes. Se uma máquina some, as outras ficam esperando dado que nunca chega. O timeout dispara, a operação coletiva falha e todos os processos encerram. Uma máquina derruba o job inteiro.
A correção padrão mora fora do código de treino. Um escalonador (Slurm, Kubernetes, Ray) percebe que o job caiu, realoca e relança tudo do zero. Você paga pods novos, containers e processos Python novos, reconexão com os aceleradores e aquecimento do carregador de dados. E perde cada passo desde o último checkpoint.
Um processo Python só, rodando na CPU
Na maioria dos lançadores de treino distribuído sobe um processo Python por nó, cada um com cópia idêntica do script, todos coordenando como iguais. É o padrão SPMD.
Com o Pathways existe exatamente um processo Python, numa máquina CPU comum, que enxerga todos os chips TPU do cluster como se fossem locais. Uma chamada a jax.devices() devolve todos. As máquinas TPU rodam apenas um binário fino de trabalhador, que recebe programas compilados e executa.
Quando uma máquina TPU morre, ainda há um processo Python vivo na CPU capaz de reagir. A falha de hardware vira exceção capturável em vez de processo encerrado. É disso que depende todo o resto.
Como a falha chega ao código
O Google descreve três componentes que precisam cooperar.
O Pathways detecta a queda por dois caminhos. Quando há operação em voo para o trabalhador morto, ela falha e o Pathways devolve DATA_LOSS. Quando não há nada em voo, o gerenciador de recursos, um container que roda ao lado do script no nó de CPU, percebe que o trabalhador parou de mandar heartbeat e devolve DEADLINE_EXCEEDED depois de cerca de 10 segundos. Nos dois casos o erro chega ao passo de treino como jax.errors.JaxRuntimeError.
O decorador elastic_retry, da biblioteca pathwaysutils, captura essa exceção específica. O MaxText já o aplica em volta da própria função de treino. Ele registra no log um aviso de slice caída, limpa estado parcial, restaura o último checkpoint viável e chama a função de treino de novo, dentro do mesmo processo e do mesmo PID.
O terceiro componente é o Orbax, que decide o que é seguro restaurar. Ele coordena a gravação a partir do controlador enquanto cada host TPU escreve sua fatia do estado direto no Cloud Storage, em paralelo.
O que a recuperação economiza, e o que ela não economiza
O próprio Google delimita o ganho.
Chamar a função de treino de novo refaz o setup do modelo, o carregador de dados e a restauração do checkpoint. São custos que um restart completo também pagaria. O agendamento de pod também não é de graça: o trabalhador que caiu ainda precisa de um substituto agendado na slice afetada, e essa espera domina o relógio.
O que a recuperação elástica corta é o entorno. Um restart completo derruba e reagenda a carga inteira: pod do controlador, todos os pods de trabalhador saudáveis e o processo Python novo que vem junto. A elástica deixa tudo isso de pé e troca só a slice que morreu.
Compilação não entra na economia. O Pathways mantém cache persistente de compilação no Cloud Storage, ligado por padrão, então um restart completo recarrega os executáveis XLA do cache em vez de recompilar do zero, e a recuperação elástica paga custo comparável ao reentrar na função de treino. A diferença entre os dois caminhos é o teardown, não o compile.
O restart completo derruba, e a recuperação elástica preserva:
- Pods de controlador e de trabalhadores saudáveis. O restart reagenda todos. A recuperação elástica mantém de pé.
- Processo Python do controlador. O restart sobe um novo. A recuperação elástica segue no mesmo, com o mesmo PID.
Os dois caminhos pagam igual:
- Setup do modelo, carregador de dados e restauração do checkpoint.
- Agendamento do pod substituto na slice que caiu.
- Carga dos executáveis XLA a partir do cache de compilação.
O número comparativo é frouxo
O post não publica tabela de benchmark. Descreve a diferença entre restart completo e recuperação elástica como sendo entre "centenas de segundos" e "vários minutos". As duas faixas se sobrepõem, e ninguém calcula economia em cima disso.
O único número fechado é o da corrida do Google: menos de dois minutos entre o kill da TPU e o passo de treino seguinte, com o agendamento do pod substituto consumindo a maior parte desse intervalo. É uma medição, de uma corrida só, feita por quem vende a plataforma.
Pausar ou encolher
Duas estratégias estão no MaxText hoje.
Pause and resume, a que o tutorial percorre, captura a exceção, espera a slice falha ser substituída, recarrega o último checkpoint viável e continua na malha completa.
Replica resize recarrega o checkpoint imediatamente nas slices sobreviventes e segue treinando com throughput menor enquanto o substituto sobe, voltando ao tamanho cheio depois. Oferece o que restart nenhum oferece: o treino continua mesmo que as TPUs perdidas nunca voltem.
Suspend-resume e treinamento elástico não são a mesma coisa
Suspend-resume e treinamento elástico resolvem problemas diferentes. O primeiro trata interrupção avisada; o segundo, falha que chega sem aviso nenhum.
Em Spot TPU, o suspend-resume do Pathways cuida da preempção planejada: escuta o aviso, salva o estado do acelerador no Cloud Storage e retoma quando o Kubernetes reagenda o pod, sem exigir código do usuário. Os nomes se confundem porque uma das modalidades elásticas também se chama pause and resume.
As lacunas do material
Uma fonte só, e ela é parte interessada. O relato é do Google sobre a pilha do Google. Nenhum veículo independente repercutiu ou reproduziu o teste até agora.
O material apurado também não traz o tamanho do modelo treinado, quantos chips estavam na corrida, custo, nem se existe mecanismo equivalente fora de Cloud TPU. A data de publicação não veio no dossiê. O trecho sobre o Orbax está incompleto no material que chegou.
MaxText e JAX são abertos. O controlador único que sustenta o mecanismo é Pathways rodando em TPU no GKE, com Cloud Storage guardando checkpoint e cache de compilação. São quatro hosts por slice na configuração descrita.
O que muda para quem decide aqui
Para equipe brasileira que treina ou faz fine-tune em acelerador alugado, queda de nó é linha de custo recorrente. Quanto maior a corrida, maior a chance de uma máquina cair, e cada queda vira hora de acelerador ocioso somada aos passos perdidos desde o último checkpoint.
O Google desloca o item que se mede nessa conta. Compilação já sai resolvida por cache nos dois caminhos. Sobra o teardown: quanto da carga sobrevive à morte de uma máquina.
A pergunta é de contrato antes de ser de engenharia. Na escolha entre TPU no Google Cloud, GPU em nuvem e cluster próprio, pergunte ao fornecedor o que acontece com o job quando um nó morre no meio do treino. Se a resposta for relançar tudo, isso tem preço, e ele aparece na fatura de acelerador ocioso.
--- Apurado em 1 veículo de 1 país. A única fonte é o blog de desenvolvedores do Google, parte interessada no assunto. O dossiê registrava 2 veículos; só 1 trazia texto sobre esta história.