From 64f661a99d9c14091bc33e499ebd69a01041042d Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:14:29 +0200 Subject: [PATCH 1/2] fix(loss): compute causal-LM cross-entropy in float32 CLMCrossEntropyLoss passed the model's logits straight into CrossEntropyLoss. Under mixed precision those are bfloat16, and nothing promotes them: FSDP2's MixedPrecisionPolicy casts parameters, it does not install torch.autocast, so the operator-level float32 autocast list never applies. The log-softmax, its backward, and -- because the reduction is "mean" -- the accumulation of the per-token losses therefore all ran in bfloat16 over a 131k-entry vocabulary. Both reference implementations up-cast at exactly this point: TorchTitan does it on every one of its cross-entropy paths (torchtitan/components/loss.py, e.g. `pred.flatten(0, 1).float()`), and HF transformers opens ForCausalLMLoss with `logits = logits.float()` (transformers/loss/loss_utils.py). Measured against a float64 evaluation of identical bfloat16-quantized logits at vocab 131072, the loss value was off by ~4.9e-2 and is now off by ~2.4e-6. The accumulator dominates that figure; the log-softmax alone accounts for ~3e-3. On a real 8B checkpoint and 262,144 held-out tokens the gradient into lm_head differs by 0.19% between the two paths, and the up-cast recovers about 1.8x of that -- the remainder is imposed by the bfloat16 logit boundary itself and is not reachable from the loss. After the change the loss and its gradient are bit-identical to TorchTitan's formulation for bfloat16, float16 and float32 inputs at vocab 1024 and 131072. Tests cover the three symptoms: the returned dtype, accuracy against a float64 reference, and agreement with TorchTitan's sum-reduction-over-valid-tokens formulation. All three fail without the fix. Co-Authored-By: Claude Opus 5 (1M context) --- src/modalities/loss_functions.py | 8 +++- tests/test_loss_functions.py | 64 +++++++++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/src/modalities/loss_functions.py b/src/modalities/loss_functions.py index e3be6100d..39277a8a6 100644 --- a/src/modalities/loss_functions.py +++ b/src/modalities/loss_functions.py @@ -48,7 +48,13 @@ def __call__(self, *args, **kwargs) -> torch.Tensor: shift_logits = lm_logits.contiguous() shift_labels = labels.contiguous().long() # Flatten the tokens. We compute here, the loss per token. - loss = self.loss_fun(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + # The up-cast to float32 is deliberate and must not be removed. Without it the log-softmax, + # its backward, and the mean reduction all run in the logits' own dtype -- and under mixed + # precision the model hands us bfloat16. FSDP2's MixedPrecisionPolicy casts *parameters*; it + # does not install torch.autocast, so nothing else promotes them. TorchTitan up-casts on + # every one of its cross-entropy paths for the same reason (torchtitan/components/loss.py), + # as does HF transformers (transformers/loss/loss_utils.py). + loss = self.loss_fun(shift_logits.view(-1, shift_logits.size(-1)).float(), shift_labels.view(-1)) return loss def _parse_arguments( diff --git a/tests/test_loss_functions.py b/tests/test_loss_functions.py index 8825f15c3..0fdfca78f 100644 --- a/tests/test_loss_functions.py +++ b/tests/test_loss_functions.py @@ -2,7 +2,7 @@ import torch from modalities.batch import InferenceResultBatch -from modalities.loss_functions import NCELoss, nce_loss +from modalities.loss_functions import CLMCrossEntropyLoss, NCELoss, nce_loss @pytest.fixture @@ -36,3 +36,65 @@ def test_nce_loss_correctness(embedding1, embedding2): bidirectional_loss = nce_loss(embedding1, embedding2, device="cpu", is_asymmetric=False, temperature=1.0) assert unidirectional_loss == pytest.approx(1.1300, 0.0001) assert bidirectional_loss == pytest.approx(2.2577, 0.0001) + + +# --------------------------------------------------------------------------- +# Causal-LM cross-entropy must be computed in float32 even when the model hands +# it bfloat16 logits. FSDP2's MixedPrecisionPolicy casts parameters but does not +# install torch.autocast, so nothing promotes them on the way into the loss. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def clm_loss() -> CLMCrossEntropyLoss: + return CLMCrossEntropyLoss(target_key="target", prediction_key="logits") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_clm_cross_entropy_returns_float32_for_half_precision_logits(clm_loss, dtype): + """The returned dtype is the giveaway: cross-entropy returns its input's dtype.""" + torch.manual_seed(0) + logits = torch.randn(2, 16, 512, dtype=dtype) + labels = torch.randint(0, 512, (2, 16)) + assert clm_loss(logits, labels).dtype == torch.float32 + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_clm_cross_entropy_is_accurate_for_half_precision_logits(clm_loss, dtype): + """Half-precision logits must not drag the log-softmax or the reduction down with them. + + Two faults are guarded at once. Without the up-cast the log-softmax runs in the logits' own + dtype, and because the reduction is ``"mean"`` the per-token losses are accumulated in that + dtype too -- over a 131k-entry vocabulary the accumulator error dominates. Against a float64 + evaluation of the identical quantized logits, the un-upcast path lands around 5e-2 and the + float32 path around 2e-6. + """ + torch.manual_seed(0) + vocab_size = 131072 + logits = (torch.randn(2, 64, vocab_size) * 8.0).to(dtype) + labels = torch.randint(0, vocab_size, (2, 64)) + + reference = clm_loss(logits.double(), labels) + assert abs(clm_loss(logits, labels).double() - reference) < 1e-4 + + +def test_clm_cross_entropy_matches_torchtitan_formulation(clm_loss): + """Mean reduction over valid tokens == sum reduction / valid-token count. + + Mirrors ``torchtitan/components/loss.py::cross_entropy_loss``, which up-casts the logits and + sum-reduces for token-based normalization. Kept inline so the test carries no dependency on + torchtitan. + """ + torch.manual_seed(0) + vocab_size = 1024 + logits = torch.randn(2, 64, vocab_size, dtype=torch.bfloat16) + labels = torch.randint(0, vocab_size, (2, 64)) + labels[0, :7] = -100 # ignored tokens must leave both sides unchanged + + flat_labels = labels.view(-1).long() + summed = torch.nn.functional.cross_entropy( + logits.view(-1, vocab_size).float(), flat_labels, reduction="sum", ignore_index=-100 + ) + torchtitan_style = summed / (flat_labels != -100).sum() + + torch.testing.assert_close(clm_loss(logits, labels), torchtitan_style, rtol=1e-6, atol=1e-6) From 70cbb83cd8468f04fbafbf6fbcb327a8fa7dd964 Mon Sep 17 00:00:00 2001 From: Max Luebbering <2804731+le1nux@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:49:17 +0200 Subject: [PATCH 2/2] fix(loss): correct the accuracy test's reference and the precision rationale Two review findings, both correct. The accuracy test compared the implementation with itself. It built its "float64 reference" by calling the loss with a float64 tensor, but the loss casts its input to float32 -- so once the fix is in place both sides evaluated the identical float32 tensor and the difference was exactly 0.0. It still went red on the bug, but it was vacuous with the fix applied. The reference is now computed directly via F.cross_entropy on float64 logits, and the assertion measures a real 2.36e-6. The rationale conflated tensor precision with accumulation. PyTorch's kernels already accumulate in float32 for half-precision inputs, so the claim that the log-softmax, its backward and the mean reduction "run in the logits' own dtype" was wrong. What the cast preserves is the precision of the stored tensors: the log-softmax output, the tensors its backward reads, and the returned loss scalar. The scalar dominates -- measured, the ~5e-2 error before this cast is exactly the bfloat16 quantum at a loss magnitude of ~36 (4.89e-2 on CPU, 3.18e-2 on an A100), not an accumulator artefact. Co-Authored-By: Claude Opus 5 (1M context) --- src/modalities/loss_functions.py | 19 +++++++++++++------ tests/test_loss_functions.py | 17 ++++++++++------- 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/src/modalities/loss_functions.py b/src/modalities/loss_functions.py index 39277a8a6..ac636565b 100644 --- a/src/modalities/loss_functions.py +++ b/src/modalities/loss_functions.py @@ -48,12 +48,19 @@ def __call__(self, *args, **kwargs) -> torch.Tensor: shift_logits = lm_logits.contiguous() shift_labels = labels.contiguous().long() # Flatten the tokens. We compute here, the loss per token. - # The up-cast to float32 is deliberate and must not be removed. Without it the log-softmax, - # its backward, and the mean reduction all run in the logits' own dtype -- and under mixed - # precision the model hands us bfloat16. FSDP2's MixedPrecisionPolicy casts *parameters*; it - # does not install torch.autocast, so nothing else promotes them. TorchTitan up-casts on - # every one of its cross-entropy paths for the same reason (torchtitan/components/loss.py), - # as does HF transformers (transformers/loss/loss_utils.py). + # The up-cast to float32 is deliberate and must not be removed. Under mixed precision the + # model hands us bfloat16 logits: FSDP2's MixedPrecisionPolicy casts *parameters*, it does + # not install torch.autocast, so nothing else promotes them on the way in. + # + # This is about *tensor* precision, not accumulation -- PyTorch's kernels already accumulate + # in float32 for half-precision inputs. What the cast preserves is the log-softmax output, + # the tensors its backward reads, and the returned loss scalar, each of which would + # otherwise be stored in the logits' dtype. The scalar matters most: at a loss magnitude of + # ~36 the bfloat16 grid is 0.25 wide, which is the whole of the ~5e-2 error measured before + # this cast was added. + # + # TorchTitan up-casts on every one of its cross-entropy paths for the same reason + # (torchtitan/components/loss.py), as does HF transformers (transformers/loss/loss_utils.py). loss = self.loss_fun(shift_logits.view(-1, shift_logits.size(-1)).float(), shift_labels.view(-1)) return loss diff --git a/tests/test_loss_functions.py b/tests/test_loss_functions.py index 0fdfca78f..e1eecb784 100644 --- a/tests/test_loss_functions.py +++ b/tests/test_loss_functions.py @@ -61,20 +61,23 @@ def test_clm_cross_entropy_returns_float32_for_half_precision_logits(clm_loss, d @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_clm_cross_entropy_is_accurate_for_half_precision_logits(clm_loss, dtype): - """Half-precision logits must not drag the log-softmax or the reduction down with them. + """Half-precision logits must not drag the loss down with them. - Two faults are guarded at once. Without the up-cast the log-softmax runs in the logits' own - dtype, and because the reduction is ``"mean"`` the per-token losses are accumulated in that - dtype too -- over a 131k-entry vocabulary the accumulator error dominates. Against a float64 - evaluation of the identical quantized logits, the un-upcast path lands around 5e-2 and the - float32 path around 2e-6. + Not about accumulation -- the kernels already accumulate in float32 for half-precision inputs. + What the up-cast preserves is tensor precision: the log-softmax output and the returned scalar, + which would otherwise be stored in the logits' dtype. The scalar dominates; at a loss magnitude + of ~36 the bfloat16 grid is 0.25 wide. Against a float64 reference the un-upcast path lands + around 5e-2 and the float32 path around 2e-6. + + The reference is computed directly rather than through ``clm_loss``: the implementation casts + its input to float32, so passing it a float64 tensor would silently compare the fix with itself. """ torch.manual_seed(0) vocab_size = 131072 logits = (torch.randn(2, 64, vocab_size) * 8.0).to(dtype) labels = torch.randint(0, vocab_size, (2, 64)) - reference = clm_loss(logits.double(), labels) + reference = torch.nn.functional.cross_entropy(logits.double().reshape(-1, vocab_size), labels.reshape(-1)) assert abs(clm_loss(logits, labels).double() - reference) < 1e-4