Entrenador GRPO asíncrono

Resumen

experimental.async_grpo.AsyncGRPOTrainer implementa el mismo algoritmo GRPO pero desacopla la generación de rollout del entrenamiento. Un worker en segundo plano transmite continuamente completions desde un servidor vLLM mientras el bucle de entrenamiento los consume, de modo que la generación y las actualizaciones de gradiente se solapan en lugar de alternarse.

Este trainer fue aportado por Amine Dirhoussi y Quentin Gallouédec.

Inicio rápido

Este ejemplo demuestra cómo entrenar un modelo utilizando el método GRPO asíncrono. Entrenamos un modelo Qwen3.5 2B con los prompts del conjunto de datos DeepMath-103K. Puede ver los datos del conjunto de datos aquí:

A continuación se muestra el script para entrenar el modelo. El trainer necesita vLLM y kernels, y accuracy_reward necesita math-verify, que puedes instalar con pip install trl[vllm,kernels,math_verify].

# train_async_grpo.py
from datasets import load_dataset
from trl.experimental.async_grpo import AsyncGRPOTrainer
from trl.rewards import accuracy_reward

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

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3.5-2B",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)
trainer.train()

La generación se ejecuta en un servidor vLLM, en GPUs diferentes a las del entrenador. Inícialo en la GPU 0:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3.5-2B \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096

Luego, ejecute el script en la GPU 1 usando el siguiente comando:

CUDA_VISIBLE_DEVICES=1 python train_async_grpo.py

Compatibilidad y restricciones

Algunas combinaciones de argumentos están restringidas intencionalmente en la implementación actual de experimental.async_grpo.AsyncGRPOTrainer:

• La generación siempre se ejecuta en un servidor vLLM separado, iniciado con vllm serve. No existe un modo de colocación.
• Prompts from the dataset must be in conversational format. With an environment_factory, the dataset may have no prompt column, and reset() provides it.
• Las funciones de recompensa deben ser funciones Python invocables. Los modelos de recompensa proporcionados como un identificador de modelo o un PreTrainedModel no son compatibles, consulte Uso de una función de recompensa personalizada.
• El trainer necesita transformers>=5.2.0.
• El entrenamiento fragmentado utiliza FSDP2. DeepSpeed no es compatible, y tampoco lo son el paralelismo de contexto y el paralelismo de secuencia (cp_size > 1 o sp_size > 1 en parallelism_config).
• use_liger_kernel=True no es compatible.
• Los adaptadores PEFT que apuntan a lm_head no son compatibles.
• No se admiten modelos con attention sinks. Las únicas capas no de atención admitidas son las capas DeltaNet con compuerta de los modelos híbridos (Qwen3.5, Qwen3.6, Qwen3-Next, …), que requieren transformers>=5.15.0.
• No se admiten imágenes, consulte Entrenamiento de modelos de lenguaje visual.

El entrenamiento carga su kernel de atención desde el Hub a través de la biblioteca kernels: kernels-community/aiter-flash-attn en GPUs AMD (ROCm), kernels-community/flash-attn4 en GPUs NVIDIA Blackwell y kernels-community/flash-attn3 en otras GPUs NVIDIA. Los modelos que limitan sus logits de atención (attn_logit_softcapping), como Gemma 2, no son compatibles con GPUs AMD ni con GPUs NVIDIA Blackwell.

Mirando más a fondo el método GRPO asíncrono

En GRPOTrainer, la generación y el entrenamiento son secuenciales: se genera un lote, se calcula la pérdida, se actualizan los pesos y se repite. Incluso en modo colocate, una fase debe terminar antes de que comience la otra. En experimental.async_grpo.AsyncGRPOTrainer, un trabajador de rollout genera y puntúa las finalizaciones en un proceso en segundo plano mientras el bucle de entrenamiento calcula la pérdida y actualiza el modelo. Para entender cómo funciona, se puede dividir en tres pasos principales: generar finalizaciones, calcular la ventaja y calcular la pérdida. No hay un modelo de referencia, por lo que no hay un término de divergencia KL (el valor predeterminado de GRPO, beta=0.0).

Generar completiones

El worker de rollout mantiene hasta max_concurrent_rollouts solicitudes en curso hacia el servidor de vLLM. Para cada prompt, genera un grupo de $G$ complementos (num_generations) y, una vez que todo el grupo está terminado, lo puntúa y calcula sus ventajas. Las muestras de entrenamiento resultantes se envían a una cola, de la cual el bucle de entrenamiento las extrae. La generación solo espera cuando está max_staleness pasos por delante del entrenamiento (ver Concurrencia), y un paso de entrenamiento solo espera cuando la cola está vacía.

Cada weight_sync_steps pasos de entrenamiento, el trainer envía sus pesos actualizados al servidor vLLM a través de NCCL, de modo que las próximas completions provengan de la política más reciente. Por lo tanto, una completion puede provenir de una política unos pocos pasos más antigua que la que se está entrenando: una muestra que esté más de max_staleness actualizaciones de pesos por detrás cuando el trainer la recupera se descarta.

Las métricas y el resto de esta página utilizan los siguientes términos para las etapas por las que pasa una finalización:

• a grupo es un prompt y sus $G$ rollouts. La ventaja se calcula dentro de él.
• Un rollout es una conversación completa generada para un prompt, que incluye cada llamada a una herramienta. Está compuesto por uno o más turns, cada uno correspondiente a una solicitud de vLLM.
• una muestra es sobre lo que se calcula la pérdida: los tokens de un rollout, con su ventaja. Un rollout de múltiples turnos puede producir más de una muestra (ver Rollouts de múltiples turnos).
• una fila es lo que un proceso envía en un micro-lote: varias muestras empaquetadas juntas (ver Tree packing). Un micro-lote contiene una fila por proceso, y gradient_accumulation_steps micro-lotes forman un paso, el paso del optimizador.

Las cuentas de tokens se dividen en tres tipos: los tokens generados son los que produjo el modelo, los tokens procesados son todos los tokens de una muestra (incluidos el prompt y los resultados de las herramientas) y los tokens entrenados son aquellos sobre los que se calcula la pérdida.

Cálculo de la ventaja

Para cada una de las $G$ completions de un grupo, la recompensa es la suma de las recompensas de cada función de recompensa. La ventaja se normaliza dentro del grupo, como en GRPO: $$\hat{A}_{i,t} = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}$$

No existe la opción scale_rewards: la ventaja siempre se escala por la desviación estándar del grupo. Una finalización para la que todas las funciones de recompensa devolvieron None se excluye de la media y la desviación estándar del grupo y recibe una ventaja de cero.

Se descarta un grupo cuyos recompensas son todas iguales: con la ventaja predeterminada, esta es cero en todas partes, por lo que entrenar con él costaría un paso hacia adelante y uno hacia atrás sin mover nada. La fracción de tales grupos se registra como frac_reward_zero_std. El paso se completa con los grupos que se conservan, por lo que muchos grupos descartados significa más generación por paso. Un grupo descartado aún cuenta hacia la época, que permanece como un solo paso sobre los prompts.

Cálculo de la pérdida

Las log-probabilidades del servidor de vLLM nunca coinciden exactamente con las del entrenador: kernels diferentes, una precisión distinta y una política unos pasos por detrás (véase Cómo abordar la discrepancia entre el entrenamiento y la inferencia). El entrenador utiliza las log-probabilidades del muestreador como la política anterior, por lo que una única relación de importancia cubre tanto la discrepancia como la obsolescencia: $$r_{i,t}(\theta) = \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\text{sampler}}(o_{i,t} \mid q, o_{i,< t})}.$$

Los tokens cuya proporción queda fuera de [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] se descartan de la pérdida. Para obtener más información, consulte Masked Importance Sampling.

Por defecto (score_centering=True), el trainer utiliza Score Centering. Los pesos de importancia por sí solos no eliminan la discrepancia: la puntuación $\nabla_\theta \log \pi_\theta$ de un token muestreado no promedia a cero bajo el muestreador $\mu$, por lo que cada actualización también tira de la política hacia el muestreador. El Score Centering resta esa expectativa en cada posición: $$\nabla_\theta \mathcal{L} = - \mathbb{E}_{o \sim \mu} \left[ \hat{A} \left( f(r_t) \nabla_\theta \log \pi_\theta(o_t) - \mathbb{E}_{v \sim \mu} \left[ f(r_v) \nabla_\theta \log \pi_\theta(v) \right] \right) \right], \quad r_v = \frac{\pi_\theta(v)}{\mu(v)},$$

