TinyZero

veRL is a flexible, efficient, and production-ready Reinforcement Learning (RL) training framework designed specifically for Large Language Models (LLMs). It is the open-source implementation of the HybridFlow paper.

Key capabilities include:

• Diverse RL Algorithms: Uses a Hybrid programming model to enable flexible representation and efficient execution of complex post-training dataflows.
• Modular Infrastructure Integration: Decouples computation and data dependencies, allowing seamless integration with PyTorch FSDP, Megatron-LM, vLLM, and TGI.
• Efficient Resource Utilization: Supports flexible device mapping for scaling across different cluster sizes and uses a 3D-HybridEngine to minimize memory redundancy and communication overhead during transitions between training and generation phases.
• Scalability: Capable of scaling up to 70B models and hundreds of GPUs.
• Supported Tasks: Supervised fine-tuning (SFT), Reward model training, and Reinforcement Learning from Human Feedback (RLHF) using PPO.

veRL is a flexible, efficient, and production-ready Reinforcement Learning (RL) training framework designed for Large Language Model (LLM) post-training. It is an open-source implementation of the HybridFlow paper.

Key Features

Flexibility and Ease of Use

• Diverse RL Algorithm Extension: Uses a Hybrid programming model that combines single-controller and multi-controller paradigms, allowing users to build complex RL dataflows with minimal code.
• Modular API Integration: Decouples computation and data dependencies to integrate seamlessly with existing LLM infrastructure like PyTorch FSDP, Megatron-LM, and vLLM.
• Flexible Device Mapping: Supports various model placements across different GPU sets for efficient resource utilization and scalability.
• HuggingFace Integration: Readily integrates with popular HuggingFace models.

Performance and Speed

• High Throughput: Achieves state-of-the-art generation and training throughput by leveraging existing SOTA LLM training and inference frameworks.
• 3D-HybridEngine: Features efficient actor model resharding to eliminate memory redundancy and reduce communication overhead during transitions between training and generation phases.

When implementing make_map_fn, each data item in the resulting Parquet file must contain exactly these 5 fields to be compatible with the training pipeline:

1. data_source: (String) The name of the dataset. This is used to index the correct reward function in the RewardModule.
2. prompt: (List of Dicts) The prompt formatted using the Hugging Face chat template (e.g., [{"role": "user", "content": "..."}]). The RLHFDataset tokenizer will apply the template and tokenize this field.
3. ability: (String) The task category (e.g., "math").
4. reward_model: (Dict) Contains the ground truth for evaluation. It must follow the structure: {"style": "rule", "ground_truth": <extracted_solution>}. Note that your extract_solution logic must align with how the reward function expects this value.
5. extra_info: (Dict) Metadata about the prompt (e.g., split name or index). Currently unused by the core logic but required for schema completeness.

# Example of the required dictionary structure returned by make_map_fn
data = {
    "data_source": "openai/gsm8k",
    "prompt": [{
        "role": "user",
        "content": "What is 2+2?"
    }],
    "ability": "math",
    "reward_model": {
        "style": "rule",
        "ground_truth": "4"
    },
    "extra_info": {
        "split": "train",
        "index": 0
    }
}

To save redundant CUDA and distributed context overhead, you can merge different roles (e.g., critic, ref, rm, actor_rollout) into the same process using create_colocated_worker_cls.

Important Constraints:

• Parallelism: For the Megatron backend, if you colocate worker groups into the same process, all roles will share the same 3D parallel size. If you require different 3D parallel sizes for different roles, do not use create_colocated_worker_cls; instead, pass different resource pools to each worker group separately.
• Memory Management: It is recommended to initialize the actor_rollout worker group last so that vLLM can more accurately estimate KV cache memory requirements.

# Example of colocating multiple roles into the same processes
all_wg = {}
for resource_pool, class_dict in self.resource_pool_to_cls.items():
    worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
    wg_dict = self.ray_worker_group_cls(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)
    spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
    all_wg.update(spawn_wg)

