[bhargavchintam]

Fixes #1618

Summary

With label_smoothing > 0, LigerVocabParallelCEFunction.forward took torch.log(exp_buf) over the whole [BT, V_local] buffer to get the log-probabilities. exp(x - max) is 0 in fp32 once a logit is about 104 below the row max, so the loss came out as inf where F.cross_entropy is finite. The same line made two [BT, V_local] fp32 temporaries.

Since log_softmax_i = (x_i - max) - log(sum_exp), its mean over K classes is sum_i(x_i - max) / K - log(sum_exp). With label smoothing on, the forward kernel now also sums the shifted logits per row (sum_shifted_ptr, None otherwise), and the smoothing term is taken from that sum. In "global" mode the AllReduce carries this sum instead of the summed log-probabilities; "partition" mode uses the local sum. Ignored rows no longer need their exp_buf row filled with ones, so that loop is removed.

The diff is this change in ops/vocab_parallel_cross_entropy.py (the forward kernel, the smoothing block, docstrings), one docstring line in transformers/vocab_parallel_cross_entropy.py, and the test below. The loss line, the backward kernel and the path without label smoothing are not changed.

Details

#1555 also adds an output to this forward kernel, so I placed my lines where the two do not overlap. A merge of this branch with #1555's head (2e68201) is clean in this file and in its test file. The one conflict is in a file this PR does not touch:

$ git merge-tree --write-tree --name-only HEAD 2e68201 | grep -E "vocab_parallel|megatron/test_cross_entropy|CONFLICT"
Auto-merging src/liger_kernel/ops/vocab_parallel_cross_entropy.py
Auto-merging test/megatron/test_cross_entropy.py
CONFLICT (content): Merge conflict in test/transformers/test_cross_entropy.py

#1331 wraps the kernel launch in device_context, which re-indents the lines this PR edits there, so those two conflict in this file:

$ git merge-tree --write-tree --name-only HEAD 8bad3cf | grep -E "CONFLICT.*vocab_parallel"
CONFLICT (content): Merge conflict in src/liger_kernel/ops/vocab_parallel_cross_entropy.py

Megatron-LM's own loss takes torch.log(exp_logits) at this step (megatron/core/ops/loss/cross_entropy.py). I read that in its source and did not run it. If it behaves as written, LigerMegatronCrossEntropy now differs from it in these cases: a finite loss instead of inf.

Testing Done

• Hardware Type: 2x NVIDIA L40S (NGC PyTorch 26.10 container: torch 2.15.0a0, Triton 3.8.0, CUDA 13.4)
• run make test to ensure correctness (68 failed, 5337 passed; the same 68 fail on main on this machine, shown below)
• run make checkstyle to ensure code style (I ran its two check-only steps, shown at the end, because the other two steps rewrite files)
• run make test-convergence to ensure convergence

I did not run make test-convergence.

New test test_label_smoothing_with_logits_far_below_the_max: both classes, 64 of 512 logits moved 120 or 1e4 below the rest, fp32 and bf16, a quarter of the rows ignored. The loss must be finite and match F.cross_entropy, and so must the gradient.

$ python -m pytest --disable-warnings -q --tb=line -p no:logging -rs test/megatron/test_cross_entropy.py -k far_below

# main (42dc23c) with only the new test applied; this line eight times, once per case
/w/base_t/test/megatron/test_cross_entropy.py:564: AssertionError: assert tensor(False, device='cuda:0')
8 failed, 125 deselected, 3 warnings in 1.40s

# this PR
8 passed, 125 deselected, 2 warnings in 1.23s

The whole file, on two GPUs so that the TP=2 cases run:

$ python -m pytest --disable-warnings -q --tb=line -p no:logging -rs test/megatron/test_cross_entropy.py

# main (42dc23c) with only the new test applied
SKIPPED [16] test/megatron/test_cross_entropy.py:734: requires >= 4 GPUs
SKIPPED [16] test/megatron/test_cross_entropy.py:734: requires >= 8 GPUs
8 failed, 93 passed, 32 skipped, 2 warnings in 127.78s (0:02:07)

# this PR
SKIPPED [16] test/megatron/test_cross_entropy.py:734: requires >= 4 GPUs
SKIPPED [16] test/megatron/test_cross_entropy.py:734: requires >= 8 GPUs
101 passed, 32 skipped, 2 warnings in 128.72s (0:02:08)

The snippet from the issue, saved as issue_snippet.py, on this branch (on main it prints the values in the issue):

$ python issue_snippet.py
label_smoothing=0.0: liger 1.407606  torch 1.407606
label_smoothing=0.1: liger 6.432606  torch 6.432606

Two ranks on two GPUs over NCCL, with my own script (not in this PR): three inputs (ordinary logits, low logits on one rank, low logits on both), both formulas, both modes, label_smoothing 0 and 0.1, each rank compared with an fp64 log_softmax reference. It checks the forward value in "partition" mode too.

