Modelo de custo do LHS

tldr;

Esta página descreve os aspectos internos do modelo de custo usado pelo Latency Hiding Scheduler. Se você quiser ajustar o modelo, vá direto para a seção de ajuste.

O Latency Hiding Scheduler (LHS) é uma transmissão do compilador que agenda um DAG de HLO de maneira a minimizar o tempo decorrido.

As decisões são orientadas pelo modelo de custo unificado, que usa uma combinação de tabelas de desempenho e modelos analíticos. Em especial, o XLA incorpora tabelas de desempenho para GEMMs e coletivos de interconexão rápida, além de usar um modelo analítico de rede e custo de fusão para outros casos. O restante do documento descreve o funcionamento interno delas em um nível mais detalhado.

Tabelas de performance: coletivos da ICI

A tabela de performance consiste em dois componentes principais: um coletor e um interpolador.

Coletor

O coletor é uma ferramenta em C++ responsável por gerar as tabelas de desempenho para operações coletivas. Ele mede o desempenho de operações individuais de HLO (por exemplo, all-gather, all-reduce) em um espaço de parâmetros definido estaticamente.

Como funciona

A ferramenta faz uma varredura em uma variedade de operações coletivas, tamanhos de transferência e esquemas de transferência para um determinado cluster. Ele usa a infraestrutura de execução de HLO multihost e os dados ExecutionProfile para executar o HLO gerado e coletar métricas de desempenho.

Parâmetros de coleta de dados

As tabelas de latência são coletadas para um produto cruzado dos seguintes parâmetros:

  • Tipo coletivo:
    • all-reduce
    • all-gather
    • reduce-scatter
  • Tamanho da transferência:
    • Escala logarítmica de 1024B até 2GiB (por exemplo, 1024B, 2048B, 4096B, ...)
  • Esquema de transferência:
    • rail-aligned
    • non-rail-aligned

Essa varredura é executada para clusters intranós com 2, 4 e 8 dispositivos.

Saída

O resultado de uma execução de coleta é uma tabela de latência no formato .pbtxt (aproximadamente 116 KB por plataforma).

Interpolator

O interpolador é o componente do compilador que consome as tabelas de desempenho geradas para fornecer estimativas de tempo de execução durante a compilação.

Estrutura de dados interna

Na inicialização, o interpolador processa a tabela de performance em um mapa. Esse mapa usa uma tupla de (collective_type, transfer_scheme) como chave.

O valor associado a cada chave é um plano euclidiano 2D. Esse plano indexa a taxa de transferência de rede (medida pelo coletor) com base em dois eixos:

  1. Tamanho da transferência.
  2. Número de dispositivos envolvidos.

Pesquisa e interpolação

Quando o compilador encontra uma operação coletiva, o interpolador executa as seguintes etapas:

  1. Ele identifica o plano de taxa de transferência 2D correto usando o (collective_type, transfer_scheme) da operação como a chave do mapa.
  2. Em seguida, ele usa uma recuperação de média ponderada (com base na distância euclidiana) nesse plano 2D, usando o (transfer_size, num_devices) da operação como o ponto de consulta.
  3. O resultado dessa pesquisa é um único valor exclusivo de taxa de transferência de rede.

Justificativa: capacidade de processamento e extrapolação

O sistema foi projetado para armazenar o throughput da rede, e não a latência bruta. Essa escolha de design simplifica significativamente a extrapolação do desempenho para tamanhos de transferência não explicitamente presentes na tabela.

Se as tabelas de latência capturarem a saturação da largura de banda da rede em um tamanho coletivo S, a capacidade T nesse ponto será considerada o máximo. Para qualquer novo coletivo de tamanho S' > S, o tempo de execução pode ser estimado como:

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

Isso permite que o modelo estime o desempenho de coletivos de qualquer tamanho, mesmo aqueles maiores que o máximo de 2 GiB medido pelo coletor.

  • Subestime a capacidade máxima de processamento.
  • Consequentemente, superestimam o tempo de execução para transferências grandes.

Em geral, as equipes do XLA:GPU mantêm tabelas de performance, mas, caso o usuário decida fornecer as próprias tabelas, é responsabilidade dele garantir que elas sejam representativas e incluam medições na região saturada por largura de banda para o hardware de destino.

Tabelas de performance – GEMMs

Semelhante ao sistema de coletivos, as tabelas de latência GEMM são compatíveis com dois componentes: um coletor e um interpolador.

Coletor

O coletor é uma ferramenta em C++ que calcula tabelas de desempenho para multiplicações de matrizes gerais (GEMMs, na sigla em inglês). Ele mede o desempenho das multiplicações de matriz no nível da operação dot do HLO.

Como funciona

A ferramenta faz uma varredura em um espaço estático de dimensões da GEMM (lote, duas dimensões não contráteis e uma contrátil) e tipos de dados.

  • Tipos de dados padrão:LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32.
  • Infraestrutura:reutiliza o criador de perfis de operações HLO.

Parâmetros de coleta

As tabelas de latência são coletadas para um produto cruzado das seguintes dimensões:

  • batch: {1, 2, 4}
  • m (não contrátil): {256, 512, ..., 4096}
  • n (não contraente): {256, 512, ..., 4096}
  • k (contração): {256, 512, ..., 4096}

Saída e armazenamento

Uma varredura completa gera uma tabela de latência .pbtxt, pronta para ser consumida pelo interpolador.

Interpolator

O interpolador é o componente do compilador que usa as tabelas geradas para estimar o desempenho da GEMM.

Justificativa: saturação de FLOPS

As tabelas de latência coletadas permitem que o interpolador reconstrua FLOPS para cada entrada:

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

Um insight importante é que os FLOPS saturam em um determinado ponto. Ou seja, o hardware atinge o pico de FLOPS além de um determinado formato de matriz. Essa saturação permite o uso do mesmo método de extrapolação empregado para coletivos.

Pesquisa e interpolação

O interpolador cria um espaço euclidiano 4D com base nos dados da tabela. Para fornecer uma estimativa de performance, ele realiza uma interpolação de média ponderada nesse espaço 4D. Se não houver uma tabela para um determinado tipo de dados, como uma heurística, cada dimensão será normalizada para o número de bytes.

Modelo de custo analítico: DCN

Modelo de custo coletivo de curva em S

O modelo de curva S é um modelo de teto de rede totalmente analítico.

Visão geral

Ele foi projetado para estimar a performance de operações coletivas com base em um conjunto de propriedades de rede fixas.

Entradas do modelo

O modelo requer duas categorias de entradas:

  1. Propriedades de rede fixa (definidas pelo usuário):

    • Sobrecarga de inicialização coletiva
    • Velocidade da placa de rede
    • RTT (tempo de retorno)

    Por padrão, o XLA detecta automaticamente uma plataforma e usa valores para as arquiteturas mais comuns. Essas propriedades podem ser configuradas pelo usuário. Consulte a seção de ajuste para mais detalhes.

  2. Entradas por coletivo:

    • Tipo coletivo (por exemplo, AllGather, ReduceScatter)
    • Tamanho da transferência
    • Número de nós envolvidos na comunicação

Integração

O modelo de curva em S está integrado ao XLA:GPU e sendo usado no Hopper e no Blackwell.

Modelo de custo analítico: fusões

Para outros kernels, usamos o modelo de custo de desempenho da GPU para estimar os tempos de execução corretos. Saiba mais neste artigo.

Ajuste

O modelo de curva S pode ser ajustado emitindo flags XLA corretas. A configuração padrão é suficiente na maioria dos casos, mas o controle do modelo é exposto em outros casos.

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"