From 3a4b2bce097d94ef9523637a73734e664d35796d Mon Sep 17 00:00:00 2001 From: Bindu Bhargava R Chintam Date: Wed, 11 Nov 2026 23:00:30 -0800 Subject: [PATCH] fix(vocab_parallel_ce): take label smoothing from the summed shifted logits With label smoothing, the forward took log(exp(x - max)) over the whole [BT, V_local] buffer in Python. exp(x - max) underflows to 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 gradient was not affected. The forward kernel now also sums the shifted logits, and the smoothing term is sum(x - max) / K - log(sum_exp). This needs no [BT, V_local] temporaries, and ignored rows no longer have to be filled with ones. --- .../ops/vocab_parallel_cross_entropy.py | 42 +++++++++++------ .../vocab_parallel_cross_entropy.py | 2 +- test/megatron/test_cross_entropy.py | 47 +++++++++++++++++++ 3 files changed, 76 insertions(+), 15 deletions(-) diff --git a/src/liger_kernel/ops/vocab_parallel_cross_entropy.py b/src/liger_kernel/ops/vocab_parallel_cross_entropy.py index 4624a4a52..736234235 100644 --- a/src/liger_kernel/ops/vocab_parallel_cross_entropy.py +++ b/src/liger_kernel/ops/vocab_parallel_cross_entropy.py @@ -5,6 +5,7 @@ Kernel 1 (forward stats): raw logits + global max -> exp(shifted) buffer + sum_exp_local + predicted_logit_local + (+ sum_shifted_local with label smoothing) Python: AllReduce(SUM) x2 (no-op at TP=1) loss = log(sum_exp_global) - predicted_global optional label-smoothing adjustment @@ -62,6 +63,8 @@ def liger_vocab_parallel_ce_forward_kernel( Y_ptr, # [BT] int64 GLOBAL target indices in [0, V_global) pred_ptr, # [BT] OUT: predicted_logit_local (= X[masked_target] - max, # 0 if off-rank / ignored) + sum_shifted_ptr, # [BT] OUT: sum over V_local of (X - max); written only with + # label smoothing (None otherwise) sum_exp_ptr, # [BT] OUT: sum_exp_local vocab_start, # int — first global vocab index owned by this rank n_cols, # V_local @@ -84,9 +87,7 @@ def liger_vocab_parallel_ce_forward_kernel( tl.store(pred_ptr + program_id, 0.0) tl.store(sum_exp_ptr + program_id, 1.0) if HAS_LABEL_SMOOTHING: - for i in range(0, n_cols, BLOCK_SIZE): - X_offsets = i + tl.arange(0, BLOCK_SIZE) - tl.store(EXP_ptr + X_offsets, 1.0, mask=X_offsets < n_cols) + tl.store(sum_shifted_ptr + program_id, 0.0) return m = tl.load(logits_max_ptr + program_id).cast(tl.float32) @@ -101,6 +102,11 @@ def liger_vocab_parallel_ce_forward_kernel( else: pred_local = tl.load(X_ptr + y_local).cast(tl.float32) - m + # With label smoothing the shifted logits are summed too: they are the + # log-probabilities up to the log(sum_exp_global) term, and unlike + # log(exp(shifted)) they stay finite when exp() underflows to 0. + sum_shifted = 0.0 + # Single pass over V_local: shift inside, exp, accumulate sum_exp, # store exp(shifted) to the fp32 output buffer. Shifted logits are <= 0 so # exp() is in (0, 1] — no overflow possible. @@ -109,12 +115,17 @@ def liger_vocab_parallel_ce_forward_kernel( X_offsets = i + tl.arange(0, BLOCK_SIZE) col_mask = X_offsets < n_cols X_block = tl.load(X_ptr + X_offsets, mask=col_mask, other=0.0).cast(tl.float32) - exp_block = tl.exp(X_block - m) + shifted_block = X_block - m + if HAS_LABEL_SMOOTHING: + sum_shifted += tl.sum(tl.where(col_mask, shifted_block, 0.0)) + exp_block = tl.exp(shifted_block) exp_block = tl.where(col_mask, exp_block, 0.0) sum_exp += tl.sum(exp_block) tl.store(EXP_ptr + X_offsets, exp_block, mask=col_mask) tl.store(pred_ptr + program_id, pred_local) + if HAS_LABEL_SMOOTHING: + tl.store(sum_shifted_ptr + program_id, sum_shifted) tl.store(sum_exp_ptr + program_id, sum_exp) @@ -197,9 +208,11 @@ class LigerVocabParallelCEFunction(torch.autograd.Function): 1. Per-token global max (cheap reduction on input dtype, cast result to fp32) 2. AllReduce(MAX) (no-op at TP=1) 3. Kernel 1: shift internally, write exp(shifted) to a fp32 output buffer, - compute sum_exp_local and predicted_logit_local in one pass + compute sum_exp_local and predicted_logit_local in one pass (and the sum + of the shifted logits when label smoothing is on) 4. AllReduce(SUM) × 2 5. Loss = log(sum_exp_global) - predicted_global, then optional label smoothing + from the summed shifted logits (one more AllReduce(SUM) in "global" mode) 6. Save the exp buffer for backward Backward layout: @@ -264,6 +277,8 @@ def forward( block_size = _select_block_size(v_local) num_warps = _get_num_warps(block_size) + # Kernel 1 writes the per-token sum of shifted logits here, only with label smoothing. + sum_shifted_local = torch.empty(BT, device=x.device, dtype=torch.float32) if label_smoothing > 0 else None x_2d = x.view(BT, v_local) liger_vocab_parallel_ce_forward_kernel[(BT,)]( @@ -274,6 +289,7 @@ def forward( logits_max_ptr=logits_max, Y_ptr=flat_target, pred_ptr=predicted_local, + sum_shifted_ptr=sum_shifted_local, sum_exp_ptr=sum_exp_local, vocab_start=vocab_start, n_cols=v_local, @@ -300,15 +316,13 @@ def forward( alpha_eff = label_smoothing eps_eff = alpha_eff / v_for_K - # log_softmax_i = log(exp_i / sum_exp_global) = log(exp_i) - log(sum_exp_global) - log_probs_local = torch.log(exp_buf) - torch.log(sum_exp_local).unsqueeze(-1) - if label_smoothing_mode == "global": - sum_log_probs = log_probs_local.sum(dim=-1) - if tp_distributed: - dist.all_reduce(sum_log_probs, op=dist.ReduceOp.SUM, group=tp_group) - mean_log_probs = sum_log_probs / v_for_K - else: # "partition" - mean_log_probs = log_probs_local.mean(dim=-1) + # log_softmax_i = (x_i - max) - log(sum_exp_global), so its mean over K classes is + # sum_i(x_i - max) / K - log(sum_exp_global). The kernel already summed the shifted + # logits: no [BT, V_local] temporaries, and no log(0) where exp(x_i - max) underflowed. + # "global" averages over V_global (one AllReduce); "partition" over this rank's V_local. + if label_smoothing_mode == "global" and tp_distributed: + dist.all_reduce(sum_shifted_local, op=dist.ReduceOp.SUM, group=tp_group) + mean_log_probs = sum_shifted_local / v_for_K - torch.log(sum_exp_local) loss = (1.0 - alpha_eff) * loss + alpha_eff * (-mean_log_probs) else: diff --git a/src/liger_kernel/transformers/vocab_parallel_cross_entropy.py b/src/liger_kernel/transformers/vocab_parallel_cross_entropy.py index 9b7ed79f2..0a0979bdf 100644 --- a/src/liger_kernel/transformers/vocab_parallel_cross_entropy.py +++ b/src/liger_kernel/transformers/vocab_parallel_cross_entropy.py @@ -18,7 +18,7 @@ Averaging scope at TP>1 (irrelevant at TP=1 since V_local == V_global): - ``"global"`` (default): ``H(u, p)`` averaged over V_global. Adds one AllReduce on - ``sum(log_softmax_local, dim=-1)``. Exact result — matches a single-rank reference + the per-token sum of the shifted logits. Exact result — matches a single-rank reference computed on the full logits. - ``"partition"``: ``H(u, p)`` averaged over V_local on each rank. No extra AllReduce. Loss tensor differs per rank under this mode — this mirrors Megatron's diff --git a/test/megatron/test_cross_entropy.py b/test/megatron/test_cross_entropy.py index 4cfa0a192..5e762b49d 100644 --- a/test/megatron/test_cross_entropy.py +++ b/test/megatron/test_cross_entropy.py @@ -522,6 +522,53 @@ def test_vp_class_correctness_with_label_smoothing_and_ignore_index( assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=atol, rtol=rtol) +@pytest.mark.parametrize( + "ce_cls, reference", + [ + (LigerVocabParallelCrossEntropy, _reference_pytorch_loss), + (LigerMegatronCrossEntropy, _reference_megatron_loss), + ], +) +@pytest.mark.parametrize("gap", [120.0, 1e4]) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 1e-5, 1e-4), + pytest.param( + torch.bfloat16, + 1e-2, + 5e-2, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported"), + ), + ], +) +def test_label_smoothing_with_logits_far_below_the_max(ce_cls, reference, gap, dtype, atol, rtol): + """exp(x - max) underflows to 0 in fp32 once a logit is ~104 below the row max (a masked + vocabulary, for example). The smoothing term must stay finite there, like F.cross_entropy.""" + s, b, v = 8, 2, 512 + n_low = 64 + ignore_index = -100 + label_smoothing = 0.1 + ce = ce_cls(ignore_index=ignore_index, label_smoothing=label_smoothing) + + base = torch.randn(s, b, v, device=device, dtype=dtype) + base[..., v - n_low :] -= gap + target = torch.randint(0, v - n_low, (s, b), device=device, dtype=torch.long) + _assign_ignore_index(target, ignore_index, frac=0.25) + + h_ref = base.detach().clone().requires_grad_(True) + h_got = base.detach().clone().requires_grad_(True) + + ref = reference(h_ref, target, ignore_index=ignore_index, label_smoothing=label_smoothing) + got = ce(h_got, target) + assert torch.isfinite(got).all() + assert_verbose_allclose(got.float(), ref.float(), atol=atol, rtol=rtol) + + ref.sum().backward() + got.sum().backward() + assert_verbose_allclose(h_got.grad.float(), h_ref.grad.float(), atol=atol, rtol=rtol) + + @pytest.mark.parametrize( "s, b, v", [