Implementation:Hpcaitech ColossalAI RL Example Zero Bubble

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Reinforcement Learning, GRPO, Distributed Training, Ray
Last Updated	2026-02-09 00:00 GMT

Overview

rl_example_zero_bubble.py is an entry-point script for launching GRPO (Group Relative Policy Optimization) or DAPO (Dynamic Advantage Policy Optimization) training with a zero-bubble distributed architecture using Ray for producer-consumer coordination.

Description

This script configures and launches a distributed reinforcement learning training pipeline that separates inference (generation) and training into producer and consumer processes. Producers generate rollout samples using either a Transformers or vLLM backend, while consumers train the policy model using GRPO or DAPO algorithms with ColossalAI's parallelism plugins. The zero-bubble architecture minimizes idle time by overlapping inference and training through a buffer-based data exchange mechanism. The script supports configurable batch sizes at multiple levels (inference, train, mini-batch, micro-batch), tensor/pipeline parallelism, and multiple reward types including think-answer tags and boxed format for math reasoning tasks.

Usage

Use this script for large-scale RLHF training with GRPO or DAPO algorithms, particularly for math reasoning tasks that benefit from structured response formats. It is designed for multi-GPU or multi-node setups where inference and training can be distributed across separate GPU groups. Launch directly with Python (Ray handles distribution internally).

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalChat/rl_example_zero_bubble.py
• Lines: 1-378

Signature

# Entry point script; main logic is in __main__ block
# Calls: launch_distributed(...) from coati.distributed.launch_zero_bubble

Import

from coati.distributed.launch_zero_bubble import launch_distributed

I/O Contract

Inputs

Name	Type	Required	Description
-m, --model	str	No	Model path or name (default: Qwen/Qwen2.5-7B)
-d, --dataset	str	No	Training dataset path (default: data.jsonl)
-ed, --eval-dataset	str	No	Evaluation dataset config in JSON format (task:path pairs)
-a, --algo	str	No	Algorithm: GRPO or DAPO (default: GRPO)
-t, --num-trainers	int	No	Number of trainer (consumer) processes (default: 2)
-i, --num-inferencer	int	No	Number of inferencer (producer) processes (default: 2)
-g, --num-generations	int	No	Number of generations per prompt (default: 8)
-ibs, --inference-batch-size	int	No	Prompts per inference step (default: 64)
-tbs, --train-batch-size	int	No	Unique prompts per training step per DP group (default: 32)
-tMbs, --train-minibatch-size	int	No	Unique prompts per training mini-batch per DP group (default: 8)
-tmbs, --train-microbatch-size	int	No	Effective batch size for forward/backward (default: 2)
-b, --backend	str	No	Inference backend: transformers or vllm (default: transformers)
-tp, --tensor-parallel-size	int	No	TP size for trainer (default: 1)
-pp, --pipeline-parallel-size	int	No	PP size for trainer (default: 1)
-zero, --zero-stage	int	No	ZeRO stage for trainer (default: 0)
-lr, --learning-rate	float	No	Learning rate (default: 1e-6)
-kl, --kl-coeff	float	No	KL penalty coefficient (default: 0.01)
-rt, --reward-type	str	No	Reward type: think_answer_tags, boxed, code (default: think_answer_tags)
-temp, --temperature	float	No	Sampling temperature (default: 1.0)
-mnt, --max-new-tokens	int	No	Maximum new tokens to generate (default: 3584)
-mpt, --max-prompt-tokens	int	No	Maximum prompt tokens (default: 512)
-si, --save-interval	int	No	Checkpoint save interval in training steps (default: 100)
-sd, --save-dir	str	No	Checkpoint save directory (default: ./model)

Outputs

Name	Type	Description
checkpoint	directory	Model checkpoints saved at --save-dir every --save-interval steps
eval_results	directory	Evaluation results saved at --eval-save-dir
rollout_logs	directory	Rollout logging saved at --rollout-save-dir

Usage Examples

# GRPO training with vLLM backend for math reasoning:
# python rl_example_zero_bubble.py \
#     -m Qwen/Qwen2.5-7B \
#     -d ./math_train.jsonl \
#     -a GRPO \
#     -b vllm \
#     -t 4 -i 2 \
#     -g 8 \
#     -tbs 32 -tMbs 8 -tmbs 2 \
#     -rt think_answer_tags \
#     -lr 1e-6 \
#     -sd ./grpo_checkpoint

# DAPO training with dynamic batching:
# python rl_example_zero_bubble.py \
#     -m Qwen/Qwen2.5-7B \
#     -d ./math_train.jsonl \
#     -a DAPO \
#     -b vllm \
#     -t 4 -i 2 \
#     -rt boxed

Key Features

• Zero-Bubble Architecture - Overlaps inference and training to minimize GPU idle time via producer-consumer coordination
• GRPO and DAPO Algorithms - Supports both standard GRPO and the DAPO variant with dynamic batching, asymmetric clipping, and soft over-length punishment
• Dual Backend Support - Inference can use either Transformers (with flash attention) or vLLM (with chunked prefill)
• Structured Response Formats - Configurable think/answer tags and boxed format for math reasoning
• Multi-Level Batch Sizes - Fine-grained control over inference, training, mini-batch, and micro-batch sizes
• Default System Prompts - Built-in prompts for think_answer_tags, boxed, and code reward types
• Buffer Size Control - Configurable consumer buffer size limit to balance producer and consumer throughput

Related Pages

• Environment:Hpcaitech_ColossalAI_CUDA_GPU_Environment