fix: detach probability weights in DFT loss

[taking-lying-flat]

PR type

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

PR information

DFT requires a stop-gradient probability weight: -stop_gradient(p_y) * log(p_y), as specified in the paper's Equation 9. Differentiating the weight instead multiplies the intended gradient by 1 + log(p_y), reversing its direction when p_y < exp(-1).

Detach the probability weight at all four expressions: standard CE, chunked CE forward, chunked backward recomputation, and the chunked precomputed-logps path. This changes four lines across two production files. Forward loss values and ordinary CE behavior remain unchanged.

Experiment results

For p_y=0.1, one SGD step with learning rate 0.1 gives the same results across both losses and their logits/logps input paths:

	Before	After
Forward loss	0.230259	0.230259
Target-logit gradient	+0.117233	-0.090000
Updated target probability	0.097910	0.101632

Validated locally with PyTorch 2.13.0 on CPU:

• An out-of-tree harness checked forward values and full analytic gradients for the baseline and edited implementations: 384 checks covering float32/float64, mean/sum, masked and all-ignored labels, three chunk sizes, upstream gradient scaling, and ordinary CE.
• tests/loss/test_ce_mse.py: 18 passed.
• uvx pre-commit run --all-files: passed.