AgentGym-RL

https://github.com/woooodyy/agentgym-rl
Admin
AgentGym-RL is a framework for training LLM agents in multi-turn interactive decision-making through...
Tokens:16,684
Snippets:97
Trust Score:8.9
License:MIT
Update:1 year ago
Show doc for...
Tokens:
Raw
### Setup Megatron-LM for AgentGym-RL (Patch Method)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/install.rst

This bash script demonstrates how to set up the Megatron-LM backend for AgentGym-RL by cloning the repository, applying a specific patch, and installing it. This is an optional step for enhanced scalability during training.

```bash
cd ..
git clone -b core_v0.4.0 https://github.com/NVIDIA/Megatron-LM.git
cp verl/patches/megatron_v4.patch Megatron-LM/
cd Megatron-LM && git apply megatron_v4.patch
pip3 install -e .
export PYTHONPATH=$PYTHONPATH:$(pwd)
```

--------------------------------

### Install AgentGym-RL via Docker

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/install.rst

This snippet shows how to run a pre-built Docker image for AgentGym-RL and install the stable version of the library within the container. Ensure you have NVIDIA Container Toolkit installed for GPU support.

```bash
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN -v <image:tag> # Replace <image:tag> with the actual image tag

# Inside the container:
pip3 install verl
```

--------------------------------

### Download and Test Qwen2.5-0.5B-Instruct Model

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/quickstart.rst

This Python command downloads and initializes the Qwen2.5-0.5B-Instruct model using the transformers library. It's a prerequisite step to obtain a base model for post-training. Ensure transformers is installed.

```python
import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-0.5B-Instruct')
```

--------------------------------

### Install AgentGym-RL from Source with Custom Environment

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/install.rst

This snippet details how to set up a Python environment using conda, install PyTorch with CUDA support, and then install AgentGym-RL from its source code. This method allows for code customization.

```bash
conda create -n verl python==3.9
conda activate verl

pip3 install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
pip3 install flash-attn --no-build-isolation
pip3 install verl
```

--------------------------------

### Setup Megatron-LM for AgentGym-RL (Patched Clone Method)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/install.rst

An alternative method to set up Megatron-LM for AgentGym-RL by cloning a pre-patched version of the repository. This simplifies the process by avoiding manual patching.

```bash
git clone -b core_v0.4.0_verl https://github.com/eric-haibin-lin/Megatron-LM
export PYTHONPATH=$PYTHONPATH:$(pwd)/Megatron-LM
```

--------------------------------

### Install Megatron Dependencies (Apex, Transformer Engine)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/install.rst

These commands install the necessary dependencies for the Megatron-LM backend, including Apex and Transformer Engine, from their respective GitHub repositories. These are prerequisites for certain advanced training configurations.

```bash
# Install Apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" \
    git+https://github.com/NVIDIA/apex

# Install Transformer Engine
pip install git+https://github.com/NVIDIA/TransformerEngine.git@v1.7
```

--------------------------------

### Environment Setup for AgentGym-RL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/README.md

This script sets up the development environment for AgentGym-RL. It creates a conda environment, installs specific versions of PyTorch with CUDA support, and installs necessary libraries like flash-attention and the agentgym-rl package itself. It also prepares the agentenv package.

```shell
echo "Preparing environment for agentgym-rl..."
conda create -n agentgym-rl python==3.10 -y
conda activate agentgym-rl
pip3 install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
# install flash-atten
FLASH_ATTENTION_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
FLASH_ATTENTION_NAME="flash_attn-2.7.3+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
wget -q $FLASH_ATTENTION_URL -O $FLASH_ATTENTION_NAME
pip3 install $FLASH_ATTENTION_NAME
rm -f $FLASH_ATTENTION_NAME
# for RL
cd AgentGym-RL
pip3 install -e .
# for agentgym
echo "Preparing environment for agentenv..."
cd AgentGym/agentenv
pip3 install -e .
pip3 install transformers==4.51.3
```

--------------------------------

### Enable Wandb for Experiment Tracking

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/quickstart.rst

This configuration enables the use of Weights & Biases (wandb) for experiment tracking by setting the logger to include 'wandb' and specifying project and experiment names. Ensure that wandb is installed and configured in your environment.

```bash
trainer.logger=['console','wandb'] 
trainer.project_name=$YOUR_PROJECT_NAME 
trainer.experiment_name=$YOUR_RUN_NAME 

```

--------------------------------

### Launch AgentGym-RL Training Script (Bash)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/README.md

This command initiates the training process for AgentGym-RL using a predefined shell script. It's a common starting point for RL training tasks within the project.

```bash
bash webarena_train.sh
```

--------------------------------

### Configure Actor Reference and FSDP Settings

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Specifies settings for the actor reference, including micro-batch size per GPU and FSDP parameter offloading. This configuration impacts how references are handled during the reinforcement learning process.

```shell
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 
actor_rollout_ref.ref.fsdp_config.param_offload=True
```

