Fix fused_linear_cross_entropy_loss adding the z-loss term to the loss when compute_z_loss=False

[ReinforcedKnowledge]

liger's lse_square_scale folds m * lse^2 into the returned loss whenever it's non-zero. return_z_loss only controls the separate return. Since LMHead.forward passes z_loss_multiplier or 1e-4, fused-loss runs with z-loss off return ce + 1e-4 * lse^2. The compute_z_loss=True path is unaffected.

Try this with:

import torch
import torch.nn.functional as F
from olmo_core.nn.functional.cross_entropy_loss import fused_linear_cross_entropy_loss

torch.manual_seed(1234)
h = torch.randn(4096, 1024, device="cuda")
w = torch.randn(151936, 1024, device="cuda") * 0.02
labels = torch.randint(0, 151936, (4096,), device="cuda")
labels[torch.rand(4096, device="cuda") < 0.4] = -100

logits = h @ w.t()
ce = F.cross_entropy(logits, labels, ignore_index=-100)
mask = labels != -100
z_term = 1e-4 * (logits.logsumexp(-1).pow(2) * mask).sum() / mask.sum()

loss, _ = fused_linear_cross_entropy_loss(
    h, w, labels, ignore_index=-100, reduction="mean",
    compute_z_loss=False, z_loss_multiplier=1e-4, accum_dtype=torch.float32,
)

print(f"plain CE         = {ce.item():.6f}")
print(f"fused, z off     = {loss.item():.6f}")
print(f"1e-4*mean(lse^2) = {z_term.item():.6f}")

if abs(loss.item() - ce.item()) > 1e-5 * ce.item():
    print("FAIL: the no-z fused loss includes the z-loss term")
else:
    print("PASS")

I tried it on a cluster on 1 A100:

Before:
plain CE         = 12.137142
fused, z off     = 12.151872
1e-4*mean(lse^2) = 0.014729
FAIL: the no-z fused loss includes the z-loss term
After:
plain CE         = 12.137142
fused, z off     = 12.137142
1e-4*mean(lse^2) = 0.014729
PASS