Merge pull request #3093 from bghira/agent/dataset-batch-resume-guard

Guard batch-size changes when resuming sampler state

Author: bghira

documentation/OPTIONS.es.md

@@ -1798,6 +1798,7 @@ Las opciones anteriores aplican en su mayor parte a `config.json`, pero algunas
 - `TRAINING_NUM_PROCESSES` debe configurarse al número de GPUs del sistema. Para la mayoría de los casos de uso, esto es suficiente para habilitar entrenamiento DistributedDataParallel (DDP). Usa `num_processes` dentro de `config.json` si prefieres no usar `config.env`.
 - `TRAINING_DYNAMO_BACKEND` por defecto es `no` pero puede configurarse a cualquier backend de torch.compile soportado (p. ej., `inductor`, `aot_eager`, `cudagraphs`) y combinarse con `--dynamo_mode`, `--dynamo_fullgraph` o `--dynamo_use_regional_compilation` para un ajuste más fino
 - `SIMPLETUNER_LOG_LEVEL` por defecto es `INFO` pero puede configurarse a `DEBUG` para añadir más información para reportes de problemas en `debug.log`
+- `SIMPLETUNER_ALLOW_MODIFYING_BSZ=1` omite la comprobación de discrepancia de batch_size del sampler al reanudar desde un checkpoint. Úsalo solo si has cambiado intencionalmente `train_batch_size` entre ejecuciones y aceptas el riesgo de estado inconsistente del sampler. Equivalente a `--i_know_what_i_am_doing`.
 - `VENV_PATH` puede configurarse a la ubicación de tu entorno virtual de python si no está en la ubicación típica `.venv`
 - `ACCELERATE_EXTRA_ARGS` puede dejarse sin configurar o contener argumentos extra como `--multi_gpu` o flags específicos de FSDP
 
@@ -2410,7 +2411,9 @@ options:
                         Source device used to generate validation seeds
   --i_know_what_i_am_doing [I_KNOW_WHAT_I_AM_DOING]
                         Unlock experimental overrides and bypass built-in
-                        safety limits.
+                        safety limits. Also bypasses the sampler batch-size
+                        mismatch check on resume (same effect as
+                        SIMPLETUNER_ALLOW_MODIFYING_BSZ=1).
   --flow_sigmoid_scale FLOW_SIGMOID_SCALE
                         Scale factor for sigmoid timestep sampling for flow-
                         matching models.

documentation/OPTIONS.hi.md

@@ -1796,6 +1796,7 @@ Upstream option mapping (LayerSync → SimpleTuner):
 - `TRAINING_NUM_PROCESSES` को सिस्टम में GPUs की संख्या पर सेट करें। अधिकांश उपयोग‑मामलों में इससे DistributedDataParallel (DDP) training सक्षम हो जाती है। यदि आप `config.env` उपयोग नहीं करना चाहते, तो `config.json` में `num_processes` उपयोग करें।
 - `TRAINING_DYNAMO_BACKEND` डिफ़ॉल्ट रूप से `no` है, लेकिन इसे किसी भी समर्थित torch.compile backend (उदा. `inductor`, `aot_eager`, `cudagraphs`) पर सेट किया जा सकता है और `--dynamo_mode`, `--dynamo_fullgraph`, या `--dynamo_use_regional_compilation` के साथ finer tuning के लिए जोड़ा जा सकता है
 - `SIMPLETUNER_LOG_LEVEL` डिफ़ॉल्ट रूप से `INFO` है, लेकिन issue reports के लिए `debug.log` में अधिक जानकारी जोड़ने हेतु इसे `DEBUG` पर सेट किया जा सकता है
