fix(dataloader): Prevent RuntimeError from DataloaderStopIteration (#1627)

The DataloaderStopIteration exception inherited from StopIteration.
According to PEP 479, raising a StopIteration subclass from a generator
causes a RuntimeError in Python 3.7+.

This change modifies the base class to `Exception` to ensure it can be
caught correctly by user code without triggering this behavior.

Fixes ISSUE #1626

Author: Lain810

torchtitan/components/dataloader.py

@@ -16,8 +16,14 @@
 from torchdata.stateful_dataloader import StatefulDataLoader
 from torchtitan.tools.logging import logger
 
-
-class DataloaderStopIteration(StopIteration):
+# NOTE: This class deliberately inherits from `Exception` and not `StopIteration`.
+# According to PEP 479, raising a `StopIteration` or its subclass from within a
+# generator will wrap it in a `RuntimeError`. Since this exception is designed
+# to be raised from a generator-based dataloader and caught by the training loop,
+# inheriting from `StopIteration` would make it uncatchable and would crash the
+# program.
+# See: https://peps.python.org/pep-0479/
+class DataloaderExhaustedError(Exception):
     """An exception that indicates dataloader exhaustion."""
 
     pass

torchtitan/experiments/forge/example_train.py

@@ -13,7 +13,7 @@
 from torch.distributed.elastic.multiprocessing.errors import record
 
 import torchtitan.protocols.train_spec as train_spec_module
-from torchtitan.components.dataloader import DataloaderStopIteration
+from torchtitan.components.dataloader import DataloaderExhaustedError
 from torchtitan.components.metrics import build_metrics_processor
 from torchtitan.components.tokenizer import build_hf_tokenizer
 from torchtitan.components.validate import build_validator
@@ -135,7 +135,7 @@ def batch_generator(
             except StopIteration as ex:
                 # If data runs out during gradient accumulation, that
                 # entire step will not be executed.
-                raise DataloaderStopIteration() from ex
+                raise DataloaderExhaustedError() from ex
             data_load_start = time.perf_counter()
             input_dict, labels = batch
             self.metrics_processor.ntokens_since_last_log += labels.numel()
@@ -292,7 +292,7 @@ def train(self):
                 self.gc_handler.run(self.step)
                 try:
                     self.train_step(data_iterator)
-                except DataloaderStopIteration:
+                except DataloaderExhaustedError:
                     logger.warning("Ran out of data; last step was canceled.")
                     break
 

torchtitan/train.py

@@ -15,7 +15,7 @@
 
 import torchtitan.protocols.train_spec as train_spec_module
 from torchtitan.components.checkpoint import CheckpointManager
-from torchtitan.components.dataloader import DataloaderStopIteration
+from torchtitan.components.dataloader import DataloaderExhaustedError
 from torchtitan.components.ft import FTManager, maybe_semi_sync_training
 from torchtitan.components.loss import rescale_accumulated_loss
 from torchtitan.components.metrics import (
@@ -386,7 +386,7 @@ def batch_generator(
             except StopIteration as ex:
                 # If data runs out during gradient accumulation, that
                 # entire step will not be executed.
-                raise DataloaderStopIteration() from ex
+                raise DataloaderExhaustedError() from ex
             input_dict, labels = batch
             ntokens_batch = labels.numel()
             self.ntokens_seen += ntokens_batch
@@ -583,7 +583,7 @@ def train(self):
                 self.gc_handler.run(self.step)
                 try:
                     self.train_step(data_iterator)
-                except DataloaderStopIteration:
+                except DataloaderExhaustedError:
                     logger.warning("Ran out of data; last step was canceled.")
                     break