--------------------------------

### Run PPO Training for GSM8K

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/quickstart.rst

This command initiates PPO training for the GSM8K dataset using the verl trainer. It configures dataset paths, batch sizes, model paths, learning rates, and other training parameters. The output logs are redirected to 'verl_demo.log'. Ensure all specified paths and models are accessible.

```bash
PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.train_batch_size=256 \
    data.val_batch_size=1312 \
    data.max_prompt_length=512 \
    data.max_response_length=256 \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.actor.ppo_mini_batch_size=64 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
    critic.optim.lr=1e-5 \
    critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \
    critic.ppo_micro_batch_size_per_gpu=4 \
    algorithm.kl_ctrl.kl_coef=0.001 \
    trainer.logger=['console'] \
    +trainer.val_before_train=False \
    trainer.default_hdfs_dir=null \
    trainer.n_gpus_per_node=1 \
    trainer.nnodes=1 \
    trainer.save_freq=10 \
    trainer.test_freq=10 \
    trainer.total_epochs=15 2>&1 | tee verl_demo.log
```

--------------------------------

### Preprocess GSM8K Dataset for RL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/quickstart.rst

This command preprocesses the GSM8K dataset into a parquet format suitable for RL reward computation and faster reading. It requires Python 3 and the verl library. The output is stored in the specified local directory.

```bash
python3 examples/data_preprocess/gsm8k.py --local_dir ~/data/gsm8k
```

--------------------------------

### PPO Training Configuration - Bash

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

This bash script configures and runs the main PPO training process. It specifies training and validation file paths, batch sizes, sequence lengths, actor model paths, learning rates, FSDP configurations, and gradient checkpointing. It is designed for use with the `verl.trainer.main_ppo` module.

```bash
set -x

python3 -m verl.trainer.main_ppo \
   data.train_files=$HOME/data/gsm8k/train.parquet \
   data.val_files=$HOME/data/gsm8k/test.parquet \
   data.train_batch_size=1024 \
   data.val_batch_size=1312 \
   data.max_prompt_length=512 \
   data.max_response_length=512 \
   actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \
   actor_rollout_ref.actor.optim.lr=1e-6 \
   actor_rollout_ref.model.use_remove_padding=True \
   actor_rollout_ref.actor.ppo_mini_batch_size=256 \
   actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
   actor_rollout_ref.actor.fsdp_config.param_offload=False \
   actor_rollout_ref.actor.fsdp_config.grad_offload=False \
   actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
   actor_rollout_ref.model.enable_gradient_checkpointing=True \

```

--------------------------------

### Prepare GSM8K Dataset - Bash

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

This command downloads and preprocesses the GSM8K dataset, preparing it for model training. It requires navigating to the data preprocessing directory and executing a Python script with specified local directories for storage.

```bash
cd examples/data_preprocess
python3 gsm8k.py --local_dir ~/data/gsm8k
```

--------------------------------

### PPO Trainer Example for verL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/VERL_README.md

This example demonstrates how to use the PPO trainer within the verl library. It includes configurations for sequence balancing and long context support via DeepSpeed Ulysses. Dependencies include specific verl components and potentially HuggingFace models.

```shell
# Example for sequence balancing with Qwen2-7B
./examples/ppo_trainer/run_qwen2-7b_seq_balance.sh

# Example for long context support with DeepSeek-7B
./examples/ppo_trainer/run_deepseek7b_llm_sp2.sh
```

--------------------------------

### Install verl and vLLM

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/README_vllm0.7.md

Commands to set up a conda environment, clone the verl repository, install verl and its dependencies including vLLM (>=0.7) and flash-attn.

```bash
conda create -n verl python==3.10
conda activate verl

git clone https://github.com/volcengine/verl.git
cd verl
pip3 install -e .

# Install vLLM>=0.7
# (Option1) pip3 install vllm --pre --extra-index-url https://wheels.vllm.ai/nightly
# (Option2) pip3 install "vllm>=0.7.0" 

# Install flash-attn
pip3 install flash-attn --no-build-isolation
```

--------------------------------

### Enable Sequence Packing with Transformers

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/perf/perf_tuning.rst

This snippet demonstrates how to enable sequence packing using the transformers library for supported models like Llama, Mistral, Gemma1, and Qwen. It involves setting a configuration flag and testing the functionality.

```bash
pytest -s tests/model/test_transformer.py
```

--------------------------------

### Initialize and Run Ray PPO Trainer

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/ppo_code_architecture.rst

This snippet initializes and runs the Ray PPO Trainer. It takes the configuration, tokenizer, worker mappings, resource pool manager, worker class, and reward functions as input. The process involves initializing workers and then fitting the trainer for PPO training. Dependencies include 'RayPPOTrainer', 'config', 'tokenizer', 'role_worker_mapping', 'resource_pool_manager', 'ray_worker_group_cls', 'reward_fn', and 'val_reward_fn'.

