From 0a8f44867b6a3020ab954b18dcaef13b1bf79752 Mon Sep 17 00:00:00 2001
From: lilei <799812479@qq.com>
Date: Mon, 20 Apr 2026 05:57:33 +0000
Subject: [PATCH] fix entropy bug and update code

---
 slime/backends/megatron_utils/loss.py | 36 +++++++++++----------------
 1 file changed, 14 insertions(+), 22 deletions(-)

diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py
index bc053fe76e..5707c57cca 100644
--- a/slime/backends/megatron_utils/loss.py
+++ b/slime/backends/megatron_utils/loss.py
@@ -292,7 +292,6 @@ def _build_shifted_tokens(
 def _extract_per_sample(
     log_prob_full: torch.Tensor,
     entropy_full: torch.Tensor | None,
-    unconcat_tokens: list[torch.Tensor],
     total_lengths: list[int],
     response_lengths: list[int],
     qkv_format: str,
@@ -302,22 +301,12 @@ def _extract_per_sample(
     """Slice per-sample response log-probs/entropy from full-length 1-D tensors."""
     cp_size = mpu.get_context_parallel_world_size()
     log_probs_list: list[torch.Tensor] = []
-    entropy_list: list[torch.Tensor | None] = []
-
-    def _append(lp: torch.Tensor) -> None:
-        log_probs_list.append(lp)
-        entropy_list.append(None)
-
-    def _append_with_entropy(lp: torch.Tensor, start: int, end: int) -> None:
-        log_probs_list.append(lp)
-        entropy_list.append(entropy_full[start:end] if entropy_full is not None else None)
+    entropy_list: list[torch.Tensor] = []
 
     if cp_size > 1 and not allgather_cp:
         # zigzag CP
         pos = 0
-        for i, (_tokens, total_length, response_length) in enumerate(
-            zip(unconcat_tokens, total_lengths, response_lengths, strict=False)
-        ):
+        for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)):
             max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None
             chunk_size_cp, chunks_offset, logits_offset, _tokens_offset = get_logits_and_tokens_offset_with_cp(
                 total_length, response_length, qkv_format, max_seq_len
@@ -344,8 +333,6 @@ def _append_with_entropy(lp: torch.Tensor, start: int, end: int) -> None:
                     dim=0,
                 )
                 entropy_list.append(ent)
-            else:
-                entropy_list.append(None)
             pos += 2 * chunk_size_cp
 
     elif allgather_cp:
@@ -363,11 +350,13 @@ def _append_with_entropy(lp: torch.Tensor, start: int, end: int) -> None:
             s = max(logit_global_start, chunk_start)
             e = min(logit_global_end, chunk_end)
             if e <= s:
-                _append(log_prob_full[0:0])
+                log_probs_list.append(torch.zeros((0,), dtype=log_prob_full.dtype, device=log_prob_full.device))
+                if entropy_full is not None:
+                    entropy_list.append(torch.zeros((0,), dtype=entropy_full.dtype, device=entropy_full.device))
             else:
-                _append_with_entropy(
-                    log_prob_full[s - chunk_start : e - chunk_start], s - chunk_start, e - chunk_start
-                )
+                log_probs_list.append(log_prob_full[s - chunk_start : e - chunk_start])
+                if entropy_full is not None:
+                    entropy_list.append(entropy_full[s - chunk_start : e - chunk_start])
             seq_start += total_length
 
     else:
@@ -377,13 +366,17 @@ def _append_with_entropy(lp: torch.Tensor, start: int, end: int) -> None:
             for total_length, response_length in zip(total_lengths, response_lengths, strict=False):
                 end = offset + total_length
                 start = end - response_length
-                _append_with_entropy(log_prob_full[start - 1 : end - 1], start - 1, end - 1)
+                log_probs_list.append(log_prob_full[start - 1 : end - 1])
+                if entropy_full is not None:
+                    entropy_list.append(entropy_full[start - 1 : end - 1])
                 offset += total_length
         else:  # bshd
             for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)):
                 end = max_seq_lens[i] * i + total_length
                 start = end - response_length
-                _append_with_entropy(log_prob_full[start - 1 : end - 1], start - 1, end - 1)
+                log_probs_list.append(log_prob_full[start - 1 : end - 1])
+                if entropy_full is not None:
+                    entropy_list.append(entropy_full[start - 1 : end - 1])
 
     return log_probs_list, entropy_list
 
@@ -450,7 +443,6 @@ def get_log_probs_and_entropy(
     log_probs_list, entropy_list = _extract_per_sample(
         log_prob_full,
         entropy_full,
-        unconcat_tokens,
         total_lengths,
         response_lengths,
         qkv_format,