verl - verl核心代码详解（与令人纠结的batch size）

在RL训练中，我们会遇到各种各样的batch, 眼花缭乱。在这里，我们细致梳理一下这些batch 代表了什么，和它们之间的关系。这个blog里面，主要涉及如下文件中的代码：
verl/verl/trainer/ppo/ray_trainer.py
verl/verl/workers/fsdp_workers.py

1. 背景

采用FSDP 进行GRPO 训练。GRPO是DeepSeek提出的PPO的高效变体。相较于PPO：

• GRPO省略了Reward Model，直接用Rule-based 方式计算reward。
• GRPO省略了Critic Model（评论家模型）, 不再额外计算$V_i$
• 计算advantage直接基于Reward $R_i$， 而原始PPO(等算法)基于$V_i$. 这表明 GRPO直接用rollout example的reward 度量其Value (价值）.

2. batch size 相关参数

在verl/verl/trainer/config/ppo_trainer.yaml中，我们会遇到若干与batch size相关的参数，让我们一一拆解分析。
Warning： verl/verl/trainer/config/ppo_trainer.yaml中的配置参数可能会被运行脚本中的配置参数覆盖掉。

2.1 General

data.train_batch_size=60
trainer.n_gpus_per_node=6 
trainer.nnodes=1 

这里我们有1个节点，共6张卡。
注意:

• data.train_batch_size 必须可以整除 trainer.n_gpus_per_node， 否则报错。

2.2 actor_rollout_ref

## actor
actor_rollout_ref.actor.ppo_mini_batch_size=60
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8  # 这个似乎没什么用
actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 
### actor.fsdp_config
actor_rollout_ref.actor.fsdp_config.param_offload=False 
actor_rollout_ref.actor.fsdp_config.optimizer_offload=False 
## fsdp_config
actor_rollout_ref.fsdp_config.fsdp_size=-1 
## rollout
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 
actor_rollout_ref.rollout.n=12
actor_rollout_ref.rollout.tensor_model_parallel_size=2
## ref
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 

我们可以在ray_trainer.py 的 fit() 函数中粗略理解一下batch

with _timer('step', timing_raw): # 记录整个步骤时间
	# 整个步骤 = 计算old policy + 计算ref policy +  计算adv(包括计算value + 计算reward) 
	#			+ 更新critic model + 更新actor model + sometimes validation + sometimes ckpt saving
     with _timer('gen', timing_raw): # 记录生成序列时间
         # 使用actor模型生成文本序列
         # gen_batch：包含input_ids等生成所需数据
         # gen_batch_output：生成的序列及其概率等信息
         print('gen_batch shape: ', gen_batch.batch['input_ids'].shape)
         '''actor 进行rollout之前，
         gen_batch shape:  torch.Size([60, 8192]), 
         data.train_batch_size = 60
         	->  gen_batch.batch['input_ids'].shape[0] = 60
         '''
         gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)
         print("gen_batch_output.batch['prompt_token_ids'].shape: ", gen_batch_output.batch['prompts'].shape)
         '''actor 进行rollout之后，
         gen_batch_output.batch['prompt_token_ids'].shape:  torch.Size([720, 8192])
         data.train_batch_size = 60, actor_rollout_ref.rollout.n=12
         	->  gen_batch_output.batch['prompt_token_ids'].shape.shape[0] = 60 * 12 =720
         '''

TL,DR:
在 verl/verl/workers/fsdp_workers.py的class ActorRolloutRefWorker(Worker)中，我们会与这些参数打上交道

# 对于Actor
actor_rollout_ref.actor.ppo_mini_batch_size 和 GPU个数将共同决定 
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu

# 在rollout过程中计算 rollout sample 的 log_prob, 每个GPU的处理样例数直接由如下配置决定
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 

# 计算reference model的 log_prob， 每个GPU的处理样例数直接由如下配置决定
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 

3. fsdp_workers 与上述配置

3.1 ActorRolloutRefWorker 的配置函数

class ActorRolloutRefWorker(Worker):
    """
    This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy
    or a hybrid engine based on the config.rollout
    """

    def __init__(self, config: DictConfig, role: str):
    	'''
    		DictConfig: actor_rollout_ref 相关的配置
    		role: 该worker的角色，一般情况下 role = actor_rollout
    	'''
        super().__init__()
        self.config = config
        import torch.distributed
        if not torch.distributed.is_initialized():
            torch.distributed.init_process_group()

        # build device mesh for FSDP
        world_size = torch.distributed.get_world_size()
        # TODO(sgm): support FSDP hybrid shard for larger model
        '''
        	trainer.n_gpus_per_node=6  -> world_size = 6
        	actor_rollout_ref.fsdp_config.fsdp_size=-1 -> fsdp_size = -1
        '''

        self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=self.config.actor.fsdp_config.fsdp_size)

        # build device mesh for Ulysses Sequence Parallel
        self.ulysses_device_mesh = None
        self.ulysses_sequence_parallel_size = self.config.actor.get('ulysses_sequence_parallel_size', 1)
        dp = world_size // self.ulysses_sequence_parallel_size
        if self.ulysses_sequence_parallel_size > 1:
            # 创建一个二维 GPU 网格，同时支持数据并行（DP）和序列并行（SP）。
            print('self.ulysses_sequence_parallel_size: ', self.ulysses_sequence_parallel_size)
            self.ulysses_device_mesh = init_device_mesh('cuda',
                                                        mesh_shape=(dp, self.ulysses_sequence_parallel_size),
                                                        mesh_dim_names=['dp', 'sp'])

        self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh)

        self.role = role
        '''
        	正如前面的注释， self.role = actor_rollout
        '''
        assert self.role in ['actor', 'rollout', 'ref', 'actor_rollout', 'actor_rollout_ref']

        self._is_actor = self.role in ['actor', 'actor_rollout', 'actor_rollout_ref']
        self._is_rollout = self.role in ['rollout', 'actor_rollout', 'actor_rollout_ref']
        self._is_ref = self.role in ['ref', 'actor_rollout_ref']
		'''
			那么
			self._is_actor = True，
			self._is_rollout = True,  
			self._is_ref = Flase
		'''
        self._is_offload_param = False
        self._is_offload_optimizer = False
        if self._is_actor:
            self._is_offload_param = self.config.actor.fsdp_config.get('param_offload', False)
            self._is_offload_optimizer = self.config.actor.fsdp_config.get('optimizer_offload', False)
        '''
        	根据actor_rollout_ref.actor的配置信息, actor使用FSDP时，不进行 param 和 optimizer的offloading
        	actor_rollout_ref.actor.fsdp_config.param_offload=False 
        		-> self._is_offload_param = False
			actor_rollout_ref.actor.fsdp_config.optimizer_offload=False 
				-> self._is_offload_optimizer = False
        '''
        elif self._is_ref:
            # TODO: it seems that manual offload is slowly than FSDP offload
            self._is_offload_param = self