[Bug] Vocab-parallel label smoothing uses the local vocabulary size and the local mean log-probability

Summary

With nonzero label smoothing, vocab_parallel_cross_entropy gives a different loss and gradient when the same vocabulary is split across tensor-parallel ranks.

Two places use the local partition as if it were the whole vocabulary: vocab_size = exp_logits.size(-1) feeds the smoothing coefficient, and mean_log_probs = log_probs.mean(dim=-1) averages over the local slice only. The saved ctx.vocab_size carries the same error into backward.

Related reports are Issue #737 and the earlier PR #5522. The reproduction below measures the loss and full-gradient difference directly between TP1 and TP2.

System Info

Environment: Megatron-LM commit d4550898, PyTorch 2.13.0+cu130, RTX 5090, FP32.

Steps/Code to reproduce bug

1. Build the same logits for six positions and a global vocabulary of 8, with the same target labels.
2. Evaluate vocab_parallel_cross_entropy and its backward with TP1, then with TP2 (two partitions of width 4).
3. Use label_smoothing=0.1 and compare the gathered full gradient.
4. Optionally apply one SGD step at learning rate 0.25 to the logits treated as parameters.

With global-vocabulary statistics, the loss-stage discrepancy falls to about 1.7e-6 of its original value. The loss and gradient match an independent float64 reference within FP32 tolerance. label_smoothing=0 is unaffected. Reproduced on three hosts.

This tests the CE API and its backward, not a full training job.

Expected behavior

TP1 and TP2 should give the same loss and gradient for the same logits and labels, with or without label smoothing.

Root cause

vocab_size is taken from the local partition, and the uniform log-probability term is averaged over that partition. sum_exp_logits is already all-reduced over the TP group, so only the shifted-logit sum and the vocabulary size are missing their global versions.

Proposed fix

Use the global vocabulary size in the smoothing coefficient and in the saved backward metadata, and compute the uniform term over the same global vocabulary:

mean_global_log_prob = global_shifted_logit_sum / global_vocab_size
                       - log(global_sum_exp_logits)

Sum the shifted logits across the TP group before the helper exponentiates them in place. This also avoids taking log of already-rounded tiny probabilities. The extra all-reduce only runs when smoothing is nonzero.

Additional context

Related: Issue #737, PR #5522.