fix(preflight): use physical batch for VRAM check

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

Author: binaryaaron

src/nemo_safe_synthesizer/preflight/checks/environment.py

@@ -270,9 +270,9 @@ def activation_memory_gib(
 ) -> float:
     r"""Rough activation VRAM on one device given micro-batch geometry.
 
-    Uses ``training.batch_size`` (HF ``per_device_train_batch_size``), not
-    ``gradient_accumulation_steps``. Matches bf16-ish training tensors at
-    2 bytes/element:
+    Uses the resolved physical microbatch (HF
+    ``per_device_train_batch_size``), not the logical batch or accumulation
+    count. Matches bf16-ish training tensors at 2 bytes/element:
 
     \[
         M_\text{act} \approx B \cdot S \cdot H \cdot L \cdot 2\text{ bytes}
@@ -397,11 +397,13 @@ def check(self, ctx: MetadataView, collector: IssueCollector) -> None:
                 n_params / 1e9,
             )
 
+        batching = config.training.resolve_batching()
+        per_device_batch_size = batching.per_device_train_batch_size
         seq_len = getattr(ctx.metadata, "max_seq_length", None)
         comp = estimate_training_vram_components(
             n_params=n_params,
             training_cfg=config.training,
-            batch_size=config.training.batch_size,
+            batch_size=per_device_batch_size,
             seq_len=seq_len,
             hidden_size=getattr(autoconfig, "hidden_size", None),
             num_hidden_layers=getattr(autoconfig, "num_hidden_layers", None),
@@ -429,7 +431,7 @@ def check(self, ctx: MetadataView, collector: IssueCollector) -> None:
                         f"~{comp.overhead_gib:.1f} GiB reserved){qualifier} "
                         f"exceeds available ~{max_free_gib:.1f} GiB "
                         f"(training.max_vram_fraction={config.training.max_vram_fraction:.2g}). "
-                        f"Per-device batch_size={config.training.batch_size}. {oom_risk}. "
+                        f"Per-device batch_size={per_device_batch_size}. {oom_risk}. "
                         "This remains an estimate -- attention blocks, adapters, optimizer state, and "
                         "checkpointing materially affect real usage."
                     ),

tests/preflight/test_preflight.py

@@ -126,7 +126,12 @@ def _metadata(autoconfig=None):
         return md
 
     @staticmethod
-    def _estimated_vram_gib(config: SafeSynthesizerParameters, autoconfig: PretrainedConfig) -> float:
+    def _estimated_vram_gib(
+        config: SafeSynthesizerParameters,
+        autoconfig: PretrainedConfig,
+        *,
+        batch_size: int | None = None,
+    ) -> float:
         from nemo_safe_synthesizer.preflight.checks.environment import (
             estimate_base_model_params,
             estimate_training_vram_components,
@@ -138,7 +143,7 @@ def _estimated_vram_gib(config: SafeSynthesizerParameters, autoconfig: Pretraine
         comp = estimate_training_vram_components(
             n_params=n_params,
             training_cfg=config.training,
-            batch_size=config.training.batch_size,
+            batch_size=batch_size if batch_size is not None else config.training.batch_size,
             seq_len=DEFAULT_MAX_SEQ_LENGTH,
             hidden_size=getattr(autoconfig, "hidden_size", None),
             num_hidden_layers=getattr(autoconfig, "num_hidden_layers", None),
@@ -175,6 +180,49 @@ def test_ample_vram_is_silent(self, default_config):
         assert not any(i.code == "low_vram" for i in issues)
         assert not any(i.code == "vram_exceeds_capacity" for i in issues)
 
+    def test_automatic_accumulation_uses_physical_batch_for_vram_estimate(self, default_config):
+        """A physical microbatch cap prevents a false VRAM issue for the logical batch."""
+        config = default_config.model_copy(deep=True)
+        config.training.batch_size = 8
+        config.training.gradient_accumulation_steps = "auto"
+        config.training.max_physical_batch_size = 4
+        autoconfig = self._autoconfig()
+        metadata = self._metadata(autoconfig=autoconfig)
+        batch_4_gib = self._estimated_vram_gib(config, autoconfig, batch_size=4)
+        batch_8_gib = self._estimated_vram_gib(config, autoconfig, batch_size=8)
+        available_gib = (batch_4_gib + batch_8_gib) / 2
+        fake_props = MagicMock(total_memory=100 * 1024**3)
+
+        with (
+            patch("torch.cuda.is_available", return_value=True),
+            patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: available_gib / 100}),
+            patch("torch.cuda.get_device_properties", return_value=fake_props),
+        ):
+            issues = VRAMHeadroomCheck().run(make_ctx(config=config, metadata=metadata))
+
+        assert not any(issue.code in {"low_vram", "vram_exceeds_capacity"} for issue in issues)
+
+    def test_automatic_accumulation_reports_physical_batch_when_vram_is_low(self, default_config):
+        """The VRAM diagnostic reports the capped physical per-device batch."""
+        config = default_config.model_copy(deep=True)
+        config.training.batch_size = 8
+        config.training.gradient_accumulation_steps = "auto"
+        config.training.max_physical_batch_size = 4
+        autoconfig = self._autoconfig()
+        metadata = self._metadata(autoconfig=autoconfig)
+        batch_4_gib = self._estimated_vram_gib(config, autoconfig, batch_size=4)
+        fake_props = MagicMock(total_memory=100 * 1024**3)
+
+        with (
+            patch("torch.cuda.is_available", return_value=True),
+            patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: batch_4_gib / 200}),
+            patch("torch.cuda.get_device_properties", return_value=fake_props),
+        ):
+            issues = VRAMHeadroomCheck().run(make_ctx(config=config, metadata=metadata))
+
+        issue = next(issue for issue in issues if issue.code in {"low_vram", "vram_exceeds_capacity"})
+        assert "Per-device batch_size=4" in issue.message
+
     def test_absurd_batch_errors(self, default_config):
         """Per-device batch_size far too large must fail preflight."""
         default_config.training.batch_size = 100_000