Transferência de Pesos para Pós-Treinamento de RL em menos de 2s

Recentemente, alcançamos atualizações de parâmetros entre máquinas em 1,3 segundos para o Kimi-K2 (1T parâmetros), transferindo pesos de 256 GPUs de treinamento (BF16) para 128 GPUs de inferência (FP8).

No ajuste fino de aprendizado por reforço assíncrono, o treinamento e a inferência rodam em GPUs separadas. Após cada etapa de treinamento, novos pesos devem ser enviados para os nós de inferência. Muitas estruturas existentes levam vários segundos — ou até minutos — para modelos de um trilhão de parâmetros.

Ao aproveitar a comunicação ponto a ponto RDMA, conseguimos tornar a transferência de pesos extremamente rápida, sem alterar o motor de inferência, e tornar o código mais fácil de escrever e manter.

Nossa solução é baseada em RDMA WRITE, uma primitiva unilateral onde a origem grava diretamente na memória da GPU do destino.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

O lado de destino nem sequer será notificado sobre a transferência. Isso nos proporciona transferências de baixa latência, alto rendimento e zero-copy, impulsionadas pelos nós de treinamento, sem qualquer lógica de controle nos nós de inferência.

Coleta de metadados – O controlador reúne metadados de parâmetros de todas as GPUs de treinamento e inferência.

Computação de agendamento – O controlador calcula um agendamento estático de transferência de pesos, mapeando qual GPU de treinamento envia qual parâmetro para qual GPU de inferência e em que ordem.

Distribuição de agendamento – O controlador envia o agendamento para todas as GPUs de treinamento.

Execução – Após cada etapa de treinamento, o controlador sinaliza para as GPUs de treinamento iniciarem as transferências.

Com o fluxo de trabalho de alto nível definido, o principal desafio é como executar transferências de pesos de forma eficiente em escala de trilhões de parâmetros. Aqui, descrevemos os detalhes do caminho de execução.

Os parâmetros no treinamento são distribuídos de acordo com os posicionamentos FSDP. Usando full_tensor(), todas as GPUs em um DeviceMesh podem reconstruir o parâmetro completo, servindo assim como origem para a transferência de pesos.

Vários DeviceMeshes disjuntos formam um mesh group. Como os DeviceMeshes no mesmo grupo são disjuntos, suas transferências não interferem e podem rodar totalmente em paralelo. Entre os mesh groups, inserimos uma barreira global para impor a ordenação.

O pipeline de transferência de pesos sobrepõe cópias do host para o dispositivo, montagem de tensores, fusão e quantização, e transferências RDMA entre MeshGroups.

Tratamos a transferência de cada tensor de parâmetro como uma tarefa. O processo de transferência de pesos utiliza vários tipos de fontes de hardware, portanto dividimos uma tarefa de transferência de pesos em diferentes estágios de pipeline que se sobrepõem no tempo:

Cópia de host para dispositivo (memcpy) — Se o FSDP descarregar o peso para a CPU

Preparação de parâmetros — Reconstruir o peso completo com full_tensor(), aplicar fusão de projeção e quantizar, se necessário.

Transferência RDMA — Gravação zero-copy na memória remota da GPU de inferência

Barreira global — Após a conclusão de todas as chamadas de full_tensor(), sincronize entre os mesh groups usando GLOO via Ethernet.

Na implementação, mantemos uma fila FIFO de tarefas para cada estágio de pipeline. Sempre que a tarefa na cabeça da fila conclui o estágio, ela é movida para a cauda da fila do próximo estágio.

full_tensor() e outras operações de GPU introduzem uso extra de memória da GPU. Para evitar erros de falta de memória (OOM), iniciamos a execução de uma tarefa apenas se as tarefas atuais em andamento ocuparem menos memória temporária da GPU do que um limite configurável.

Várias escolhas de design tornam nosso sistema significativamente mais rápido para executar e mais fácil de manter do que soluções comuns de código aberto.

A transferência de pesos baseada em Rank0 reúne os pesos antes de distribuí-los para o nó de inferência.

A comunicação ponto a ponto transfere os pesos de ambos os nós de treino para as interfaces de rede do nó de inferência.

Um padrão comum é canalizar todos os parâmetros através de GPUs de rank-0: reunir no rank-0 de treinamento, enviar para o rank-0 de inferência e, em seguida, espalhar novamente. Isso rapidamente se torna um gargalo, limitado pela largura de banda PCIe de uma única GPU e pela placa de rede (ex.: 400 Gbps ≈ 50 GB/s).

Em contraste, nossa configuração ponto a ponto permite que cada GPU de treinamento envie diretamente para cada GPU de inferência, saturando toda a estrutura de rede em vez de um único link.

A transferência baseada em RPC usa um buffer no lado da inferência e chamadas do worker, enquanto a transferência unidirecional grava os pesos diretamente sem a participação do worker de inferência.

Alguns sistemas dependem da chamada do método update_weight() do motor de inferência para cada tensor. Isso significa alterações intrusivas no código de inferência, além de sobrecarga de RPCs, serialização e coordenação do plano de controle.

Com a primitiva RDMA WRITE, atualizamos os pesos silenciosamente na memória da GPU de inferência, sem cópias extras. Nenhuma mensagem do plano de controle e nenhuma lógica de controle de CPU estão envolvidas. Nenhuma modificação no motor de inferência é necessária.

O processo de transferência de pesos pode aproveitar quatro tipos de recursos de hardware: (1) Movimentação de dados entre host e dispositivo (2) Computação de GPU para fusão de projeção e quantização (3) Rede RDMA para o plano de dados (4) Ethernet para o plano de controle.

Nosso design divide as tarefas de transferência de pesos em estágios de pipeline, permitindo fácil sobreposição entre diferentes recursos de hardware.

Algumas implementações recalculam um agendamento de transferência a cada etapa de treinamento, coletando metadados e distribuindo instruções repetidamente. Isso adiciona latência desnecessária ao plano de controle.

Nosso agendamento é calculado uma única vez na inicialização. Cada iteração de treinamento simplesmente reproduz o plano: o controlador emite um sinal de "ir" e as GPUs seguem suas rotas pré-atribuídas. A execução é previsível e leve.

É tentador envolver todo o processo de atualização de pesos em uma única função monolítica: coleta de metadados, correspondência de nomes, agrupamento intra-nó, fusão de projeção, quantização, sub-fatiamento do mundo de comunicação, transferência de rede inter-nó. É difícil programar corretamente e ainda mais difícil de otimizar.

Em nossa engenharia, filtramos e separamos essas etapas como componentes individuais. Cada componente pode ser testado unitariamente, analisado e otimizado de forma isolada.

A transferência de pesos rápida e confiável é um bloco de construção essencial para o ajuste fino de RL em larga escala. Ao combinar a primitiva RDMA WRITE, um agendamento de transferência estático e execução em pipeline, reduzimos atualizações de trilhões de parâmetros para apenas 1,3 segundos no Kimi-K2. A abordagem é simples de entender, fácil de manter e evita os gargalos dos designs tradicionais.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: transferências unilaterais

Fluxo de trabalho de alto nível

Execução da transferência de pesos

DeviceMesh e Mesh Groups

Pipeline de tarefas

Controle de uso de memória da GPU

Por que é rápido e simples

Comunicação ponto a ponto

Transferência de dados unilateral

Criação de pipeline

Agendamento estático

Separação limpa

Conclusão

Citação