Transferência de pesos para pós-treinamento RL em

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

No ajuste fino por reforço assíncrono, o treinamento e a inferência são executados em GPUs separadas. Após cada etapa de treinamento, os 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 trilhões de parâmetros.

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

Nossa solução é baseada em RDMA WRITE, uma primitiva unidirecional onde a origem grava diretamente na memória da GPU de 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 do destino nem sequer será notificado da transferência. Isso nos proporciona transferências de baixa latência, alto rendimento e zero cópia, controladas 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.

Cálculo do cronograma – O controlador calcula um cronograma 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 do cronograma – O controlador envia o cronograma para todas as GPUs de treinamento.

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

Com o fluxo de trabalho de alto nível definido, o desafio principal é como executar transferências de pesos de forma eficiente na 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 uma DeviceMesh podem reconstruir o parâmetro completo; portanto, todas podem servir como origem para a transferência de pesos.

Múltiplas DeviceMeshes disjuntas formam um grupo de malha. Como as DeviceMeshes no mesmo grupo são disjuntas, suas transferências não interferem e podem ser executadas totalmente em paralelo. Entre os grupos de malha, inserimos uma barreira global para garantir a ordenação.

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

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

Memcpy de host para dispositivo — Se o FSDP descarrega o peso para a CPU

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

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

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

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

full_tensor() e outras operações de GPU introduzem um uso extra de memória da GPU. Para evitar erro de memória insuficiente, iniciamos a execução de uma tarefa apenas se as tarefas atuais em processamento 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 as soluções comuns de código aberto.

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

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

Um padrão comum é canalizar todos os parâmetros através de GPUs de classificação 0: reunir na classificação de treinamento 0, enviar para a classificação de inferência 0 e, em seguida, dispersar novamente. Isso rapidamente se torna um ponto de estrangulamento, limitado pela largura de banda PCIe e pela NIC de uma única GPU (por exemplo, 400 Gbps ≈ 50 GB/s).

Em contraste, nossa configuração ponto a ponto permite que cada GPU de treinamento envie dados 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 do lado da inferência e chamadas de worker, enquanto a transferência unilateral grava os pesos diretamente, sem a participação do worker de inferência.

Alguns sistemas dependem da chamada do método update_weight() do mecanismo de inferência para cada tensor. Isso significa mudanças intrusivas no código de inferência, além da 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 mecanismo 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 plano de dados (4) Ethernet para plano de controle.

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

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

Nosso cronograma é calculado uma vez na inicialização. Cada iteração de treinamento simplesmente repete o plano: o controlador emite um sinal de início 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: coletar metadados, correspondência de nomes, reunião intra-nó, fusão de projeção, quantização, subslicing do ambiente de comunicação, transferência de rede inter-nó. É difícil programar corretamente e ainda mais difícil de otimizar.

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

A transferência de pesos rápida e confiável é um bloco de construção fundamental para o ajuste fino de RL em larga escala. Ao combinar a primitiva RDMA WRITE, um cronograma de transferência estático e execução em pipeline, reduzimos as 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 unidirecionais

Fluxo de trabalho de alto nível

Execução da transferência de pesos

DeviceMesh e grupos de malha

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 unidirecional

Pipelining

Cronograma estático

Separação clara

Conclusão

Citação