Repository: volcengine/verl
Pull request: [misc] feat: spport rmpad/data-packing in FSDP with transformers
- Use `actor_rollout_ref.model.use_rmpad=True` + `critic.model.use_rmpad=True \` + `reward_model.model.use_rmpad=True ` to enable rmpad for different models. Default set to False
- Using `AutoModelForTokenClassification` for Value and Reward Model. Instead of using SeqenceClassification
- Compute logprob convert to `log_probs_from_logits_response_rmpad`

Resolve: https://github.com/volcengine/verl/issues/53

**Comparison using DeepSeek7b and GSM8k:**
About 1.7x speedup compare to no rmpad (original cases)


Problem statement:
Do we have plans for data packing?
Thank you for creating this excellent project. The design is elegant and the implementation is very efficient.

I'm wondering if there are plans to support sequence packing? 
OpenRLHF has demonstrated significant speedup after supporting sequence packing. Given VERL's already impressive performance, I believe adding packing support could bring even greater efficiency gains.

Hints from issue discussion:
Hi @YixinSong-e, thanks for your positive feedback about veRL!
Do you mean data packing in the FSDP backend?
We already implemented data packing (i.e. remove padding?) in the Megatron-LM backend.
For the FSDP backend, we do have some monkey-patch to implement it upon transformers but it's not general for all models in transformers. So we haven't released them yet.
And we also noticed that only some important models (e.g., llama, qwen...) support data packing naturally in transformers.

We're thinking about some elegant ways to support data packing for all models in our fsdp workers. Do you have any suggestions or references?
Hey @PeterSH6! 

+1 for such an awesome project. My favourite in the new wave of open-source RLHF frameworks.

> For the FSDP backend, we do have some monkey-patch to implement it upon transformers but it's not general for all models in transformers. So we haven't released them yet.

I'm curious what kind of monkey patch you mean here? I'd be interested in contributing if you'd like some help, I'm working on something similar in another repo.
+1 This seems to be widely used and speed up training, and I am happy to PR as well.
Hi @SalmanMohammadi and @fzyzcjy , sorry for the late response.

The monkey patch we apply in transformers==4.44 is something like that:

```python
def flash_attn2_rmpad_forward(
        self: LlamaFlashAttention2,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs  # for compatibility
):
    hidden_states = hidden_states.squeeze(0)  # (total_nnz, hidden_size)
    # this matches with hidden_states
    indices, cu_seqlens, max_seqlen_in_batch = _get_unpad_data(attention_mask)

    hidden_states = pad_input(hidden_states,
                              indices=indices,
                              batch=attention_mask.shape[0],
                              seqlen=attention_mask.shape[1])

    position_ids = compute_position_id_with_mask(attention_mask)

    attn_output, _, _ = llama_flash_attn2_orig_forward(self,
                                                       hidden_states=hidden_states,
                                                       attention_mask=attention_mask,
                                                       position_ids=position_ids,
                                                       past_key_value=past_key_value,
                                                       output_attentions=output_attentions,
                                                       use_cache=use_cache,
                                                       cache_position=cache_position,
                                                       **kwargs)

    attn_output = index_first_axis(rearrange(attn_output, "b s ... -> (b s) ..."), indices)

    return attn_output, None, None

