[data] feat: TransferQueue - remove redundant data collect for both TQ and DataProto (verl-project#4618)

### What does this PR do?

This PR optimizes the dataflow between the single-controller
(RayPPOTrainer) and worker processes. By applying the result filtering
(de-duplication) logic from the controller to the individual workers in
advance, we significantly reduce redundant data transmission.

In the current architecture, worker functions use the @register
decorator to manage data dispatching and collection.



```python3
    # in megatron_workers.py 
    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor"))
    @GPUMemoryLogger(role="update_actor", logger=logger)
    @DistProfiler.annotate(color="red", role="actor_update")
    def update_actor(self, data: DataProto):
        assert self._is_actor
        if self._is_offload_param:
            load_megatron_model_to_gpu(self.actor_module)
            log_gpu_memory_usage("After load actor params and grad during update_actor", logger=logger)
        if self._is_offload_optimizer:
            load_megatron_optimizer(self.actor_optimizer)
            log_gpu_memory_usage("After load actor optimizer during update_actor", logger=logger)

```

Specifically, `make_nd_compute_dataproto_dispatch_fn` is used to

1. Dispatch: Shard and distribute input data from the controller to
workers.
2. Collect: Gather results back from all workers and de-duplicate them
(usually keeping only one copy per DP group) using a collect_mask.

```python3
def make_nd_compute_dataproto_dispatch_fn(mesh_name):
    return {
        "dispatch_fn": partial(dispatch_lazy_compute_data_proto, mesh_name),
        "collect_fn": partial(collect_lazy_compute_data_proto, mesh_name),
    }
```


```python3
def collect_nd_compute(collect_mask: list[bool], worker_group, output):
    from verl.single_controller.base.worker_group import WorkerGroup

    assert isinstance(worker_group, WorkerGroup)
    assert len(output) == worker_group.world_size

    output_in_dp = []
    for global_rank in range(worker_group.world_size):
        collect_dp_rank = collect_mask[global_rank]
        if collect_dp_rank:
            output_in_dp.append(output[global_rank])
    return output_in_dp
```

Previously, the de-duplication process happened entirely on the
controller side.

- Redundant Transmission: Every worker sent its full output back to the
controller, regardless of whether that data would be kept or discarded.
- Overhead: For large-scale LLM training (large DP groups), this led to
massive, unnecessary network traffic or repeated put operations to the
TransferQueue.
- Bottleneck: The controller became a bottleneck as it had to receive
and process redundant data chunks before applying the mask.

In this PR, we shifts the filtering logic "left" (to the worker side).
Now each worker automatically determine whether to return real data or
return an empty obj to the single-controller according to the
`collect_mask` in advance, thus reducing data transfer overhead.

### Checklist Before Starting

- [x] Search for similar PRs. Paste at least one query link here: ...
- [x] Format the PR title as `[{modules}] {type}: {description}` (This
will be checked by the CI)
- `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`,
`trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`,
`ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`,
`env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`
- If this PR involves multiple modules, separate them with `,` like
`[megatron, fsdp, doc]`
  - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`
- If this PR breaks any API (CLI arguments, config, function signature,
etc.), add `[BREAKING]` to the beginning of the title.
  - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching`

### Test

> For changes that can not be tested by CI (e.g., algorithm
implementation, new model support), validate by experiment(s) and show
results like training curve plots, evaluation results, etc.

### API and Usage Example

> Demonstrate how the API changes if any, and provide usage example(s)
if possible.

```python
# Add code snippet or script demonstrating how to use this
```

### Design & Code Changes

> Demonstrate the high-level design if this PR is complex, and list the
specific changes.

### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [x] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [x] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [x] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [x] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [ ] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)

---------

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>
Signed-off-by: jianjunzhong <jianjunzhong@foxmail.com>
Co-authored-by: Jianjun Zhong <87791082+jianjunzhong@users.noreply.github.com>
Co-authored-by: jianjunzhong <jianjunzhong@foxmail.com>

Author: 0oshowero0

verl/single_controller/base/decorator.py

@@ -447,7 +447,7 @@ def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocki
     _check_execute_mode(execute_mode=execute_mode)
 
     def decorator(func):
-        func = tqbridge()(func)
+        func = tqbridge(dispatch_mode=dispatch_mode)(func)
 
         @wraps(func)
         def inner(*args, **kwargs):

verl/single_controller/base/worker.py

@@ -117,6 +117,9 @@ def _query_dispatch_info(self, mesh_name: str):
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL)
     def _query_collect_info(self, mesh_name: str):
