fix(loss): compute causal-LM cross-entropy in float32

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(

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)