fix(gkd): promote low-precision logits before loss computation

[taking-lying-flat]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

GKD can produce negative JSD and inaccurate student-logit gradients when student logits are FP32 and teacher logits are BF16. Temperature scaling and probability calculations round the teacher values before they are combined with the student values.

Promote FP16/BF16 logits to FP32 in the shared loss pipeline before vocabulary alignment and temperature scaling, preserving existing FP32/FP64 inputs. The top-k path gathers the selected student scores before promotion to avoid an additional full-vocabulary FP32 buffer.

The loss change adds 7 net lines. Update one existing empty-mask test assertion to expect FP32 for FP16/BF16 inputs while preserving FP64. The test continues to check zero loss, zero gradients, and vocabulary alignment.

Fixes #10156.

Experiment results

Validated on CPU and one NVIDIA RTX A1000 Laptop GPU with PyTorch 2.13.0+cu130:

• A local validation harness, kept outside this PR, passed 144 loss/gradient comparisons against FP64, covering full-vocabulary and top-k inputs, mixed FP32/BF16, BF16, FP16, FP64, beta values 0/0.5/1, non-unit temperatures, and masked positions.
• The same precision checks against the unpatched main loss reproduce 108 failing low-precision cases; the 36 FP64 cases pass.
• The local harness also verified zero loss and zero gradients for empty masks, expecting FP32 loss for low-precision inputs.
• python -m pytest tests/train/test_gkd_loss.py -q: 5 tests and 24 subtests passed.
• pre-commit run --files swift/rlhf_trainers/gkd_loss.py tests/train/test_gkd_loss.py: all applicable hooks passed.
• git diff --check: passed.

Multi-rank TP/CP and end-to-end training were not run.