```python
   trainer = RayPPOTrainer(config=config,
                           tokenizer=tokenizer,
                           role_worker_mapping=role_worker_mapping,
                           resource_pool_manager=resource_pool_manager,
                           ray_worker_group_cls=ray_worker_group_cls,
                           reward_fn=reward_fn,
                           val_reward_fn=val_reward_fn)
   trainer.init_workers()
   trainer.fit()
```

--------------------------------

### Gradient Checkpointing Configuration

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/perf/perf_tuning.rst

This example shows how to enable gradient checkpointing for both the actor and critic models. This technique can allow for larger micro-batch sizes, which is beneficial for training with large mini-batches.

```python
actor_rollout_ref.model.enable_gradient_checkpointing=True
critic.model.enable_gradient_checkpointing=True
```

--------------------------------

### Download DeepSeek Math Model - Bash

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

This command downloads a pre-trained DeepSeek math model from Hugging Face to a local directory. It uses the `huggingface-cli` tool and specifies the model name and local directory for storage, with an option to disable symbolic links.

```bash
huggingface-cli download deepseek-ai/deepseek-math-7b-instruct --local-dir ~/models/deepseek-math-7b-instruct --local-dir-use-symlinks False
```

--------------------------------

### Configure Critic Model and FSDP Settings

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Defines parameters for the critic model, including learning rate, padding removal, model path, gradient checkpointing, and FSDP (Fully Sharded Data Parallel) configurations such as parameter offloading. These settings are vital for training the critic network.

```shell
critic.optim.lr=1e-5 
critic.model.use_remove_padding=True 
critic.model.path=deepseek-ai/deepseek-llm-7b-chat 
critic.model.enable_gradient_checkpointing=True 
critic.model.fsdp_config.param_offload=False 
critic.model.fsdp_config.grad_offload=False 
critic.model.fsdp_config.optimizer_offload=False
```

--------------------------------

### PPO Training Script Execution - Bash

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

This bash script executes the Proximal Policy Optimization (PPO) training for the model on the GSM8K dataset. It requires navigating to the PPO trainer directory and running a Python module with extensive configuration parameters for data loading, model paths, and optimization settings.

```bash
cd ../ppo_trainer
bash run_deepseek7b_llm.sh
```

--------------------------------

### Configure Actor Rollout Settings

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Sets parameters for actor rollouts, including micro-batch size, tensor model parallelism, logging method, and GPU memory utilization. These settings influence how the actor interacts with the environment and processes data.

```shell
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 
actor_rollout_ref.rollout.tensor_model_parallel_size=4 
actor_rollout_ref.rollout.name=vllm 
actor_rollout_ref.rollout.gpu_memory_utilization=0.5
```

--------------------------------

### Main Training Loop with Ray Actors

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/advance/dpo_extension.rst

Provides an example of a main training loop implemented as a Ray remote process. It demonstrates constructing various Ray actors (SampleGenerator, ReferencePolicy, DPOActor) and orchestrating data generation, scoring, and policy updates.

```python
@ray.remote(num_cpus=1)
def main_task(config):
    # construct SampleGenerator
    resource_pool = RayResourcePool(process_on_nodes=[8] * 2)  # 16 GPUs
    ray_cls = RayClassWithInitArgs(SampleGenerator, config=config) 
    # put SampleGenerator onto resource pool
    sample_gen = RayWorkerGroup(resource_pool, ray_cls)
    
    # construct reference policy
    ray_cls = RayClassWithInitArgs(ReferencePolicy)
    ref_policy = RayWorkerGroup(resource_pool, ray_cls)
    
    # construct actor
    ray_cls = RayClassWithInitArgs(DPOActor)  
    dpo_policy = RayWorkerGroup(resource_pool, ray_cls)
    
    dataloader = DataLoader()
    
    for data in dataloader:
        # generate data
        data = sample_gen.generate_sequences(data)
        # generate scores for each data 
        data = generate_scores(data)
        # generate pairwise data using scores
        data = generate_pairwise_data(data)
        # generate ref_log_prob
        data.batch['ref_log_prob'] = ref_policy.infer(data)
        # update using dpo
        dpo_policy.update(data)
        # logging
```

--------------------------------

### Configure Trainer and Logging Settings

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Sets up trainer parameters, including logger types, project and experiment names, GPU allocation per node, number of nodes, saving frequency, testing frequency, and total epochs. These control the overall training process and experiment management.

```shell
trainer.logger=['console','wandb'] 
trainer.project_name='verl_example_gsm8k' 
trainer.experiment_name='deepseek_llm_7b_function_rm' 
trainer.n_gpus_per_node=8 
trainer.nnodes=1 
trainer.save_freq=-1 
trainer.test_freq=1 
trainer.total_epochs=15
```

--------------------------------

### Install yapf for Code Formatting

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/VERL_README.md

Installs the latest version of yapf, a code formatter, used for enforcing strict code formatting in the project's pull requests. Ensure you have pip3 installed.

