Destilação Assíncrona

Este trainer requer vllm, transformers>=5.2.0 e kernels (o treinamento carrega o FlashAttention do Hub por meio dele). Para treinamento distribuído, apenas o FSDP2 é suportado (o DeepSpeed ZeRO não é).

Atualmente, vllm e transformers possuem restrições de dependência conflitantes. Para contornar isso, instale o vLLM primeiro, em seguida, force a instalação do transformers e, depois, adicione kernels por meio do extra do TRL, que fixa a versão que o transformers aceita:

Copiado

pip install vllm
pip install 'transformers>=5.2.0' --no-deps
pip install 'trl[kernels]'

Visão geral

experimental.async_distillation.AsyncDistillationTrainer é o equivalente assíncrono do DistillationTrainer, com arquitetura semelhante à do experimental.async_grpo.AsyncGRPOTrainer: um worker de rollout em segundo plano gera as próprias conclusões on-policy do aluno e as pontua em relação a um professor, enquanto o treinamento prossegue de forma concorrente, em vez de alternar entre geração e atualizações de gradiente. Diferentemente do trainer síncrono, o professor nunca é carregado localmente — apenas uma URL de servidor vLLM é necessária, de modo que o professor pode ser executado em hardware totalmente separado do aluno e do trainer, ou até mesmo ser um modelo muito maior do que caberia ao lado do aluno.

compute_loss minimiza uma Divergência de Jensen-Shannon generalizada entre as distribuições de tokens por posição do aluno e do professor (beta=0.0 é KL direta, beta=1.0 é KL reversa, valores intermediários interpolam), o mesmo objetivo usado por DistillationTrainer e ServerDistillationTrainer. O treinamento é sempre on-policy: o aluno gera todas as conclusões nas quais ele é treinado.

Destilação on-policy com múltiplos professores (MOPD)

O MOPD não faz parte do artigo 2306.13649, no qual se baseia o objetivo central deste trainer (que descreve um cenário síncrono com um único professor); trata-se de um método separado, descrito em MOPD: Multi-Teacher On-Policy Distillation for Capability Integration in LLM Post-Training. Lá, o MOPD é a terceira de três etapas — SFT geral, em seguida treinamento de RL independente por domínio de um especialista por domínio, e então o MOPD funde esses especialistas congelados em um único aluno. O AsyncDistillationTrainer implementa apenas essa terceira etapa de fusão: os professores especialistas por domínio devem já existir (por exemplo, treinados separadamente com GRPOTrainer/RLOOTrainer) e estar sendo servidos via HTTP antes de você apontar teacher_server_urls para eles. A própria Etapa 3 do artigo usa KL reversa (beta=1.0), e não o padrão deste trainer, beta=0.0 (KL direta).

teacher_server_urls aceita mais de uma entrada. Com uma única entrada, cada amostra é pontuada por aquele único professor (destilação on-policy simples). Com múltiplas entradas, a coluna teacher_id de cada amostra de treinamento seleciona qual professor a pontua — por exemplo, roteando prompts de matemática para um professor especialista em matemática e prompts de código para um professor especialista em código, cada um servido independentemente. Cada amostra é despachada para o seu professor correspondente, nunca sendo média ou ensemble entre professores. Uma amostra com teacher_id ausente ou não mapeado gera um erro em vez de cair silenciosamente para o professor errado. Veja examples/async_distillation_mopd/async_distillation_mopd.py para um exemplo executável que funde os especialistas em matemática, código e seguimento de instruções liberados com Open-MOPD em seu checkpoint SFT compartilhado, e examples/async_distillation_mopd/launcher.py para executá-lo em 4 GPUs com um único comando do Hugging Face Jobs.

Cada professor deve compartilhar o tokenizer do aluno. As conclusões são enviadas aos professores como ids de tokens brutos, e os ids de candidatos que um professor relata de volta indexam diretamente o vocabulário do próprio aluno em compute_loss. Professores da mesma família de modelos que o aluno (como no exemplo de MOPD, em que um aluno SmolLM3 é fundido a partir de especialistas SmolLM3 treinados a partir do mesmo checkpoint) satisfazem isso; um professor com um vocabulário diferente treina o aluno contra os tokens errados, silenciosamente, a menos que seu vocabulário seja maior que o do aluno.

