Merge pull request #3097 from bghira/agent/dataset-batch-mps-limit

Validate per-dataset batch sizes on MPS

Author: bghira

simpletuner/helpers/configuration/cmd_args.py

@@ -17,6 +17,7 @@
 from accelerate.utils import ProjectConfiguration
 
 from simpletuner.helpers.configuration.cli_utils import mapping_to_cli_args, normalize_lr_scheduler_value
+from simpletuner.helpers.configuration.platform_validation import validate_mps_train_batch_size
 from simpletuner.helpers.configuration.template_vars import render_modelspec_comment
 from simpletuner.helpers.distillation.common import validate_distillation_text_encoder_training
 from simpletuner.helpers.logging import get_logger
@@ -963,12 +964,7 @@ def _normalize_input_args(raw_args):
     if torch.backends.mps.is_available():
         if args.model_family.lower() not in ["sd3", "flux", "legacy"] and not args.unet_attention_slice:
             warning_log("MPS may benefit from the use of --unet_attention_slice for memory savings at the cost of speed.")
-        if args.train_batch_size > 16:
-            raise ValueError(
-                "An M3 Max 128G will use 12 seconds per step at a batch size of 1 and 65 seconds per step at a batch size of 12."
-                " Any higher values will result in NDArray size errors or other unstable training results and crashes."
-                "\nPlease reduce the batch size to 12 or lower."
-            )
+        validate_mps_train_batch_size(args.train_batch_size)
 
         if args.quantize_via == "accelerator":
             args.quantize_via = "cpu"

simpletuner/helpers/configuration/platform_validation.py

@@ -0,0 +1,11 @@
+import torch
+
+
+def validate_mps_train_batch_size(train_batch_size: int) -> None:
+    """Reject training batch sizes that exceed the existing MPS safety limit."""
+    if torch.backends.mps.is_available() and train_batch_size > 16:
+        raise ValueError(
+            "An M3 Max 128G will use 12 seconds per step at a batch size of 1 and 65 seconds per step at a batch size of 12."
+            " Any higher values will result in NDArray size errors or other unstable training results and crashes."
+            "\nPlease reduce the batch size to 12 or lower."
+        )

simpletuner/helpers/data_backend/factory.py

@@ -96,6 +96,7 @@ def _coerce_bucket_keys(indices: Dict[Any, Iterable]) -> Dict[Any, list]:
 from simpletuner.helpers.caching.image_embed import ImageEmbedCache
 from simpletuner.helpers.caching.text_embeds import TextEmbeddingCache
 from simpletuner.helpers.caching.vae import VAECache
+from simpletuner.helpers.configuration.platform_validation import validate_mps_train_batch_size
 from simpletuner.helpers.configuration.template_vars import resolve_value_placeholders
 from simpletuner.helpers.data_backend.aws import S3DataBackend
 from simpletuner.helpers.data_backend.base import BaseDataBackend
@@ -312,6 +313,8 @@ def init_backend_config(backend: dict, args: dict, accelerator) -> dict:
     dataset_train_batch_size = (
         resolve_dataset_train_batch_size(backend, args, dataset_type) if has_training_batch_size else None
     )
+    if dataset_train_batch_size is not None:
+        validate_mps_train_batch_size(dataset_train_batch_size)
 
     start_epoch = normalize_start_epoch(backend.get("start_epoch", 1))
     start_step = normalize_start_step(backend.get("start_step", 0))

tests/test_factory_edge_cases.py

@@ -453,6 +453,71 @@ def test_init_backend_config_uses_dataset_train_batch_size(self):
         self.assertEqual(result["config"]["train_batch_size"], 2)
         self.assertEqual(result["bucket_report"].constraints["train_batch_size"], 2)
 
+    def test_init_backend_config_rejects_unsafe_dataset_train_batch_size_on_mps(self):
+        from simpletuner.helpers.data_backend.factory import init_backend_config
+
+        self.args.train_batch_size = 1
+        training_dataset_types = (
+            DatasetType.IMAGE,
+            DatasetType.VIDEO,
+            DatasetType.AUDIO,
+            DatasetType.CONDITIONING,
+            DatasetType.CAPTION,
+            DatasetType.GROUNDING,
+        )
+
+        with patch(
+            "simpletuner.helpers.configuration.platform_validation.torch.backends.mps.is_available", return_value=True
+        ):
+            for dataset_type in training_dataset_types:
+                with self.subTest(dataset_type=dataset_type.value):
+                    backend = {
+                        "id": f"{dataset_type.value}-unsafe-batch",
+                        "type": "local",
+                        "dataset_type": dataset_type.value,
+                        "train_batch_size": 17,
+                        "instance_data_dir": self.temp_dir,
+                    }
+
+                    with self.assertRaisesRegex(ValueError, "Please reduce the batch size to 12 or lower"):
+                        init_backend_config(backend, self.args, self.accelerator)
+
+    def test_init_backend_config_allows_train_batch_size_at_mps_limit(self):
+        from simpletuner.helpers.data_backend.factory import init_backend_config
+
+        backend = {
+            "id": "image-safe-mps-batch",
+            "type": "local",
+            "dataset_type": "image",
+            "train_batch_size": 16,
+            "instance_data_dir": self.temp_dir,
+        }
+
+        with patch(
+            "simpletuner.helpers.configuration.platform_validation.torch.backends.mps.is_available", return_value=True
+        ):
+            result = init_backend_config(backend, self.args, self.accelerator)
+
+        self.assertEqual(result["config"]["train_batch_size"], 16)
+
+    def test_init_backend_config_allows_large_train_batch_size_without_mps(self):
+        from simpletuner.helpers.data_backend.factory import init_backend_config
+
+        backend = {
+            "id": "image-non-mps-batch",
+            "type": "local",
+            "dataset_type": "image",
+            "train_batch_size": 17,
+            "instance_data_dir": self.temp_dir,
+        }
+
+        with patch(
+            "simpletuner.helpers.configuration.platform_validation.torch.backends.mps.is_available", return_value=False
+        ):
+            result = init_backend_config(backend, self.args, self.accelerator)
+
+        self.assertEqual(result["config"]["train_batch_size"], 17)
+
     def test_eval_backend_config_forces_train_batch_size_one(self):
         from simpletuner.helpers.data_backend.factory import init_backend_config
 
@@ -461,11 +526,14 @@ def test_eval_backend_config_forces_train_batch_size_one(self):
             "id": "eval-custom-batch",
             "type": "local",
             "dataset_type": "eval",
-            "train_batch_size": 3,
+            "train_batch_size": 17,
             "instance_data_dir": self.temp_dir,
         }
 
-        result = init_backend_config(backend, self.args, self.accelerator)
+        with patch(
+            "simpletuner.helpers.configuration.platform_validation.torch.backends.mps.is_available", return_value=True
+        ):
+            result = init_backend_config(backend, self.args, self.accelerator)
 
         self.assertEqual(result["config"]["train_batch_size"], 1)