```bash
pip3 install yapf --upgrade
```

--------------------------------

### Configure Algorithm KL Control Coefficient

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Sets the KL divergence control coefficient for the algorithm. This parameter helps to regularize the policy updates in reinforcement learning.

```shell
algorithm.kl_ctrl.kl_coef=0.001
```

--------------------------------

### Configure Critic PPO Micro Batch Size

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Sets the PPO micro batch size per GPU for the critic. This parameter is important for the Proximal Policy Optimization algorithm's training efficiency.

```shell
critic.ppo_micro_batch_size_per_gpu=32
```

--------------------------------

### SFT Model Training Script - Shell

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

This shell script initiates the Supervised Fine-Tuning (SFT) process for a model on the GSM8K dataset using PyTorch FSDP. It configures training and validation file paths, prompt/response keys, batch sizes, model paths, and experiment tracking.

```shell
set -x

torchrun -m verl.trainer.fsdp_sft_trainer \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.prompt_key=question \
    data.response_key=answer \
    data.micro_batch_size_per_gpu=8 \
    model.partial_pretrain=deepseek-ai/deepseek-coder-6.7b-instruct \
    trainer.default_hdfs_dir=hdfs://user/verl/experiments/gsm8k/deepseek-coder-6.7b-instruct/ \
    trainer.project_name=gsm8k-sft \
    trainer.experiment_name=gsm8k-sft-deepseek-coder-6.7b-instruct \
    trainer.total_epochs=4 \
    trainer.logger=['console','wandb']
```

--------------------------------

### Environment Interaction with EnvClient

Source: https://github.com/woooodyy/agentgym-rl/blob/main/README.md

Provides an EnvClient for facilitating interactions with the environment during rollouts. Key methods include observation() to get the current state, available_actions() to list possible actions, step() to execute an action, and reset() to reinitialize the environment. The framework supports parallel environment initialization and trajectory collection for improved efficiency.

```python
class EnvClient:
    def __init__(self, env_config):
        # Initializes environment(s)
        pass

    def observation(self):
        # Returns current observation from the environment
        pass

    def available_actions(self):
        # Returns currently available actions
        pass

    def step(self, action):
        # Performs an action in the environment
        pass

    def reset(self):
        # Resets the environment
        pass

# Example of parallel initialization and collection
# env_clients = [EnvClient(config) for _ in range(num_parallel_envs)]
# trajectories = collect_trajectories(env_clients)
```

--------------------------------

### Reinforcement Learning Algorithms in verL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/VERL_README.md

Illustrates the usage of different reinforcement learning algorithms within the verl library, including PPO, GRPO, and ReMax. These examples showcase the flexibility in choosing RL algorithms and the support for various reward models.

```shell
# PPO trainer example
./examples/ppo_trainer

# GRPO trainer example
./examples/grpo_trainer

# ReMax trainer example
./examples/remax_trainer
```

--------------------------------

### Initialize HybridEngine, Actor, and Rollout Models (Python)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/workers/fsdp_workers.rst

Initializes the HybridEngine, Actor, and Rollout models using FSDP. The `init_model` function is decorated with `Dispatch.ONE_TO_ALL`, meaning it's executed on each worker process when called from the driver process. This setup is crucial for distributed training scenarios.

```python
@register(dispatch_mode=Dispatch.ONE_TO_ALL)
def init_model(self):

```

--------------------------------

### Configure Trainer Critic Warmup

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/gsm8k_example.rst

Sets the warmup period for the critic during training. This parameter can affect the initial stages of the learning process.

```shell
trainer.critic_warmup=0
```

--------------------------------

### Define Roles for Distributed Training Workers (Python Enum)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/ppo_code_architecture.rst

This Python Enum defines distinct roles for workers in a distributed training setup. Each role represents a specific set of functionalities that a worker can perform, such as acting as an Actor, a Critic, a Rollout engine, or a combination thereof. This allows for flexible configuration of hybrid engines and specialized workers.

```python
class Role(Enum):
    """
    To create more roles dynamically, you can subclass Role and add new members
    """
    Actor = 0  # This worker only has Actor
    Rollout = 1 # This worker only has Rollout
    ActorRollout = 2 # This worker has both actor and rollout, it's a HybridEngine
    Critic = 3 # This worker only has critic
    RefPolicy = 4 # This worker only has reference policy
    RewardModel = 5 # This worker only has reward model
    ActorRolloutRef = 6 # This worker contains actor, rollout and reference policy simultaneously 
```

--------------------------------

### Enable CUDA Graph for vLLM

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/README_vllm0.7.md

Bash script snippet to disable eager execution and cache freeing in vLLM, enabling CUDA graphs for performance improvements during rollout generation.

```bash
actor_rollout_ref.rollout.enforce_eager=False \
actor_rollout_ref.rollout.free_cache_engine=False \
```

--------------------------------

### Configure Reward Model Worker and Initialize Reward Function

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/ppo_code_architecture.rst