+- `SIMPLETUNER_ALLOW_MODIFYING_BSZ=1` checkpoint से resume करते समय sampler batch_size मेल न खाने की जांच को छोड़ देता है। इसे केवल तभी उपयोग करें जब आपने जानबूझकर `train_batch_size` बदला हो और असंगत sampler state के जोखिम को स्वीकार करते हों। `--i_know_what_i_am_doing` के समकक्ष है।
 - `VENV_PATH` को आपके python virtual env की लोकेशन पर सेट किया जा सकता है यदि वह सामान्य `.venv` लोकेशन में नहीं है
 - `ACCELERATE_EXTRA_ARGS` को unset छोड़ा जा सकता है, या इसमें `--multi_gpu` या FSDP‑specific flags जैसे अतिरिक्त arguments जोड़े जा सकते हैं
 
@@ -2408,7 +2409,9 @@ options:
                         Source device used to generate validation seeds
   --i_know_what_i_am_doing [I_KNOW_WHAT_I_AM_DOING]
                         Unlock experimental overrides and bypass built-in
-                        safety limits.
+                        safety limits. Also bypasses the sampler batch-size
+                        mismatch check on resume (same effect as
+                        SIMPLETUNER_ALLOW_MODIFYING_BSZ=1).
   --flow_sigmoid_scale FLOW_SIGMOID_SCALE
                         Scale factor for sigmoid timestep sampling for flow-
                         matching models.

documentation/OPTIONS.ja.md

@@ -1799,6 +1799,7 @@ LayerSync は同一 Transformer 内の「学生」レイヤーを、より強い
 - `TRAINING_NUM_PROCESSES` はシステム内の GPU 数に設定します。ほとんどの用途で DistributedDataParallel（DDP）を有効化するには十分です。`config.env` を使いたくない場合は、`config.json` の `num_processes` を使用してください。
 - `TRAINING_DYNAMO_BACKEND` は既定で `no` ですが、対応する torch.compile バックエンド（例: `inductor`, `aot_eager`, `cudagraphs`）に設定でき、`--dynamo_mode`、`--dynamo_fullgraph`、`--dynamo_use_regional_compilation` と組み合わせて微調整できます。
 - `SIMPLETUNER_LOG_LEVEL` は既定で `INFO` ですが、`DEBUG` にすると `debug.log` に問題報告向けの詳細情報を追加できます。
+- `SIMPLETUNER_ALLOW_MODIFYING_BSZ=1` を設定すると、チェックポイントから再開する際のサンプラー batch_size 不一致チェックをスキップします。意図的に `train_batch_size` を変更し、サンプラー状態の不整合リスクを許容する場合のみ使用してください。`--i_know_what_i_am_doing` と同等です。
 - `VENV_PATH` は Python 仮想環境の場所を指定できます（標準の `.venv` 以外の場合）。
 - `ACCELERATE_EXTRA_ARGS` は未設定でも構いませんが、`--multi_gpu` や FSDP 特有のフラグなど追加引数を含められます。
 
@@ -2410,7 +2411,9 @@ options:
                         Source device used to generate validation seeds
   --i_know_what_i_am_doing [I_KNOW_WHAT_I_AM_DOING]
                         Unlock experimental overrides and bypass built-in
-                        safety limits.
+                        safety limits. Also bypasses the sampler batch-size
+                        mismatch check on resume (same effect as
+                        SIMPLETUNER_ALLOW_MODIFYING_BSZ=1).
   --flow_sigmoid_scale FLOW_SIGMOID_SCALE
                         Scale factor for sigmoid timestep sampling for flow-
                         matching models.

documentation/OPTIONS.md

@@ -1802,6 +1802,7 @@ The above options apply for the most part, to `config.json` - but some entries m
 - `TRAINING_NUM_PROCESSES` should be set to the number of GPUs in the system. For most use-cases, this is enough to enable DistributedDataParallel (DDP) training. Use `num_processes` inside `config.json` if you prefer to not use `config.env`.
 - `TRAINING_DYNAMO_BACKEND` defaults to `no` but can be set to any supported torch.compile backend (e.g. `inductor`, `aot_eager`, `cudagraphs`) and combined with `--dynamo_mode`, `--dynamo_fullgraph`, or `--dynamo_use_regional_compilation` for finer tuning
 - `SIMPLETUNER_LOG_LEVEL` defaults to `INFO` but can be set to `DEBUG` to add more information for issue reports into `debug.log`
