fix chunk and mlp_factor

Author: root

verl/single_controller/base/decorator.py

@@ -71,7 +71,7 @@ def init_predefined_execute_mode():
 init_predefined_execute_mode()
 
 
-def _balance_data_proto(data_proto, chunks, model_config=None):
+def _balance_data_proto(data_proto, chunks):
     if (
         not isinstance(data_proto, DataProtoFuture)
         and data_proto.batch is not None
@@ -121,7 +121,7 @@ def _padding_and_split_data(obj, chunks):
                     f"expecting all arg share same length of {data_proto_len}, but got {len(obj)}"
                 )
             obj.padding(padding_size=padding_size)
-        return _balance_data_proto(obj, chunks)
+        return _balance_data_proto(obj, chunks).chunk(chunks=chunks)
 
     splitted_args = [_padding_and_split_data(arg, chunks) for arg in args]
     splitted_kwargs = {key: _padding_and_split_data(val, chunks) for key, val in kwargs.items()}