$ TP_BACKEND=nccl python tp2_check.py

# main (42dc23c): each bad case is an inf loss
rank0 SUMMARY bad=6 worst_rel=1.275e-07
rank1 SUMMARY bad=8 worst_rel=1.033e-07
TOTAL bad cases: 14 of 48

# this PR
rank0 SUMMARY bad=0 worst_rel=1.623e-07
rank1 SUMMARY bad=0 worst_rel=1.824e-07
TOTAL bad cases: 0 of 48

The 4-GPU and 8-GPU cases of the repo's test did not run on this machine (the 32 skips above). The same script with 4 and with 8 CPU ranks over gloo, under the Triton interpreter:

$ TRITON_INTERPRET=1 TP_SIZE=4 python tp2_check.py

# main (42dc23c)
TOTAL bad cases: 22 of 96

# this PR
TOTAL bad cases: 0 of 96

$ TRITON_INTERPRET=1 TP_SIZE=8 python tp2_check.py

# main (42dc23c)
TOTAL bad cases: 40 of 192

# this PR
TOTAL bad cases: 0 of 192

Memory and time at TP=1 on one L40S, main -> this PR, from my own script (not in this PR). "peak fwd MB" is what one forward allocates above its inputs; times are the median of 30 calls after 10 warm-up calls.

$ python bench_vpce.py --tag before --out before.jsonl   # on main (42dc23c)
$ python bench_vpce.py --tag after --out after.jsonl     # on this branch
$ python bench_compare.py before.jsonl after.jsonl
dtype        BT       V   ls |        peak fwd MB |             fwd ms |         fwd+bwd ms
bfloat16   4096   32000  0.0 |         500 -> 500 |       1.45 -> 1.45 |       4.15 -> 4.14
bfloat16   4096   32000  0.1 |        1500 -> 500 |       5.43 -> 1.46 |       8.21 -> 4.16
bfloat16   4096  131072  0.0 |       2048 -> 2048 |       5.92 -> 5.92 |     17.14 -> 17.16
bfloat16   4096  131072  0.1 |       6144 -> 2048 |      22.20 -> 5.92 |     33.32 -> 17.15
bfloat16   8192  131072  0.0 |       4096 -> 4096 |     11.79 -> 11.78 |     34.34 -> 34.21
bfloat16   8192  131072  0.1 |      12288 -> 4096 |     44.42 -> 11.79 |     66.91 -> 34.33
bfloat16  16384   32000  0.0 |       2000 -> 2000 |       5.99 -> 5.86 |     16.33 -> 16.49
bfloat16  16384   32000  0.1 |       6000 -> 2001 |      21.47 -> 5.69 |     32.25 -> 16.36
float32    4096   32000  0.0 |         500 -> 500 |       2.24 -> 2.24 |       3.84 -> 3.66
float32    4096   32000  0.1 |        1500 -> 500 |       6.18 -> 2.18 |       7.72 -> 3.70
float32    4096  131072  0.0 |       2048 -> 2048 |       8.75 -> 8.76 |     15.09 -> 15.10
float32    4096  131072  0.1 |       6144 -> 2048 |      24.98 -> 9.02 |     31.19 -> 15.37
float32    8192  131072  0.0 |       4096 -> 4096 |     17.46 -> 17.49 |     30.13 -> 30.14
float32    8192  131072  0.1 |      12288 -> 4096 |     49.82 -> 17.99 |     62.70 -> 30.67
float32   16384   32000  0.0 |       2000 -> 2000 |       8.93 -> 8.91 |     14.45 -> 14.44
float32   16384   32000  0.1 |       6000 -> 2001 |      24.16 -> 8.56 |     30.21 -> 14.56

make test on this branch. None of the 68 failures is in test/megatron/. They are in ten other files: test_grpo_loss.py (24), test_mlp.py (18), test_fused_linear_score_centering.py (10), test_fused_neighborhood_attention.py (6), test_swiglu.py (3), test_rms_norm.py (2), test_flex_attention.py (2), and one each in test_muse_glimmer.py, test_geglu.py and test_fused_linear_selected_log_probs.py. I re-ran those 68 ids on main (42dc23c) and on this branch. The same 68 fail on both, and the two outputs are identical apart from the timing. I did not look into why they fail on an L40S.

$ make test
= 68 failed, 5337 passed, 1023 skipped, 20 xfailed, 88 warnings in 3812.90s (1:03:32) =

$ python -m pytest --disable-warnings -q --tb=line -p no:logging <the 68 failed ids>

# main (42dc23c)
68 failed, 7 warnings in 811.23s (0:13:31)

# this PR
68 failed, 7 warnings in 13.24s

$ ruff check --no-cache --output-format=concise .
All checks passed!

$ ruff format --no-cache --check .
500 files already formatted