This snippet configures the reward model worker based on configuration and initializes reward functions for training and validation. It maps the RewardModel role to a specific worker and pool, and sets up RewardManager instances. Dependencies include 'verl.workers.fsdp_workers' and 'RewardManager'.

```python
   if config.reward_model.enable:
       from verl.workers.fsdp_workers import RewardModelWorker
       role_worker_mapping[Role.RewardModel] = RewardModelWorker
       mapping[Role.RewardModel] = global_pool_id
    
   reward_fn = RewardManager(tokenizer=tokenizer, num_examine=0)

   # Note that we always use function-based RM for validation
   val_reward_fn = RewardManager(tokenizer=tokenizer, num_examine=1)

   resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
```

--------------------------------

### Supervised Fine-Tuning with PEFT/Liger-kernel in verL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/VERL_README.md

Demonstrates supervised fine-tuning (SFT) using PEFT (Parameter-Efficient Fine-Tuning) and Liger-kernel within the verl framework. These examples focus on specific models like Qwen and datasets like GSM8k, showcasing efficient training techniques.

```shell
# SFT with PEFT on Qwen for GSM8k dataset
./examples/sft/gsm8k/run_qwen_05_peft.sh

# SFT with Liger-kernel on Qwen for GSM8k dataset with sequence packing
./examples/sft/gsm8k/run_qwen_05_sp2_liger.sh
```

--------------------------------

### Initialize Multi-Process Worker Group with Ray for RL Training

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/advance/dpo_extension.rst

Demonstrates how to initialize a multi-process worker group using Ray for distributed RL training. This setup allows a single control process to manage and drive computations across multiple worker processes, specifying resource allocation per node.

```python
@ray.remote(num_cpus=1)
def main_task(config):
    # construct SampleGenerator
    resource_pool = RayResourcePool(process_on_nodes=[8] * 2)  # 16 GPUs
    ray_cls = RayClassWithInitArgs(SampleGenerator, config=config)
```

--------------------------------

### Enable LigerKernel for SFT Training

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/perf/perf_tuning.rst

This configuration snippet shows how to enable LigerKernel for Supervised Fine-Tuning (SFT) in AgentGym RL. LigerKernel is a high-performance kernel designed to improve SFT training efficiency. Set `use_liger` to `True` in your SFT configuration file.

```yaml
model:
  use_liger: True  # Enable LigerKernel for SFT
```

--------------------------------

### Optimize for Low HBM Environments

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/start/quickstart.rst

These settings are designed to reduce High Bandwidth Memory (HBM) usage during training, which is beneficial for systems with less than 32GB of HBM. By setting the PPO micro batch size per GPU to 1 for both actor and critic, the memory footprint per GPU is minimized.

```bash
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 
critic.ppo_micro_batch_size_per_gpu=1 

```

--------------------------------

### Format Code with yapf

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/index.rst

This snippet demonstrates how to install and use yapf, a code formatting tool, to enforce strict code formatting according to Google style within the verl repository. It's used for code reviews and maintaining code consistency.

```bash
pip3 install yapf
yapf -ir -vv --style ./.style.yapf verl examples tests
```

--------------------------------

### vLLM Patch for older versions

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/README_vllm0.7.md

Manual code modifications required for stable vLLM versions (prior to necessary merges) to ensure compatibility with verl. These patches address assertion errors and rank assignments.

```python
# vllm/distributed/parallel_state.py: Remove the assertion below:
if (world_size
        != tensor_model_parallel_size * pipeline_model_parallel_size):
    raise RuntimeError(
        f"world_size ({world_size}) is not equal to "
        f"tensor_model_parallel_size ({tensor_model_parallel_size}) x "
        f"pipeline_model_parallel_size ({pipeline_model_parallel_size})")

# vllm/executor/uniproc_executor.py: change `local_rank = rank` to `local_rank = int(os.environ["LOCAL_RANK"])`

# vllm/model_executor/model_loader/weight_utils.py: remove the `torch.cuda.empty_cache()` in `pt_weights_iterator`
```

--------------------------------

### Select Reward Score Function Based on Data Source

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/ppo_code_architecture.rst

This function selects the appropriate reward scoring function based on the provided data source. It handles specific datasets like 'openai/gsm8k' and 'lighteval/MATH', raising an error for unsupported sources. This is crucial for function-based reward models.

```python
   def _select_rm_score_fn(data_source):
       if data_source == 'openai/gsm8k':
           return gsm8k.compute_score
       elif data_source == 'lighteval/MATH':
           return math.compute_score
       else:
           raise NotImplementedError
```

--------------------------------

### Initialize Ray Resource Pool and Actor Rollout Worker Group

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/workers/ray_trainer.rst

This code initializes a Ray Resource Pool and an Actor Rollout Worker Group. The resource pool is configured based on the number of GPUs per node and whether to use GPUs. The actor rollout worker group is then defined using the resource pool and a class that can be initialized remotely. This setup is for distributing actor rollout tasks across Ray workers.

