feat: account for batch size in preflight

docs/user-guide/running.md

@@ -290,7 +290,7 @@ execute in order (`config` → `dataframe` → `metadata` → `advisory`).
 | `columns.pseudo` | dataframe | Input does not use the reserved `__nss_sequence_id` column name |
 | `columns.constant` | dataframe | No column is constant (warning only) |
 | `timeseries.timestamp` | dataframe | Timestamp column is present and has no nulls (time-series mode) |
-| `gpu.vram` | metadata | Free VRAM headroom for the chosen model + PEFT mode (warning only) |
+| `gpu.vram` | metadata | Free VRAM headroom for the chosen model, quantization load mode, and per-device batch size; emits `low_vram` as a warning and `vram_exceeds_capacity` as an error when the estimate is far above capacity |
 | `token_budget` | metadata | Schema prompt, sampled records, and top groups each fit in the model's context window |
 | `dataset.row_count` | advisory | Training split is above a comfort threshold (warning only) |
 | `training.oversampling` | advisory | Sampling fraction is not extreme (warning only) |
@@ -327,8 +327,11 @@ the best-effort caveat below.
       can be shorter or longer than the original and shift token budgets.
     - Token-budget checks sample rows and top groups -- a long-tail
       outlier outside the sample can still exceed the budget at assembly.
-    - VRAM headroom is a lower bound (LoRA adapters, optimizer state, and
-      activations are not modeled); full training may still OOM.
+    - VRAM headroom includes a coarse bf16 compute activation term
+      (`batch_size` x `metadata.max_seq_length` x width x depth) plus base
+      weights using the configured quantization mode; LoRA adapters,
+      optimizer state, and attention blocks beyond that shape are not modeled
+      tightly. Full training may still OOM.
 
     Treat `--validate` as a quick fail-fast gate, not a full-run guarantee.
 

docs/user-guide/troubleshooting.md

@@ -101,8 +101,8 @@ path. Common causes:
 
 - Local cached tokenizer is missing `tokenizer.json` — re-download
   with `huggingface-cli download <model>`
-- Model ships only a SentencePiece vocab (``tokenizer.model``) with no Rust
-  ``tokenizer.json`` — common on older checkpoints; fast conversion may land upstream.
+- Model ships only a SentencePiece vocab (`tokenizer.model`) with no Rust
+  `tokenizer.json` — common on older checkpoints; fast conversion may land upstream.
 - `trust_remote_code=True` model with a custom slow tokenizer class.
 
 The warning is informational. To suppress it, switch to a model with a
@@ -513,6 +513,7 @@ check of its own.
 | `torch_missing` | error | `gpu.cuda` | PyTorch not installed; cannot verify GPU availability |
 | `no_gpu` | error | `gpu.cuda` | No CUDA GPU detected (required for training or generation) |
 | `low_vram` | warning | `gpu.vram` | Free GPU VRAM may be insufficient |
+| `vram_exceeds_capacity` | error | `gpu.vram` | Estimated training VRAM is far above available GPU memory |
 | `inference_key_missing` | warning | `env.inference` | `NSS_INFERENCE_KEY` not set; PII classification degraded |
 | `inference_model_blank` | warning | `env.inference` | `NSS_INFERENCE_MODEL` set but empty; the blank value is ignored and the default model id is used |
 | `inference_endpoint_invalid` | error | `env.inference` | `NSS_INFERENCE_ENDPOINT` set but not a valid http(s) URL; classification requests will fail |

src/nemo_safe_synthesizer/configurator/pydantic_click_options.py

@@ -197,6 +197,17 @@ def _is_auto_only_literal_union(args: set) -> bool:
     return string_values == {AUTO_STR}
 
 
