doc updated

Author: BalaBalaYi

dlrover/python/rl/trainer/example/realhf/__init__.py

@@ -113,61 +113,61 @@ def _split_args_kwargs_data_proto_with_auto_padding(chunks, *args, **kwargs):
     return splitted_args, splitted_kwargs
 
 
-def dispatch_one_to_all(worker_group, *args, **kwargs):
-    args = tuple([arg] * worker_group.world_size for arg in args)
-    kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()}
+def dispatch_one_to_all(role_group, *args, **kwargs):
+    args = tuple([arg] * role_group.world_size for arg in args)
+    kwargs = {k: [v] * role_group.world_size for k, v in kwargs.items()}
     return args, kwargs
 
 
-def dummy_direct_rollout_call(worker_group, *args, **kwargs):
+def dummy_direct_rollout_call(role_group, *args, **kwargs):
     raise NotImplementedError("Direct rollout call is forbidden.")
 
 
-def dispatch_all_to_all(worker_group, *args, **kwargs):
+def dispatch_all_to_all(role_group, *args, **kwargs):
     return args, kwargs
 
 
-def collect_all_to_all(worker_group, output):
+def collect_all_to_all(role_group, output):
     return output
 
 
-def dispatch_megatron_compute(worker_group, *args, **kwargs):
+def dispatch_megatron_compute(role_group, *args, **kwargs):
     """
     User passes in dp data. The data is dispatched to all tp/pp ranks with the same dp
     """
 
     all_args = []
     for arg in args:
         assert (
-            isinstance(arg, (Tuple, List)) and len(arg) == worker_group.dp_size
+            isinstance(arg, (Tuple, List)) and len(arg) == role_group.dp_size()
         )
         transformed_args = []
-        for i in range(worker_group.world_size):
-            local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank
+        for i in range(role_group.world_size):
+            local_dp_rank = role_group.get_megatron_rank_info()[i].dp_rank
             transformed_args.append(arg[local_dp_rank])
         all_args.append(transformed_args)
     all_args = tuple(all_args)
 
     all_kwargs = {}
     for k, v in kwargs.items():
-        assert isinstance(v, (Tuple, List)) and len(v) == worker_group.dp_size
+        assert isinstance(v, (Tuple, List)) and len(v) == role_group.dp_size
         transformed_v = []
-        for i in range(worker_group.world_size):
-            local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank
+        for i in range(role_group.world_size):
+            local_dp_rank = role_group.get_megatron_rank_info()[i].dp_rank
             transformed_v.append(v[local_dp_rank])
         all_kwargs[k] = transformed_v
     return all_args, all_kwargs
 
 
-def collect_megatron_compute(worker_group, output):
+def collect_megatron_compute(role_group, output):
     """
     Only collect the data from the tp=0 and pp=last and every dp ranks
     """
 
     output_in_dp = []
