Principle:Hpcaitech ColossalAI GRPO Consumer Setup

Knowledge Sources	ColossalAI DeepSeekMath: Pushing the Limits of Mathematical Reasoning
Domains	Reinforcement_Learning, Distributed_Computing
Last Updated	2026-02-09 00:00 GMT

Overview

A training worker pattern that receives experiences from producers and updates the policy model using the GRPO objective with ColossalAI distributed training.

Description

The GRPO Consumer is a training worker that receives experience batches (generated responses, log probabilities, rewards, advantages) from producer actors and updates the policy model. It uses ColossalAI's Booster for distributed training, supporting ZeRO and hybrid parallelism. After each update, it broadcasts updated weights back to producers.

Usage

Consumers are automatically created by launch_distributed(). Configure the number of consumer GPUs based on model size and available memory.

Theoretical Basis

The consumer minimizes the GRPO policy loss with importance sampling:

${ℒ}=-\mathbb{𝔼}[\min (\frac{{\pi }_{\theta }}{{\pi }_{old}}A,\text{clip}(\frac{{\pi }_{\theta }}{{\pi }_{old}},1-ϵ,1+ϵ)A)]+\beta \cdot KL({\pi }_{\theta }||{\pi }_{ref})$

Related Pages

Implemented By

• Implementation:Hpcaitech_ColossalAI_GRPOConsumer

Heuristic Links

• Heuristic:Hpcaitech_ColossalAI_Empty_Cache_Between_Phases