[Bug] Negative JSD loss in Megatron GKD at beta=0.5; FP32 gives a positive result on the same inputs

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

We observed negative JSD loss while running Megatron GKD with ms-swift 4.5.3.

With beta=0.5 and sft_alpha=0, our understanding is that the distillation objective is:

M = (P_student + P_teacher) / 2
JSD = 0.5 * KL(P_student || M) + 0.5 * KL(P_teacher || M) >= 0

JSD should be nonnegative for normalized probability distributions. Floating-point arithmetic can introduce small numerical errors, but the negative value we observed differs substantially from a higher-precision reference.

What happened

In a two-step training diagnostic, the second step reported a loss of -0.00016957933743.

We captured the logits actually used by the loss and recomputed it on the GPU. The input tensors, top64 token IDs, masks, beta, and temperature were unchanged. Only the precision of the loss computation changed:

Same training inputs, averaged over 8 DP ranks	JSD loss
Original computation: student FP32 / teacher BF16	-0.00016957933743
Both inputs converted to FP32 before temperature scaling	+0.00050891250673
Both inputs converted to FP64 before temperature scaling, as a reference	+0.00050891054394

FP32 closely matches FP64, while the original computation differs in both magnitude and sign. This is a comparison on the same inputs, not a comparison between separate training runs.

Checks already performed

• The official loss was not replaced. Backward still used the original official loss during the diagnostic.
• Teacher/student token-ID vocabularies, input IDs, and shifted labels were checked for alignment; top64 indices were valid.
• Across 16 captured cases (2 steps * 8 ranks), original-precision CUDA replay matched each original loss within 1e-8. FP32 and FP64 results were positive in all cases.
• The synthetic reproduction below also produces negative JSD without a tokenizer, dataset, model weights, or distributed communication.

Expected behavior: JSD should be nonnegative within a reasonable numerical tolerance, and the loss and gradients should closely match a higher-precision reference.

How to Reproduce / 如何复现

Environment and training configuration

• ms-swift 4.5.3
• PyTorch 2.10.0+cu128
• Megatron-Core 0.18.0
• mcore-bridge 1.6.3
• Transformers 5.12.1
• Python 3.12, NVIDIA H100 80GB
• Local frozen Qwen3.5-9B teacher and Qwen3.5-0.8B student
• Megatron TP1/DP8, colocated student vLLM, BF16 model precision
• gkd_logits_topk=64, beta=0.5, temperature=0.9, sft_alpha=0

Minimal reproduction

Save the following as repro.py and run python repro.py in the environment above.

It loads the installed official gkd_loss.py directly and calls jsd_loss, using the same temperature-scaling order as its caller. It requires only synthetic tensors and one CUDA GPU; no model loading or training is needed.

import importlib.util
import sys
from pathlib import Path
import torch

source = Path(importlib.util.find_spec('swift').origin).parent / 'rlhf_trainers/gkd_loss.py'
spec = importlib.util.spec_from_file_location('gkd_repro', source)
gkd = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = gkd
spec.loader.exec_module(gkd)

def evaluate(student, teacher, dtype=None):
    s = (student if dtype is None else student.to(dtype)).detach().requires_grad_(True)
    t = teacher if dtype is None else teacher.to(dtype)
    loss = gkd.jsd_loss(s / 0.9, t / 0.9, beta=0.5) / s.shape[0]
    grad, = torch.autograd.grad(loss, s)
    return loss.item(), grad.double()

rng = torch.Generator().manual_seed(0)
# BF16 model output upcast by the student wrapper.
student = (torch.randn(3, 64, generator=rng) * 2).bfloat16().float()
teacher = (student + torch.randn(3, 64, generator=rng) * 0.03).bfloat16()
student, teacher = student.cuda(), teacher.cuda()

native, gn = evaluate(student, teacher)
fp32, g32 = evaluate(student, teacher, torch.float32)
fp64, g64 = evaluate(student, teacher, torch.float64)
print('native_mixed', native, 'fp32', fp32, 'fp64', fp64)
print('native_grad_relative_l2', ((gn-g64).norm()/g64.norm()).item())
print('fp32_grad_relative_l2', ((g32-g64).norm()/g64.norm()).item())

Observed output on H100 / PyTorch 2.10.0+cu128:

native_mixed -0.0012147885281592607
fp32 7.973663741722703e-05
fp64 7.975079402659209e-05
native_grad_relative_l2 0.22682357533754788
fp32_grad_relative_l2 1.167498132765438e-05

The gradient comparison measures the relative L2 error of dLoss/dLogits for these small student-logit tensors against FP64. It does not measure full-model parameter-gradient error or downstream quality degradation. CPU BF16 kernels may produce different numbers; the values above are CUDA results.

Additional Information / 补充信息

Likely cause

The observed loss inputs were student FP32 and teacher BF16, with CUDA autocast disabled during loss computation.

The student is wrapped by Megatron's Float16Module, which promotes its final-stage output to FP32 by default. The separate frozen local teacher does not use the same output wrapper and retains BF16 logits. Therefore, configuring both models for BF16 does not necessarily produce matching output dtypes at the loss boundary.

The current loss path does not explicitly promote both inputs before temperature scaling and probability normalization. BF16 temperature division, log_softmax, and exp on the teacher side introduce rounding errors; mixing the results with FP32 student values later cannot recover the lost precision. The same-input comparisons support this diagnosis.

Different student/teacher output dtypes need not be prohibited; the loss should handle this combination safely.

Suggested fix

Please consider promoting low-precision logits used by the loss to FP32 before temperature scaling, log-softmax, and JSD calculation. Model weights and forward passes can remain BF16.

For the top64 / TP1 path tested here, a candidate is:

# Select the student's scores at the teacher's top64 token IDs, then upcast.
s_logits = gather_fn(s_active, t_active.topk_indices).float()
t_logits = t_active.topk_logprobs.float()

# Keep the existing temperature scaling and JSD formula.
s_logits = s_logits / temperature
t_logits = t_logits / temperature

Default gather only selects values by index. Gathering before conversion avoids an unnecessary full-vocabulary FP32 allocation when the source is low precision. Our actual student output was already FP32, so this is not a claim of measured speed or memory improvement in this run. TP-aware/custom gather needs separate validation, and a production implementation should preserve existing FP64 inputs rather than unconditionally downcasting them.

Please avoid clamp(min=0): it hides the negative value without correcting the probability calculations and changes gradients.

Suggested regression coverage:

• Mixed FP32/BF16 inputs with identical or nearly identical distributions.
• beta=0/0.5/1 and positive temperatures other than 1.
• Loss and gradient comparisons against FP64 with numerical tolerances.
• Top-k, full-vocabulary, empty-mask, and TP/CP paths.

The same-input FP32 calculation has been validated; a patched end-to-end multi-step training run has not yet been validated.

Related upstream work

When checked on 2026-09-16, PR #10102 proposed opt-in FP32 loss through SWIFT_GKD_JSD_FP32=1. It was an unmerged draft, with the option disabled by default. We inspected its diff but have not executed that PR. This report provides additional independent reproduction and regression evidence for that work.

The main source at commit 6dd7774, inspected on the same date, still lacked the explicit promotion described above. Main was source-inspected only; the executable reproduction was tested on the installed 4.5.3 release.

Read swift/rlhf_trainers/gkd_loss.py, focusing on jsd_loss and the caller’s temperature-scaling path. Run the supplied repro.py first and compare native mixed-precision results with FP32 and FP64, then add regression coverage for the listed mixed-dtype, beta, top-k, full-vocabulary, masking, and TP/CP cases. Done means nonnegative loss and gradients that remain within defined tolerances against FP64.