Checkpoint e retomada: ignore_data_skip tem valor padrão True; o loop de pulo e reprocessamento do Trainer base não se aplica a uma fila de rollout ao vivo. Em vez disso, o índice do primeiro prompt ainda não treinado é salvo em rollout_state.json junto com cada checkpoint e restaurado na retomada, de modo que o worker avança rapidamente até esse prompt sem reprocessar amostras. É a posição treinada, não a do gerador: o worker opera à frente do treinamento pela profundidade da fila de rollout, e essas amostras em buffer são perdidas quando a execução termina, portanto, retomar a partir da posição do gerador pularia prompts que foram gerados, mas nunca treinados. Conjuntos de dados em streaming (IterableDataset) não podem ser reposicionados; seu worker reinicia do prompt 0 na retomada. Com FSDP e bf16=True (o padrão) ou fp16=True, model.safetensors é gravado nesse dtype, mesmo que os pesos mestres sejam fp32, de modo que um checkpoint publicado tem o tamanho do modelo servido, e não o dobro; os pesos mestres fp32 e o estado do otimizador ainda vão para os arquivos de retomada (pytorch_model_fsdp.bin, optimizer.bin), que são os lidos na retomada, portanto, a retomada não é afetada. Sem FSDP, o checkpoint também é o arquivo de retomada, então ele mantém o dtype mestre.

Como difere do DistillationTrainer

Em DistillationTrainer, o modelo professor é carregado localmente: a geração, a passagem direta do professor e a atualização do gradiente ocorrem sequencialmente no mesmo processo. experimental.async_distillation.AsyncDistillationTrainer separa essas responsabilidades da mesma forma que experimental.async_grpo.AsyncGRPOTrainer separa o rollout do GRPO de sua atualização:

• Trabalhador de rollout (processo em segundo plano) — gera conclusões a partir do servidor vLLM do aluno, envia a sequência completa de volta para o /v1/completions do professor resolvido com prompt_logprobs para pontuação forçada pelo professor (sem novos tokens gerados pelo professor) e insere amostras prontas para treinamento em uma fila.
• Loop de treinamento (processo principal) — extrai amostras da fila, calcula a perda JSD generalizada e atualiza os pesos do aluno.

Como o professor é pontuado via HTTP em vez de uma passagem local de forward, apenas uma fatia esparsa, top-k, de sua distribuição é transmitida (teacher_top_k), e não o vocabulário completo — consulte a documentação de beta e teacher_top_k em experimental.async_distillation.AsyncDistillationConfig para saber exatamente quais candidatos o protocolo de rede garante um logprob do professor para cada regime de beta.

A cada weight_sync_steps passos de treinamento, os pesos atualizados do aluno são transferidos para seu servidor vLLM via NCCL. Assim como no experimental.async_grpo.AsyncGRPOTrainer, a geração fica à frente do treinamento, portanto as amostras podem refletir uma política ligeiramente desatualizada; max_staleness controla quantas atualizações de pesos uma amostra pode ficar para trás antes de ser descartada.

Início rápido

# train_async_distillation.py
from datasets import load_dataset
from trl.experimental.async_distillation import AsyncDistillationTrainer

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

O servidor do professor, o servidor vLLM do aluno e o treinador devem ser executados em GPUs separadas. Ambos são instâncias simples de vllm serve, mas precisam de flags diferentes: o professor é apenas pontuado (--logprobs-mode processed_logprobs para que teacher_temperature alcance os logprobs retornados, --max-logprobs -1 para que teacher_top_k possa exceder o limite padrão de 20 do vLLM), enquanto o aluno é apenas gerado e precisa ter a transferência de pesos NCCL habilitada para que o treinador possa enviar os pesos atualizados para ele. A pontuação envia a sequência inteira do aluno e pede ao professor mais um token, portanto o --max-model-len do professor (ou seu padrão) deve ser maior que o do aluno:

# Terminal 1: teacher server on GPU 0 (static, never updated)
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen2.5-1.5B-Instruct \
    --port 8001 \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

# Terminal 2: student's vLLM server on GPU 1 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}'

# Terminal 3: training on GPU 2
CUDA_VISIBLE_DEVICES=2 python train_async_distillation.py

LoRA

Passe um peft_config para treinar um adaptador LoRA no aluno em vez do modelo completo:

from peft import LoraConfig

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)

Cada sincronização de pesos precisa enviar o aluno atualizado para o servidor. Por padrão, o trainer mescla o adaptador no modelo base e envia os pesos completos mesclados, o que corresponde a cerca de 15 GB para um modelo de 7B em cada sincronização. Inicie o servidor vLLM do aluno com --enable-lora e ele enviará apenas o adaptador, cerca de 160 MB em r=32:

CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}' \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 expõe o endpoint para o qual o trainer envia cada novo adapter. Adicione VLLM_WORKER_MULTIPROC_METHOD=spawn se você estiver servindo com tensor-parallel e mantenha --weight-transfer-config em qualquer caso: o trainer só escolhe um modo de sincronização quando inicia, momento em que o servidor já está ativo, e a sincronização mesclada é o fallback. O servidor teacher não é afetado: ele nunca é atualizado.

--max-lora-rank deve ser um dos valores 1, 8, 16, 32, 64, 128, 256, 320, 512. Ele define o rank máximo que o servidor pode atender, e não o rank que ele efetivamente servirá, portanto um adaptador r=4 funciona perfeitamente sob 8. --max-loras deve ser pelo menos max_staleness + 2: o trainer mantém max_staleness + 1 versões de adaptador registradas para que um rollout iniciado sob uma política mais antiga possa ser concluído sob ela, em vez de trocar de política no meio da geração, e cada sincronização carrega a próxima versão antes de descarregar a mais antiga. O trainer verifica ambos os valores ao iniciar e informa com quais valores reiniciar o servidor.

O vLLM carrega adaptadores do disco em vez da rede, portanto o trainer grava cada versão em <output_dir>/.vllm_lora/ e fornece o caminho ao servidor. Se o servidor estiver em uma máquina diferente, esse diretório precisa estar em um sistema de arquivos visível para ambos.

<output_dir>/.vllm_lora/ é um cache de serviço, não o local onde o seu adaptador treinado é armazenado. Ele apenas mantém as versões mais recentes, e cada uma é excluída pouco depois de sair da janela de obsolescência. Os checkpoints são o que preservam o adaptador que você treinou, portanto, não os desative: com save_strategy="no", uma execução de várias horas termina e não deixa nada para trás. O output_dir também precisa ser legível pelo servidor e precisa sobreviver ao término do job, portanto, um armazenamento temporário local do nó não funcionará, mesmo que isso mantivesse as sincronizações fora de um sistema de arquivos compartilhado lento.

A sincronização de adaptadores nem sempre é mais rápida. O vLLM gera tokens aproximadamente 1,33 a 1,39 vezes mais devagar ao servir um adaptador LoRA do que ao servir pesos mesclados, e a geração geralmente é o que domina o tempo real de execução. Em contrapartida, você obtém uma pausa muito mais curta em cada sincronização e cerca de 1% da largura de banda de sincronização. Se o throughput de geração for o seu gargalo, a sincronização mesclada ainda costuma ser a melhor escolha.

Métricas registradas

Um rollout passa por várias etapas antes de se tornar um gradiente. Primeiro, vamos dar uma olhada nas diferentes etapas, de um prompt para a composição final do nosso lote:

Dataset row = PROMPT (message list + optional `teacher_id`)
 └─ ROLLOUT            (1 prompt -> 1 student completion -> 1 teacher scoring call)
     ├─ generate       (student's vLLM /v1/completions, sampled)
     └─ score          (the routed teacher's /v1/completions with `prompt_logprobs`, teacher-forced)
         └─ SAMPLE     (prompt + completion + the teacher's sparse per-position distribution)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

         └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
             └─ ROW            (Planner assigns it to one of `dp` rows, Σ Lᵢ²-balanced)
                 └─ MICRO-BATCH (`dp` rows, one per rank)
                     └─ PACKED ROW (1 concat sequence, `position_ids` reset per sample)
                         └─ FORWARD (`compute_loss`, bs=1, inter-rank padding stripped)
             └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

Lendo de cima para baixo:

• um rollout é um prompt gerado uma vez e pontuado uma vez. Não há grupo: a destilação não possui uma linha de base de vantagem para calcular entre gerações, portanto, um prompt não é repetido e um rollout gera exatamente uma amostra de treinamento.
• um sample é aquela conclusão mais os teacher_top_k principais candidatos do professor para cada posição da conclusão. É isso que atravessa a fronteira do processo.
• além do limite, o trainer planeja as amostras em linhas (uma por rank de DP), empacota cada linha em uma única sequência concatenada e acumula grad_accum micro-lotes em um único passo do otimizador.

Portanto, dividimos nossas métricas com base na etapa (ou entidade) sobre a qual elas se referem:

Tudo o que a própria perda mede é registrado sem prefixo (jsd, entropy, teacher_entropy), com uma divisão teacher_*/<teacher_id> sob MOPD — veja a divergência abaixo.

Vamos também definir alguns termos para tornar as métricas de tokens inequívocas:

• Os tokens gerados são aqueles produzidos pelo aluno (completions/*).
• Os tokens forwarded são todos os tokens que a passagem forward processa: o prompt e os tokens gerados.
• Os tokens treinados são o subconjunto sobre o qual a perda é calculada, onde completion_mask == 1. Treinado ≠ gerado: uma posição de conclusão para a qual o professor não pontuou nenhum candidato é mascarada na divergência, mas ainda assim é encaminhada.

Como os números de lote se compõem

Uma linha é o que um rank de DP faz o forward em um micro-lote: várias amostras concatenadas em uma única sequência, com position_ids reiniciando em cada fronteira de amostra. O planejador decide em qual linha cada amostra será alocada (equilibrando Σ Lᵢ² para que nenhum rank fique para trás), e um data collator realiza a concatenação.

Em todas as nossas métricas, um “step” sempre significa um passo COMPLETO do otimizador, nunca um micro-lote. gradient_accumulation_steps micro-lotes formam um passo, e toda métrica chamada _per_step é por passo do otimizador, correspondendo a global_step, logging_steps e ao restante do vocabulário do Trainer. Quando uma métrica é por micro-lote, isso é indicado (batch/microbatches_per_step) ou ela é uma quantidade por linha (batch/row_*, batch/samples_per_row).

Portanto, uma etapa contém um número fixo de row-slots:

row-slots per step = gradient_accumulation_steps x world_size

o que faz com que as métricas do lote se componham em uma escada que você pode verificar entre si:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

Portanto, batch/samples_per_step ≈ row-slots x batch/samples_per_row, e o mesmo vale para tokens. Os dois lados só coincidem até a variação entre micro-lotes dentro do passo — as métricas por passo são somas, enquanto as por linha são médias —, então espere uma diferença de fração de um percentual, não uma correspondência exata.

batch/row_fill_frac é o que deve ser monitorado quando as amostras são longas: uma amostra de 10k tokens preenche mal um orçamento de 32k (três cabem, quatro nunca cabem, então o empacotador frequentemente obtém apenas duas e a linha fica ~77% cheia), enquanto amostras de 1k token preenchem-no quase perfeitamente. Isso é quantização, não um bug, e token_budget é o parâmetro de ajuste — tendo em mente que a atenção é O(L²) por sequência, de modo que uma linha mais cheia de sequências longas não custa linearmente mais memória.

A fila de rollout

O worker envia amostras pontuadas para uma fila e o trainer as retira dela. Quatro métricas descrevem essa única fila e respondem a perguntas diferentes:

As duas últimas métricas (perf/rollout_wait_s e rollout/backpressure_s) são imagens espelhadas e nunca são grandes ao mesmo tempo! Lidas em conjunto com o tamanho da fila, elas indicam qual lado é o gargalo:

• fila quase vazia, perf/rollout_wait_s alto → limitado pela geração. O treinador está com falta de dados; verifique rollout/generated_tok_s, rollout/concurrency e rollout/score_s.
• fila quase cheia, rollout/backpressure_s alto → limitado pelo trainer. A geração está sendo limitada e sua saída está envelhecendo na fila, então observe sample/staleness_mean subir.
• ambos próximos de zero → equilibrado!

Completions

Uma completion é o que o servidor vLLM do aluno gerou para um prompt.

Rollouts

Um rollout é um prompt levado até o fim: gerado pelo aluno e, em seguida, processado pelo professor com forçamento de professor. Como a chamada ao professor está no caminho crítico de cada rollout (diferentemente do GRPO, em que a pontuação é um loop separado sobre grupos concluídos), um professor lento aparece diretamente como rollout/duration_s.

Amostras chegando da fila

Cada amostra é um RolloutSample: um prompt, a conclusão do aluno e a pontuação esparsa do professor para ela. O worker as insere em rollout_buffer; tudo abaixo é medido no lado do trainer, conforme ele as extrai para o treinamento.

Lotes

Um micro-lote é world_size linhas (lembre-se de que o packing aplanar em 1 sequência), então um por rank de DP, de modo que o rank i faz o forward da linha i. gradient_accumulation_steps desses formam um passo do otimizador. Um passo então cobre gradient_accumulation_steps × world_size linhas no total. Cada métrica _per_step abaixo é uma soma sobre o passo inteiro e entre todos os ranks; as métricas batch/row_* são médias sobre as linhas.

A divergência

O que a própria métrica de objetivo mede, em média, sobre os tokens treinados da janela. Estes são os números a serem observados para o aprendizado, em oposição aos de throughput.

Não existe entropy por professor: a entropia do aluno é uma propriedade da própria política dele, não de qual professor pontuou a amostra, então a métrica combinada já a cobre.

Desempenho

Throughput mede quantos tokens o trainer processa por segundo. Model FLOPs Utilization (MFU) estima a porcentagem da capacidade teórica de computação dos dispositivos de treinamento utilizada para as passadas de forward e backward do modelo. A capacidade utilizada para essa estimativa depende do hardware e da precisão do treinamento.

O throughput e o MFU são reportados em duas bases de tempo:

• _fwd_bwd mede o desempenho durante as passadas de forward e backward.
• _wall_clock mede o desempenho ao longo da etapa completa do otimizador, incluindo as atualizações do otimizador, a sincronização de pesos e o tempo gasto aguardando os rollouts.

Essa diferença ajuda a distinguir o desempenho de computação do custo de espera e coordenação. Ambas as métricas de MFU descrevem apenas os dispositivos de treinamento: elas não incluem a computação ou a capacidade de hardware de servidores externos de rollout ou de teacher. A MFU de tempo real reflete o tempo que o treinador gasta aguardando esses serviços, e não a própria utilização desses serviços.

O MFU é reportado apenas quando a capacidade de computação de pico é conhecida para o dispositivo de treinamento local na precisão selecionada. Caso contrário, o MFU é omitido; o treinamento continua e as métricas de throughput e tempo permanecem disponíveis.

Filosofia de design

Este trainer é mantido intencionalmente minimalista e não foi projetado para se tornar uma solução de uso geral. Se você precisar de um recurso que não é suportado, recomendamos clonar o repositório e adaptar o trainer diretamente às suas necessidades. Novos recursos só serão considerados quando houver demanda significativa da comunidade.

AsyncDistillationConfig

AsyncDistillationTrainer

>>> from trl.experimental.async_distillation import AsyncDistillationTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

>>> trainer = AsyncDistillationTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()