+- `SIMPLETUNER_ALLOW_MODIFYING_BSZ=1` bypasses the sampler batch-size mismatch check when resuming from a checkpoint. Use this only if you intentionally changed `train_batch_size` between runs and accept the risk of inconsistent sampler state. Equivalent to passing `--i_know_what_i_am_doing`.
 - `VENV_PATH` can be set to the location of your python virtual env, if it is not in the typical `.venv` location
 - `ACCELERATE_EXTRA_ARGS` can be left unset, or, contain extra arguments to add like `--multi_gpu` or FSDP-specific flags
 
@@ -2413,7 +2414,9 @@ options:
                         Source device used to generate validation seeds
   --i_know_what_i_am_doing [I_KNOW_WHAT_I_AM_DOING]
                         Unlock experimental overrides and bypass built-in
-                        safety limits.
+                        safety limits. Also bypasses the sampler batch-size
+                        mismatch check on resume (same effect as
+                        SIMPLETUNER_ALLOW_MODIFYING_BSZ=1).
   --flow_sigmoid_scale FLOW_SIGMOID_SCALE
                         Scale factor for sigmoid timestep sampling for flow-
                         matching models.

documentation/OPTIONS.pt-BR.md

@@ -1794,6 +1794,7 @@ As opcoes acima se aplicam em grande parte ao `config.json` — mas algumas entr
 - `TRAINING_NUM_PROCESSES` deve ser definido para o numero de GPUs no sistema. Para a maioria dos casos, isso basta para habilitar treino DDP. Use `num_processes` no `config.json` se preferir nao usar `config.env`.
 - `TRAINING_DYNAMO_BACKEND` padrao e `no`, mas pode ser definido para qualquer backend suportado do torch.compile (ex.: `inductor`, `aot_eager`, `cudagraphs`) e combinado com `--dynamo_mode`, `--dynamo_fullgraph` ou `--dynamo_use_regional_compilation` para ajuste fino.
 - `SIMPLETUNER_LOG_LEVEL` padrao e `INFO`, mas pode ser definido para `DEBUG` para adicionar mais informacoes de issues no `debug.log`.
+- `SIMPLETUNER_ALLOW_MODIFYING_BSZ=1` ignora a verificação de incompatibilidade de batch_size do sampler ao retomar de um checkpoint. Use apenas se intencionalmente alterou `train_batch_size` entre execuções e aceita o risco de estado inconsistente do sampler. Equivalente a `--i_know_what_i_am_doing`.
 - `VENV_PATH` pode ser definido para o caminho do seu virtual env python, se nao estiver no local tipico `.venv`.
 - `ACCELERATE_EXTRA_ARGS` pode ficar vazio ou conter argumentos extras como `--multi_gpu` ou flags especificas do FSDP.
 
@@ -2405,7 +2406,9 @@ options:
                         Source device used to generate validation seeds
   --i_know_what_i_am_doing [I_KNOW_WHAT_I_AM_DOING]
                         Unlock experimental overrides and bypass built-in
-                        safety limits.
+                        safety limits. Also bypasses the sampler batch-size
+                        mismatch check on resume (same effect as
+                        SIMPLETUNER_ALLOW_MODIFYING_BSZ=1).
   --flow_sigmoid_scale FLOW_SIGMOID_SCALE
                         Scale factor for sigmoid timestep sampling for flow-
                         matching models.

documentation/OPTIONS.zh.md

@@ -1801,6 +1801,7 @@ LayerSync 通过在同一 Transformer 内让“学生”层对齐更强的“教
 - `TRAINING_NUM_PROCESSES` 应设置为系统中的 GPU 数量。多数场景下已足够启用 DistributedDataParallel（DDP）。如果不想使用 `config.env`，可在 `config.json` 中设置 `num_processes`。
 - `TRAINING_DYNAMO_BACKEND` 默认 `no`，但可设置为任意支持的 torch.compile 后端（例如 `inductor`, `aot_eager`, `cudagraphs`），并与 `--dynamo_mode`、`--dynamo_fullgraph` 或 `--dynamo_use_regional_compilation` 配合微调。
 - `SIMPLETUNER_LOG_LEVEL` 默认 `INFO`，设置为 `DEBUG` 可在 `debug.log` 中记录更多问题报告信息。