```python
# max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool
# For FSDP backend, we recommend using max_colocate_count=1 that merge all WorkerGroups into one.
# For Megatron backend, we recommend using max_colocate_count>1 that can utilize different WorkerGroup for differnt models
resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes,
                                   use_gpu=True,
                                   max_colocate_count=1)
# define actor_rollout cls to be init on remote
actor_rollout_cls = RayClassWithInitArgs(cls=ActorRolloutWorker)
# define actor_rollout worker group
actor_rollout_worker_group = MegatronRayWorkerGroup(resource_pool=resource_pool,
                                                       ray_cls_with_init=actor_rollout_cls,
                                                       default_megatron_kwargs=config.actor_rollout.megatron)
```

--------------------------------

### Define Worker Classes based on Backend Strategy (Python)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/ppo_code_architecture.rst

This code snippet defines how to select and import appropriate worker classes (ActorRolloutRefWorker, CriticWorker) and Ray worker group classes based on the chosen backend strategy (FSDP or Megatron). It also sets up a mapping between different roles and their corresponding worker implementations, along with resource pool configurations for distributed training.

```python
if config.actor_rollout_ref.actor.strategy == 'fsdp': # for FSDP backend
    assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
    from verl.workers.fsdp_workers import ActorRolloutRefWorker, CriticWorker
    from verl.single_controller.ray import RayWorkerGroup
    ray_worker_group_cls = RayWorkerGroup

elif config.actor_rollout_ref.actor.strategy == 'megatron': # for Megatron backend
    assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
    from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker
    from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup
    ray_worker_group_cls = NVMegatronRayWorkerGroup # Ray worker class for Megatron-LM

else:
    raise NotImplementedError

from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role

role_worker_mapping = {
    Role.ActorRollout: ActorRolloutRefWorker,
    Role.Critic: CriticWorker,
    Role.RefPolicy: ActorRolloutRefWorker
}

global_pool_id = 'global_pool'
resource_pool_spec = {
    global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
}
mapping = {
    Role.ActorRollout: global_pool_id,
    Role.Critic: global_pool_id,
    Role.RefPolicy: global_pool_id,
}
```

--------------------------------

### Data Mapping Function for GSM8K Dataset (Python)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/preparation/prepare_data.rst

This factory function generates a data processing function (`process_fn`) tailored for the GSM8K dataset. It takes a split ('train' or 'test') and processes each example to create a structured dictionary containing data source, formatted prompt, ability, reward model details (including the extracted solution), and extra information. It relies on the `extract_solution` function.

```python
instruction_following = "Let's think step by step and output the final answer after \"####\"."

# add a row to each data item that represents a unique id
def make_map_fn(split):

    def process_fn(example, idx):
        question = example.pop('question')

        question = question + ' ' + instruction_following

        answer = example.pop('answer')
        solution = extract_solution(answer)
        data = {
            "data_source": data_source,
            "prompt": [
                {
                    "role": "user",
                    "content": question
                }
            ],
            "ability": "math",
            "reward_model": {
                "style": "rule",
                "ground_truth": solution
            },
            "extra_info": {
                'split': split,
                'index': idx
            }
        }
        return data

    return process_fn
```

--------------------------------

### Reference Model Configuration

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/config.rst

Configuration for the reference model, including FSDP settings and micro-batch size for log probability calculations. Recommends offloading for larger models.

```yaml
ref:
  log_prob_micro_batch_size_per_gpu: 1 # The value represent the local num per gpu
  # For models larger than 7B, it's recommended to turn on offload for ref by default
```

--------------------------------

### Launch AgentGym-RL Evaluation Script (Bash)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/README.md

This command initiates the evaluation process for AgentGym-RL using a predefined shell script. It's the standard way to run evaluations for the project's RL tasks.

```bash
bash webarena_eval.sh
```

--------------------------------

### Initialize RLHFDataset for Data Loading

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/workers/ray_trainer.rst

This snippet demonstrates initializing the RLHFDataset class for loading and preprocessing training data. It configures parameters like parquet file paths, tokenizer, prompt key, max prompt length, and truncation settings. This is crucial for preparing a complete batch of samples for distributed training.

```python
self.train_dataset = RLHFDataset(parquet_files=self.config.data.train_files,
                                       tokenizer=self.tokenizer,
                                       prompt_key=self.config.data.prompt_key,
                                       max_prompt_length=self.config.data.max_prompt_length,
                                       filter_prompts=True,
                                       return_raw_chat=self.config.data.get('return_raw_chat', False),
                                       truncation='error')
```

--------------------------------

### Trainer Configuration (YAML)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/config.rst

Sets up the training environment and process, including the total number of epochs, project and experiment names for logging, hardware specifications (nodes and GPUs), and frequencies for saving checkpoints and testing.

