feat: remove left-right padding mode in FSDP engine

1. add transformation layer in actor/critic worker
2. support nested tensor in megatron engine
3. pass test_engine.py test for megatron/fsdp engines

Author: houminz

tests/models/test_engine.py

@@ -84,6 +84,8 @@ def test_actor_engine(strategy):
     # init model
     wg.init_model()
 
+    print(f"test_actor_engine strategy: {strategy}, config: {config} after init_model")
+
     batch_size = 8
     seqlen = 32
 
@@ -100,8 +102,6 @@ def test_actor_engine(strategy):
 
     global_token_num = torch.sum(attention_mask, dim=-1).tolist()
 
-    print(input_ids.float().mean(), attention_mask.float().mean())
-
     responses = input_ids[:, response_length:]
     response_mask = attention_mask[:, response_length:]
 
@@ -129,6 +129,7 @@ def test_actor_engine(strategy):
     hf_logprobs = logprobs_from_logits_naive(
         hf_output.logits[:, -response_length - 1 : -1, :].float(), input_ids[:, -response_length:]
     )
+
     hf_logprobs_mean = torch.mean(hf_logprobs * response_mask)
     mcore_logprobs_mean = torch.mean(output.batch["old_log_probs"] * response_mask)
 
@@ -173,6 +174,8 @@ def create_model():
 def test_critic_engine(strategy):
     ray.init()
 
+    torch.autograd.set_detect_anomaly(True)
+
     path = create_model()
     model_config = HFModelConfig(path=path, load_tokenizer=False)
 
@@ -209,6 +212,8 @@ def test_critic_engine(strategy):
     # init model
     wg.init_model()
 
+    print(f"test_critic_engine strategy: {strategy}, config: {config}")
+
     batch_size = 8
     seqlen = 32
 
@@ -257,6 +262,9 @@ def test_critic_engine(strategy):
 
     engine_values = torch.mean(output.batch["values"] * response_mask)
 
+    print(f"engine_values: {output.batch['values']}")
+    print(f"hf_values_mean: {hf_values_mean}, engine_values: {engine_values}")
+
     torch.testing.assert_close(hf_values_mean, engine_values, atol=1e-2, rtol=1e-2)
 
     data = data.union(output)
@@ -265,6 +273,7 @@ def test_critic_engine(strategy):
     data.batch["values"] = torch.rand_like(responses, dtype=torch.float32)
     data.batch["returns"] = torch.rand_like(responses, dtype=torch.float32)
 
+    print(f"before update critic: {data}")
     # update again
     ppo_metrics = wg.update_critic(data)
     print(ppo_metrics)
@@ -354,6 +363,8 @@ def test_per_tensor_generator(world_size, tmp_path, config, strategy):
     os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)
     # create a model
     model_path = create_actor_model(tmp_path, config)
+
+    print(f"test_per_tensor_generator world_size: {world_size}, strategy: {strategy}, config: {config}")
     # spawn workers
     mp.spawn(
         fn=_worker,

tests/special_e2e/sft/run_sft_engine_gsm8k.sh

@@ -29,7 +29,7 @@ PP_SIZE=${PP_SIZE:-1}
 VPP_SIZE=${VPP_SIZE:-null}
 CP_SIZE=${CP_SIZE:-1}
 
-PAD_MODE=${PAD_MODE:-left_right}
+PAD_MODE=${PAD_MODE:-no_padding}
 
 USE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True}
 
@@ -80,8 +80,6 @@ torchrun --standalone --nnodes=1 --nproc_per_node=${NUM_GPUS} ${ENTRYPOINT} \
     data.train_files="${TRAIN_FILES}" \
     data.val_files="${VAL_FILES}" \
     data.train_batch_size=256 \
-    data.max_prompt_length=1024 \
-    data.max_response_length=1024 \
     data.pad_mode=${PAD_MODE} \
     data.truncation=error \
     data.use_dynamic_bsz=True \

tests/special_e2e/sft/test_sft_engine_all.sh

@@ -9,15 +9,6 @@ echo "run with single gpu as golden"
 BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp VERL_FILE_LOGGER_PATH=~/verl/test/log/golden.jsonl bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
 
 # test with fsdp 1
-echo "run with sp1 fsdp_size2 num_gpus8 fsdp_strategy fsdp pad_mode left_right"
-BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=left_right bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
-echo "run with sp1 fsdp_size-1 num_gpus8 fsdp_strategy fsdp pad_mode left_right"
-BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=left_right bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
-echo "run with sp2 fsdp_size-1 num_gpus8 fsdp_strategy fsdp pad_mode left_right"
-BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=left_right bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
-echo "run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode left_right"
-BACKEND=fsdp SP_SIZE=4 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=left_right bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh 
-
 echo "run with sp1 fsdp_size2 num_gpus8 fsdp_strategy fsdp pad_mode no_padding"
 BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
 echo "run with sp1 fsdp_size-1 num_gpus8 fsdp_strategy fsdp pad_mode no_padding"
