AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers

Resumo Técnico: AoiZora

Definição do Problema
A inferência de difusão de vídeo, particularmente para modelos como o Wan 2.1, tornou-se uma carga de trabalho generativa crítica. No entanto, o processo de inferência é computacionalmente intensivo, exigindo centenas de milhares de tokens espaço-temporais e etapas de denoising repetidas sobre grandes tensores de ativação. Isso necessita de inferência distribuída através de múltiplos aceleradores. Os atuais sistemas de auto-paralelismo (por exemplo, baseados em JAX/XLA) buscam efetivamente malhas de dispositivos lógicos otimizadas (determinando como os tensores são fragmentados), mas ignoram amplamente o posicionamento físico desses ranks lógicos no hardware.

Em sub-fatias de TPU (subconjuntos retangulares contíguos de um pod de TPU), a interconexão entre chips (ICI) é uma malha de vizinhos próximos sem links de contorno (wrap-around) em todas as dimensões. Consequentemente, a mesma estratégia de fragmentação lógica pode resultar em latências finais significativamente diferentes dependendo de como os ranks lógicos são mapeados para os chips físicos. O artigo observa que este "efeito de posicionamento" sozinho pode causar uma variação de até 16,6% na latência de denoising de um passo em uma fatia v5e-8, primariamente devido à contenção de links compartilhados, em vez de saltos (hop count). Os sistemas existentes vinculam os ranks lógicos aos dispositivos físicos precocemente no fluxo de trabalho (frequentamente via ordenação padrão) antes que o compilador exponha o grafo concreto de comunicação coletiva, deixando o desempenho dependente da topologia não otimizado.

Metodologia: Um Planejador Mediado pelo Compilador em Dois Estágios
O AoiZora aborda essa lacuna posicionando o planejamento consciente da topologia na fronteira entre a análise do compilador e a execução. Ele aproveita a "assimetria de custo/fidelidade" do pipeline de compilação JAX/XLA: representações intermediárias (IRs) iniciais são baratas, mas carecem de detalhes concretos de comunicação, enquanto HLOs totalmente compilados são caros, mas expõem o grafo de comunicação e grupos de dispositivos exatos.

O planejador opera em dois estágios distintos:

1. Estágio 1: Poda de Fragmentação Lógica Pré-Compilação

   • Objetivo: Eliminar eficientemente candidatos de fragmentação lógica ineficientes.
   • Mecanismo: O AoiZora enumera atribuições de eixos de carga de trabalho legais (ex: paralelismo de tensor, sequência, dados) e as reduz para as IRs StableHLO/Shardy.
   • Análise: Ele constrói um DAG de predição que inclui comunicações explícitas e nós de comunicação implícita sintéticos (ex: para resharding, reduções fragmentadas e fronteiras de kernels opacos) para estimar o volume de comunicação lógica.
   • Pontuação: Utiliza um limite de vazão independente de posicionamento ($Q_1$) para pontuar candidatos. Este estágio filtra fragmentações com baixa pegada de memória ou comunicação lógica excessiva, retendo apenas um pequeno conjunto "Top-K" de sobreviventes.

2. Estágio 2: Classificação de Topologia Pós-Compilação

   • Objetivo: Selecionar o posicionamento físico ideal para as fragmentações lógicas sobreviventes.
   • Mecanismo: Os Top-K candidatos lógicos são totalmente compilados pelo XLA para gerar HLOs concretos.
   • Análise: O sistema enumera distintas ordens de eixos físicos (mapeando eixos lógicos para coordenadas de sub-fatia de TPU). Ele analisa o HLO compilado para extrair o grafo de comunicação exato, incluindo grupos de réplica, tamanhos de carga útil (payload) e dependências.
   • Pontuação: Um modelo de custo consciente da topologia ($Q_2$) classifica cada par fragmentação-posicionamento. Este modelo reconstrói rotas físicas para coletivos, estima cargas de link e penaliza posicionamentos que causem contenção em links de ICI compartilhados.
   • Seleção: O par com o menor custo estimado é selecionado.

Principais Contribuições

• Caracterização da Sensibilidade de Posicionamento: Os autores demonstram que, em sub-fatias de TPU, o posicionamento físico é uma variável de otimização de primeira classe. Eles mostram que fragmentações lógicas idênticas podem gerar até 16,6% de diferença de latência devido à contenção de links compartilhados.
• Busca Consciente da Fidelidade do Compilador: O artigo identifica e explora a assimetria entre a IR de pré-compilação barata (adequada para poda de alto recall) e o HLO compilado caro (necessário para classificação de topologia precisa).
• Design do Sistema AoiZora: Um planejador mediado por compilador que separa a poda lógica da classificação física. Crucialmente, ele se integra ao stack padrão JAX/XLA sem modificar o código do modelo, as regras de redução, as implementações de kernels coletivos ou o roteamento de rede. Ele produz uma estratégia de fragmentação e uma ordem de eixo físico que o compilador padrão então execza.
• Avaliação: Experimentos extensos no Wan 2.1 através de sub-fatias de TPU v5e (4, 8 e 16 chips) e várias resoluções de vídeo (480p/720p) e contagens de frames.

Resultados Experimentais
Avaliado no Wan 2.1, o AoiZora supera consistentemente duas linhas de base:

1. Implantação Existente: Uma estratégia de paralelismo configurada por especialistas e ajustada manualmente (usando FSDP, paralelismo de sequência e de tensor) com posicionamento físico fixo.
2. Linha de Base de Busca Lógica: Um sistema inspirado no Alpa que otimiza a fragmentação lógica usando um modelo de custo abstrato, mas que não realiza a classificação de posicionamento físico pós-compilação.

• Ganhos de Desempenho: O AoiZora alcança acelerações de até 1,42× sobre a solução existente e 1,37× sobre a linha de base de busca lógica. A média geométrica de aceleração é de 1,09× e 1,11×, respectivamente.
• Fidelidade: A etapa de pré-compilação retém com sucesso a melhor fragmentação lógica medida em todos os cenários testados (100% de recall top-1), enquanto a etapa pós-compilação atinge uma alta correlação de rank ($\rho_s$ até 0,985) entre a latência prevista e a medida.
• Sobrecarga (Overhead): O custo de planejamento offline é prático, variando de ~47 segundos (v5e-4) a ~376 segundos (v5e-16), o que é amortizado ao longo de solicitações de inferência repetidas da mesma forma.

Significância e Alegações
O artigo afirma que o AoiZora preenche a lacuna entre o paralelismo lógico e a topologia física da TPU sem exigir mudanças no stack de serviço subjacente. Ao adiar as decisões de posicionamento físico até que o compilador materialize o grafo de comunicação concreto, o AoiZora evita a "omissão" dos atuais sistemas de auto-paralelismo que tratam malhas lógicas como abstrações completas.

Os autores posicionam o AoiZora como uma solução prática para as restrições específicas de sub-fatias de TPU, onde a perda de links de contorno torna o posicionamento consciente da topologia crítico. Eles afirmam que sua abordagem permite que desenvolvedores de modelos permaneçam agnósticos aos detalhes de topologia enquanto alcançam latência de inferência próxima da ótima para cargas de trabalho de difusão de vídeo. O trabalho sugere que, para arquiteturas de DiT de vídeo densas com formas de tensores estáticas, um passo de planejamento offline único pode melhorar significativamente o desempenho de serviço ao resolver os problemas de contenção inerentes aos posicionamentos físicos fixos.