```yaml
trainer:
  total_epochs: 30
  project_name: verl_examples
  experiment_name: gsm8k
  logger: ['console', 'wandb']
  nnodes: 1
  n_gpus_per_node: 8
  save_freq: -1
  test_freq: 2
  critic_warmup: 0
  default_hdfs_dir: ~/experiments/gsm8k/ppo/${trainer.experiment_name} # hdfs checkpoint path
  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name} # local checkpoint path
```

--------------------------------

### YAML Configuration for Actor, Rollout, and Reference Policies with FSDP

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/config.rst

This extensive configuration block defines settings for the actor, rollout, and reference policies, heavily utilizing the FSDP (Fully Sharded Data Parallel) backend. It includes model paths, gradient checkpointing, FSDP-specific parameters like sharding and offloading, optimizer settings (learning rate, warmup), and rollout configurations for inference, including dynamic batching and tensor parallelism. It also details parameters for reward model log probabilities.

```yaml
actor_rollout_ref:
  hybrid_engine: True
  model:
    path: ~/models/deepseek-llm-7b-chat
    external_lib: null
    override_config: { }
    enable_gradient_checkpointing: False
    use_remove_padding: False
  actor:
    strategy: fsdp  # This is for backward-compatibility
    ppo_mini_batch_size: 256
    ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu
    ppo_micro_batch_size_per_gpu: 8
    use_dynamic_bsz: False
    ppo_max_token_len_per_gpu: 16384 # n * ${data.max_prompt_length} + ${data.max_response_length}
    grad_clip: 1.0
    clip_ratio: 0.2
    entropy_coeff: 0.001
    use_kl_loss: False # True for GRPO
    kl_loss_coef: 0.001 # for grpo
    kl_loss_type: low_var_kl # for grpo
    ppo_epochs: 1
    shuffle: False
    ulysses_sequence_parallel_size: 1 # sp size
    optim:
      lr: 1e-6
      lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime
      min_lr_ratio: null   # only useful for warmup with cosine
      warmup_style: constant  # select from constant/cosine
      total_training_steps: -1  # must be override by program
    fsdp_config:
      wrap_policy:
        # transformer_layer_cls_to_wrap: None
        min_num_params: 0
      param_offload: False
      grad_offload: False
      optimizer_offload: False
      fsdp_size: -1
  ref:
    fsdp_config:
      param_offload: False
      wrap_policy:
        # transformer_layer_cls_to_wrap: None
        min_num_params: 0
    log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu
    log_prob_micro_batch_size_per_gpu: 16
    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}
    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}
    ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size} # sp size
  rollout:
    name: vllm
    temperature: 1.0
    top_k: -1 # 0 for hf rollout, -1 for vllm rollout
    top_p: 1
    prompt_length: ${data.max_prompt_length}  # not use for opensource
    response_length: ${data.max_response_length}
    # for vllm rollout
    dtype: bfloat16 # should align with FSDP
    gpu_memory_utilization: 0.5
    ignore_eos: False
    enforce_eager: True
    free_cache_engine: True
    load_format: dummy_dtensor
    tensor_model_parallel_size: 2
    max_num_batched_tokens: 8192
    max_num_seqs: 1024
    log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu
    log_prob_micro_batch_size_per_gpu: 16
    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}
```

--------------------------------

### Load, Preprocess, and Save Datasets to Parquet (Python)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/preparation/prepare_data.rst

This script demonstrates the common data preprocessing pipeline. It loads a dataset from Hugging Face, applies a mapping function (to be defined separately), and saves the processed data into parquet files locally and to HDFS. It requires the `datasets`, `re`, `os`, and `argparse` libraries.

```python
import re
import os
import datasets

from verl.utils.hdfs_io import copy, makedirs
import argparse

# To extract the solution for each prompts in the dataset
# def extract_solution(solution_str):
# ...


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_dir', default='/opt/tiger/gsm8k')
    parser.add_argument('--hdfs_dir', default=None)

    args = parser.parse_args()

    num_few_shot = 5
    data_source = 'openai/gsm8k'

    dataset = datasets.load_dataset(data_source, 'main')

    train_dataset = dataset['train']
    test_dataset = dataset['test']

        # Construct a `def make_map_fn(split)` for the corresponding datasets.
    # ...
        
    train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
    test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)

    local_dir = args.local_dir
    hdfs_dir = args.hdfs_dir

    train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet'))
    test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet'))

    makedirs(hdfs_dir)

    copy(src=local_dir, dst=hdfs_dir)
```

--------------------------------

### Repository Structure of AgentGym-RL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/hybrid_flow.rst

This is a bash snippet outlining the directory structure of the AgentGym-RL project. It highlights key packages and files involved in training RL algorithms, configuration, worker implementations, model definitions, and utility functions.