@@ -27,18 +18,14 @@ BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_pa
 echo "run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode no_padding"
 BACKEND=fsdp SP_SIZE=4 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
 
-# test use_remove_padding and pad_mode left_right/no_padding
-echo "run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode left_right use_remove_padding False"
-BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=left_right USE_REMOVE_PADDING=False bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
+# test use_remove_padding and pad_mode no_padding
 echo "run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode no_padding use_remove_padding False"
 BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding USE_REMOVE_PADDING=False bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
 
 
 # test with fsdp 2
-echo "run with sp1 fsdp_size1 num_gpus1 fsdp_strategy fsdp2 pad_mode left_right"
-BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp2 PAD_MODE=left_right bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
-echo "run with sp1 fsdp_size1 num_gpus1 fsdp_strategy fsdp2 pad_mode no_padding"
-BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp2 PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
+echo "run with sp1 fsdp_size1 num_gpus1 fsdp_strategy fsdp2"
+BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
 
 echo "run with sp1 fsdp_size-1 num_gpus8 fsdp_strategy fsdp2"
 BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
@@ -50,6 +37,7 @@ BACKEND=fsdp SP_SIZE=4 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/spe
 # test with megatron
 echo "run with tp1 pp1 cp1 num_gpus1"
 BACKEND=megatron TP_SIZE=1 PP_SIZE=1 CP_SIZE=1 NUM_GPUS=1 bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
+
 echo "run with tp2 pp2 vpp2 cp1 num_gpus8"
 BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=2 CP_SIZE=1 NUM_GPUS=8 bash tests/special_e2e/sft/run_sft_engine_gsm8k.sh
 

tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py

@@ -177,28 +177,26 @@ def test_multiturn_sft_dataset():
     assert torch.all(padded_item["attention_mask"][actual_length:] == 0), "Attention mask not set correctly for padding"
     assert torch.all(padded_item["loss_mask"][actual_length:] == 0), "Loss mask not set correctly for padding"
 
-    # test left right padding
+    # test no-padding
     config = {
         "max_length": 512,
         "truncation": "error",
         "multiturn": {"messages_key": "messages"},
-        "pad_mode": "left_right",
-        "max_prompt_length": 64,
-        "max_response_length": 64,
+        "pad_mode": "no_padding",
     }
     dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, config=config)
 
     item0 = dataset[0]
 
-    # make sure all the input_ids with attention_mask == 0 are all padding
-    assert torch.all(item0["input_ids"][item0["attention_mask"] == 0] == tokenizer.pad_token_id)
+    # Verify that the output contains expected keys for no-padding mode
+    required_keys = ["input_ids", "position_ids", "loss_mask"]
+    for key in required_keys:
+        assert key in item0, f"Missing key {key} in no-padding mode dataset item"
+        assert isinstance(item0[key], torch.Tensor), f"Expected torch.Tensor for {key} in no-padding mode"
 
     # make sure assistant_text matches with expected
-    assistant_text = tokenizer.decode(item0["responses"][item0["response_mask"] == 1])
+    assistant_text = tokenizer.decode(item0["input_ids"][item0["loss_mask"] == 1])
     assert assistant_text == "2+2 equals 4.<|im_end|>\n4+4 equals 8.<|im_end|>\n"
 
-    # make sure responses are part of input_ids
-    assert torch.all(item0["input_ids"][-item0["responses"].shape[0] :] == item0["responses"])
-
     print("All tests passed!")
     print("Starting test...")

verl/models/mcore/__init__.py

@@ -16,6 +16,7 @@
 from .registry import (
     get_mcore_forward_fn,
     get_mcore_forward_fused_fn,
+    get_mcore_forward_no_padding_fn,
     get_mcore_weight_converter,
     hf_to_mcore_config,
     init_mcore_model,
@@ -27,4 +28,5 @@
     "get_mcore_forward_fn",
     "get_mcore_weight_converter",
     "get_mcore_forward_fused_fn",
+    "get_mcore_forward_no_padding_fn",
 ]

verl/models/mcore/model_forward.py

@@ -16,7 +16,14 @@
 
 from verl.utils.megatron_utils import unwrap_model
 
