Este artículo describe el proceso de construir, desde cero, un bucle de entrenamiento de aprendizaje por refuerzo (RL) para modelos de lenguaje grandes, aplicándolo al algoritmo Dr. GRPO sobre el modelo Qwen2.5-0.5B-Instruct y la tarea GSM8K, con una sola GPU A10G. El punto de partida es una observación estructural: la fase de rollout, que ejecuta la generación con model.generate, domina el tiempo total porque realiza un forward pass por cada token generado de forma secuencial, mientras que la fase de update se reduce a tres llamadas masivas a la GPU.
El autor primero monta el bucle como un esqueleto en PyTorch (unas 300 líneas en grpo.py) y, antes de escribir ningún kernel, consigue acelerar la fase de rollout 4,8×. Después identifica la ruta que aún domina y escribe un kernel de atención de decodificación fusionado con CuteDSL. En microbenchmark, el kernel es 2,2× más rápido que la ruta SDPA que reemplaza. Sin embargo, al integrarlo en la función generate de HuggingFace, el paso de decodificación se vuelve casi 3× más lento. El motivo no está en el kernel, que hacía exactamente lo que el microbench indicaba, sino en que la integración rompió una ruta de auto-compilación de la que la implementación base se beneficiaba en silencio.
El texto también repasa la evolución de PPO a GRPO y a Dr. GRPO. GRPO elimina la red de valor de PPO y emplea la media del grupo de completions como línea base. Dr. GRPO corrige dos sesgos del GRPO original: el de longitud, donde las respuestas largas recibían una penalización por token más débil, y el de dificultad, donde dividir por la desviación estándar amplificaba los gradientes en grupos muy fáciles o muy difíciles. Para ello suprime ambos denominadores y agrega por suma de tokens en lugar de media por respuesta. El artículo cierra con detalles prácticos: la máscara de completion con la trampa del argmax cuando no hay EOS, y el cálculo de log-probs con su correspondiente trampa off-by-one.
El proyecto es de código abierto. La lección implícita es que optimizar un kernel sin entender las rutas de compilación e integración del framework anfitrión puede producir regresiones mucho mayores que las ganancias medidas en aislación.