Explicit error when the number of training samples < global batch size (#3464)

Co-authored-by: Chen Cui <chcui@nvidia.com>

Author: OlegSudakov

src/megatron/bridge/data/loaders.py

@@ -228,6 +228,13 @@ def build_train_valid_test_data_loaders(
         cfg=cfg, build_train_valid_test_datasets_provider=build_train_valid_test_datasets_provider
     )
 
+    # Check that the train dataset has at least one global batch of samples.
+    if cfg.dataset.dataloader_type != "external" and len(train_ds) < cfg.train.global_batch_size:
+        raise RuntimeError(
+            f"Not enough train samples for a single global batch: "
+            f"train dataset size ({len(train_ds)}) < global batch size ({cfg.train.global_batch_size})."
+        )
+
     exit_signal = cfg.train.exit_signal
 
     def worker_init_fn(_):

tests/functional_tests/test_groups/data/test_loaders.py

@@ -231,7 +231,8 @@ def test_build_train_valid_test_data_loaders_specialized_dispatch_updates_flags(
         dp_group = object()
         dataset_provider = mock.Mock()
 
-        fake_train_ds = mock.Mock()
+        fake_train_ds = mock.MagicMock()
+        fake_train_ds.__len__.return_value = cfg.train.global_batch_size
         fake_train_ds.collate_fn = None
         mock_build_datasets.return_value = (fake_train_ds, None, None)
         fake_train_loader = object()