Skip to content

fix entropy bug and update code #1846

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhuzilin merged 1 commit into from
Apr 20, 2026
Merged

fix entropy bug and update code #1846

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 14 additions & 22 deletions slime/backends/megatron_utils/loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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))
Comment on lines +353 to +355

Copilot AI Apr 20, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the allgather-CP path, the empty-slice case allocates new zero-length tensors via torch.zeros((0,), ...). Since these tensors are always length-0 placeholders, consider using log_prob_full[:0] / entropy_full[:0] (or new_empty((0,))) to avoid allocations and keep dtype/device/grad semantics consistent with the source tensors.

Suggested change
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))
log_probs_list.append(log_prob_full[:0])
if entropy_full is not None:
entropy_list.append(entropy_full[:0])

Copilot uses AI. Check for mistakes.
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:
Expand All @@ -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

Expand Down Expand Up @@ -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,
Expand Down
Loading