+- `SIMPLETUNER_ALLOW_MODIFYING_BSZ=1` 在从检查点恢复时跳过采样器 batch_size 不匹配检查。仅在您有意在不同训练轮次间更改 `train_batch_size` 并接受采样器状态不一致的风险时使用。等效于 `--i_know_what_i_am_doing`。
 - `VENV_PATH` 可设置为 Python 虚拟环境的位置（如果不在常见 `.venv` 目录）。
 - `ACCELERATE_EXTRA_ARGS` 可留空，或包含额外参数，例如 `--multi_gpu` 或 FSDP 专用标志。
 
@@ -2411,7 +2412,9 @@ options:
                         Source device used to generate validation seeds
   --i_know_what_i_am_doing [I_KNOW_WHAT_I_AM_DOING]
                         Unlock experimental overrides and bypass built-in
-                        safety limits.
+                        safety limits. Also bypasses the sampler batch-size
+                        mismatch check on resume (same effect as
+                        SIMPLETUNER_ALLOW_MODIFYING_BSZ=1).
   --flow_sigmoid_scale FLOW_SIGMOID_SCALE
                         Scale factor for sigmoid timestep sampling for flow-
                         matching models.

simpletuner/helpers/multiaspect/sampler.py

@@ -176,6 +176,24 @@ def load_states(self, state_path: str):
         except Exception as e:
             raise e
 
+        saved_batch_size = previous_state.get("batch_size")
+        if "batch_size" in previous_state and saved_batch_size != self.batch_size:
+            _allow = os.environ.get("SIMPLETUNER_ALLOW_MODIFYING_BSZ", "").lower() in (
+                "1",
+                "true",
+                "yes",
+            ) or getattr(StateTracker.get_args(), "i_know_what_i_am_doing", False)
+            if not _allow:
+                raise ValueError(
+                    f"Dataset '{self.id}' checkpoint batch_size={saved_batch_size} does not match "
+                    f"current batch_size={self.batch_size}. Resume with the same per-dataset train_batch_size, "
+                    f"or set --i_know_what_i_am_doing / SIMPLETUNER_ALLOW_MODIFYING_BSZ=1 to bypass."
+                )
+            self.logger.warning(
+                f"Dataset '{self.id}': resuming with batch_size={self.batch_size} but checkpoint recorded "
+                f"batch_size={saved_batch_size}. Proceeding because override is set."
+            )
+
         # Checkpoints contain the rank-local schedule. Restore it before seen
         # state so legacy boolean flags can be expanded to all occurrences.
         saved_schedule = previous_state.get("aspect_ratio_bucket_indices")

tests/test_sampler.py

@@ -270,10 +270,11 @@ def read_image(path):
         self.assertEqual(conditioning_sample.image_path(), "/conditioning/11.png")
         self.assertEqual(conditioning_sample.caption, "caption")
 
