Pular para o conteúdo principal

Sistemas Distribuídos e TPU

O sucesso do AlphaGo não foi apenas uma vitória de algoritmos, mas também uma vitória de engenharia. Para treinar uma IA de Go que supera os humanos em tempo razoável, foi necessário um sistema distribuído cuidadosamente projetado e suporte de hardware especializado.

Este artigo analisará profundamente a arquitetura de sistemas por trás do AlphaGo, incluindo fluxo de treinamento, arquitetura de inferência, MCTS paralelo e o papel crucial das TPUs.


Visão Geral da Arquitetura de Treinamento

Arquitetura de Treinamento do AlphaGo Original

A arquitetura de treinamento do AlphaGo original (a versão que derrotou Lee Sedol) foi dividida em múltiplas fases, cada uma usando uma configuração de recursos diferente:

Arquitetura de Treinamento do AlphaGo Zero

O AlphaGo Zero simplificou significativamente o processo de treinamento, usando um único ciclo de treinamento de ponta a ponta:

As vantagens desta arquitetura:

  1. Aprendizado contínuo: Self-play e Training acontecem simultaneamente, sem necessidade de esperar
  2. Eficiência de recursos: Todos os recursos estão fazendo trabalho útil
  3. Iteração rápida: A rede atualizada é imediatamente usada para gerar novos dados

Workers de Auto-Jogo (Self-play Workers)

Distribuição de Tarefas

Os Self-play Workers são responsáveis por realizar auto-jogo com a rede mais forte atual, gerando dados de treinamento.

ConfiguraçãoAlphaGo Zero
Número de WorkersDezenas
Por Worker1-4 TPU
MCTS por jogo1600 simulações
Produção diária~100.000 jogos

Fluxo de Trabalho

O fluxo de trabalho de cada Self-play Worker:

while True:
# 1. Baixar os pesos da rede mais recente
network = download_latest_checkpoint()

# 2. Realizar múltiplos jogos de auto-jogo
for game in range(batch_size):
positions = []
board = EmptyBoard()

while not board.is_terminal():
# Executar MCTS
mcts = MCTS(network, board)
policy = mcts.search(num_simulations=1600)

# Selecionar jogada
action = sample(policy)

# Registrar
positions.append((board.state, policy))

# Jogar
board = board.play(action)

# 3. Obter resultado do jogo
result = board.get_result()

# 4. Fazer upload dos dados
upload_to_replay_buffer(positions, result)

Balanceamento de Carga

Múltiplos Workers precisam de balanceamento de carga:

  • Sincronização de rede: Todos os Workers usam a mesma versão da rede
  • Balanceamento de dados: Garantir que os dados de diferentes Workers sejam todos utilizados
  • Tratamento de falhas: Falha de um único Worker não afeta o treinamento geral

Workers de Treinamento (Training Workers)

Distribuição de Tarefas

Os Training Workers são responsáveis por amostrar dados do Replay Buffer e treinar a rede neural.

ConfiguraçãoAlphaGo Zero
Número de Workers1-4
Por Worker4 TPU
Batch Size2048 (512 por TPU)
Passos de treinamentoDezenas de milhares por dia

Treinamento Distribuído

O treinamento em larga escala usa Paralelismo de Dados (Data Parallelism):

Cada TPU processa um mini-batch diferente, calcula o gradiente local e então agrega para atualizar os parâmetros globais.

Atualização Síncrona vs. Assíncrona

Modo de atualizaçãoVantagensDesvantagens
SíncronoEstável, reproduzívelWorkers precisam esperar pelo mais lento
AssíncronoAlto throughputGradientes podem estar desatualizados

O AlphaGo Zero usa atualização síncrona para garantir a estabilidade do treinamento.


O Papel das TPUs

O Que é uma TPU?

TPU (Tensor Processing Unit) é um acelerador projetado especificamente pela Google para deep learning:

CaracterísticaTPUGPUCPU
Objetivo de designOperações matriciaisParalelismo geralComputação geral
PrecisãoOtimizado para FP16/BF16FP32/FP16FP64/FP32
Consumo de energiaRelativamente baixoMais altoMais alto
LatênciaBaixaMédiaAlta

Arquitetura da TPU

O núcleo da TPU é a MXU (Matrix Multiply Unit):

A MXU pode executar 16K operações de multiplicação-acumulação por ciclo, o que é crucial para a multiplicação de matrizes em redes neurais.

