update

Author: Williamren97

slime/backends/fsdp_utils/actor.py

@@ -246,20 +246,9 @@ def setup_context_parallelism(self):
         print(f"Ring attention rank: {ring_attn_rank}")
 
     def _update_cp_cu_seqlens(self, packed_batch):
-        cu_seqlens = packed_batch["cu_seqlens"].to(device=torch.cuda.current_device(), dtype=torch.int32)
+        cu_seqlens = packed_batch["cu_seqlens"]
 
-        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
-            logits = self.model(
-                input_ids=packed_batch["tokens"].unsqueeze(0),
-                attention_mask=None,
-                position_ids=packed_batch["position_ids"].unsqueeze(0),
-            ).logits
-        # Handle packed sequences
-        log_probs = gather_log_probs_packed(logits, packed_batch["tokens"], packed_batch["cu_seqlens"])
-        packed_batch["cur_log_probs"] = log_probs
-        unpacked_batches = unpack_sequences(packed_batch)
-
-        # Sync ring flash attention parameters
+        # Update the ring attention parameters
         update_ring_flash_attn_params(cu_seqlens, self.cp_group)