donde $f(r) = r$ dentro de $[C_{\min}, C_{\max}]$ y $0$ fuera. La expectativa interna se toma sobre los num_sampler_logprobs tokens más probables en cada posición, que el trabajador de rollout solicita al servidor.

Con score_centering=False, el trainer utiliza en su lugar el objetivo sustituto recortado de GRPO, con la misma proporción enmascarada: $$\mathcal{L}(\theta) = - \frac{1}{\sum_{i} |o_i|} \sum_{i} \sum_{t=1}^{|o_i|} \mathbb{1}\left[ C_{\min} \le r_{i,t}(\theta) \le C_{\max} \right] \min \left( r_{i,t}(\theta) \hat{A}_{i,t}, \, \text{clip}\left( r_{i,t}(\theta), 1 - \epsilon_{\text{low}}, 1 + \epsilon_{\text{high}} \right) \hat{A}_{i,t} \right),$$

donde $\epsilon_{\text{low}}$ es epsilon y $\epsilon_{\text{high}}$ es epsilon_high.

En ambos casos, la pérdida se normaliza por el número de tokens entrenados en el micro-lote, contados en todos los procesos, y se promedia sobre los gradient_accumulation_steps micro-lotes de un paso.

Métricas registradas

Durante el entrenamiento, registramos las siguientes métricas:

• num_tokens: El número total de tokens reenviados hasta el momento, tanto prompts como completions. Con tree packing, un prefijo compartido por varias muestras se cuenta una sola vez.
• completions/mean_length: La longitud media de las completaciones generadas. Al usar herramientas, solo se cuentan los tokens que no son de herramientas.
• completions/min_length: La longitud mínima de las completaciones generadas. Al usar herramientas, solo se cuentan los tokens que no son de herramientas.
• completions/max_length: La longitud máxima de las completaciones generadas. Al usar herramientas, solo se cuentan los tokens que no son de herramientas.
• completions/mean_terminated_length: La longitud media de las completaciones generadas que terminan con EOS. Al usar herramientas, solo se cuentan los tokens que no son de herramientas.
• completions/min_terminated_length: La longitud mínima de las completaciones generadas que terminan con EOS. Al usar herramientas, solo se cuentan los tokens que no son de herramientas.
• completions/max_terminated_length: La longitud máxima de las completaciones generadas que terminan con EOS. Al usar herramientas, solo se cuentan los tokens que no son de herramientas.
• completions/clipped_ratio: La proporción de completados truncados (recortados), cuyo último turno no termina con un token EOS.
• completions/duplicate_frac: La fracción de rollouts que repiten un rollout anterior de su grupo token por token. Consulte Saturación y duplicados.
• tools/call_frequency: El número medio de llamadas a herramientas por finalización, en todos los despliegues, incluidos los grupos que luego se descartan. Se registra solo cuando el entrenador tiene herramientas.
• tools/failure_frequency: La proporción de llamadas a herramientas que fallaron (la herramienta no se encontró o lanzó una excepción), sobre todas las llamadas a herramientas, incluidos los grupos que luego se descartan. Es NaN en una ventana de registro sin ninguna llamada a herramienta, donde GRPO registra 0.0. Se registra solo cuando el entrenador tiene herramientas.
• rewards/{reward_func_name}/mean: La recompensa media de una función de recompensa específica, sobre las muestras entrenadas. Cuando un entorno posee la recompensa a través de get_reward, {reward_func_name} es el nombre de la clase del entorno.
• rewards/{reward_func_name}/group_std: La desviación estándar de la recompensa de una función de recompensa específica, calculada dentro de cada grupo y promediada sobre las muestras entrenadas.
• reward: La recompensa media general tras sumar las recompensas de todas las funciones, sobre todas las finalizaciones puntuadas, incluidos los grupos que luego se descartan. Es la recompensa de la política, por lo que es la que hay que seguir: reward_kept excluye los grupos cuyas recompensas son todas iguales.
• reward_group_std: La desviación estándar de las recompensas sumadas entre funciones, calculada dentro de cada grupo y promediada sobre las muestras entrenadas. Es el valor por el que se divide la ventaja, a diferencia de reward_std de GRPO, que se toma sobre el lote completo.
• reward_kept: La recompensa media de las muestras entrenadas, que excluye los grupos descartados.
• reward_dropped: La recompensa media de los grupos que se descartaron, la recompensa en la que se quedan atascados. Consulta Saturación y duplicados para leerla junto con frac_reward_zero_std y reward.
• frac_reward_zero_std: La fracción de grupos con una desviación estándar de recompensa igual a cero (todas las respuestas son correctas o incorrectas) o con menos de dos completados puntuables. Estos grupos se descartan.
• sampling/sampling_logp_difference/mean: La diferencia absoluta media entre las probabilidades logarítmicas devueltas por el muestreador (vLLM) y las recalculadas por el modelo de entrenamiento, sobre los tokens entrenados.
• sampling/sampling_logp_difference/max: La mayor diferencia absoluta de este tipo.
• sampling/importance_sampling_ratio/min: La proporción de muestreo por importancia más pequeña $r_{i,t}(\theta)$ entre el modelo y el muestreador, sobre los tokens entrenados, antes de descartar los tokens fuera de [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] .
• sampling/importance_sampling_ratio/mean: La media de dicha relación de muestreo por importancia.
• sampling/importance_sampling_ratio/max: La mayor proporción de muestreo por importancia.
• sampling/masked_mean: La fracción de tokens entrenados que se excluyen de la pérdida porque su ratio de muestreo de importancia cae fuera de [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max].
• sampling/sampler_mass: La masa de probabilidad promedio que cubren los num_sampler_logprobs tokens más probables del muestreador. Se registra solo cuando score_centering=True.
• entropy: Entropía media de las predicciones de tokens en los tokens entrenados.
• aux_loss: La pérdida auxiliar de equilibrio de carga de un modelo Mixture-of-Experts, antes de que se escale por router_aux_loss_coef y se sume a la pérdida. Se registra solo cuando el modelo es un modelo MoE y router_aux_loss_coef es distinto de cero.
• clip_ratio/region_mean: La proporción de probabilidades de tokens donde el objetivo de GRPO se recorta para mantenerse dentro de la región de confianza: $\text{clip}\left( r_{i,t}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)$. Con score_centering=True la pérdida no se recorta, por lo que las clip_ratio/* métricas cuentan los tokens fuera de la región de confianza.
• clip_ratio/low_mean: La proporción media de probabilidades de tokens que se recortaron en el límite inferior de la región de confianza: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: La fracción más pequeña por finalización de tokens recortados en el límite inferior de la región de confianza: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: La proporción media de probabilidades de tokens que se recortaron en el límite superior de la región de confianza: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: La mayor fracción de tokens recortados en el límite superior de la región de confianza por finalización: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

Las otras métricas describen el pipeline asíncrono.

Las métricas se registran una vez por ventana de registro: todo lo registrado desde el registro anterior, que abarca los últimos logging_steps pasos del optimizador. En la ventana:

• Una media medida por rollout, muestra o llamada, como reward, rollout/turns_mean y tools/latency_s, se promedia sobre todos los rollouts, muestras o llamadas de la ventana. Una tasa como rollout/generated_tokens_per_s es el total de tokens dividido por el total de segundos.
• En caso contrario, un nombre cuya última parte contiene la palabra count se suma: sample/dropped_stale_count es el número de muestras descartadas desde el último registro, no desde el inicio del entrenamiento.
• Un nombre cuya última parte contiene la palabra max o min, como sample/staleness_max o completions/max_length, conserva el valor más grande o más pequeño.
• Cualquier otra métrica es la media de los valores registrados en la ventana.

Unidades:

• _s: segundos.
• tokens_per_s: tokens por segundo.
• _frac: fracción entre 0 y 1.
• _per_step: cantidad en un paso del optimizador, promediada sobre los pasos de la ventana.

Métricas de despliegue

• rollout/duration_s: El tiempo real de un rollout, desde el despacho hasta que el bucle de rollout lo recoge.
• rollout/generated_tokens_per_s: El rendimiento de generación, en tokens generados por segundo, desde el informe anterior.
• rollout/concurrency: El número de rollouts en curso hacia el servidor de vLLM.
• rollout/concurrency_cap: El límite actual de rollouts en curso, que sigue la carga del servidor. Consulte Concurrencia.
• rollout/turns_mean, rollout/turns_max: El número medio y máximo de turnos por rollout.
• rollout/samples_per_rollout: El número de muestras de entrenamiento por rollout. Por encima de 1.0, algunos rollouts se bifurcaron; por debajo de 1.0, algunos no produjeron ningún token entrenado. Consulte Multi-turn rollouts.
• rollout/fork_frac: La fracción de transiciones de turno cuyo prompt no comienza con los tokens retenidos, por lo que el turno inició una nueva muestra.
• rollout/drift_tokens_mean, rollout/drift_tokens_max: El número de tokens retenidos que un fork volvió a renderizar, promediado sobre las transiciones de turno (uno que no hace fork cuenta como 0). La nueva muestra comparte los tokens anteriores a la divergencia, por lo que esto es lo que cada fork añade al forward.
• rollout/score_queue_size: El número de grupos completados que esperan para ser puntuados.
• rollout/queue_size: El número de muestras puntuadas que esperan en la cola de rollout, medido en el punto donde el trainer las extrae. Consulte La cola de rollout.
• rollout/score_s: El tiempo para puntuar un grupo.
• rollout/score_wait_s: El tiempo que un grupo completado esperó antes de ser puntuado.
• rollout/score_block_s: El tiempo durante el cual la generación se bloqueó porque la cola de puntuación estaba llena.
• rollout/backpressure_s: El tiempo que la generación esperó para iniciar un rollout porque estaba max_staleness pasos por delante del entrenamiento, promediado sobre las esperas en la ventana de registro. Consulte Concurrencia.
• rollout/vllm_retry_count: El número de solicitudes de vLLM reintentadas. Un servidor degradado, de lo contrario, parece una lentitud inexplicable.
• rollout/vllm_kv_cache_usage_frac: El uso de la caché KV de vLLM (máximo entre los motores), promediado sobre las consultas en la ventana de registro.
• rollout/vllm_running_requests: El número de solicitudes que los motores del servidor están decodificando (el num_requests_running de vLLM, sumado sobre todos los motores). La diferencia con rollout/concurrency corresponde a lo que se encuentra en el proceso del servidor de API, construyendo las respuestas.
• rollout/vllm_queued_requests: El número de solicitudes que el servidor ha aceptado pero aún no ha programado (el num_requests_waiting de vLLM, sumado sobre los motores).
• rollout/vllm_preemption_count: El número de solicitudes que el servidor interrumpió (el num_preemptions_total de vLLM, sumado sobre los motores). Una sincronización de pesos restablece las cachés del servidor, lo que interrumpe todas las solicitudes en ejecución, por lo que esto no es una señal de presión.

Métricas de la herramienta

• tools/latency_s, tools/{tool_name}/latency_s: El tiempo empleado en ejecutar una llamada a una herramienta, en total y por herramienta.
• tools/{tool_name}/call_count, tools/{tool_name}/failure_count: El número de llamadas y llamadas fallidas por herramienta.
• tools/unknown_name_count: El número de llamadas a una herramienta que no existe. A diferencia de una herramienta que se ejecutó y lanzó una excepción, esto es un error de política.
• tools/parallel_calls_mean: El número medio de llamadas a herramientas en un mensaje del asistente que invoca herramientas.
• tools/loop_exhausted_frac: La fracción de rollouts que se cortan en max_tool_calling_iterations mientras aún están llamando a herramientas, sobre todos los rollouts, incluidos los rollouts sin llamadas y los grupos que luego se descartan. Se registra solo cuando el trainer tiene herramientas.

HarnessAsyncGRPOTrainer registra la mayoría de estos con los mismos nombres. Consulte Entrenamiento en harnesses para ver las diferencias.

Métricas del entorno

• env/reset_s: El tiempo que se invierte en el reset() del entorno antes de que comience un rollout. Se registra solo cuando se proporciona environment_factory.
• env/reward_s: El tiempo transcurrido en get_reward() del entorno tras finalizar un rollout. Se registra solo cuando el entorno define get_reward.

Métricas de ejemplo

• sample/forwarded_tokens_mean, sample/forwarded_tokens_max: El número medio y máximo de tokens reenviados por muestra.
• sample/trained_tokens_mean: El número medio de tokens entrenados por muestra.
• sample/time_in_queue_s: El tiempo que una muestra esperó en la cola de rollout antes de ser utilizada para el entrenamiento.
• sample/staleness_mean, sample/staleness_max: El número de actualizaciones de pesos entre la política que generó una muestra y la que la está entrenando. Las métricas sampling/importance_sampling_ratio/* y sampling/sampling_logp_difference/* combinan esta obsolescencia con la discrepancia entre entrenamiento e inferencia: una razón alejada de 1.0 o una diferencia de log-probabilidad alejada de 0.0 con obsolescencia cero indica la discrepancia, por ejemplo, una diferencia de dtype con el servidor.
• sample/dropped_stale_count: El número de muestras descartadas por exceder max_staleness.

Métricas por lote

• batch/forwarded_tokens_per_step, batch/trained_tokens_per_step: El número de tokens reenviados y entrenados por paso, en todos los procesos.
• batch/samples_per_step, batch/groups_per_step: El número de muestras y de grupos distintos por paso. Difieren en el número de muestras por grupo.
• batch/microbatches_per_step: El número de micro-lotes por paso.
• batch/masked_token_frac: La fracción de tokens reenviados que no se entrenan.
• batch/samples_per_row, batch/row_tokens_mean, batch/row_tokens_max: El número de muestras por fila, y el número medio y máximo de tokens reenviados por fila.
• batch/row_fill_frac: La fracción de token_budget que usan las filas.
• batch/row_imbalance: La proporción entre el coste de atención de la fila más costosa y el coste medio. El proceso con la fila más costosa detiene a los demás en el all-reduce de gradientes; 1.0 es un equilibrio perfecto.
• batch/packing_ratio: El número de tokens antes del empaquetado por token reenviado: el factor por el cual el empaquetado reduce el paso hacia adelante, 1.0 cuando no se comparte nada.
• batch/pad_frac: La fracción de relleno que se añade para igualar la longitud de las filas. Solo tiene un coste de transferencia y se elimina antes de la propagación hacia adelante.
• batch/dropped_oversize_count: El número de muestras descartadas por exceder token_budget.

Métricas de rendimiento

• perf/step_s: El tiempo real de un paso, incluyendo el optimizador, la sincronización de pesos y el tiempo de espera de muestras.
• perf/fwd_bwd_s, perf/fwd_s: El tiempo empleado en las pasadas hacia adelante y hacia atrás de un paso, y solo en sus pasadas hacia adelante.
• perf/optimizer_s: El tiempo empleado en el paso del optimizador.
• perf/rollout_wait_s: El tiempo que el entrenamiento esperó porque la cola de rollout estaba vacía. Consulte La cola de rollout.
• perf/weight_sync_s: El tiempo de una sincronización de pesos, dividido en perf/weight_sync_pause_s (esperando a que vLLM se pause), perf/weight_sync_barrier_s (esperando a los otros procesos) y perf/weight_sync_transfer_s (enviando los pesos).
• perf/fwd_bwd_forwarded_tokens_per_s, perf/wall_clock_forwarded_tokens_per_s: El número de tokens reenviados por segundo, durante las pasadas de adelante y atrás o durante todo el paso.
• perf/wall_clock_trained_tokens_per_s: El número de tokens entrenados por segundo, durante todo el paso.
• perf/fwd_bwd_mfu, perf/wall_clock_mfu: La utilización de FLOPs del modelo en los dispositivos de entrenamiento, durante las pasadas hacia adelante y hacia atrás o durante todo el paso. Excluyen el servidor vLLM. Se registran solo cuando se conoce el cómputo máximo del dispositivo de entrenamiento a la precisión de entrenamiento y cuando la estimación de FLOPs es compatible con la arquitectura del modelo (como Llama, Qwen2, Qwen3, Qwen3-MoE o Qwen3-VL).

Personalización

Configuración del servidor vLLM

El trainer genera con un servidor vLLM, al que accede en vllm_server_base_url (por defecto, http://localhost:8000):

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096

• VLLM_SERVER_DEV_MODE=1 y --weight-transfer-config exponen los puntos finales por los que el entrenador envía sus pesos.
• --logprobs-mode processed_logprobs hace que el servidor devuelva las log-probabilidades de la distribución de la que muestreó, después de aplicar la temperatura y top-k/top-p, que es lo que compara la razón de importancia.
• --api-server-count 8 mantiene las log-probabilidades superiores para que las solicitudes de centrado de puntuaciones sean baratas de servir (vea la advertencia a continuación).
• --max-model-len es la longitud total máxima de la secuencia (prompt y completion) que esperas. Un valor más bajo reduce el uso de memoria GPU en el servidor, liberando más memoria para la caché KV y aumentando el throughput. Un buen punto de partida es la longitud del prompt más max_completion_length de tu configuración. Con herramientas, cada turno suma hasta max_completion_length tokens más los resultados de las herramientas, así que establécelo en la conversación más larga que esperes. Si el prompt de un turno más max_completion_length es más largo que --max-model-len, el servidor rechaza la solicitud y la ejecución se detiene con el error del servidor. Con top_k o top_p, el servidor en su lugar detiene ese turno cuando alcanza --max-model-len, y solo rechaza un prompt que ya sea de esa longitud. También es el valor predeterminado de token_budget.

Asegúrate de que el servidor utilice GPUs diferentes a las del entrenador, de lo contrario es posible que aparezcan errores de NCCL. Puedes especificar las GPUs a utilizar con la variable de entorno CUDA_VISIBLE_DEVICES.

La serialización de las log-probabilidades superiores es trabajo de CPU en el servidor de API de vLLM, y con su valor predeterminado de un solo servidor de API, num_sampler_logprobs=8 reduce a la mitad el rendimiento de generación. Inicie el servidor con --api-server-count 8 en vLLM 0.32 o posterior: medido en una H100 con modelos de 1.5B a 7B, las log-probabilidades superiores de 8 cuestan menos del 3% del rendimiento de generación. En vLLM 0.31, el mismo parámetro aún deja una pérdida del 20 al 40%, corregida aguas arriba por vllm-project/vllm#58181.

Concurrencia

Tres límites deciden cuántos rollouts se ejecutan a la vez. Los dos primeros se adaptan por sí solos; el tercero lo defines tú.

Antigüedad. max_staleness establece cuántos pasos de entrenamiento puede adelantarse la generación al entrenador. El worker rastrea esto en tokens, porque el número de muestras por paso cambia a medida que las completaciones se vuelven más largas o más cortas. Después de cada paso, el entrenador cuenta los tokens que utilizó y permite que el worker mantenga en cola o en proceso hasta max(max_staleness, 1) veces esa cantidad de tokens. Antes del primer paso, asume que cada fila está completa. El worker inicia un nuevo grupo de rollouts solo mientras esté por debajo de ese límite. Al principio, no sabe aún cuántos tokens produce un rollout, por lo que inicia un grupo a la vez hasta que el primer grupo termina. Una muestra que al final resulta ser más antigua que max_staleness pasos se descarta.

Carga de vLLM. Enviar demasiadas solicitudes a vLLM a la vez lo hace más lento: la caché KV se llena y vLLM tiene que pausar y repetir algunas solicitudes. Por lo tanto, el worker también mantiene un límite en vivo sobre cuántos rollouts se ejecutan a la vez, registrado como rollout/concurrency_cap, y comprueba los /metrics del servidor cada 5 segundos para ajustarlo. El límite comienza en el número de solicitudes que caben en la caché KV a longitud completa (KV cache tokens / max_model_len), de modo que el servidor no se sobrecargue al inicio. Mientras el uso de la caché KV se mantenga por debajo del 60% y no haya solicitudes en espera, el límite aumenta un 25%, como máximo una vez por cada cantidad de rollouts finalizados equivalente al límite. Cuando el uso supera el 80%, el límite disminuye un 30%: las solicitudes en ejecución continúan y las nuevas esperan. Si el servidor no tiene un endpoint /metrics, el límite es max_concurrent_rollouts.

Recursos. max_concurrent_rollouts (por defecto, 512) es el límite superior de rollouts que se ejecutan simultáneamente. Es el único límite que configuras tú, porque el trainer no puede ver lo que tus entornos y herramientas pueden soportar. Cada rollout en ejecución utiliza un entorno, puede ejecutar llamadas de herramientas síncronas en un thread pool de este tamaño y mantiene una conexión abierta con el servidor vLLM. Establece este valor en lo que esos componentes puedan soportar, por ejemplo, el número de sesiones que acepta tu servidor de entornos:

from trl.experimental.async_grpo import AsyncGRPOConfig

training_args = AsyncGRPOConfig(
    ...,
    max_concurrent_rollouts=64,
)

Si es demasiado bajo, la generación se ralentiza. Un rollout mantiene su slot mientras ejecuta herramientas o espera entre turnos, no solo mientras vLLM genera, por lo que los rollouts que llaman a muchas herramientas necesitan más slots para mantener vLLM ocupado. Si es mucho mayor que el valor predeterminado, el servidor de vLLM puede comenzar a descartar conexiones.

Las métricas registradas muestran qué límite está frenando la generación:

• rollout/backpressure_s por encima de cero: el límite de obsolescencia. La generación va por delante del entrenamiento, lo cual es esperado cuando la inferencia es más rápida que el entrenamiento.
• rollout/concurrency en rollout/concurrency_cap, por debajo de max_concurrent_rollouts: la carga del servidor de vLLM.
• rollout/concurrency en max_concurrent_rollouts, con un rollout/vllm_kv_cache_usage_frac bajo y un perf/rollout_wait_s alto: el límite. Aumenta max_concurrent_rollouts si tus entornos y herramientas lo permiten.

La cola de despliegue

El worker envía las muestras puntuadas a una cola y el trainer las extrae de ella. Cuatro métricas describen esa cola:

• rollout/queue_size: cuántas muestras están esperando.
• sample/time_in_queue_s: cuánto tiempo esperó una sola muestra antes de ser utilizada para el entrenamiento, la parte temporal de su obsolescencia.
• perf/rollout_wait_s: cuánto tiempo esperó el entrenamiento porque la cola estaba vacía.
• rollout/backpressure_s: cuánto tiempo esperó la generación porque iba max_staleness pasos por delante del entrenamiento.

Los dos últimos son imágenes especulares. Leídos junto con el tamaño de la cola, te indican qué lado es el cuello de botella:

• Cola casi vacía, perf/rollout_wait_s alto: limitado por la generación. El entrenador se queda sin trabajo; revisa rollout/generated_tokens_per_s, rollout/concurrency y rollout/score_queue_size.
• cola que mantiene aproximadamente max_staleness pasos de tokens, rollout/backpressure_s alto: limitado por el trainer. La generación se retrasa y sample/staleness_mean se sitúa cerca de su límite.
• ambos altos: el propio límite de max_staleness es demasiado estricto. Los rollouts en curso lo llenan antes de que sus muestras lleguen a la cola, por lo que el entrenador espera mientras la generación se retiene. Aumenta max_staleness.
• ambos cercanos a cero: equilibrado.

Saturación y duplicados

La política aprende de las diferencias entre los rollouts de un grupo: cada ventaja se calcula en relación con la recompedia media del grupo. Cuando todos los rollouts de un grupo obtienen la misma recompensa, todos reciben una ventaja cero, por lo que el grupo se descarta (ver Métricas registradas) y no enseña nada a la política. A medida que una ejecución se satura, cada vez más grupos terminan así, ya sea porque todos los rollouts resuelven el prompt o porque ninguno lo hace. Cuatro métricas distinguen estos casos, y cada una responde a su propia pregunta, por lo que puedes leerlas directamente de sus gráficos:

Lee las tres primeras juntas, en relación con el rango de tu recompensa:

completions/duplicate_frac es la proporción de rollouts que repiten un rollout anterior de su grupo, token por token, a lo largo de todos sus turnos. Los rollouts idénticos reciben recompensas idénticas, por lo que los duplicados también hacen subir frac_reward_zero_std.

Estas métricas se registran cuando el worker puntúa un grupo, no cuando se entrena con sus muestras, por lo que se ejecutan ligeramente por delante del entrenamiento e incluyen grupos cuyas muestras se descartan posteriormente por exceder max_staleness. Un grupo en el que no se pudo puntuar ningún rollout, porque todas las funciones de recompensa devolvieron None, se descarta y se cuenta en frac_reward_zero_std, pero como no tiene recompensa, reward_dropped lo excluye.

Empaquetado de árboles

Las muestras de un micro-lote se empaquetan en una fila por proceso, sin relleno. Una fila es un prefix trie, no una concatenación: un token compartido por varias muestras se propaga una sola vez, y la atención sigue mostrando a cada muestra exactamente su propio historial. Los rollouts de un grupo repiten todos su prompt, y cada turno de un rollout de varios turnos repite la conversación hasta ese momento, por lo que la propagación se reduce en la proporción de empaquetado (batch/packing_ratio). Es solo un cambio de disposición: cada token entrenado conserva su propio objetivo, ventaja y log-probabilidad anterior, y su log-probabilidad coincide con la de propagar la muestra por separado.

En dos despliegues de un solo prompt, donde p es el prompt y a, b las dos completaciones:

1. Rollouts         r0:  p1 p2 p3 a1 a2                10 tokens
                    r1:  p1 p2 p3 b1 b2

2. Trie             S0 [p1 p2] ─┬─ S1 [p3 a1 a2]       p3 is not shared: a position predicts
                                └─ S2 [p3 b1 b2]       one token, and p3 predicts a1 or b1

3. Packed row       index         0  1 |  2  3  4 |  5  6  7
   (depth-first)    input_ids    p1 p2 | p3 a1 a2 | p3 b1 b2    8 tokens: everything but attention
                    position_ids  0  1 |  2  3  4 |  2  3  4    runs on 8 instead of 10
                    segment      S0    | S1       | S2

4. Attention        segment   queries    keys and values: the segment's path in the trie
   (one call)       S0        p1 p2      p1 p2              index 0 1
                    S1        p3 a1 a2   p1 p2 p3 a1 a2     index 0 1 2 3 4
                    S2        p3 b1 b2   p1 p2 p3 b1 b2     index 0 1 5 6 7

5. Loss             r0: a1 a2 predicted at index 2 3   each trained token is predicted at the index
                    r1: b1 b2 predicted at index 5 6   of the token before it in its rollout

Algunas consecuencias:

• Cualquier prefijo común se comparte: los despliegues de un prompt, y entre prompts, un prompt de sistema o un esquema de herramienta común. Una muestra se envía a la fila que ya contiene su grupo, donde ya está su prompt. Cuando esa fila está llena, la muestra se envía a otra fila si un proceso aún tiene una vacía, o si eso no aumenta la fila más cargada, que determina el tiempo del micro-lote; de lo contrario, el micro-lote se cierra.
• Nothing shared is plain packing: cada muestra es entonces su propio segmento, lo que equivale a un empaquetado sin relleno ordinario.
• Solo se duplica K/V: un prefijo compartido se recopila una vez por segmento descendiente, lo cual es aproximadamente lo que retendría el entrenamiento sin empaquetar, y es pequeño con atención de consultas agrupadas.
• Se admiten los híbridos Gated DeltaNet (Qwen3.5, Qwen3.6, Qwen3-Next, …): su convolución lee los predecesores propios de cada token, y su regla delta por bloques inicia cada segmento desde el estado final de su padre.

token_budget limita los tokens de una fila, contando un token compartido una sola vez. Por defecto, es el max_model_len del servidor, por lo que ninguna muestra puede superarlo. Sin nada compartido, el tree packing se ejecuta a una velocidad similar a la del padding-free packing convencional, y la aceleración crece aproximadamente en proporción a la ratio de empaquetado (entre 2,8 y 3x con una ratio de 3,4 en Qwen3-0.6B y 1.7B, en H100). La ratio es el prefijo compartido por el número de muestras por fila, por lo que batch/samples_per_row es el indicador principal: cerca de 1, no se puede compartir nada, y la palanca es token_budget. batch/row_fill_frac importa cuando las muestras son largas: una muestra de 10k tokens llena mal un presupuesto de 32k (caben tres, cuatro nunca, por lo que el planificador a menudo obtiene dos y la fila queda llena al 77% aproximadamente), mientras que las muestras de 1k tokens lo llenan casi perfectamente.

Evalúe el empaquetado del árbol de juicio en función del rendimiento, batch/trained_tokens_per_step / perf/fwd_bwd_s. Un micro-lote empaquetado contiene más muestras, por lo que perf/fwd_bwd_s aumenta y perf/fwd_bwd_mfu, que divide entre los tokens procesados, disminuye: ninguno de los dos es una regresión. Más muestras por fila también significa más despliegues extraídos por paso, por lo que una ejecución limitada por la generación lo será aún más.

Las métricas por lote se componen de la siguiente manera: un paso contiene gradient_accumulation_steps × num_processes filas, por lo que batch/samples_per_step es aproximadamente ese número de veces batch/samples_per_row, y lo mismo ocurre con los tokens. Los dos lados solo coinciden hasta la variación entre los micro-lotes de un paso, ya que las métricas por paso son sumas y las por fila son medias.

Cualquier cosa que haga que las filas sean más largas o que haya menos filas elimina el uso compartido. Una conversación de varios turnos que nunca se ramifica es una sola muestra larga, y rollout/samples_per_rollout indica cuántas muestras produce una conversación.

Muestreo top-k y top-p

Top-k y top-p renormalizan la distribución de muestreo sobre un conjunto restringido de candidatos en lugar de sobre todo el vocabulario. Las log-probabilidades del servidor reflejan esto (--logprobs-mode processed_logprobs), por lo que el entrenador debe normalizar sobre el mismo conjunto; de lo contrario, todas las proporciones de importancia están sesgadas: la probabilidad de un token bajo todo el vocabulario es menor que bajo el conjunto de candidatos del que fue extraído (DeepSeek-V3.2, sección 3.1). Con top_p < 1 o top_k > 0, el worker solicita al servidor el conjunto de candidatos de cada token muestreado y el entrenador renormaliza sus log-probabilidades sobre dicho conjunto. Esto requiere dos banderas adicionales del servidor, y top_k debe estar establecido, ya que vLLM limita el tamaño del conjunto de candidatos con él:

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096 \
    --return-sampling-mask \
    --enable-scale-out

from trl.experimental.async_grpo import AsyncGRPOConfig

training_args = AsyncGRPOConfig(
    ...,
    top_p=0.95,
    top_k=512,
)

--return-sampling-mask hace que vLLM registre el conjunto de candidatos (vLLM 0.28 o posterior, en su ejecutor de modelos V2: VLLM_USE_V2_MODEL_RUNNER=1 si el servidor rechaza la bandera), y --enable-scale-out expone el punto de entrada de entrada-salida de tokens que lo devuelve. El worker comprueba ambos en el servidor al iniciar. Una versión anterior de vLLM no puede devolver el conjunto, por lo que el entrenador sigue normalizando sobre todo el vocabulario y emite una advertencia sobre la proporción sesgada.

Entrenar adaptadores con PEFT

Soportamos una integración estrecha con la biblioteca 🤗 PEFT, lo que te permite entrenar adaptadores y compartirlos en el Hub en lugar de entrenar el modelo completo.

from datasets import load_dataset
from peft import LoraConfig
from trl.experimental.async_grpo import AsyncGRPOTrainer
from trl.rewards import accuracy_reward

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

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3.5-2B",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)
trainer.train()

Cada sincronización de pesos debe enviar la política actualizada al servidor. De forma predeterminada, el entrenador fusiona el adaptador con el modelo base y envía los pesos fusionados completos, lo que equivale a unos 15 GB para un modelo de 7B en cada sincronización. Inicie el servidor vLLM con --enable-lora y solo enviará el adaptador, aproximadamente 160 MB con r=32:

VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096 \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 expone el endpoint al que el entrenador publica cada nuevo adaptador. Mantén --weight-transfer-config incluso con --enable-lora: el entrenador solo elige un modo de sincronización al iniciar, momento en el que el servidor ya está en marcha, y la sincronización combinada es la opción de respaldo.

--max-lora-rank debe ser uno de 1, 8, 16, 32, 64, 128, 256, 320, 512. Establece el rango máximo que el servidor puede servir, en lugar del rango que servirá, por lo que un adaptador r=4 funciona bien bajo 8. --max-loras debe ser al menos max_staleness + 2: el entrenador mantiene registradas max_staleness + 1 versiones de adaptadores para que un rollout que comenzó bajo una política anterior pueda completarse bajo ella en lugar de cambiar de política a mitad de la generación, y cada sincronización carga la siguiente versión antes de descargar la más antigua. El entrenador comprueba ambos valores al iniciar y te indica con qué valores reiniciar el servidor.

vLLM carga los adaptadores desde el disco en lugar de a través de la red, por lo que el entrenador escribe cada versión bajo <output_dir>/.vllm_lora/ y le proporciona la ruta al servidor. Si el servidor se ejecuta en una máquina diferente, ese directorio debe estar en un sistema de archivos al que ambos puedan acceder.

<output_dir>/.vllm_lora/ es una caché de servicio, no el lugar donde se guarda el adaptador entrenado. Solo contiene las versiones más recientes y cada una se elimina poco después de salir de la ventana de obsolescencia. Los checkpoints son los que preservan el adaptador que entrenaste, así que no los desactives: con save_strategy="no", una ejecución de varias horas termina sin dejar nada. output_dir también debe ser legible por el servidor y debe sobrevivir al trabajo, por lo que un espacio temporal local al nodo no funcionará, aunque evitaría las sincronizaciones en un sistema de archivos compartido lento.

La sincronización de adaptadores no siempre es más rápida. vLLM genera aproximadamente 1,33 a 1,39 veces más lento por token cuando sirve un adaptador LoRA en comparación con cuando sirve pesos fusionados, y la generación suele ser lo que domina el tiempo real en RL. A cambio, se obtiene una pausa mucho más corta en cada sincronización y aproximadamente el 1% de la banda ancha de sincronización. Si el rendimiento de generación es su cuello de botella, la sincronización fusionada suele ser la mejor opción.

Uso de una función de recompensa personalizada

El experimental.async_grpo.AsyncGRPOTrainer admite el uso de funciones de recompensa personalizadas. Para garantizar la compatibilidad, su función de recompensa debe cumplir con los siguientes requisitos:

Las funciones de recompensa pueden ser funciones Python síncronas o corrutinas asíncronas async def. Cuando se proporcionan varias funciones de recompensa, se ejecutan de forma concurrente: las asíncronas se esperan juntas y las síncronas se ejecutan en hilos, de modo que su latencia se solapa.

1. Argumentos de entrada:

   • La función debe aceptar lo siguiente como argumentos de palabra clave:

     • prompts (contiene los prompts),
     • completions (contiene las completaciones generadas),
     • completion_ids (contiene las completas tokenizadas),
     • Todos los nombres de columnas (excepto prompt) que el conjunto de datos puede tener. Por ejemplo, si el conjunto de datos contiene una columna llamada ground_truth, la función se llamará con ground_truth como argumento de palabra clave.

     La forma más sencilla de cumplir con este requisito es usar **kwargs en la firma de la función.

   • El conjunto de datos está en formato conversacional, por lo que prompts y completions son listas de diccionarios de mensajes.

2. Valor de retorno: La función debe devolver una lista de floats. Cada float representa la recompensa correspondiente a una única finalización.

Las funciones de recompensa se ejecutan en el worker de rollout, un proceso separado, por lo que todo lo que se le pasa se serializa con pickle. Cada función de recompensa, herramienta y environment_factory (y cualquier cosa que cierren) debe poder serializarse con pickle: usa una función a nivel de módulo, functools.partial o una instancia de clase invocable. Las lambdas y los cierres provocan un TypeError en trainer.train(). Esta es una diferencia con GRPOTrainer, donde las funciones de recompensa se invocan dentro del proceso y los cierres funcionan.

El proceso de rollout también se ejecuta con CUDA_VISIBLE_DEVICES="", por lo que no puede usar la GPU. Un modelo de recompensa basado en GPU (por ejemplo, un evaluador AutoModelForSequenceClassification) se carga sin errores, pero recurre silenciosamente a la CPU. Mantenga las funciones de recompensa ligeras (verificadores como accuracy_reward, comprobaciones de formato o longitud). Si necesita un modelo de recompensa en GPU, sírvalo detrás de su propio motor de inferencia (vLLM, TGI, …) en GPUs separadas y haga que una función de recompensa serializable lo llame a través de HTTP.

Ejemplo 1: Recompensar las finalizaciones más largas

A continuación se muestra un ejemplo de una función de recompensa que premia las finalizaciones más largas:

def reward_func(completion_ids, **kwargs):
    """Reward function that assigns higher scores to longer completions (in terms of token count)."""
    return [float(len(ids)) for ids in completion_ids]

Puedes probarlo de la siguiente manera:

>>> prompts = [[{"role": "user", "content": "The sky is"}], [{"role": "user", "content": "The sun is"}]]  # not used in the reward function, but the trainer will pass it
>>> completions = [[{"role": "assistant", "content": " blue."}], [{"role": "assistant", "content": " in the sky."}]]  # not used in the reward function, but the trainer will pass it
>>> completion_ids = [[6303, 13], [304, 279, 12884, 13]]
>>> reward_func(prompts=prompts, completions=completions, completion_ids=completion_ids)
[2.0, 4.0]

Ejemplo 2: Recompensar las finalizaciones con un formato específico

A continuación se muestra un ejemplo de una función de recompensa que verifica si la finalización tiene un formato específico. Este ejemplo se inspira en la función format reward utilizada en el artículo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Está diseñada para un formato conversacional, donde los prompts y las finalizaciones consisten en mensajes estructurados.

import re

def format_reward_func(completions, **kwargs):
    """Reward function that checks if the completion has a specific format."""
    pattern = r"^<think>.*?</think><answer>.*?</answer>$"
    completion_contents = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, content) for content in completion_contents]
    return [1.0 if match else 0.0 for match in matches]

Puedes probar esta función de la siguiente manera:

>>> prompts = [
...     [{"role": "user", "content": "What is the result of (1 + 2) * 4?"}],
...     [{"role": "user", "content": "What is the result of (3 + 1) * 2?"}],
... ]
>>> completions = [
...     [{"role": "assistant", "content": "<think>The sum of 1 and 2 is 3, which we multiply by 4 to get 12.</think><answer>(1 + 2) * 4 = 12</answer>"}],
...     [{"role": "assistant", "content": "The sum of 3 and 1 is 4, which we multiply by 2 to get 8. So (3 + 1) * 2 = 8."}],
... ]
>>> format_reward_func(prompts=prompts, completions=completions)
[1.0, 0.0]

Ejemplo 3: Recompensar las finalizaciones basándose en una referencia

A continuación se muestra un ejemplo de una función de recompensa que verifica si la finalización es correcta. Este ejemplo se inspira en la función de recompensa de precisión utilizada en el artículo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. El conjunto de datos contiene una columna llamada ground_truth.

import re

def reward_func(completions, ground_truth, **kwargs):
    # Regular expression to capture content inside \boxed{}
    matches = [re.search(r"\\boxed\{(.*?)\}", completion[0]["content"]) for completion in completions]
    contents = [match.group(1) if match else "" for match in matches]
    # Reward 1 if the content is the same as the ground truth, 0 otherwise
    return [1.0 if c == gt else 0.0 for c, gt in zip(contents, ground_truth)]

Puedes probar esta función de la siguiente manera:

>>> prompts = [
...     [{"role": "user", "content": "Solve the equation $2x + 3 = 7$."}],
...     [{"role": "user", "content": "Solve the equation $3x - 5 = 10$."}],
... ]
>>> completions = [
...     [{"role": "assistant", "content": r"The solution is \boxed{2}."}],
...     [{"role": "assistant", "content": r"The solution is \boxed{6}."}],
... ]
>>> ground_truth = ["2", "5"]
>>> reward_func(prompts=prompts, completions=completions, ground_truth=ground_truth)
[1.0, 0.0]

Ejemplo 4: Funciones de recompensa multitarea

A continuación se muestra un ejemplo de uso de múltiples funciones de recompensa en experimental.async_grpo.AsyncGRPOTrainer. En este ejemplo, definimos dos funciones de recompensa específicas para cada tarea: math_reward_func y coding_reward_func. La función math_reward_func otorga recompensas a los problemas de matemáticas en función de su corrección, mientras que la función coding_reward_func otorga recompensas a los problemas de programación en función de si la solución funciona.

from datasets import Dataset
from trl.experimental.async_grpo import AsyncGRPOTrainer

# Define a dataset that contains both math and coding problems
dataset = Dataset.from_list(
    [
        {"prompt": [{"role": "user", "content": "What is 2+2?"}], "task": "math"},
        {"prompt": [{"role": "user", "content": "Write a function that returns the sum of two numbers."}], "task": "code"},
        {"prompt": [{"role": "user", "content": "What is 3*4?"}], "task": "math"},
        {"prompt": [{"role": "user", "content": "Write a function that returns the product of two numbers."}], "task": "code"},
    ]
)

# Math-specific reward function
def math_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "math":
            # Calculate math-specific reward
            correct = check_math_solution(prompt, completion)
            reward = 1.0 if correct else -1.0
            rewards.append(reward)
        else:
            # Return None for non-math tasks
            rewards.append(None)
    return rewards

# Coding-specific reward function
def coding_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "code":
            # Calculate coding-specific reward
            works = test_code_solution(prompt, completion)
            reward = 1.0 if works else -1.0
            rewards.append(reward)
        else:
            # Return None for non-coding tasks
            rewards.append(None)
    return rewards

# Use both task-specific reward functions
trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    reward_funcs=[math_reward_func, coding_reward_func],
    train_dataset=dataset,
)

trainer.train()

En este ejemplo, math_reward_func y coding_reward_func están diseñados para funcionar con un conjunto de datos mixto que contiene tanto problemas de matemáticas como de programación. La columna task del conjunto de datos se utiliza para determinar qué función de recompensa aplicar a cada problema. Si no hay una función de recompensa relevante para una muestra del conjunto de datos, la función de recompensa devolverá None y el experimental.async_grpo.AsyncGRPOTrainer continuará con las funciones y tareas válidas. Esto permite que el experimental.async_grpo.AsyncGRPOTrainer maneje múltiples funciones de recompensa con diferentes alcances de aplicación.

Tenga en cuenta que experimental.async_grpo.AsyncGRPOTrainer ignorará las recompensas None devueltas por las funciones de recompensa y solo considerará las recompensas devueltas por las funciones relevantes. Esto garantiza que el modelo se entrene en las tareas relevantes e ignore las tareas para las que no existe una función de recompensa relevante.

Ejemplo 5: Funciones de recompensa asíncronas

Las funciones de recompensa personalizadas también pueden definirse como corrutinas async def. Esto resulta útil si la recompensa depende de una E/S lenta (por ejemplo, al llamar a un servicio remoto). Al pasar varias funciones de recompensa asíncronas, experimental.async_grpo.AsyncGRPOTrainer las ejecuta de forma concurrente para que su latencia se solape.

A continuación se muestra un ejemplo mínimo de una función de recompensa asíncrona que simula una operación limitada por E/S:

import asyncio

async def async_reward_func(prompts, completions, **kwargs):
    # Simulate an I/O-bound call (e.g., HTTP request, database lookup)
    await asyncio.sleep(0.01)
    # Simple toy reward: 1.0 if the completion is non-empty, else 0.0
    return [1.0 if completion else 0.0 for completion in completions]

Pasar la función de recompensa al entrenador

Para usar su función de recompensa personalizada, pásela al experimental.async_grpo.AsyncGRPOTrainer de la siguiente manera:

from trl.experimental.async_grpo import AsyncGRPOTrainer

trainer = AsyncGRPOTrainer(
    reward_funcs=reward_func,
    ...,
)

Puedes pasar varias funciones de recompensa como una lista; esta lista puede incluir tanto funciones síncronas como asíncronas:

from trl.experimental.async_grpo import AsyncGRPOTrainer

trainer = AsyncGRPOTrainer(
    reward_funcs=[reward_func, async_reward_func1, async_reward_func2],
    ...,
)

y la recompensa se calculará como la suma de las recompensas de cada función.

Personalizar el bucle de despliegue

La generación, las recompensas, las ventajas y el filtrado se ejecutan en el proceso de rollout, en experimental.async_grpo.AsyncRolloutLoop. Para modificar uno de ellos, subclasea el bucle, sobrescribe el método de ese paso y establece la subclase como loop_cls en una subclase del trainer. El constructor del trainer no cambia.

Por ejemplo, un plan de estudios que deja de muestrear los prompts que la política ya resuelve, con su umbral en una subclase de configuración:

from dataclasses import dataclass

import numpy as np

from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer, AsyncRolloutLoop


@dataclass
class CurriculumConfig(AsyncGRPOConfig):
    pass_rate_cap: float = 0.9


class CurriculumRolloutLoop(AsyncRolloutLoop):
    def filter_prompt(self, row_index, row, epoch):
        return self.pass_rates.get(row_index, 0.0) < self.config.pass_rate_cap

    async def setup(self):
        self.pass_rates = {}

    def compute_advantages(self, group):
        scored = group.rewards[~np.isnan(group.rewards)]
        if len(scored):
            self.pass_rates[group.row_index] = float(np.mean(scored > 0))
        return super().compute_advantages(group)


class CurriculumTrainer(AsyncGRPOTrainer):
    loop_cls = CurriculumRolloutLoop

Las recompensas se suman sobre reward_funcs, los entornos que definen get_reward y la reward que un generate_one personalizado establece en su experimental.async_grpo.Rollout: para puntuar con un juez, pase una función de recompensa asíncrona, o puntúe dentro de generate_one. compute_advantages se ejecuta para cada grupo, incluido uno que luego se descarta porque todas sus recompensas puntuadas son iguales; el descarte mira las recompensas, no las ventajas. El bucle se construye en un proceso generado, así que defina la subclase a nivel de módulo. Cree los clientes en setup, que se ejecuta en el bucle de eventos del bucle, en lugar de sobrescribir __init__. Un prompt que filter_prompt rechaza se omite para esa época y cuenta hacia ella, como un grupo descartado. Si rechaza todos los prompts de una época, el bucle lanza una excepción.

Reanudar desde un punto de control

El Trainer base se reanuda omitiendo y reproduciendo los lotes en los que ya se entrenó, lo cual no se aplica a una cola de despliegue en vivo, por lo que ignore_data_skip siempre es True. En su lugar, el índice del primer prompt que aún no se ha entrenado se guarda en rollout_state.json junto con cada punto de control y se restaura al reanudar, de modo que el trabajador avanza rápidamente hasta ese prompt sin reproducir muestras. Es la posición trained, no la del generador: el trabajador va por delante del entrenamiento en la profundidad de la cola de despliegue, y esas muestras en búfer se pierden cuando termina la ejecución, por lo que reanudar desde la posición del generador omitiría prompts que se generaron pero nunca se entrenaron. Los conjuntos de datos de streaming (IterableDataset) no se pueden reposicionar, por lo que su trabajador se reinicia desde el primer prompt al reanudar.

Entrenamiento de agentes

GRPO asíncrono admite entrenamiento de agentes: el modelo llama a herramientas durante la generación y aprende del resultado.

• Una herramienta es una función Python simple (síncrona o asíncrona) expuesta al modelo. Utilice tools para llamadas sin estado (una calculadora, una búsqueda web).
• Un environment es la forma más general: un objeto con estado creado de nuevo para cada rollout, cuyos métodos públicos se exponen como herramientas, más un hook de ciclo de vida reset y un get_reward opcional que le permite poseer la recompensa. Use environment_factory cuando necesite estado por rollout, un hook de reset o una recompensa poseída por el environment.

Se combinan: puedes pasar tools independientes junto con un environment_factory.

Herramientas

El argumento tools espera una lista de funciones de Python (síncronas o asíncronas) que definen las herramientas disponibles para el agente:

from trl.experimental.async_grpo import AsyncGRPOTrainer

trainer = AsyncGRPOTrainer(
    tools=[tool1, tool2],
    ...,
)

Cada herramienta debe ser una función de Python estándar con argumentos y tipos de retorno con anotaciones de tipo, junto con una docstring en estilo Google que describa su propósito, sus argumentos y su valor de retorno. Para más detalles, consulte la guía para pasar herramientas.

El modelo invoca una herramienta por su nombre, por lo que cada herramienta se registra bajo su __name__ y debe ser una función a nivel de módulo, lo que también la mantiene serializable para el worker de rollout. Las herramientas síncronas se ejecutan en un pool de hilos dimensionado según max_concurrent_rollouts, de modo que una herramienta lenta nunca bloquea los otros rollouts en curso. Dentro de un turno asistente, las llamadas a herramientas se ejecutan en el orden en que el modelo las emitió. Por lo tanto, hasta max_concurrent_rollouts herramientas síncronas pueden ejecutarse a la vez, por lo que una herramienta que lea o escriba estado compartido debe ser segura para hilos.

El bucle de llamadas a herramientas requiere que la plantilla de chat sea de prefijo conservador (añadir un mensaje de herramienta no debe cambiar cómo se renderizan los mensajes anteriores). Para las familias de modelos conocidas (p. ej., Qwen3, DeepSeek-V3), TRL cambia automáticamente a una plantilla de entrenamiento parcheada cuando se habilitan las herramientas. Consulte Plantillas de chat para obtener la lista completa.

Use max_tool_calling_iterations en experimental.async_grpo.AsyncGRPOConfig para limitar el número de turnos de llamada a herramientas. Por defecto no hay límite, y la generación se detiene cuando el modelo produce un turno de respuesta sin llamadas a herramientas.

Ejemplo:

from trl.experimental.async_grpo import AsyncGRPOTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

async def async_add(a: int, b: int) -> int:
    """
    Asynchronously adds two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The sum of the two integers.
    """
    return a + b

trainer = AsyncGRPOTrainer(
    tools=[multiply, async_add],
    ...,
)

Entornos

También puede proporcionar herramientas a través de environment_factory. En este modo, experimental.async_grpo.AsyncGRPOTrainer crea una instancia de entorno por rollout y expone los métodos públicos del entorno como herramientas.

Cuando se utiliza un environment_factory, el entorno posee los datos: no existe un train_dataset externo. En cada rollout, reset() genera la tarea —muestreando de forma autónoma de un corpus que el entorno mantiene, o generando un estado de forma procedural— y devuelve el prompt. max_steps establece la duración del entrenamiento. Lo siguiente es un ejemplo mínimo de un entorno que muestrea de forma autónoma un objetivo y expone un método increment (expuesto como una herramienta).

import random

from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

class IncrementEnv:
    # Reserved methods
    def reset(self, **kwargs) -> str | None:  # required; called at the start of each rollout
        self.counter = 0
        self.target = random.randint(1, 6)  # self-sample the task
        return f"Increment the counter by {self.target}."  # returned string becomes the prompt

    def get_reward(self) -> float:  # optional: the environment scores itself from its own state
        return float(self.counter == self.target)

    # Public methods (exposed as tools)
    def increment(self, step: int) -> int:
        """
        Increment the internal counter.

        Args:
            step: Value to add to the counter.

        Returns:
            The updated counter value.
        """
        self.counter += step
        return self.counter

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=AsyncGRPOConfig(max_steps=1000, chat_template_kwargs={"enable_thinking": False}),
    environment_factory=IncrementEnv,
)
trainer.train()

Una clase de entorno tiene dos métodos reservados: reset y get_reward. A diferencia de cualquier otro método público, estos no se exponen al modelo como herramientas.

• reset (obligatorio) se llama cuando comienza el rollout. Puede devolver None o una cadena de caracteres. Si devuelve una cadena, esa cadena es el prompt del usuario.
• get_reward (opcional, síncrono o asíncrono) no toma argumentos y devuelve un float: el entorno puntúa el episodio que acaba de ejecutar a partir de su propio estado interno (¿el juego terminó en una victoria? ¿se adivinó la palabra?). Se llama una vez por cada rollout completado y actúa como fuente de recompensa.

Ambos se ejecutan en el pool de hilos de herramientas del worker, por lo que un entorno lento retrasa su propio despliegue, pero no a los demás en curso. Sus duraciones se registran como env/reset_s y env/reward_s.

Dado que el entorno se muestrea a sí mismo en cada reset(), los G miembros de un grupo GRPO pueden no compartir el mismo estado inicial, lo que hace que la línea base del grupo sea ligeramente más ruidosa. Para estados idénticos dentro de un grupo, haz que reset() derive el estado de forma determinista a partir de una clave compartida.

Proporcionar un conjunto de datos externo (opcional)

En lugar de dejar que el entorno muestreea por sí mismo, puedes proporcionar un train_dataset externo. Por ejemplo, cuando tus tareas ya están en un conjunto de datos, o para enrutar entre múltiples entornos. Por cada rollout, el entrenador muestrea una fila, usa su columna "prompt" como prompt y pasa las columnas de la fila a reset() como argumentos de palabra clave (de modo que el entorno lea la tarea del conjunto de datos).

from datasets import Dataset
from trl.experimental.async_grpo import AsyncGRPOTrainer

# Each row carries the prompt and the task data (here, the counter target to reach).
dataset = Dataset.from_dict(
    {
        "prompt": [[{"role": "user", "content": f"Increment the counter to {i}."}] for i in range(1, 7)],
        "target": list(range(1, 7)),
    }
)

class IncrementEnv:
    def reset(self, target, **kwargs) -> None:  # the row's `target` column arrives as a keyword argument
        self.counter = 0
        self.target = target

    def get_reward(self) -> float:
        return float(self.counter == self.target)

    def increment(self, step: int) -> int:  # exposed as a tool (see the full docstring above)
        self.counter += step
        return self.counter

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=dataset,  # the dataset owns the task; the "target" column reaches reset(target=...)
    environment_factory=IncrementEnv,
)
trainer.train()

Recompensas

Las recompensas gestionadas por el entorno y las gestionadas por el entrenador no son excluyentes. Pasa reward_funcs para recompensas que es natural calcular a partir de la finalización (p. ej., una comprobación de formato del texto) y deja que el entorno devuelva su recompensa dependiente del estado a través de get_reward. Todas las fuentes se suman con peso 1.

def format_reward(completions, **kwargs):  # trainer-owned: scores the completion text
    ...

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=AsyncGRPOConfig(max_steps=1000),
    reward_funcs=format_reward,          # trainer-owned: scores the completion
    environment_factory=IncrementEnv,    # env-owned: scores its internal state via `get_reward`
)

Varios entornos

Para entrenar con una mezcla de tareas en una sola ejecución (por ejemplo, una tarea de programación y un juego), pase un diccionario que mapee nombres de entornos a fábricas en lugar de un único invocable. Cada rollout selecciona su entorno a través de un campo environment en el conjunto de datos, y solo se exponen las herramientas de ese entorno. Esto evita filtrar herramientas irrelevantes (la herramienta move del juego a un ejemplo de programación, y viceversa). Cada entorno sigue siendo propietario de sus datos y devuelve el prompt desde reset(), por lo que el conjunto de datos solo necesita la columna environment para enrutar entre ellos.

import random

from datasets import Dataset
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer


class CodingEnv:
    def reset(self, **kwargs) -> str:
        n = random.randint(1, 20)
        return f"Compute the {n}th Fibonacci number."

    def run_code(self, code: str) -> str:  # exposed as a tool only for "coding" examples
        """Run the given Python code and return its stdout.

        Args:
            code: The Python source to execute.

        Returns:
            The captured standard output.
        """
        ...

class GameEnv:
    def reset(self, **kwargs) -> str:
        return "Reach the goal tile."

    def move(self, action: str) -> str:  # exposed as a tool only for "game" examples
        """Apply an action to the game and return the new observation.

        Args:
            action: One of "up", "down", "left", "right".

        Returns:
            The observation after the move.
        """
        ...

# The dataset only routes: each row picks the environment for that rollout
dataset = Dataset.from_dict({"environment": ["coding", "game"] * 500})

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=AsyncGRPOConfig(chat_template_kwargs={"enable_thinking": False}),
    train_dataset=dataset,
    reward_funcs=reward_func,
    environment_factory={"coding": CodingEnv, "game": GameEnv},
)
trainer.train()

Las instancias de entorno se reutilizan entre pasos (se reinician entre episodios), por lo que un __init__ costoso se paga una vez en lugar de en cada paso. El campo environment lo consume el entrenador y no se reenvía a reset.

Despliegues de varios turnos

En cada turno, se renderiza toda la conversación a través de la plantilla de chat y se genera a partir de ella. Una plantilla de chat no siempre renderiza una conversación en crecimiento de la misma manera (por ejemplo, algunas descartan el razonamiento de turnos anteriores), por lo que los tokens de entrada de un turno pueden diferir de los tokens acumulados hasta ese momento. Al final del despliegue, el trabajador reconcilia los turnos en muestras de entrenamiento: un turno cuya entrada comienza con los tokens acumulados se añade a la misma muestra, y cualquier otro turno inicia una nueva. Cada token generado se entrena en exactamente una muestra, por lo que una bifurcación tiene un costo en tokens de propagación hacia adelante, pero no en señal de entrenamiento.

Cada muestra de un rollout recibe la ventaja del rollout. rollout/fork_frac y rollout/drift_tokens_* muestran cómo se reconciliaron los turnos.

Integraciones de entorno

Para entrenar agentes que se ejecutan dentro de su propio entorno, como agentes de codificación en los entornos aislados de Harbor, consulte Entrenamiento en entornos en la guía de OpenEnv.

Entrenamiento de modelos de visión y lenguaje

Los modelos de visión y lenguaje (Qwen3.5, Qwen3.6, Qwen3-VL, …) se pueden entrenar en conjuntos de datos de solo texto. Pasa el identificador del modelo como de costumbre; no cambia nada más.

La torre de visión está congelada. Un conjunto de datos solo de texto nunca genera tokens de imagen, por lo que la torre nunca se utiliza en la pasada hacia adelante. Todo lo que está fuera del modelo de texto (torre de visión, proyector multimodal) tiene requires_grad=False: no recibe gradientes ni estado del optimizador, y la sincronización de pesos lo omite, por lo que el servidor mantiene los valores que cargó desde el punto de control. La torre sigue cargándose, por lo que consume memoria de GPU para sus pesos.

Las imágenes aún no son compatibles. Los prompts que contienen imágenes no se envían al servidor de vLLM ni a la pasada de entrenamiento. El entrenamiento multimodal también requiere la ruta de empaquetado sin relleno para construir posiciones M-RoPE 3D a partir de la cuadrícula de la imagen, lo cual no está implementado aquí.

AsyncGRPOTrainer

>>> from trl.experimental.async_grpo import AsyncGRPOTrainer
>>> from trl.rewards import accuracy_reward
>>> from datasets import load_dataset

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

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

>>> class CurriculumTrainer(AsyncGRPOTrainer):
...     loop_cls = CurriculumRolloutLoop

AsyncGRPOConfig

AsyncRolloutLoop

Rollout

TurnRecord

RolloutGroup