-    pp_size = worker_group.get_megatron_global_info().pp_size
-    for global_rank in range(worker_group.world_size):
-        local_rank_info = worker_group.get_megatron_rank_info(rank=global_rank)
+    pp_size = role_group.get_megatron_global_info().pp_size
+    for global_rank in range(role_group.world_size):
+        local_rank_info = role_group.get_megatron_rank_info()[global_rank]
         if (
             local_rank_info.tp_rank == 0
             and local_rank_info.pp_rank == pp_size - 1
@@ -177,16 +177,16 @@ def collect_megatron_compute(worker_group, output):
     return output_in_dp
 
 
-def dispatch_megatron_compute_data_proto(worker_group, *args, **kwargs):
+def dispatch_megatron_compute_data_proto(role_group, *args, **kwargs):
     """
     All the args and kwargs must be DataProto. The batch will be chunked by dp_size and passed to each rank
     """
 
     splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(
-        worker_group.dp_size, *args, **kwargs
+        role_group.dp_size(), *args, **kwargs
     )
     return dispatch_megatron_compute(
-        worker_group, *splitted_args, **splitted_kwargs
+        role_group, *splitted_args, **splitted_kwargs
     )
 
 
@@ -208,14 +208,14 @@ def _concat_data_proto_or_future(output: List):
         raise NotImplementedError
 
 
-def collect_megatron_compute_data_proto(worker_group, output):
+def collect_megatron_compute_data_proto(role_group, output):
     """
     Each output must be a DataProto. We concat the dim=0 of output
     """
     import ray
     from verl.protocol import DataProto
 
-    output = collect_megatron_compute(worker_group, output)
+    output = collect_megatron_compute(role_group, output)
     for o in output:
         assert isinstance(
             o, (DataProto, ray.ObjectRef)
@@ -224,24 +224,25 @@ def collect_megatron_compute_data_proto(worker_group, output):
     return _concat_data_proto_or_future(output)
 
 
-def dispatch_megatron_pp_as_dp(worker_group, *args, **kwargs):
+def dispatch_megatron_pp_as_dp(role_group, *args, **kwargs):
     """
     treat pp as dp.
     """
 
-    pp_size = worker_group.pp_size
-    dp_size = worker_group.dp_size
-    cp_size = worker_group.cp_size
+    pp_size = role_group.pp_size()
+    dp_size = role_group.dp_size()
+    cp_size = role_group.cp_size()
     pp_dp_cp_size = pp_size * dp_size * cp_size
 
     all_args = []
     for arg in args:
         assert isinstance(arg, (List, Tuple)) and len(arg) == pp_dp_cp_size
         transformed_args = []
-        for i in range(worker_group.world_size):
-            local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank
-            local_pp_rank = worker_group.get_megatron_rank_info(rank=i).pp_rank
-            local_cp_rank = worker_group.get_megatron_rank_info(rank=i).cp_rank
+        for i in range(role_group.world_size):
+            rank_info = role_group.get_megatron_rank_info()[i]
+            local_dp_rank = rank_info.dp_rank
+            local_pp_rank = rank_info.pp_rank
+            local_cp_rank = rank_info.cp_rank
             # compute the rank in arg. Note that the order is dp then cp then pp
             # Also note that the outputs within a pp group will be firstly allgathered, then only the output of pp0 will be collected.
             # For pp=2 dp=4, a batch of data "ABCDEFGH" should be dispatched and collected in below order:
@@ -264,10 +265,11 @@ def dispatch_megatron_pp_as_dp(worker_group, *args, **kwargs):
             isinstance(v, (List, Tuple)) and len(v) == pp_dp_cp_size
         ), f"expect len(v)=={pp_dp_cp_size}, got {len(v)}"
         transformed_v = []
-        for i in range(worker_group.world_size):
-            local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank
-            local_pp_rank = worker_group.get_megatron_rank_info(rank=i).pp_rank
-            local_cp_rank = worker_group.get_megatron_rank_info(rank=i).cp_rank
+        for i in range(role_group.world_size):
+            rank_info = role_group.get_megatron_rank_info()[i]
+            local_dp_rank = rank_info.dp_rank
+            local_pp_rank = rank_info.pp_rank
+            local_cp_rank = rank_info.cp_rank
             # compute the rank in arg. Note that the order is dp then cp then pp
             dp_cp_rank = local_cp_rank * dp_size + local_dp_rank
             arg_rank = dp_cp_rank * pp_size + local_pp_rank
@@ -276,94 +278,92 @@ def dispatch_megatron_pp_as_dp(worker_group, *args, **kwargs):
     return all_args, all_kwargs
 
 
-def collect_megatron_pp_as_dp(worker_group, output):
+def collect_megatron_pp_as_dp(role_group, output):
     """
     treat pp as dp. Only collect data on tp=0
     """
     output_in_dp = []
-    for global_rank in range(worker_group.world_size):
-        local_rank_info = worker_group.get_megatron_rank_info(rank=global_rank)
+    for global_rank in range(role_group.world_size):
+        local_rank_info = role_group.get_megatron_rank_info()[global_rank]
         if local_rank_info.tp_rank == 0:
             output_in_dp.append(output[global_rank])
     return output_in_dp
 
 
-def collect_megatron_pp_only(worker_group, output):
+def collect_megatron_pp_only(role_group, output):
     """
     Only collect output of megatron pp. This is useful when examine weight names as they are identical in tp/dp
     """
     output_in_pp = []
-    for global_rank in range(worker_group.world_size):
-        local_rank_info = worker_group.get_megatron_rank_info(rank=global_rank)
+    for global_rank in range(role_group.world_size):
+        local_rank_info = role_group.get_megatron_rank_info()[global_rank]
         if local_rank_info.tp_rank == 0 and local_rank_info.dp_rank == 0:
             output_in_pp.append(output[global_rank])
     return output_in_pp
 
 
-def dispatch_megatron_pp_as_dp_data_proto(worker_group, *args, **kwargs):
+def dispatch_megatron_pp_as_dp_data_proto(role_group, *args, **kwargs):
     pp_dp_cp_size = (
-        worker_group.dp_size * worker_group.pp_size * worker_group.cp_size
+        role_group.dp_size() * role_group.pp_size() * role_group.cp_size()
     )
     splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(
         pp_dp_cp_size, *args, **kwargs
     )
     ret = dispatch_megatron_pp_as_dp(
-        worker_group, *splitted_args, **splitted_kwargs
+        role_group, *splitted_args, **splitted_kwargs
     )
     return ret
 
 
-def collect_megatron_pp_as_dp_data_proto(worker_group, output):
-    output = collect_megatron_pp_as_dp(worker_group, output)
+def collect_megatron_pp_as_dp_data_proto(role_group, output):
+    output = collect_megatron_pp_as_dp(role_group, output)
     return _concat_data_proto_or_future(output)
 
 
-def dispatch_dp_compute(worker_group, *args, **kwargs):
+def dispatch_dp_compute(role_group, *args, **kwargs):
     for arg in args:
         assert (
             isinstance(arg, (Tuple, List))
-            and len(arg) == worker_group.world_size
+            and len(arg) == role_group.world_size
         )
     for k, v in kwargs.items():
-        assert (
-            isinstance(v, (Tuple, List)) and len(v) == worker_group.world_size
-        )
+        assert isinstance(v, (Tuple, List)) and len(v) == role_group.world_size
     return args, kwargs
 
 
-def collect_dp_compute(worker_group, output):
-    assert len(output) == worker_group.world_size
+def collect_dp_compute(role_group, output):
+    assert len(output) == role_group.world_size
     return output
 
 
-def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
+def dispatch_dp_compute_data_proto(role_group, *args, **kwargs):
     # Note: enable auto padding for dp compute DatapProto
     (
         splitted_args,
         splitted_kwargs,
     ) = _split_args_kwargs_data_proto_with_auto_padding(
-        worker_group.world_size,
+        role_group.world_size,
         *args,
         **kwargs,
     )
     return splitted_args, splitted_kwargs
 
 
-def dispatch_dp_compute_data_proto_with_func(worker_group, *args, **kwargs):
+def dispatch_dp_compute_data_proto_with_func(role_group, *args, **kwargs):
     assert isinstance(
         args[0], FunctionType
     )  # NOTE: The first one args is a function!
 
     splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(
-        worker_group.world_size, *args[1:], **kwargs
+        role_group.world_size, *args[1:], **kwargs
     )
     splitted_args_with_func = [
-        [args[0]] * worker_group.world_size
+        [args[0]] * role_group.world_size
     ] + splitted_args
     return splitted_args_with_func, splitted_kwargs
 
 
-def collect_dp_compute_data_proto(worker_group, output):
+def collect_dp_compute_data_proto(role_group, output):
     import ray
     from verl.protocol import DataProto
 
@@ -372,7 +372,7 @@ def collect_dp_compute_data_proto(worker_group, output):
             o, (DataProto, ray.ObjectRef)
         ), f"expecting {o} to be DataProto, but got {type(o)}"
 
-    output = collect_dp_compute(worker_group, output)
+    output = collect_dp_compute(role_group, output)
     return _concat_data_proto_or_future(output)
 
 

dlrover/python/rl/trainer/example/verl/base/decorator.py

@@ -183,6 +183,7 @@ def get_megatron_global_info(self):
         )
         return info
 
+    @trainer_invocation()
     def get_megatron_rank_info(self):
         from megatron.core import parallel_state as mpu