fix(loss): return fp32 empty log-probs from the chunked path

[yueming-yuan]

Summary

With model-precision logits (#2764/#2818), a CP rank whose contiguous half holds no response logits returns its empty log-probs as logits.new_zeros((0,)) — in bf16/fp16 — while every other rank returns fp32 CE outputs. The dtype leaks through allgather_cp_redistribute (zeros(dtype=value.dtype)) into the cross-rank dist.nn.all_reduce, and the byte-size mismatch crashes the collective:

gloo::EnforceNotMet: op.preamble.length <= op.nbytes. 24 vs 12.
Received data size doesn't match expected size.

Before #2764 the logits reaching this code were always fp32 on the Megatron backend, so the empty tensor's dtype could not diverge. Reachable with allgather-CP + bf16 whenever one rank's half contains only prompt tokens (long prompts).

The fix pins the empty log-probs/entropy of the chunked (non-true-on-policy) path to fp32, matching the dtype the fused CE produces on non-empty ranks. The true-on-policy path is untouched: its non-empty outputs stay in model precision, so its empty tensor already matches.

Tests

Covered by the 2-rank allgather-CP consistency tests in #2826 (stacked on this PR): the contiguous_empty_rank case crashes on this exact path without the fix and passes with it.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yueming-yuan]

Folded into #2818.