# Monkey patches the forward function.
from transformers.models.llama.modeling_llama import LlamaFlashAttention2
LlamaFlashAttention2.forward = flash_attn2_rmpad_forward
```
With this patch, we also need to change the input to the actor model to `input_ids =  unpad_input(input_ids, attention_mask)`. Therefore, this method is hard to generalize to all models in transformers.

We also noticed that in transformers>4.45, it naturally support data packing in many models.
We really need some help to support data packing in the latest transformers. 
I've created a Slack channel for this [feature](https://verlgroup.slack.com/archives/C086K823VPH) and invited you. We can discuss the data-packing features there!
Btw, in SFT stage (e.g. trl.SFTTrainer), I personally use https://huggingface.co/blog/packing-with-FA2, which seems to provide position_ids and remove pads. That one seems to work well as long as FA2 exists for a model and does not need patching. Have not checked into details though...
Yeah, we implement this patching in older transformers version (<=4.44). I believe it's not necessary now.
I am thinking about another kind of data packing, which I used during SFT and seem to work well, thus propose here for discussions and happy to PR.

Since I never use this in PPO (only in SFT), I would appreciate it if I could hear some thoughts, e.g. will we have some trouble for this due to special properties of RL?

Suppose we have 8 sequences to be learnt by policy/critic of length 100,100,400,100,100,100,500,300, and max seq len is 500. Suppose our cuda memory only allows to process 1000 tokens in a batch.

The packing mentioned above will have to use bs=2, because we have to deal with the worst case that everything is on the max seq len. Then we will need 4 batches in this example.

The proposal is that, we can put as many sequences as possible into a batch, until it exceeds the max seq len. For example, here we can put into two batches: [100,100,400,100,100,100], [500,300]. Then we will only do 2 batches.

A shortcoming is that, the number of *examples* it learns in a batch is changed. I am not sure whether this will be a problem for RL. In addition, the number of *tokens* it learns is indeed more consistent than the original approach.

This is also similar to the "packing" commonly used in pretraining with some differences. In pretraining, it seems we often concat documents, and cut off whenever it reaches 2048 tokens. This is not very ideal for PPO especially for math, because it seems not very meaningful to learn "second half of a math example" without knowing the first half. The proposal is a slight variant that avoids this.

Commit message: fix util issue

Patch to review:
diff --git a/examples/ppo_trainer/run_deepseek7b_llm.sh b/examples/ppo_trainer/run_deepseek7b_llm.sh
index 108fba14d2..b54ecea803 100644
--- a/examples/ppo_trainer/run_deepseek7b_llm.sh
+++ b/examples/ppo_trainer/run_deepseek7b_llm.sh
@@ -9,6 +9,7 @@ python3 -m verl.trainer.main_ppo \
     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_rmpad=True \
     actor_rollout_ref.actor.ppo_mini_batch_size=256 \
     actor_rollout_ref.actor.ppo_micro_batch_size=32 \
     actor_rollout_ref.actor.fsdp_config.param_offload=False \
@@ -21,6 +22,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.ref.log_prob_micro_batch_size=128 \
     actor_rollout_ref.ref.fsdp_config.param_offload=True \
     critic.optim.lr=1e-5 \
+    critic.model.use_rmpad=True \
     critic.model.path=deepseek-ai/deepseek-llm-7b-chat \
     critic.model.enable_gradient_checkpointing=False \
     critic.ppo_micro_batch_size=32 \
@@ -31,7 +33,7 @@ python3 -m verl.trainer.main_ppo \
     trainer.critic_warmup=0 \
     trainer.logger=['console','wandb'] \
     trainer.project_name='verl_example_gsm8k' \
-    trainer.experiment_name='deepseek_llm_7b_function_rm' \
+    trainer.experiment_name='deepseek_llm_7b_function_rm_normpad' \
     trainer.n_gpus_per_node=8 \
     trainer.nnodes=1 \
     trainer.save_freq=-1 \
diff --git a/examples/ppo_trainer/run_gemma.sh b/examples/ppo_trainer/run_gemma.sh
index 200ebdb143..541c860dbf 100644
--- a/examples/ppo_trainer/run_gemma.sh
+++ b/examples/ppo_trainer/run_gemma.sh
@@ -9,6 +9,7 @@ python3 -m verl.trainer.main_ppo \
     data.max_response_length=512 \
     actor_rollout_ref.model.path=google/gemma-2-2b-it \
     actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_rmpad=True \
     actor_rollout_ref.actor.ppo_mini_batch_size=128 \
     actor_rollout_ref.actor.ppo_micro_batch_size=4 \
     actor_rollout_ref.actor.fsdp_config.param_offload=False \
@@ -21,6 +22,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.ref.log_prob_micro_batch_size=4 \
     actor_rollout_ref.ref.fsdp_config.param_offload=True \
     critic.optim.lr=1e-5 \
+    critic.model.use_rmpad=True \
     critic.model.path=google/gemma-2-2b-it \
     critic.model.enable_gradient_checkpointing=False \
     critic.ppo_micro_batch_size=4 \
diff --git a/examples/ppo_trainer/run_qwen2-7b.sh b/examples/ppo_trainer/run_qwen2-7b.sh
index c6ffc1b4bf..ac4bf2eb85 100644
--- a/examples/ppo_trainer/run_qwen2-7b.sh
+++ b/examples/ppo_trainer/run_qwen2-7b.sh
@@ -17,6 +17,7 @@ python3 -m verl.trainer.main_ppo \
     data.max_response_length=512 \
     actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \
     actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_rmpad=True \
     actor_rollout_ref.actor.ppo_mini_batch_size=256 \
     actor_rollout_ref.actor.ppo_micro_batch_size=16 \
     actor_rollout_ref.actor.fsdp_config.param_offload=False \
@@ -29,6 +30,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.ref.log_prob_micro_batch_size=16 \
     actor_rollout_ref.ref.fsdp_config.param_offload=True \
     critic.optim.lr=1e-5 \
+    critic.model.use_rmpad=True \
     critic.model.path=Qwen/Qwen2-7B-Instruct \
     critic.model.enable_gradient_checkpointing=False \
     critic.ppo_micro_batch_size=16 \
diff --git a/examples/ppo_trainer/run_qwen2-7b_rm.sh b/examples/ppo_trainer/run_qwen2-7b_rm.sh
index 3755b38064..4cfc5ead7d 100644
--- a/examples/ppo_trainer/run_qwen2-7b_rm.sh
+++ b/examples/ppo_trainer/run_qwen2-7b_rm.sh
@@ -18,6 +18,7 @@ python3 -m verl.trainer.main_ppo \
     data.return_raw_chat=True \
     actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \
     actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_rmpad=True \
     actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.1 \
     actor_rollout_ref.actor.ppo_mini_batch_size=256 \
     actor_rollout_ref.actor.ppo_micro_batch_size=16 \
@@ -31,6 +32,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.ref.log_prob_micro_batch_size=16 \
     actor_rollout_ref.ref.fsdp_config.param_offload=True \
     critic.optim.lr=1e-5 \
+    critic.model.use_rmpad=True \
     critic.optim.lr_warmup_steps_ratio=0.05 \
     critic.model.path=Qwen/Qwen2-7B-Instruct \
     critic.model.enable_gradient_checkpointing=False \
@@ -40,6 +42,7 @@ python3 -m verl.trainer.main_ppo \
     critic.model.fsdp_config.optimizer_offload=False \
     reward_model.enable=True \
     reward_model.model.path=sfairXC/FsfairX-Gemma2-RM-v0.1\
+    reward_model.model.use_rmpad=True \
     reward_model.model.fsdp_config.param_offload=True \
     reward_model.micro_batch_size=16 \
     algorithm.kl_ctrl.kl_coef=0.001 \
diff --git a/examples/ppo_trainer/run_qwen2.5-32b.sh b/examples/ppo_trainer/run_qwen2.5-32b.sh
index 1192f1eba9..90f22d7d41 100644
--- a/examples/ppo_trainer/run_qwen2.5-32b.sh
+++ b/examples/ppo_trainer/run_qwen2.5-32b.sh
@@ -18,6 +18,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.model.path=Qwen/Qwen2.5-32B-Instruct \
     actor_rollout_ref.model.enable_gradient_checkpointing=False \
     actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_rmpad=True \
     actor_rollout_ref.actor.ppo_mini_batch_size=256 \
     actor_rollout_ref.actor.ppo_micro_batch_size=16 \
     actor_rollout_ref.actor.fsdp_config.param_offload=False \
@@ -30,6 +31,7 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.ref.log_prob_micro_batch_size=128 \
     actor_rollout_ref.ref.fsdp_config.param_offload=True \
     critic.optim.lr=1e-5 \
+    critic.model.use_rmpad=True \
     critic.model.path=Qwen/Qwen2.5-32B-Instruct \
     critic.model.enable_gradient_checkpointing=False \
     critic.ppo_micro_batch_size=32 \
diff --git a/tests/model/test_transformer.py b/tests/model/test_transformer.py
new file mode 100644
index 0000000000..1db1dcb58a
--- /dev/null
+++ b/tests/model/test_transformer.py
@@ -0,0 +1,107 @@
+from transformers import AutoModelForCausalLM, AutoConfig, AutoModelForTokenClassification, AutoTokenizer
+
+import torch
+from verl.utils.model import create_random_mask, compute_position_id_with_mask
+from verl.utils.torch_functional import masked_mean
+from flash_attn.bert_padding import unpad_input, pad_input, index_first_axis, rearrange
+
+
+def test_hf_casual_models():
+    batch_size = 4
+    seqlen = 128
+
+    # TODO(sgm): add more models for test
+    # we only need one scale for each model
+    test_cases = ['deepseek-ai/deepseek-llm-7b-chat', 'Qwen/Qwen2-7B-Instruct']
+    for test_case in test_cases:
+        config = AutoConfig.from_pretrained(test_case)
+        model = AutoModelForCausalLM.from_pretrained(test_case,
+                                                     torch_dtype=torch.bfloat16,
+                                                     attn_implementation='flash_attention_2')
+        model = model.to(device='cuda')
+        input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device='cuda')
+        attention_mask = create_random_mask(input_ids=input_ids,
+                                            max_ratio_of_left_padding=0.1,
+                                            max_ratio_of_valid_token=0.8,
+                                            min_ratio_of_valid_token=0.5)
+        position_ids = compute_position_id_with_mask(
+            attention_mask)  # TODO(sgm): we can construct the position_ids_rmpad here
+
+        input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+            input_ids.unsqueeze(-1), attention_mask)  # input_ids_rmpad (total_nnz, ...)
+        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)
+
+        # unpad the position_ids to align the rotary
+        position_ids_rmpad = index_first_axis(rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."),
+                                              indices).transpose(0, 1)
+
+        # input with input_ids_rmpad and postition_ids to enable flash attention varlen
+        rmpad_logits = model(input_ids_rmpad, position_ids=position_ids_rmpad,
+                             use_cache=False).logits  # (1, total_nnz, vocab_size)
+        pad_logits = pad_input(rmpad_logits.squeeze(0), indices, batch_size, seqlen=seqlen)
+
+        origin_logits = model(input_ids=input_ids,
+                              attention_mask=attention_mask,
+                              position_ids=position_ids,
+                              use_cache=False).logits
+
+        torch.testing.assert_close(masked_mean(pad_logits, attention_mask[:, :, None]),
+                                   masked_mean(origin_logits, attention_mask[:, :, None]),
+                                   msg=f'{test_case} rmpad and non-rmpad logits are not equal')
+    print(f'Check pass')
+
+
+def test_hf_value_models():
+    batch_size = 4
+    seqlen = 128
+
+    # TODO(sgm): add more models for test
+    # we only need one scale for each model
+    test_cases = ['deepseek-ai/deepseek-llm-7b-chat', 'Qwen/Qwen2-7B-Instruct']
+    for test_case in test_cases:
+        config = AutoConfig.from_pretrained(test_case)
+        config.num_labels = 1
+        setattr(config, 'classifier_dropout', 0)
+        setattr(config, 'hidden_dropout', 0)
+        model = AutoModelForTokenClassification.from_pretrained(test_case,
+                                                                torch_dtype=torch.bfloat16,
+                                                                config=config,
+                                                                attn_implementation='flash_attention_2')
+        model = model.to(device='cuda')
+        input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device='cuda')
+        attention_mask = create_random_mask(input_ids=input_ids,
+                                            max_ratio_of_left_padding=0.1,
+                                            max_ratio_of_valid_token=0.8,
+                                            min_ratio_of_valid_token=0.5)
+        position_ids = compute_position_id_with_mask(
+            attention_mask)  # TODO(sgm): we can construct the position_ids_rmpad here
+
+        input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+            input_ids.unsqueeze(-1), attention_mask)  # input_ids_rmpad (total_nnz, ...)
+        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)
+
+        # unpad the position_ids to align the rotary
+        position_ids_rmpad = index_first_axis(rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."),
+                                              indices).transpose(0, 1)
+
+        origin_logits = model(input_ids=input_ids,
+                              attention_mask=attention_mask,
+                              position_ids=position_ids,
+                              use_cache=False).logits
+        print(f'origin_logits: {origin_logits.shape}')
+
+        # input with input_ids_rmpad and postition_ids to enable flash attention varlen
+        rmpad_logits = model(input_ids_rmpad, position_ids=position_ids_rmpad,
+                             use_cache=False).logits  # (1, total_nnz, vocab_size)
+        print(f'rmpad_logits: {rmpad_logits.shape}')
+        pad_logits = pad_input(rmpad_logits.squeeze(0), indices, batch_size, seqlen=seqlen)
+
+        torch.testing.assert_close(masked_mean(pad_logits, attention_mask[:, :, None]),
+                                   masked_mean(origin_logits, attention_mask[:, :, None]),
+                                   msg=f'{test_case} rmpad and non-rmpad logits are not equal')
+    print('Value model check pass')
+
+
+if __name__ == '__main__':
+    test_hf_casual_models()
+    test_hf_value_models()
diff --git a/tests/rollout/run_fsdp_vllm.py b/tests/rollout/run_fsdp_vllm.py
index d9e165a9fa..3267bee0e5 100644
--- a/tests/rollout/run_fsdp_vllm.py
+++ b/tests/rollout/run_fsdp_vllm.py
@@ -127,8 +127,9 @@ def main():
     for i in range(batch_size):
         idx_list.append(_pre_process_inputs(pad_token_id, input_ids[i]))
     print('start generation')
-    outputs = llm.generate(prompt_token_ids=idx_list, sampling_params=sampling_params, use_tqdm=False)
-    vllm_output = outputs[0].cuda()
+    while True:
+        outputs = llm.generate(prompt_token_ids=idx_list, sampling_params=sampling_params, use_tqdm=False)
+        vllm_output = outputs[0].cuda()
     if torch.distributed.get_rank() == 0:
         print(f'hf response: {tokenizer.batch_decode(response)}')
         print(f'vllm response: {tokenizer.batch_decode(vllm_output)}')
diff --git a/verl/trainer/config/ppo_trainer.yaml b/verl/trainer/config/ppo_trainer.yaml
index 774125de69..df7511bfd0 100644
--- a/verl/trainer/config/ppo_trainer.yaml
+++ b/verl/trainer/config/ppo_trainer.yaml
@@ -17,6 +17,7 @@ actor_rollout_ref:
     external_lib: null
     override_config: { }
     enable_gradient_checkpointing: False
+    use_rmpad: False
   actor:
     strategy: fsdp  # This is for backward-compatibility
     ppo_mini_batch_size: 256
@@ -83,6 +84,7 @@ critic:
     override_config: { }
     external_lib: ${actor_rollout_ref.model.external_lib}
     enable_gradient_checkpointing: False
+    use_rmpad: False
     fsdp_config:
       param_offload: False
       grad_offload: False
@@ -105,6 +107,7 @@ reward_model:
     input_tokenizer: ${actor_rollout_ref.model.path}  # set this to null if the chat template is identical
     path: ~/models/FsfairX-LLaMA3-RM-v0.1
     external_lib: ${actor_rollout_ref.model.external_lib}
+    use_rmpad: False
     fsdp_config:
       min_num_params: 0
       param_offload: False
diff --git a/verl/workers/actor/dp_actor.py b/verl/workers/actor/dp_actor.py
index e0001c9525..18c5e19c60 100644
--- a/verl/workers/actor/dp_actor.py
+++ b/verl/workers/actor/dp_actor.py
@@ -14,7 +14,7 @@
 """
 Single Process Actor
 """
-from typing import Iterable
+from typing import Iterable, Tuple
 
 import torch
 from torch import nn
@@ -24,7 +24,9 @@
 from verl.trainer.ppo import core_algos
 from verl.workers.actor import BasePPOActor
 from verl.utils.py_functional import append_to_dict
-from verl.utils.torch_functional import logprobs_from_logits
+from verl.utils.torch_functional import logprobs_from_logits, log_probs_from_logits_response_rmpad
+
+from flash_attn.bert_padding import pad_input, unpad_input, rearrange, index_first_axis
 
 __all__ = ['DataParallelPPOActor']
 
@@ -41,17 +43,44 @@ def __init__(
         super().__init__(config)
         self.actor_module = actor_module
         self.actor_optimizer = actor_optimizer
+        self.use_rmpad = self.config.get('use_rmpad', False)
+        print(f'Actor use_rmpad={self.use_rmpad}')
 
-    def _forward_micro_batch(self, micro_batch, temperature):
+    def _forward_micro_batch(self, micro_batch, temperature) -> Tuple[torch.Tensor, torch.Tensor]:
         response_length = micro_batch['responses'].size(-1)
         with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-            output = self.actor_module(input_ids=micro_batch['input_ids'],
-                                       attention_mask=micro_batch['attention_mask'],
-                                       position_ids=micro_batch['position_ids'],
-                                       use_cache=False)  # prevent model thinks we are generating
-            logits = output.logits / temperature
-            logits = logits[:, -response_length - 1:-1]
-            log_probs = logprobs_from_logits(logits, micro_batch['responses'])
+            input_ids = micro_batch['input_ids']
+            attention_mask = micro_batch['attention_mask']
+            position_ids = micro_batch['position_ids']
+
+            if self.use_rmpad:
+                input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+                    input_ids.unsqueeze(-1), attention_mask)  # input_ids_rmpad (total_nnz, ...)
+                input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)
+
+                # unpad the position_ids to align the rotary
+                position_ids_rmpad = index_first_axis(rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."),
+                                                      indices).transpose(0, 1)
+                # only pass input_ids and position_ids to enable flash_attn_varlen
+                output = self.actor_module(input_ids=input_ids_rmpad,
+                                           attention_mask=None,
+                                           position_ids=position_ids_rmpad,
+                                           use_cache=False)  # prevent model thinks we are generating
+                logits_rmpad = output.logits.squeeze(0)  # (total_nnz, vocab_size)
+                logits_rmpad = logits_rmpad / temperature
+                log_probs = log_probs_from_logits_response_rmpad(input_ids=input_ids,
+                                                                 attention_mask=attention_mask,
+                                                                 logits_rmpad=logits_rmpad,
+                                                                 response_length=response_length)  # (batch, seqlen)
+                logits = logits_rmpad
+            else:
+                output = self.actor_module(input_ids=input_ids,
+                                           attention_mask=attention_mask,
+                                           position_ids=position_ids,
+                                           use_cache=False)  # prevent model thinks we are generating
+                logits = output.logits / temperature
+                logits = logits[:, -response_length - 1:-1]
+                log_probs = logprobs_from_logits(logits, micro_batch['responses'])
             return logits, log_probs
 
     def _make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]:
@@ -145,8 +174,17 @@ def update_policy(self, data: DataProto):
                                                                               advantages=advantages,
                                                                               eos_mask=response_mask,
                                                                               cliprange=clip_ratio)
-
-                entropy_loss = core_algos.compute_entropy_loss(logits, response_mask)
+                # compute entropy loss
+                if self.use_rmpad:
+                    full_response_mask = attention_mask.clone()
+                    full_response_mask[:, :-response_length] = 0  # set the prompt part to zero
+                    full_response_mask_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+                        full_response_mask.unsqueeze(-1), attention_mask=attention_mask)
+                    full_response_mask_rmpad = full_response_mask_rmpad.squeeze(-1)  # (total_nnz)
+                    entropy_loss = core_algos.compute_entropy_loss(logits, full_response_mask_rmpad)  # (total_nnz,)
+                else:
+                    entropy_loss = core_algos.compute_entropy_loss(logits, response_mask)
+                # compute policy loss
                 policy_loss = pg_loss - entropy_loss * entropy_coeff
 
                 loss = policy_loss / self.gradient_accumulation
diff --git a/verl/workers/critic/dp_critic.py b/verl/workers/critic/dp_critic.py
index 8e845f008d..94923bf5b1 100644
--- a/verl/workers/critic/dp_critic.py
+++ b/verl/workers/critic/dp_critic.py
@@ -29,6 +29,8 @@
 from verl.utils.py_functional import append_to_dict
 from verl.utils.torch_functional import masked_mean
 
+from flash_attn.bert_padding import pad_input, unpad_input, rearrange, index_first_axis
+
 __all__ = ['DataParallelPPOCritic']
 
 
@@ -38,6 +40,8 @@ def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Opt
         super().__init__(config=config)
         self.critic_module = critic_module
         self.critic_optimizer = critic_optimizer
+        self.use_rmpad = self.config.model.get('use_rmpad', False)
+        print(f'Critic use_rmpad={self.use_rmpad}')
 
         assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size == 0
         self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size
@@ -45,12 +49,37 @@ def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Opt
     def _forward_micro_batch(self, micro_batch):
         response_length = micro_batch['responses'].size(-1)
         with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-            output = self.critic_module(input_ids=micro_batch['input_ids'],
-                                        attention_mask=micro_batch['attention_mask'],
-                                        position_ids=micro_batch['position_ids'],
-                                        use_cache=False)  # prevent model thinks we are generating
-            values = output.logits
-            values = values[:, -response_length - 1:-1]
+            input_ids = micro_batch['input_ids']
+            batch, seqlen = input_ids.shape
+            attention_mask = micro_batch['attention_mask']
+            position_ids = micro_batch['position_ids']
+
+            if self.use_rmpad:
+                input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+                    input_ids.unsqueeze(-1), attention_mask)  # input_ids_rmpad (total_nnz, ...)
+                input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)
+
+                # unpad the position_ids to align the rotary
+                position_ids_rmpad = index_first_axis(rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."),
+                                                      indices).transpose(0, 1)
+                # only pass input_ids and position_ids to enable flash_attn_varlen
+                output = self.critic_module(input_ids=input_ids_rmpad,
+                                            attention_mask=None,
+                                            position_ids=position_ids_rmpad,
+                                            use_cache=False)  # prevent model thinks we are generating
+                values_rmpad = output.logits
+                values_rmpad = values_rmpad.squeeze(0)  # (total_nnz)
+
+                # pad it back
+                values = pad_input(values_rmpad, indices=indices, batch=batch, seqlen=seqlen).squeeze(-1)
+                values = values[:, -response_length - 1:-1]
+            else:
+                output = self.critic_module(input_ids=input_ids,
+                                            attention_mask=attention_mask,
+                                            position_ids=position_ids,
+                                            use_cache=False)  # prevent model thinks we are generating
+                values = output.logits
+                values = values[:, -response_length - 1:-1].squeeze(-1)
             return values
 
     def _make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]:
diff --git a/verl/workers/fsdp_workers.py b/verl/workers/fsdp_workers.py
index 36ab6e6ec8..7f4e9ee891 100644
--- a/verl/workers/fsdp_workers.py
+++ b/verl/workers/fsdp_workers.py
@@ -23,7 +23,7 @@
 import torch.distributed
 import verl.utils.hdfs_io as hdfs_io
 import verl.utils.torch_functional as verl_F
-from omegaconf import DictConfig
+from omegaconf import DictConfig, open_dict
 from verl import DataProto
 from verl.single_controller.base import Worker
 from verl.single_controller.base.decorator import register, Dispatch
@@ -254,6 +254,8 @@ def init_model(self):
         from omegaconf import OmegaConf
         override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create()))
 
+        use_rmpad = self.config.model.get('use_rmpad', False)
+
         if self._is_actor or self._is_rollout:
             # we need the model for actor and rollout
             if self._is_actor:
@@ -283,6 +285,8 @@ def init_model(self):
         # load from checkpoint
         if self._is_actor:
             OmegaConf.set_struct(self.config.actor, True)
+            with open_dict(self.config.actor):
+                self.config.actor.use_rmpad = use_rmpad
             self.actor = DataParallelPPOActor(config=self.config.actor,
                                               actor_module=self.actor_module_fsdp,
                                               actor_optimizer=self.actor_optimizer)
@@ -301,6 +305,8 @@ def init_model(self):
                 offload_fsdp_param_and_grad(module=self.ref_module_fsdp, offload_grad=self._is_offload_grad)
 
             OmegaConf.set_struct(self.config.ref, True)
+            with open_dict(self.config.ref):
+                self.config.ref.use_rmpad = use_rmpad
             self.ref_policy = DataParallelPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp)
 
         torch.cuda.empty_cache()
@@ -463,7 +469,6 @@ def _build_critic_model_optimizer(self, config):
         local_path = copy_local_path_from_hdfs(config.model.path)
         # note that the tokenizer between actor and critic may be different. So override tokenizer info with actor info
         # using random initialized model from any architecture. May not be the same as Actor.
-        # TODO: support loading critic weights from RM. Support using AutoModelForTokenClassification
 
         tokenizer_path = copy_local_path_from_hdfs(config.model.tokenizer_path)
         self.tokenizer = hf_tokenizer(tokenizer_path, trust_remote_code=config.model.get('trust_remote_code', False))
@@ -482,22 +487,22 @@ def _build_critic_model_optimizer(self, config):
         torch_dtype = self.config.model.fsdp_config.get('model_dtype', 'fp32')
         torch_dtype = PrecisionType.to_dtype(torch_dtype)
 
-        from transformers import AutoConfig, AutoModelForCausalLM
+        from transformers import AutoConfig, AutoModelForTokenClassification
         from torch import nn
 
         trust_remote_code = False
         critic_model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code)
-
+        critic_model_config.num_labels = 1
         init_context = get_init_weight_context_manager()
         with init_context(), warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            critic_module = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=local_path,
-                                                                 torch_dtype=torch_dtype,
-                                                                 config=critic_model_config,
-                                                                 attn_implementation='flash_attention_2',
-                                                                 trust_remote_code=trust_remote_code)
-            critic_module.lm_head = nn.Sequential(nn.Linear(critic_model_config.hidden_size, 1, dtype=torch_dtype),
-                                                  LambdaLayer(fn=squeeze))
+            setattr(critic_model_config, 'classifier_dropout', 0.)
+            setattr(critic_model_config, 'hidden_dropout', '0')
+            critic_module = AutoModelForTokenClassification.from_pretrained(pretrained_model_name_or_path=local_path,
+                                                                            torch_dtype=torch_dtype,
+                                                                            config=critic_model_config,
+                                                                            attn_implementation='flash_attention_2',
+                                                                            trust_remote_code=trust_remote_code)
 
             # some parameters may not in torch_dtype
             critic_module.to(torch_dtype)
@@ -642,9 +647,10 @@ def save_checkpoint(self, local_path, hdfs_path=None):
             offload_fsdp_param_and_grad(module=self.critic_module, offload_grad=self._is_offload_grad)
 
 
+# TODO(sgm): we may need to extract it to dp_reward_model.py
 class RewardModelWorker(Worker):
     """
-    Note that we only implement the reward model that is subclass of AutoModelForSequenceClassification.
+    Note that we only implement the reward model that is subclass of AutoModelForTokenClassification.
     """
 
     def __init__(self, config):
@@ -653,12 +659,12 @@ def __init__(self, config):
         if not torch.distributed.is_initialized():
             torch.distributed.init_process_group(backend="nccl")
         self.config = config
-
+        self.use_rmpad = self.config.model.get('use_rmpad', False)
         self.config.micro_batch_size //= torch.distributed.get_world_size()
 
     def _build_model(self, config):
         # the following line is necessary
-        from transformers import AutoModelForSequenceClassification, AutoConfig
+        from transformers import AutoModelForTokenClassification, AutoConfig
         from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, CPUOffload
 
         # download the checkpoint from hdfs
@@ -675,15 +681,18 @@ def _build_model(self, config):
 
         trust_remote_code = config.model.get('trust_remote_code', False)
         model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code)
+        model_config.num_labels = 1
         # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect
         init_context = get_init_weight_context_manager(use_meta_tensor=not model_config.tie_word_embeddings)
 
         with init_context(), warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            reward_module = AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path=local_path,
-                                                                               torch_dtype=torch.bfloat16,
-                                                                               attn_implementation='flash_attention_2',
-                                                                               trust_remote_code=trust_remote_code)
+            setattr(model_config, 'classifier_dropout', 0.)
+            reward_module = AutoModelForTokenClassification.from_pretrained(pretrained_model_name_or_path=local_path,
+                                                                            config=model_config,
+                                                                            torch_dtype=torch.bfloat16,
+                                                                            attn_implementation='flash_attention_2',
+                                                                            trust_remote_code=trust_remote_code)
             reward_module.to(torch.bfloat16)
         auto_wrap_policy = get_fsdp_wrap_policy(module=reward_module, config=self.config.model.fsdp_config)
 
@@ -707,12 +716,39 @@ def init_model(self):
         torch.cuda.empty_cache()
 
     def _forward_micro_batch(self, micro_batch):
+        from verl.utils.torch_functional import prepare_input_for_rmpad
+        from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis, rearrange
+
         with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-            output = self.reward_module(input_ids=micro_batch['input_ids'],
-                                        attention_mask=micro_batch['attention_mask'],
-                                        position_ids=micro_batch['position_ids'])
-            rm_score = output.logits  # (batch_size,)
-            rm_score = rm_score.squeeze(-1)
+            input_ids = micro_batch['input_ids']
+            batch, seqlen = input_ids.shape
+            attention_mask = micro_batch['attention_mask']
+            position_ids = micro_batch['position_ids']
+
+            if self.use_rmpad:
+                input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+                    input_ids.unsqueeze(-1), attention_mask)  # input_ids_rmpad (total_nnz, ...)
+                input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)
+
+                # unpad the position_ids to align the rotary
+                position_ids_rmpad = index_first_axis(rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."),
+                                                      indices).transpose(0, 1)
+                # only pass input_ids and position_ids to enable flash_attn_varlen
+                output = self.reward_module(input_ids=input_ids_rmpad,
+                                            attention_mask=None,
+                                            position_ids=position_ids_rmpad,
+                                            use_cache=False)  # prevent model thinks we are generating
+                reward_rmpad = output.logits
+                reward_rmpad = reward_rmpad.squeeze(0)  # (total_nnz)
+
+                # pad it back
+                rm_score = pad_input(reward_rmpad, indices=indices, batch=batch, seqlen=seqlen).squeeze(-1)
+            else:
+                output = self.reward_module(input_ids=input_ids,
+                                            attention_mask=attention_mask,
+                                            position_ids=position_ids)
+                rm_score = output.logits  # (batch_size, seq_len, 2)
+                rm_score = rm_score.squeeze(-1)
             return rm_score
 
     def _expand_to_token_level(self, data: DataProto, scores: torch.Tensor):


Code review comments:
diff_hunk: @@ -41,17 +43,44 @@ def __init__(
         super().__init__(config)
         self.actor_module = actor_module
         self.actor_optimizer = actor_optimizer
+        self.use_rmpad = self.config.get('use_rmpad', False)
+        print(f'Actor use_rmpad={self.use_rmpad}')
 
-    def _forward_micro_batch(self, micro_batch, temperature):
+    def _forward_micro_batch(self, micro_batch, temperature) -> Tuple[torch.Tensor, torch.Tensor]:
         response_length = micro_batch['responses'].size(-1)
         with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-            output = self.actor_module(input_ids=micro_batch['input_ids'],
-                                       attention_mask=micro_batch['attention_mask'],
-                                       position_ids=micro_batch['position_ids'],
-                                       use_cache=False)  # prevent model thinks we are generating
-            logits = output.logits / temperature
-            logits = logits[:, -response_length - 1:-1]
-            log_probs = logprobs_from_logits(logits, micro_batch['responses'])
+            input_ids = micro_batch['input_ids']
+            attention_mask = micro_batch['attention_mask']
+            position_ids = micro_batch['position_ids']
+
+            if self.use_rmpad:
+                input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+                    input_ids.unsqueeze(-1), attention_mask)  # input_ids_rmpad (total_nnz, ...)
+                input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)
+
+                # unpad the position_ids to align the rotary
+                position_ids_rmpad = index_first_axis(rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."),
+                                                      indices).transpose(0, 1)
+                # only pass input_ids and position_ids to enable flash_attn_varlen
+                output = self.actor_module(input_ids=input_ids_rmpad,
+                                           attention_mask=None,
+                                           position_ids=position_ids_rmpad,
+                                           use_cache=False)  # prevent model thinks we are generating
+                logits_rmpad = output.logits.squeeze(0)  # (total_nnz, vocab_size)
+                logits_rmpad = logits_rmpad / temperature
path: verl/workers/actor/dp_actor.py
text: @user1:
Change this line to in-place op