# Initialize roles from the spawned worker groups
if self.use_critic:
    self.critic_wg = all_wg['critic']
    self.critic_wg.init_model()

if self.use_reference_policy:
    self.ref_policy_wg = all_wg['ref']
    self.ref_policy_wg.init_model()

if self.use_rm:
    self.rm_wg = all_wg['rm']
    self.rm_wg.init_model()

# Initialize rollout last for better vLLM KV cache estimation
self.actor_rollout_wg = all_wg['actor_rollout']
self.actor_rollout_wg.init_model()

The RLHF datasets in TinyZero are stored as single Parquet files. Data is organized using a chat-based format within the prompt field to support multi-turn conversations. To facilitate answer extraction, instruction-following text is often embedded directly into the prompt (e.g., asking the model to output the final answer after a specific delimiter like ####).

Each entry in the dataset typically includes:

• data_source: The origin of the data (e.g., openai/gsm8k).
• prompt: A list of message objects in chat format (e.g., [{"role": "user", "content": "..."}]).
• ability: The capability being tested (e.g., math).
• reward_model: A configuration object defining how the model is evaluated, including the style (e.g., rule) and the ground_truth values.

{
    "data_source": "openai/gsm8k",
    "prompt": [{"role": "user", "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Let's think step by step and output the final answer after \"####\""}],
    "ability": "math",
    "reward_model": {
        "style": "rule",
        "ground_truth": ["72"]
    }
}

Instead of manually using execute_all_sync, you can use the @register decorator on Worker methods to define how data is dispatched and collected. This allows you to call methods directly on the RayWorkerGroup object.

Supported Dispatch Modes:

• Dispatch.ONE_TO_ALL: A single input value is automatically broadcast to all workers in the group.
• Dispatch.ALL_TO_ALL: Requires a list of inputs corresponding to the world size.
• Custom Dispatch: You can provide a dispatch_fn and collect_fn for complex logic.

Supported Execution Modes:

• Execute.RANK_ZERO: The operation is executed only on the first rank (rank 0).

from verl.single_controller.base.decorator import register, Dispatch, Execute

@ray.remote
class GPUAccumulatorDecorator(Worker):
    @register(Dispatch.ONE_TO_ALL)  # Automatically broadcasts 'x' to all workers
    def add(self, x):
        self.value = self.value + x
        return self.value.cpu()

# Usage on the worker group
# If x=10, every worker receives 10
print(gpu_accumulator_decorator.add(x=10))

The PPO (Proximal Policy Optimization) training loop in verl is implemented by orchestrating calls to different WorkerGroup roles via RPC. The driver process manages the dataflow, while heavy computations are distributed across GPUs.

Key components of the loop:

• Data Transfer: Uses DataProto objects (defined in protocol.py) to move data between the driver and workers. The trainer dispatches and collects data following transfer protocols wrapped in worker functions.
• Roles: The loop interacts with several specialized worker groups:
  • actor_rollout_wg: Handles sequence generation (generate_sequences) and actor updates (update_actor).
  • ref_policy_wg: Computes reference log probabilities (compute_ref_log_prob).
  • critic_wg: Computes values (compute_values) and updates the critic (update_critic).
  • rm_wg: Computes reward model scores (compute_rm_score).
• Advantage Computation: Advantage estimation (including KL penalty application) is performed on the driver process using compute_advantage.
• Extensibility: To implement other RLHF algorithms like DPO or GRPO, refer to the DPO extension documentation.

# Conceptual flow of the PPO loop
# 1. Generate sequences via actor_rollout_wg
gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)

# 2. Compute reference log probs via ref_policy_wg
ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)

# 3. Compute values via critic_wg
values = self.critic_wg.compute_values(batch)

# 4. Compute rewards via rm_wg and reward_fn
reward_tensor = self.rm_wg.compute_rm_score(batch)

# 5. Update critic and actor
self.critic_wg.update_critic(batch)
self.actor_rollout_wg.update_actor(batch)