```bash
verl # the verl package
  trainer
    main_ppo.py  # the entrypoint for RL training
    ppo
      ray_trainer.py  # the training loop for RL algorithms such as PPO
    fsdp_sft_trainer.py  # the SFT trainer with FSDP backend
  config
    generation.yaml  # configuration template for rollout
    ppo_trainer.yaml  # configuration template for the RL trainer
  workers
    protocol.py  # the interface of DataProto
    fsdp_workers.py   # the FSDP worker interfaces: ActorRolloutRefWorker, CriticWorker, RewardModelWorker
    megatron_workers.py  # the Megatron worker interfaces: ActorRolloutRefWorker, CriticWorker, RewardModelWorker
    actor
      dp_actor.py  #  data parallel actor with FSDP backend
      megatron_actor.py  # nD parallel actor with Megatron backend
    critic
      dp_critic.py  # data parallel critic with FSDP backend
      megatron_critic.py  # nD parallel critic with FSDP backend
    reward_model
      megatron
        reward_model.py  # reward model with Megatron backend
    rollout
      vllm
        vllm_rollout.py  # rollout with vllm backend
      hf_rollout.py  # rollout with huggingface TGI backend
    sharding_manager
      fsdp_ulysses.py  # data and model resharding when using FSDP + ulysses
      fsdp_vllm.py  # data and model resharding when using FSDP + ulysses + vllm
      megatron_vllm.py  # data and model resharding when using Megatron + vllm
  utils
    dataset  # datasets for SFT/RM/RL
    reward_score  # function based reward
      gsm8k.py  # reward function for gsm8k dataset
      math.py  # reward function for math dataset
    seqlen_balancing.py  # the sequence balance optimization
  models
    llama  # Megatron implementation for llama, deepseek, mistral, etc
    transformers  # ulysses integration with transformer models such as llama, qwen, etc
    weight_loader_registery.py  # registry of weight loaders for loading hf ckpt into Megatron
  third_party
    vllm  # adaptor for vllm's usage in RL
      vllm_v_0_6_3  # vllm v0.6.3 adaptor
        llm.py  # entrypoints for generate, sync_model_weight, offload_model_weights
        parallel_state.py  # vllm related device mesh and process groups
        dtensor_weight_loaders.py  # weight loader for huggingface models with FSDP
```

--------------------------------

### Rollout Model Sampling Parameters

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/config.rst

Configuration for the rollout model, primarily using vLLM. Includes sampling parameters like temperature, top_k, and top_p, along with GPU memory utilization and tensor model parallel size.

```yaml
rollout:
  name: "hf/vllm"
  temperature: 0.7
  top_k: 50
  top_p: 0.95
  gpu_memory_utilization: 0.9
  tensor_model_parallel_size: 1
```

--------------------------------

### Algorithm Configuration (YAML)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/config.rst

Specifies the parameters for the reinforcement learning algorithm, including discount factor, GAE estimator settings, and KL divergence penalty options. These settings influence how the agent learns and optimizes its policy.

```yaml
algorithm:
  gamma: 1.0
  lam: 1.0
  adv_estimator: gae
  kl_penalty: kl  # how to estimate kl divergence
  kl_ctrl:
    type: fixed
    kl_coef: 0.005
```

--------------------------------

### YAML Configuration for Data Loading in AgentGym RL

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/examples/config.rst

This configuration section defines parameters for loading and processing training and validation data. It includes settings for tokenizers, file paths, prompt/response lengths, batch sizes, and options for handling input IDs and chat templates. It supports local and HDFS paths for data files and has constraints on dataset size.

```yaml
data:
  tokenizer: null
  train_files: ~/data/rlhf/gsm8k/train.parquet
  val_files: ~/data/rlhf/gsm8k/test.parquet
  prompt_key: prompt
  max_prompt_length: 512
  max_response_length: 512
  train_batch_size: 1024
  val_batch_size: 1312
  return_raw_input_ids: False  # This should be set to true when the tokenizer between policy and rm differs
  return_raw_chat: False
```

--------------------------------

### Configure ScalingInter-RL Training (Bash)

Source: https://github.com/woooodyy/agentgym-rl/blob/main/README.md

This command configures and launches ScalingInter-RL training, which dynamically adjusts the number of interaction turns. It sets the `algorithm.rounds_ctrl.type` to 'scaling_inter_stepwise' and specifies the initial number of rounds and the frequency for increasing them.

```bash
algorithm.rounds_ctrl.type=scaling_inter_stepwise\
algorithm.rounds_ctrl.steps_scaling_inter=100 \
algorithm.rounds_ctrl.rounds=[10,20,30] \

```

--------------------------------

### Using Pre-implemented Dispatch Modes with Verl

Source: https://github.com/woooodyy/agentgym-rl/blob/main/AgentGym-RL/docs/advance/dpo_extension.rst

Illustrates using pre-implemented dispatch modes from Verl, such as Dispatch.DP_COMPUTE_PROTO, which requires a specific data interface like DataProto.

```python
from verl.single_controller.base.decorator import register, Dispatch

@register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
def generate_sequences(self, data: DataProto) -> DataProto:
    pass
```
Tokens:
Raw