+def _literal_value_types(args: set) -> set:
+    """Replace non-string ``Literal`` annotations with their value types."""
+    normalized: set = set()
+    for arg in args:
+        if get_origin(arg) is Literal:
+            normalized.update(type(value) for value in get_args(arg))
+        else:
+            normalized.add(arg)
+    return normalized
+
+
 def _click_type(annotation: Any) -> click.ParamType:
     """Map a Pydantic field annotation to a Click type.
 
@@ -228,6 +239,7 @@ def _click_type(annotation: Any) -> click.ParamType:
                 if py_type in args:
                     return AutoParamType(click_type)
         return click.STRING
+    args = _literal_value_types(args)
     for py_type, click_type in _CLICK_TYPE_PRIORITY:
         if py_type in args:
             return click_type

src/nemo_safe_synthesizer/llm/utils.py

@@ -428,7 +428,15 @@ def round_gb(value: float) -> float:
     logger.info(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
 
 
-def get_max_vram(max_vram_fraction: float | None = None) -> dict[int, float]:
+@dataclass(frozen=True, slots=True)
+class _VRAMAllocation:
+    """Shared GPU memory calculation for runtime loaders."""
+
+    utilization: float
+    memory_bytes: int
+
+
+def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _VRAMAllocation]:
     """Calculate maximum memory allocation for each available GPU.
 
     Reserves a 2 GiB safety buffer on each device, then applies
@@ -439,27 +447,39 @@ def get_max_vram(max_vram_fraction: float | None = None) -> dict[int, float]:
             Defaults to ``0.8`` (80 %).
 
     Returns:
-        Mapping of CUDA device index to the usable memory fraction.
+        Mapping of CUDA device index to utilization and byte limit.
     """
     import torch
 
     if max_vram_fraction is None:
         max_vram_fraction = 0.8
-    max_memory = {}
+    allocations = {}
 
     if torch.cuda.is_available():
         num_gpus = torch.cuda.device_count()
         for i in range(num_gpus):
             free, total = torch.cuda.mem_get_info(device=i)
-            safe_free = free - (2 * 1024**3)
-            gpu_memory_utilization = min(max_vram_fraction, safe_free / total)
-            memory_gib = gpu_memory_utilization * total / (1024**3)
-            max_memory[i] = gpu_memory_utilization
+            safe_free = max(free - (2 * 1024**3), 0)
+            gpu_memory_utilization = min(max_vram_fraction, safe_free / total) if total > 0 else 0.0
+            memory_bytes = int(gpu_memory_utilization * total)
+            memory_gib = memory_bytes / (1024**3)
+            allocations[i] = _VRAMAllocation(utilization=gpu_memory_utilization, memory_bytes=memory_bytes)
             logger.info(
-                f"GPU {i}: Will allocate {memory_gib:.2f}GiB ({max_vram_fraction * 100}% of {total / (1024**3):.2f}GiB)"
+                f"GPU {i}: Will allocate {memory_gib:.2f}GiB "
+                f"({gpu_memory_utilization * 100:.1f}% of {total / (1024**3):.2f}GiB)"
             )
 
-    return max_memory
+    return allocations
+
+
+def get_max_vram(max_vram_fraction: float | None = None) -> dict[int, float]:
+    """Return vLLM-style GPU utilization fractions for each available GPU."""
+    return {device: allocation.utilization for device, allocation in _get_vram_allocations(max_vram_fraction).items()}
+
+
+def get_max_memory_map(max_vram_fraction: float | None = None) -> dict[int, int]:
+    """Return Hugging Face ``max_memory`` byte limits for each available GPU."""
+    return {device: allocation.memory_bytes for device, allocation in _get_vram_allocations(max_vram_fraction).items()}
 
 
 def add_bos_eos_tokens_to_tokenizer(tokenizer: PreTrainedTokenizer) -> PreTrainedTokenizer:
@@ -583,11 +603,15 @@ def load_fast_tokenizer(model_name_or_path: Path | str, **kwargs: Any) -> PreTra
 
 def get_device_name() -> str:
     """Get the name of the current device (first index). Returns 'undefined' if the device is not available."""
+    # torch may be absent (CPU-only install); CUDA/driver problems surface as
+    # RuntimeError/AssertionError from get_device_properties. Anything else is
+    # unexpected and should propagate rather than masquerade as 'undefined'.
     try:
         import torch
 
         return torch.cuda.get_device_properties(0).name
-    except Exception:
+    except (ImportError, RuntimeError, AssertionError):
+        logger.debug("Could not resolve CUDA device name; reporting 'undefined'.", exc_info=True)
         return "undefined"