rebase main

Author: RangiLyu

xtuner/v1/ray/environment/single_turn_env.py

@@ -64,16 +64,13 @@ async def generate(
         if self.rollout_controller:
             # 在env中对输入的数据进行转换，是为了支持rollout_controller单独作为rollout engine使用，使各个模块进行解耦
             # 每个模块返回独立的data item, 在env中进行更新
-            response_futures = []
-            for sample in group_data_items:
-                extra_info = sample.data.extra_info if hasattr(sample.data, "extra_info") else {}
-                extra_info.update({"action_id": sample.uid.action_id})
-                response_future = self.rollout_controller.rollout.remote(
+            response_future = [
+                self.rollout_controller.rollout.remote(
                     prompt=sample.data.messages,
                     input_ids=sample.data.input_ids,
                     sample_params=sample_params,
                     extra_params=extra_params,
-                    extra_info=extra_info,
+                    extra_info=sample.data.extra_info,
                 )
                 for sample in group_data_items
             ]
@@ -129,4 +126,4 @@ async def run(
                     for _ in group_data_items
                 ]
             group_data_items = update_dataflow_item(group_data_items, "env.judger", judger_responses)
-        return group_data_items
+        return group_data_items

xtuner/v1/rl/base/rollout_is.py

@@ -41,6 +41,7 @@
 import numpy as np
 import torch
 from pydantic import BaseModel, ConfigDict
+import torch.distributed as dist
 
 from xtuner.v1.data_proto.utils import convert_packed_to_padded, convert_padded_to_packed, masked_mean, masked_sum
 
@@ -521,14 +522,20 @@ def compute_mismatch_metrics(
     return metrics
 
 
-def merge_rollout_is_metrics(rollout_is_metrics: list[dict[str, float]]) -> dict[str, float]:
+def merge_rollout_is_metrics(rollout_is_metrics: list[dict[str, float]], device="cuda") -> dict[str, float]:
     metrics = {}
     for key in rollout_is_metrics[0].keys():
         all_values = [m[key] for m in rollout_is_metrics]
         if "max" in key:
-            metrics[key] = np.max(all_values)
+            max_value = torch.tensor(all_values).max().to(torch.float32).to(device)
+            dist.all_reduce(max_value, op=dist.ReduceOp.MAX)
+            metrics[key] = max_value.item()
         elif "min" in key:
-            metrics[key] = np.min(all_values)
+            min_value = torch.tensor(all_values).min().to(torch.float32).to(device)
+            dist.all_reduce(min_value, op=dist.ReduceOp.MIN)
+            metrics[key] = min_value.item()
         else:
-            metrics[key] = np.mean(all_values)
+            mean_value = torch.tensor(all_values).mean().to(torch.float32).to(device)
+            dist.all_reduce(mean_value, op=dist.ReduceOp.AVG)
+            metrics[key] = mean_value.item()
     return metrics

xtuner/v1/rl/base/worker.py

@@ -408,7 +408,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int):
         logger_msg = f"Rollout {rollout_idx}: "
 
         if len(all_rollout_is_metrics) > 0:
-            rollout_is_metrics = merge_rollout_is_metrics(all_rollout_is_metrics)
+            rollout_is_metrics = merge_rollout_is_metrics(all_rollout_is_metrics, DEVICE)
             logger_msg += f"\n\nrollout importance sampling metrics:\n{json.dumps(rollout_is_metrics, indent=4)}"
 
         sum_entropy = cast(torch.Tensor, sum_entropy)