Por Que o AlphaGo Precisa de TPUs?

O gargalo computacional da IA de Go está na inferência da rede neural:

OperaçãoProporção
Forward pass da rede neural~95%
Operações da árvore MCTS~4%
Outros~1%

Cada passo do MCTS requer 1600 inferências da rede neural. O alto throughput das TPUs torna isso possível.

Evolução do Uso de TPUs

VersãoTPU para TreinamentoTPU para Inferência
AlphaGo Lee50 GPU48 TPU (v1)
AlphaGo Master4 TPU (v2)4 TPU (v2)
AlphaGo Zero4 TPU (v2)4 TPU (v2) (escalável)

O número de TPUs usadas pelo AlphaGo Zero diminuiu significativamente, graças a uma arquitetura mais eficiente e versões mais recentes das TPUs.


MCTS Paralelo e Perda Virtual

O Desafio da Paralelização

A implementação padrão do MCTS é serial:

for i in range(num_simulations):
1. Selection: Selecionar da raiz para baixo
2. Expansion: Expandir nó folha
3. Evaluation: Avaliação da rede neural
4. Backup: Propagar atualização para cima

Mas a avaliação da rede neural é uma operação em batch amigável para GPU/TPU. Como fazer múltiplas simulações ocorrerem simultaneamente?

Paralelização de Folhas (Leaf Parallelization)

A forma mais simples de paralelização: executar múltiplas simulações completas simultaneamente, depois mesclar os resultados.

Problema: Cada simulação começa da raiz, e pode explorar repetidamente os mesmos caminhos.

Perda Virtual (Virtual Loss)

A DeepMind adotou a técnica de perda virtual para implementar Paralelização de Árvore (Tree Parallelization).

Conceito Básico

Quando uma thread está explorando um determinado nó, reduz temporariamente o valor desse nó, fazendo com que outras threads escolham outros caminhos.

UCB normal: Q(s,a) + c * P(s,a) * sqrt(N(s)) / (1 + N(s,a))

Com perda virtual:
(Q(s,a) * N(s,a) - v * n_virtual) / (N(s,a) + n_virtual) + c * P(s,a) * sqrt(N(s)) / (1 + N(s,a) + n_virtual)

Onde:

  • n_virtual é o número de threads explorando esse nó
  • v é o valor da perda virtual (geralmente 1 ou o valor correspondente à taxa de vitória)

Fluxo de Operação

Tempo T1:
Thread 1 seleciona caminho A → B → C
Nó C recebe perda virtual -1

Tempo T2:
Thread 2 seleciona caminho A → B → D (porque C foi "penalizado")
Nó D recebe perda virtual -1

Tempo T3:
Thread 1 completa avaliação, atualiza valor real de C, remove perda virtual
Thread 3 agora pode escolher C (se o valor real for bom o suficiente)

Efeitos da Perda Virtual

AspectoEfeito
Diversidade de exploraçãoForça exploração de diferentes caminhos
Eficiência de batchPode avaliar múltiplos nós folha simultaneamente
ConvergênciaPerda virtual é eventualmente substituída pelo valor real, não afeta convergência

Avaliação de Rede Neural em Batch

Através da perda virtual, múltiplos nós folha a serem avaliados podem ser coletados para inferência em batch:

A eficiência da inferência em batch da TPU é muito maior que a inferência individual, tornando o MCTS paralelo possível.


Arquitetura de Inferência

Configuração Durante Partidas

A arquitetura de inferência do AlphaGo durante partidas oficiais:

VersãoConfiguração de Hardware
AlphaGo Fan176 GPU
AlphaGo Lee48 TPU + múltiplos servidores
AlphaGo Master4 TPU
AlphaGo Zero4 TPU (escalável)

Fluxo de Inferência Distribuída

Fluxo de inferência durante partidas (usando AlphaGo Lee como exemplo):

Gerenciamento de Tempo de Reflexão

Estratégia de gerenciamento de tempo do AlphaGo:

PosiçãoTempo de ReflexãoSimulações MCTS
Abertura (tem joseki)Mais curto~10.000
Meio-jogo (complexo)Mais longo~100.000
Posição simplesMais curto~5.000
ByoyomiFixo~1.600

Mais simulações MCTS geralmente significam jogadas de melhor qualidade.


Comunicação e Sincronização

Formato de Dados

Formato de transmissão de dados de treinamento:

message TrainingExample {
// Estado do tabuleiro (17 × 19 × 19)
repeated float board_planes = 1;

// Resultado da busca MCTS (362)
repeated float mcts_policy = 2;

// Resultado do jogo (1 = jogador atual vence, -1 = jogador atual perde)
float game_result = 3;
}

Requisitos de Largura de Banda de Rede

Fluxo de dadosTamanhoFrequência
Amostras de treinamento~10 KB/amostraMilhares de amostras/segundo
Pesos da rede~200 MBVárias vezes/hora
Mensagens de controle< 1 KBContínuo

Requisito total de largura de banda: ~100 Mbps (rede interna é suficiente)

Tratamento de Falhas

Tratamento de falhas em sistemas distribuídos:

Tipo de falhaForma de tratamento
Worker caiuReinicia, continua usando checkpoint mais recente
Desconexão de redeBuffer de dados, continua transmissão após reconexão
Falha de TPUTroca automática para TPU reserva
Dados corrompidosDescarta após verificação, regenera

Análise de Custos

Estimativa de Custos de Hardware

Estimativa de custos de treinamento do AlphaGo Zero baseada nos preços de TPU do Google Cloud:

RecursoQuantidadePreço/horaTotal/dia
TPU v2 Pod4~$32~$3.000
VM alta memóriaVários~$5~$500
Espaço de armazenamento10 TB~$0.02/GB~$200
Rede-Incluído-

Aproximadamente 3.700/dia,treinamentocompleto(40dias)aproximadamente3.700/dia**, treinamento completo (40 dias) aproximadamente **150.000.

Nota: Esta é uma estimativa de 2017, a DeepMind como subsidiária do Google pode ter descontos internos.

Comparação com Treinamento Humano

AspectoAlphaGo ZeroJogador profissional humano
Tempo para nível profissional2 dias10-15 anos
Custo de treinamento~$7.500Milhões (mensalidades, custo de vida, custo de oportunidade)
Custos contínuosEletricidadeCusto de vida
ReplicabilidadeRéplica perfeitaNão replicável

Claro, esta comparação não é totalmente justa — os humanos aprendem muito mais do que apenas Go durante seu treinamento.

Custos de Inferência

Custos de inferência em partidas oficiais:

ConfiguraçãoCusto por jogo
48 TPU (AlphaGo Lee)~$500
4 TPU (AlphaGo Zero)~$50
GPU única (KataGo)~$1

Os custos de inferência diminuíram significativamente com o avanço tecnológico.


Evolução Tecnológica

De AlphaGo a AlphaZero

AspectoAlphaGo LeeAlphaGo ZeroAlphaZero
TPU de treinamento50+ GPU → TPU4 TPU4 TPU
TPU de inferência48 TPU4 TPU4 TPU
MCTS/jogada~100.000~1.600~800
Tempo de treinamentoMeses40 diasHoras a dias

Melhoria de eficiência de aproximadamente 100 vezes.

Impacto na Comunidade Open Source

A arquitetura do AlphaGo inspirou múltiplos projetos open source:

ProjetoCaracterísticas
Leela ZeroTreinamento distribuído comunitário, replica AlphaGo Zero
KataGoTreinamento eficiente em GPU única, supera AlphaGo Zero
ELF OpenGoOpen source do Facebook, usa PyTorch
MinigoOpen source do Google, usa TensorFlow

Esses projetos permitem que pesquisadores comuns também possam treinar IAs de Go poderosas.


Correspondência com Animações

Os conceitos principais deste artigo e números de animação correspondentes:

NúmeroConceitoCorrespondência Física/Matemática
🎬 C9MCTS paraleloProblema de N-corpos
🎬 E9Treinamento distribuídoComputação distribuída
🎬 C5Perda virtualPotencial de repulsão
🎬 D15Inferência em batchComputação vetorizada

Leitura Adicional


Referências

  1. Silver, D., et al. (2017). "Mastering the game of Go without human knowledge." Nature, 550, 354-359.
  2. Jouppi, N., et al. (2017). "In-Datacenter Performance Analysis of a Tensor Processing Unit." ISCA 2017.
  3. Dean, J., et al. (2012). "Large Scale Distributed Deep Networks." NeurIPS 2012.
  4. Chaslot, G., et al. (2008). "Parallel Monte-Carlo Tree Search." CIG 2008.
  5. Segal, R. (2010). "On the Scalability of Parallel UCT." CIG 2010.