fix precommit

Author: zyzshishui

slime/backends/megatron_utils/cp_utils.py

@@ -98,6 +98,7 @@ def sum_of_token(x: torch.Tensor):
 
     return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token
 
+
 def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int):
     """
     Gather tensors across all ranks in the context parallel group.
@@ -140,4 +141,4 @@ def slice_with_cp(tokens: torch.Tensor, pad_value):
     # get 2 chunk for thd cp
     start_1, end_1 = chunk_size * cp_rank, chunk_size * (cp_rank + 1)
     start_2, end_2 = chunk_size * (2 * cp_size - cp_rank - 1), chunk_size * (2 * cp_size - cp_rank)
-    return torch.cat([tokens[start_1:end_1], tokens[start_2:end_2]])
+    return torch.cat([tokens[start_1:end_1], tokens[start_2:end_2]])

slime/backends/megatron_utils/data.py

@@ -161,7 +161,15 @@ def log_rollout_data(rollout_id, args, rollout_data):
         log_dict = {}
         response_lengths = rollout_data["response_lengths"]
         for key, val in rollout_data.items():
-            if key == "tokens" or key == "loss_masks" or key == "sample_indices"or key == "rollout_time" or key == "completion_tokens_stats" or key == "partial_samples" or key == "total_off_policy_tokens":
+            if (
+                key == "tokens"
+                or key == "loss_masks"
+                or key == "sample_indices"
+                or key == "rollout_time"
+                or key == "completion_tokens_stats"
+                or key == "partial_samples"
+                or key == "total_off_policy_tokens"
+            ):
                 continue
             # Upload per sample mean for each rollout value
             # There are the following assumptions:
@@ -235,7 +243,9 @@ def log_partial_rollout_data(rollout_id, args, rollout_data):
         total_off_policy_tokens = rollout_data["total_off_policy_tokens"]
         if total_off_policy_tokens is not None:
             log_dict["total_off_policy_tokens"] = total_off_policy_tokens
-            log_dict["off_policy_ratio"] = total_off_policy_tokens / (log_dict["total_tokens"] + total_off_policy_tokens)
+            log_dict["off_policy_ratio"] = total_off_policy_tokens / (
+                log_dict["total_tokens"] + total_off_policy_tokens
+            )
 
         response_lengths = rollout_data["response_lengths"]
 
@@ -285,7 +295,6 @@ def log_partial_rollout_data(rollout_id, args, rollout_data):
                 dst=mpu.get_data_parallel_src_rank(with_context_parallel=True),
                 group=mpu.get_data_parallel_group(with_context_parallel=True),
             )
-            
 
 
 def log_multi_turn_data(rollout_id, args, rollout_data):

slime/backends/utils/data.py

@@ -167,4 +167,4 @@ def get_partition(val):
     if "partial_samples" in data:
         rollout_data["partial_samples"] = data["partial_samples"]
     if "total_off_policy_tokens" in data:
-        rollout_data["total_off_policy_tokens"] = data["total_off_policy_tokens"]
+        rollout_data["total_off_policy_tokens"] = data["total_off_policy_tokens"]

slime/rollout/sglang_example.py

@@ -184,7 +184,7 @@ async def abort(args, rollout_id: int, data_buffer):
         print(f"Abort request for {url}", flush=True)
         # await post(f"{url}/abort_request", {"abort_all": True}, use_http2=False)
         # based on https://github.com/THUDM/slime/pull/63/files
-        await post(f"{url}/abort_request", {"rid":"", "abort_all": True}, use_http2=False)
+        await post(f"{url}/abort_request", {"rid": "", "abort_all": True}, use_http2=False)
 
     # make sure all the pending tasks are finished
     count = 0
@@ -281,26 +281,28 @@ async def generate_rollout_async(args, rollout_id: int, data_buffer) -> list[lis
 
     assert len(data) == args.rollout_batch_size, f"Got {len(data)} samples, expected {args.rollout_batch_size}"
     data = sorted(data, key=lambda group: group[0].index)
-    
+
     rollout_time = time.time() - state.rollout_start_time
 
     completion_tokens_stats = {}
     if state.completion_tokens_list:
         completion_tokens_array = np.array(state.completion_tokens_list)
         completion_tokens_stats = {
-            'total_completion_tokens': np.sum(completion_tokens_array).item(),
-            'completion_tokens_mean': np.mean(completion_tokens_array).item(),
-            'completion_tokens_std': np.std(completion_tokens_array).item(),
-            'completion_tokens_count': len(completion_tokens_array),
+            "total_completion_tokens": np.sum(completion_tokens_array).item(),
+            "completion_tokens_mean": np.mean(completion_tokens_array).item(),
+            "completion_tokens_std": np.std(completion_tokens_array).item(),
+            "completion_tokens_count": len(completion_tokens_array),
         }
 
     if len(data) > 0:
-        data[0][0].metadata.update({
-            'rollout_time': rollout_time,
-            'completion_tokens_stats': completion_tokens_stats,
-            'partial_samples': state.partial_samples_count,
-            'total_off_policy_tokens': state.total_off_policy_tokens,
-        })
+        data[0][0].metadata.update(
+            {
+                "rollout_time": rollout_time,
+                "completion_tokens_stats": completion_tokens_stats,
+                "partial_samples": state.partial_samples_count,
+                "total_off_policy_tokens": state.total_off_policy_tokens,
+            }
+        )
     if completion_tokens_stats:
         print(f"[DEBUG] Rollout {rollout_id}: Completion tokens stats: {completion_tokens_stats}", flush=True)