-    def test_load_states_restores_schedule_before_normalizing_legacy_seen_flags(self):
+    def test_load_states_with_matching_batch_size_restores_schedule_before_normalizing_legacy_seen_flags(self):
         self.sampler.state_manager.load_state.return_value = {
             "aspect_ratio_bucket_indices": {"1.0": ["same.jpg", "same.jpg", "other.jpg"]},
             "buckets": ["1.0"],
+            "batch_size": self.batch_size,
             "current_bucket": 0,
             "exhausted_buckets": ["old"],
             "seen_images": {"same.jpg": True, "other.jpg": False, "legacy.jpg": True},
@@ -296,7 +297,71 @@ def test_load_states_restores_schedule_before_normalizing_legacy_seen_flags(self
         self.assertEqual(self.metadata_backend.seen_images["other.jpg"], 0)
         self.assertTrue(self.metadata_backend.seen_images["legacy.jpg"])
 
-    def test_load_states_keeps_the_fresh_split_when_the_checkpoint_records_no_layout(self):
+    def test_load_states_rejects_batch_size_mismatch_before_mutation(self):
+        self.sampler.state_manager.load_state.return_value = {
+            "aspect_ratio_bucket_indices": {"saved": ["saved.jpg"]},
+            "buckets": ["saved"],
+            "batch_size": 3,
+            "current_bucket": "saved",
+            "exhausted_buckets": ["saved-exhausted"],
+            "seen_images": {"saved.jpg": 1},
+            "current_epoch": 9,
+            "dp_size": 1,
+            "dp_rank": 0,
+        }
+
+        self.metadata_backend.aspect_ratio_bucket_indices = {"fresh": ["fresh.jpg"]}
+        self.metadata_backend.seen_images = {"fresh.jpg": 1}
+        self.sampler._val_master_list = ["fresh.jpg"]
+        self.sampler.buckets = ["fresh"]
+        self.sampler.current_bucket = "fresh"
+        self.sampler.exhausted_buckets = ["fresh-exhausted"]
+        self.sampler.current_epoch = 7
+
+        with (
+            patch.dict(os.environ, {"SIMPLETUNER_ALLOW_MODIFYING_BSZ": ""}),
+            patch.object(StateTracker, "get_args", return_value=SimpleNamespace(i_know_what_i_am_doing=False)),
+            self.assertRaises(ValueError) as error,
+        ):
+            self.sampler.load_states(self.state_path)
+
+        self.assertEqual(
+            str(error.exception),
+            "Dataset 'foo' checkpoint batch_size=3 does not match current batch_size=2. "
+            "Resume with the same per-dataset train_batch_size, "
+            "or set --i_know_what_i_am_doing / SIMPLETUNER_ALLOW_MODIFYING_BSZ=1 to bypass.",
+        )
+        self.assertEqual(self.metadata_backend.aspect_ratio_bucket_indices, {"fresh": ["fresh.jpg"]})
+        self.assertEqual(self.metadata_backend.seen_images, {"fresh.jpg": 1})
+        self.assertEqual(self.sampler._val_master_list, ["fresh.jpg"])
+        self.assertEqual(self.sampler.buckets, ["fresh"])
+        self.assertEqual(self.sampler.current_bucket, "fresh")
+        self.assertEqual(self.sampler.exhausted_buckets, ["fresh-exhausted"])
+        self.assertEqual(self.sampler.current_epoch, 7)
+
+    def test_load_states_batch_size_mismatch_allowed_via_env_var(self):
+        self.sampler.state_manager.load_state.return_value = {
+            "aspect_ratio_bucket_indices": {"saved": ["saved.jpg"]},
+            "buckets": ["saved"],
+            "batch_size": 3,
+        }
+        self.metadata_backend.aspect_ratio_bucket_indices = {"fresh": ["fresh.jpg"]}
+
+        with patch.dict(os.environ, {"SIMPLETUNER_ALLOW_MODIFYING_BSZ": "1"}):
+            with patch.object(StateTracker, "get_args", return_value=SimpleNamespace(i_know_what_i_am_doing=False)):
+                self.sampler.load_states(self.state_path)
+
+    def test_load_states_batch_size_mismatch_allowed_via_i_know_flag(self):
+        self.sampler.state_manager.load_state.return_value = {
+            "aspect_ratio_bucket_indices": {"saved": ["saved.jpg"]},
+            "buckets": ["saved"],
+            "batch_size": 3,
+        }
+        self.metadata_backend.aspect_ratio_bucket_indices = {"fresh": ["fresh.jpg"]}
+        with patch.object(StateTracker, "get_args", return_value=SimpleNamespace(i_know_what_i_am_doing=True)):
+            self.sampler.load_states(self.state_path)
+
+    def test_load_states_legacy_state_without_batch_size_keeps_fresh_split_when_no_layout(self):
         # Checkpoints written before the layout was recorded cannot be attributed to a rank, so
         # the schedule is left alone. Seen state still loads.
         self.sampler.state_manager.load_state.return_value = {