-from .util import postprocess_packed_seqs, preprocess_packed_seqs, recover_left_padding, remove_left_padding
+from .util import (
+    postprocess_packed_seqs,
+    postprocess_packed_seqs_no_padding,
+    preprocess_packed_seqs,
+    preprocess_packed_seqs_no_padding,
+    recover_left_padding,
+    remove_left_padding,
+)
 
 
 def gptmodel_forward(
@@ -37,13 +44,16 @@ def gptmodel_forward(
     if pack_seqs:
         batch_size, seq_len = attention_mask.shape[:2]
         input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process)
+        print(f"input_ids_rmpad shape: {input_ids_rmpad.shape}, packed_seq_params: {packed_seq_params}")
         input_ids_rmpad = input_ids_rmpad.contiguous()
         output_orig = model(
             input_ids=input_ids_rmpad,
             attention_mask=None,
             position_ids=position_ids,
             packed_seq_params=packed_seq_params,
         )
+        print(f"output_orig: {output_orig}")
+
         if post_process and logits_processor is not None:
             args = {
                 k: preprocess_packed_seqs(v, attention_mask, pre_process=True)[0]
@@ -146,3 +156,55 @@ def gptmodel_forward_qwen2_5_vl(
     if value_model and post_process:
         output = output[..., 0]
     return output
+
+
+def gptmodel_forward_no_padding(
+    model,
+    input_ids,
+    value_model=False,
+    pack_seqs=True,
+    logits_processor=None,
+    logits_processor_args: dict = None,
+    **kwargs,
+):
+    """Default forward pass for GPT models with optional sequence packing."""
+    pre_process = unwrap_model(model).pre_process
+    post_process = unwrap_model(model).post_process
+    if pack_seqs:
+        batch_size = input_ids.shape[0]
+        input_ids_rmpad, packed_seq_params = preprocess_packed_seqs_no_padding(input_ids, pre_process=pre_process)
+        input_ids_rmpad = input_ids_rmpad.contiguous()
+        output_orig = model(
+            input_ids=input_ids_rmpad,
+            attention_mask=None,
+            position_ids=None,
+            packed_seq_params=packed_seq_params,
+        )
+
+        if post_process and logits_processor is not None:
+            args = {
+                k: preprocess_packed_seqs_no_padding(v, pre_process=True)[0] for k, v in logits_processor_args.items()
+            }
+            output_dict = logits_processor(output_orig, **args)
+            # print(f'gptmodel_forward_no_padding: {output_dict=}')
+            output = {
+                k: postprocess_packed_seqs_no_padding(
+                    v, packed_seq_params, input_ids, batch_size, post_process=post_process
+                )
+                for k, v in output_dict.items()
+            }
+        else:
+            output = postprocess_packed_seqs_no_padding(
+                output_orig, packed_seq_params, input_ids, batch_size, post_process=post_process
+            )
+    else:
+        raise NotImplementedError("gptmodel_forward_no_padding only supports packed sequences")
+
+    if value_model and post_process:
+        # output = output[..., 0]
+        # while using nested tensor, the advanced indexing operation above will result in an error at backward, i.e.
+        # ValueError: NestedTensor _nested_select_backward_default(grad_output: t, self: jt_all, dim: any, index: any)
+        # so we use `squeeze` to remove the last dimension
+        output = output.squeeze(-1)
+
+    return output

verl/models/mcore/registry.py

@@ -35,6 +35,7 @@
 )
 from .model_forward import (
     gptmodel_forward,
+    gptmodel_forward_no_padding,
     gptmodel_forward_qwen2_5_vl,
 )
 from .model_forward_fused import (
@@ -122,6 +123,23 @@ class SupportedModel(Enum):
     SupportedModel.QWEN3_TOKEN_CLASSIFICATION: gptmodel_forward,
 }
 
+# Registry for model forward functions
+MODEL_FORWARD_NOPAD_REGISTRY: dict[SupportedModel, Callable] = {
+    SupportedModel.LLAMA: gptmodel_forward_no_padding,
+    SupportedModel.QWEN2: gptmodel_forward_no_padding,
+    SupportedModel.QWEN2_MOE: gptmodel_forward_no_padding,
+    SupportedModel.MIXTRAL: gptmodel_forward_no_padding,
+    SupportedModel.DEEPSEEK_V3: gptmodel_forward_no_padding,
+    SupportedModel.QWEN2_5_VL: gptmodel_forward_no_padding,
+    SupportedModel.LLAMA4: gptmodel_forward_no_padding,
+    SupportedModel.QWEN3: gptmodel_forward_no_padding,
+    SupportedModel.QWEN3_MOE: gptmodel_forward_no_padding,
+    # SupportedModel.QWEN2_5_VL: gptmodel_forward_qwen2_5_vl,
+    SupportedModel.DEEPSEEK_V3: gptmodel_forward_no_padding,
+    SupportedModel.GLM4_MOE: gptmodel_forward_no_padding,
+    SupportedModel.QWEN3_TOKEN_CLASSIFICATION: gptmodel_forward_no_padding,
+}
+
 # Registry for model forward functions
 MODEL_FORWARD_FUSED_REGISTRY: dict[SupportedModel, Callable] = {
     SupportedModel.LLAMA: fused_forward_gptmodel,
@@ -227,6 +245,15 @@ def get_mcore_forward_fn(hf_config: PretrainedConfig) -> Callable:
     return MODEL_FORWARD_REGISTRY[model]
 
 
+def get_mcore_forward_no_padding_fn(hf_config: PretrainedConfig) -> Callable:
+    """
+    Get the forward function for given model architecture.
+    """
+    assert len(hf_config.architectures) == 1, "Only one architecture is supported for now"
+    model = get_supported_model(hf_config.architectures[0])
+    return MODEL_FORWARD_NOPAD_REGISTRY[model]
+
+
 def get_mcore_forward_fused_fn(hf_config: PretrainedConfig) -> Callable:
     """
     Get the forward function for given model architecture.