Avoid per-microbatch D2H caused by isfinite (#3873)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* __->__ #3873

#3863 introduces a good
feature to raise if loss is NaN but also introduces per-micro-batch D2H
sync. This PR fixes the issue with the penalty that the actual batch
that causes the NaN is unkown (within a certain range).

#2728 is a potential solution
to remove the penalty.

Another mitigation from @mlazos: you could also copy N loss values on
GPU to keep a record of which of the N steps had the NaN.

Author: fegin

tests/unit_tests/test_invalid_loss.py

@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from torchtitan.components.loss import IGNORE_INDEX
+from torchtitan.trainer import Trainer
+
+
+class TestInvalidLoss(unittest.TestCase):
+    """Trainer.train_step crashes on a non-finite loss, aligned with logging.
+
+    The finiteness check reuses the device-to-host copy that logging already
+    performs, so it only runs on steps where metrics are logged.
+    """
+
+    def _make_trainer(self, loss_value: float, should_log: bool) -> Trainer:
+        # Build a bare Trainer and inject only the collaborators train_step
+        # touches on the non-distributed (single-rank) path.
+        trainer = object.__new__(Trainer)
+
+        trainer.optimizers = MagicMock()
+        trainer.lr_schedulers = MagicMock()
+        trainer.lr_schedulers.get_metrics.return_value = {}
+        trainer.checkpointer = MagicMock()
+        trainer.model_parts = []
+        trainer.config = MagicMock()
+        trainer.config.training.max_norm = 1.0
+        trainer.device = torch.device("cpu")
+        trainer.gradient_accumulation_steps = 1
+        trainer.step = 1
+        trainer.ntokens_seen = 0
+
+        parallel_dims = MagicMock()
+        parallel_dims.dp_enabled = False
+        parallel_dims.dp_cp_enabled = False
+        parallel_dims.ep_enabled = False
+        parallel_dims.get_optional_mesh.return_value = None
+        trainer.parallel_dims = parallel_dims
+
+        trainer.metrics_processor = MagicMock()
+        trainer.metrics_processor.should_log.return_value = should_log
+
+        # Shadow the bound method so forward/backward returns a canned loss.
+        trainer.forward_backward_step = MagicMock(return_value=torch.tensor(loss_value))
+        return trainer
+
+    def _data_iterator(self):
+        labels = torch.tensor([1, 2, IGNORE_INDEX])
+        input_dict = {"input": torch.tensor([1, 2, 3])}
+        while True:
+            yield input_dict, labels
+
+    def _run_step(self, loss_value: float, should_log: bool) -> None:
+        trainer = self._make_trainer(loss_value, should_log)
+        # sl.* are logging side effects; clip_grad_norm_ needs real params.
+        with patch("torchtitan.trainer.sl", MagicMock()), patch(
+            "torchtitan.trainer.dist_utils.clip_grad_norm_",
+            return_value=torch.tensor(1.0),
+        ):
+            trainer.train_step(self._data_iterator())
+
+    def test_nan_loss_raises_on_log_step(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._run_step(float("nan"), should_log=True)
+        self.assertIn("not finite", str(ctx.exception))
+
+    def test_inf_loss_raises_on_log_step(self):
+        with self.assertRaises(RuntimeError) as ctx:
+            self._run_step(float("inf"), should_log=True)
+        self.assertIn("not finite", str(ctx.exception))
+
+    def test_finite_loss_does_not_raise(self):
+        self._run_step(1.5, should_log=True)
+
+    def test_nan_loss_ignored_when_not_logging(self):
+        # Detection is gated on logging, so a non-log step must not raise even
+        # when the loss is NaN.
+        self._run_step(float("nan"), should_log=False)
+
+
+if __name__ == "__main__":
+    unittest.main()

torchtitan/trainer.py

@@ -6,6 +6,7 @@
 
 import dataclasses
 import json
+import math
 import os
 import time
 from collections.abc import Iterable, Iterator
@@ -799,15 +800,6 @@ def train_step(
                 labels=labels,
                 global_valid_tokens=global_valid_tokens,
             )
-
-            if not loss.isfinite():
-                raise RuntimeError(
-                    f"Loss is not finite (loss={loss.item()}) at step "
-                    f"{self.step}. Dumping the offending microbatch for "
-                    f"debugging -- input_dict: {input_dict}, labels: {labels}, "
-                    f"global_valid_tokens: {global_valid_tokens}"
-                )
-
             accumulated_losses.append(loss.detach())
 
         with sl.log_trace_span("optim"):
@@ -861,6 +853,16 @@ def train_step(
                 global_avg_loss = global_max_loss = float(loss.detach().item())
                 global_ntokens_seen = self.ntokens_seen
 
+        # Crash on invalid loss. global_avg_loss is a SUM reduction, so a infinite
+        # loss on any rank propagates here. This reuses the D2H copy already done
+        # for logging, so it adds no extra sync.
+        # TODO: make this step work even logging is off.
+        if not math.isfinite(global_avg_loss):
+            raise RuntimeError(
+                f"Loss is not finite (global_avg_loss={global_avg_loss}) at "
+                f"step {self.step}. Stopping training."
+            )
+
         extra_metrics = {
             "n_tokens_seen": global_ntokens_seen,
             **lr_metrics,