+        return self.query_collect_info(mesh_name)
+
+    def query_collect_info(self, mesh_name: str):
         """Query the collect info for a given mesh name.
 
         Args:

verl/utils/transferqueue_utils.py

@@ -13,12 +13,16 @@
 # limitations under the License.
 
 import asyncio
+import functools
 import inspect
 import logging
 import os
 import threading
 from functools import wraps
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    from verl.single_controller.base.decorator import Dispatch
 
 from tensordict import TensorDict
 
@@ -144,7 +148,95 @@ def _update_batchmeta_with_output(output: DataProto, batchmeta: "BatchMeta", fun
     return updated_batch_meta
 
 
-def tqbridge(put_data: bool = True):
+def _compute_need_collect(dispatch_mode: "dict | Dispatch", args: list) -> bool:
+    """Compute whether data collection is needed for the current worker.
+
+    This function determines whether the current worker should collect data based on
+    the dispatch mode configuration and worker parameters. It's used to optimize
+    distributed data collection by ensuring only the appropriate rank collects data.
+
+    Args:
+        dispatch_mode: Controls data collection logic for the current worker. Can be None,
+                      a Dispatch instance, or a dict with 'collect_fn' key. If None or Dispatch,
+                      always returns True (current worker should collect). If dict, checks
+                      collect_fn for lazy compute optimization.
+        args: List of arguments passed to the function. Should contain a Worker instance
+             as the first argument when using lazy compute mode.
+
+    Returns:
+        bool: True if data collection is needed, False otherwise.
+
+    Note:
+        Only checks worker attributes when dispatch_mode is a dict with 'collect_fn',
+        the collect_fn is 'collect_lazy_compute_data_proto', and args[0] is a Worker.
+        Otherwise, returns True. For the lazy compute case, checks the worker's
+        data parallel rank for the mesh specified in collect_fn.args[0] to determine
+        if this worker should collect data.
+    """
+    from verl.single_controller.base.decorator import Dispatch
+    from verl.single_controller.base.worker import Worker
+
+    if dispatch_mode is None or isinstance(dispatch_mode, Dispatch):
+        return True
+
+    assert "collect_fn" in dispatch_mode.keys(), "collect_fn should be in dispatch_mode."
+
+    collect_fn = dispatch_mode["collect_fn"]
+
+    # Check if collect_fn is a functools.partial and handle gracefully
+    if isinstance(collect_fn, functools.partial):
+        collect_fn_name = collect_fn.func.__name__
+        if collect_fn_name != "collect_lazy_compute_data_proto" or len(args) < 1 or not isinstance(args[0], Worker):
+            return True
+
+        collect_mesh_name = collect_fn.args[0] if collect_fn.args else None
+        if collect_mesh_name is None:
+            return True
+
+        return args[0].query_collect_info(collect_mesh_name)
+    else:
+        # If collect_fn is not a partial, we can't extract mesh_name information
+        # Fall back to default behavior (collect data)
+        return True
+
+
+def _postprocess_common(output, put_data, need_collect):
+    """Common post-processing logic for function outputs in TransferQueue bridge.
+
+    This function handles the final return value based on whether data should be
+    put into storage (put_data) and whether collection is needed (need_collect).
+    It ensures proper return types based on the execution context.
+
+    Args:
+        output: The original output from the decorated function. Can be any type.
+        put_data: bool, indicating whether the output should be put into TransferQueue.
+                 If True, output will be put to TQ and return the corresponding BatchMeta;
+                 if False, output will not be put into TQ.
+        need_collect: bool, indicating whether this process needs to collect data.
+                     If False, the output will be replaced by an empty BatchMeta or DataProto
+                     to avoid redundant communication.
+
+    Returns:
+        - BatchMeta.empty(): When put_data=True but need_collect=False, indicating
+          no data should be stored but BatchMeta structure is expected.
+        - DataProto(): When put_data=False, need_collect=False, and output is DataProto,
+          returning an empty DataProto.
+        - output: In all other cases, returns the original output unchanged.
+
+    Note:
+        This function is used in the tqbridge decorator to normalize return values
+        across different execution paths and avoid redundant data operations in
+        distributed scenarios.
+    """
+    if put_data and not need_collect:
+        return BatchMeta.empty()
+    elif not put_data and not need_collect and isinstance(output, DataProto):
+        return DataProto()
+    else:
+        return output
+
+
+def tqbridge(dispatch_mode: "dict | Dispatch" = None, put_data: bool = True):
     """Creates a decorator for bridging BatchMeta and DataProto.
 
     This decorator automatically handles conversions between `BatchMeta` and
@@ -155,6 +247,9 @@ def tqbridge(put_data: bool = True):
     simply calls the original function as-is).
 
     Args:
+        dispatch_mode: Controls data collection behavior for the current worker. Passed to
+                      _compute_need_collect to determine if current worker should collect data.
+                      If None, _compute_need_collect will return True to fallback default logics.
         put_data: Whether put the DataProto into Storage after func return.
                   If True, after function execution, the output result will be
                   updated to `BatchMeta` and `BatchMeta` will be returned;
@@ -181,11 +276,11 @@ def inner(*args, **kwargs):
                 args = [_batchmeta_to_dataproto(arg) if isinstance(arg, BatchMeta) else arg for arg in args]
                 kwargs = {k: _batchmeta_to_dataproto(v) if isinstance(v, BatchMeta) else v for k, v in kwargs.items()}
                 output = func(*args, **kwargs)
-                if put_data:
+                need_collect = _compute_need_collect(dispatch_mode, args)
+                if put_data and need_collect:
                     updated_batch_meta = _update_batchmeta_with_output(output, batchmeta, func.__name__)
                     return updated_batch_meta
-                else:
-                    return output
+                return _postprocess_common(output, put_data, need_collect)
 
         @wraps(func)
         async def async_inner(*args, **kwargs):
@@ -203,18 +298,21 @@ async def async_inner(*args, **kwargs):
                     for k, v in kwargs.items()
                 }
                 output = await func(*args, **kwargs)
-                if put_data:
+                need_collect = _compute_need_collect(dispatch_mode, args)
+                if put_data and need_collect:
                     updated_batchmeta = await _async_update_batchmeta_with_output(output, batchmeta, func.__name__)
                     return updated_batchmeta
-                return output
+                return _postprocess_common(output, put_data, need_collect)
 
         @wraps(func)
         def dummy_inner(*args, **kwargs):
-            return func(*args, **kwargs)
+            output = func(*args, **kwargs)
+            return output
 
         @wraps(func)
         async def dummy_async_inner(*args, **kwargs):
-            return await func(*args, **kwargs)
+            output = await func(*args, **kwargs)
+            return output
 
         wrapper_inner = inner if is_transferqueue_enabled else dummy_inner
         wrapper_async_inner = async_inner if is_transferqueue_enabled else dummy_async_inner