diff --git a/docs/user-guide/running.md b/docs/user-guide/running.md index 9aa2398f6..07f84c42e 100644 --- a/docs/user-guide/running.md +++ b/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. diff --git a/docs/user-guide/troubleshooting.md b/docs/user-guide/troubleshooting.md index ce231f6d7..662da1e8e 100644 --- a/docs/user-guide/troubleshooting.md +++ b/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 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 | diff --git a/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py b/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py index 10eda9ac7..83ab6733e 100644 --- a/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py +++ b/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 diff --git a/src/nemo_safe_synthesizer/llm/utils.py b/src/nemo_safe_synthesizer/llm/utils.py index 8fc680d17..5ac6ed0ba 100644 --- a/src/nemo_safe_synthesizer/llm/utils.py +++ b/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" diff --git a/src/nemo_safe_synthesizer/preflight/checks/environment.py b/src/nemo_safe_synthesizer/preflight/checks/environment.py index 91daffbab..2f9d9ab6b 100644 --- a/src/nemo_safe_synthesizer/preflight/checks/environment.py +++ b/src/nemo_safe_synthesizer/preflight/checks/environment.py @@ -6,6 +6,7 @@ from __future__ import annotations import os +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Literal from urllib.parse import urlparse @@ -28,10 +29,12 @@ "CUDAAvailabilityCheck", "HFModelAvailabilityCheck", "InferenceModelCheck", + "VRAMComponentEstimate", "VRAMHeadroomCheck", "bytes_per_base_weight", "estimate_base_model_params", "estimate_params_from_shape", + "estimate_training_vram_components", "param_count_from_empty_model", ] @@ -185,16 +188,21 @@ def estimate_base_model_params(autoconfig: PretrainedConfig) -> tuple[int, Liter def bytes_per_base_weight(training_cfg: TrainingHyperparams) -> float: - r"""Return expected bytes/param for the base model given PEFT mode. + r"""Return expected bytes/param for the base model load mode. - NSS always trains via LoRA or QLoRA, so the base model's storage + NSS always trains via LoRA-style adapters, so the base model's storage precision dominates VRAM (LoRA adapter params, gradients, and optimizer state are comparatively negligible). - - QLoRA: \(\text{bits}/8 + 0.1\) to cover quant state (absmax / block - scales) and dequant workspace. Yields \(\approx 0.6\) for 4-bit, + Runtime quantization is controlled by ``training.quantize_model``. The + PEFT type string alone is not enough: with ``quantize_model=False`` the + base weights are loaded as bf16 even when ``peft_implementation`` is + configured as ``"QLORA"``. + + - Quantized load: \(\text{bits}/8 + 0.1\) to cover quant state (absmax / + block scales) and dequant workspace. Yields \(\approx 0.6\) for 4-bit, \(\approx 1.1\) for 8-bit. - - LoRA (unquantized): \(2\) bytes (bf16/fp16 base weights). + - Unquantized load: \(2\) bytes (bf16/fp16 base weights). References: - Hu, E. J. et al. "LoRA: Low-Rank Adaptation of Large Language @@ -206,7 +214,7 @@ def bytes_per_base_weight(training_cfg: TrainingHyperparams) -> float: scales; the \(+0.1\) term accounts for these scales and the dequant workspace. """ - if training_cfg.peft_implementation.upper() == "QLORA": + if training_cfg.quantize_model: # Prefer the explicit scheme if set; otherwise fall back to the legacy # bits-based field. Both routes yield bits/param for memory estimation. if training_cfg.quantization_scheme is not None: @@ -217,15 +225,11 @@ def bytes_per_base_weight(training_cfg: TrainingHyperparams) -> float: return 2.0 -_VRAM_FIXED_OVERHEAD_GIB = 2.0 -r"""CUDA kernels, activations (under gradient checkpointing), and runtime fudge. +_VRAM_LEGACY_OVERHEAD_GIB = 2.0 +r"""Legacy overhead when activation memory cannot be modeled. -Intentionally a single constant: activation memory scales like -\(\mathcal{O}(B \cdot S \cdot H \cdot \sqrt{L})\) (Megatron-style, where -\(B\) is batch, \(S\) sequence length, \(H\) hidden size, \(L\) layers), -which adds another axis of estimation error for marginal gain versus the -order-of-magnitude correction this heuristic is already making over the -previous \(6 H L\) formula. Tune if we see systematic false negatives. +Covers CUDA context, unchecked activations (under gradient checkpointing), +kernels, LoRA adapters, and optimizer footprint at a coarse level. References: - Korthikanti, V. et al. "Reducing Activation Recomputation in Large @@ -234,34 +238,123 @@ def bytes_per_base_weight(training_cfg: TrainingHyperparams) -> float: """ +_VRAM_KERNEL_RESERVED_GIB = 0.5 +"""Small reservation when activation memory is modeled explicitly (kernels, graphs).""" + +_VRAM_HARD_FAIL_RATIO = 1.5 +"""Error instead of warning when the estimate is this many times available VRAM.""" + + +@dataclass(frozen=True) +class VRAMComponentEstimate: + """Per-device training VRAM components for ``gpu.vram`` preflight.""" + + base_weights_gib: float + overhead_gib: float + activation_gib: float | None + total_gib: float + + +def _positive_int_scalar(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int): + return None + return value if value > 0 else None -class VRAMHeadroomCheck(MetadataCheck): - r"""Estimate whether GPU VRAM is sufficient for training. - The estimate is intentionally a *lower bound*: +def activation_memory_gib( + *, + batch_size: int, + seq_len: int, + hidden_size: int, + num_hidden_layers: int, + bytes_per_activation_element: float = 2.0, +) -> 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: \[ - \text{VRAM}_\text{est} = N \cdot b + C + M_\text{act} \approx B \cdot S \cdot H \cdot L \cdot 2\text{ bytes} \] - where \(N\) is the base-model parameter count (see - [estimate_base_model_params][nemo_safe_synthesizer.preflight.checks.environment.estimate_base_model_params]; - exact via the meta-tensor path, or the shape-heuristic fallback), - \(b\) is the bytes-per-param for the selected PEFT mode (see - [bytes_per_base_weight][nemo_safe_synthesizer.preflight.checks.environment.bytes_per_base_weight]), - and \(C\) is a fixed overhead for CUDA kernels and checkpointed - activations. The expression excludes the fine-grained activation - term \(\mathcal{O}(B \cdot S \cdot H \cdot L)\), LoRA adapter - parameters, gradients, and optimizer state. Those are typically - small compared to the base weights for parameter-efficient - fine-tuning, but not zero. Passing this check does not guarantee - training will fit in VRAM; failing it is a strong signal that it - will OOM. + Omit attention \(O(B S^2)\) blocks and recomputation specifics; goal is + order-of-magnitude headroom versus absurd ``batch_size`` values. + + References: + - Korthikanti, V. et al. (2022) -- recomputation vs stored activations. + + """ + nbytes = batch_size * seq_len * hidden_size * num_hidden_layers * bytes_per_activation_element + return nbytes / (1024**3) + + +def estimate_training_vram_components( + *, + n_params: int, + training_cfg: TrainingHyperparams, + batch_size: int, + seq_len: int | None, + hidden_size: int | None, + num_hidden_layers: int | None, + bytes_per_activation_element: float = 2.0, +) -> VRAMComponentEstimate: + """Compose base weights, overhead, and optional activation estimate (GiB).""" + bpw = bytes_per_base_weight(training_cfg) + base_weights_gib = (n_params * bpw) / (1024**3) + + b_sz = _positive_int_scalar(batch_size) + seq = _positive_int_scalar(seq_len) if seq_len is not None else None + h_sz = _positive_int_scalar(hidden_size) if hidden_size is not None else None + n_layers = _positive_int_scalar(num_hidden_layers) if num_hidden_layers is not None else None + + activation_gib: float | None + overhead_gib: float + if b_sz is not None and seq is not None and h_sz is not None and n_layers is not None: + activation_gib = activation_memory_gib( + batch_size=b_sz, + seq_len=seq, + hidden_size=h_sz, + num_hidden_layers=n_layers, + bytes_per_activation_element=bytes_per_activation_element, + ) + overhead_gib = _VRAM_KERNEL_RESERVED_GIB + else: + activation_gib = None + overhead_gib = _VRAM_LEGACY_OVERHEAD_GIB + + total_gib = base_weights_gib + overhead_gib + (activation_gib if activation_gib is not None else 0.0) + + return VRAMComponentEstimate( + base_weights_gib=base_weights_gib, + overhead_gib=overhead_gib, + activation_gib=activation_gib, + total_gib=total_gib, + ) + + +class VRAMHeadroomCheck(MetadataCheck): + r"""Estimate whether GPU VRAM is sufficient for training. + + The estimate is intentionally *conservative/heuristic*, not worst-case-accurate. + + Parameter counts come from ``estimate_base_model_params`` via meta tensors + when possible. + + Activation memory uses ``estimate_training_vram_components`` when + ``metadata.max_seq_length`` and transformer shape fields resolve to + positive integers; missing inputs leave activations unspecified and revert + to a legacy lumped overhead. Per-device VRAM compares against + ``get_max_vram(max_vram_fraction=training.max_vram_fraction)`` headroom. + + LoRA adapters, full optimizer footprint, \(O(B S^2)\) attention material, + and quantization workspace are partially covered only by residual overhead -- + passing does not guarantee a fit; failing is a strong signal of OOM risk. References: - - EleutherAI, "Transformer Math 101" -- grounds the rule of thumb - that inference adds ~20% over raw weights; training adds - considerably more. + - EleutherAI, "Transformer Math 101". + """ name = "gpu.vram" @@ -275,7 +368,7 @@ def check(self, ctx: MetadataView, collector: IssueCollector) -> None: from ...llm.utils import get_max_vram config = ctx.config - vram_map = get_max_vram() + vram_map = get_max_vram(max_vram_fraction=config.training.max_vram_fraction) # ModelMetadata.autoconfig is populated by ``from_config`` but set to # ``None`` by ``ModelMetadata.stub`` (used on ``--validate`` when the # model is not cached / no network). Skip in that case: without the @@ -304,26 +397,53 @@ def check(self, ctx: MetadataView, collector: IssueCollector) -> None: n_params / 1e9, ) - bytes_per_param = bytes_per_base_weight(config.training) - # VRAM_est = N·b + C (EleutherAI, "Transformer Math 101" https://blog.eleuther.ai/transformer-math/) - # N·b: base-weight bytes; C: fixed overhead for kernels + checkpointed activations - estimated_gib = (n_params * bytes_per_param) / (1024**3) + _VRAM_FIXED_OVERHEAD_GIB + 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, + seq_len=seq_len, + hidden_size=getattr(autoconfig, "hidden_size", None), + num_hidden_layers=getattr(autoconfig, "num_hidden_layers", None), + ) + estimated_gib = comp.total_gib # frac is the usable fraction of device memory (gpu_memory_utilization); # take the best single-device headroom — multi-GPU tensor-parallel splits are not modelled here max_free_gib = max( frac * torch.cuda.get_device_properties(dev).total_memory / (1024**3) for dev, frac in vram_map.items() ) if max_free_gib < estimated_gib: - qualifier = "" if method == "exact" else " (approximate; shape-heuristic fallback)" - collector.warning( - "low_vram", - ( - f"Estimated required VRAM (~{estimated_gib:.1f} GiB){qualifier} " - f"exceeds available ~{max_free_gib:.1f} GiB. " - "Training may OOM. This is an estimate -- actual usage depends on " - "batch size, sequence length, and activation checkpointing." - ), - ) + qualifier = "" if method == "exact" else " (approximate base param count; shape-heuristic fallback)" + ratio = estimated_gib / max_free_gib if max_free_gib > 0 else float("inf") + hard_fail = ratio >= _VRAM_HARD_FAIL_RATIO + report_issue = collector.error if hard_fail else collector.warning + issue_code = "vram_exceeds_capacity" if hard_fail else "low_vram" + oom_risk = "Training is expected to OOM" if hard_fail else "Training may OOM" + if comp.activation_gib is not None: + report_issue( + issue_code, + ( + f"Estimated required VRAM ~{estimated_gib:.1f} GiB total" + f" (~{comp.base_weights_gib:.1f} GiB base weights, " + f"~{comp.activation_gib:.1f} GiB bf16 compute activations, " + 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}. " + "This remains an estimate -- attention blocks, adapters, optimizer state, and " + "checkpointing materially affect real usage." + ), + ) + else: + report_issue( + issue_code, + ( + f"Estimated required VRAM (~{estimated_gib:.1f} GiB){qualifier} " + f"exceeds available ~{max_free_gib:.1f} GiB. " + f"{oom_risk}. This is an estimate -- actual usage depends on " + "batch size, sequence length, activation checkpointing, and quantization." + ), + ) def _is_blank(value: str | None) -> bool: diff --git a/src/nemo_safe_synthesizer/training/huggingface_backend.py b/src/nemo_safe_synthesizer/training/huggingface_backend.py index 29b0b016a..e3573a446 100644 --- a/src/nemo_safe_synthesizer/training/huggingface_backend.py +++ b/src/nemo_safe_synthesizer/training/huggingface_backend.py @@ -55,7 +55,7 @@ add_bos_eos_tokens_to_tokenizer, cleanup_memory, get_device_map, - get_max_vram, + get_max_memory_map, get_quantization_config, load_fast_tokenizer, ) @@ -321,9 +321,15 @@ def prepare_config(self, add_max_memory: bool = True, **kwargs: Any) -> None: logger.info(f"preparing parameters for HF Automodel with model: {self.params.training.pretrained_model}") model_kwargs = self._filter_model_kwargs(kwargs) + # Pop unconditionally: max_vram_fraction is an NSS-internal kwarg that + # transformers does not accept, so it must not leak into + # framework_load_params even when add_max_memory is False. + frac = model_kwargs.pop("max_vram_fraction", None) if add_max_memory: - model_kwargs["max_memory"] = get_max_vram(max_vram_fraction=model_kwargs.pop("max_vram_fraction", None)) + if frac is None: + frac = self.params.training.max_vram_fraction + model_kwargs["max_memory"] = get_max_memory_map(max_vram_fraction=frac) framework_params = self._build_base_framework_params(model_kwargs) quant_config = self._get_quantization_config_if_enabled() diff --git a/tests/configurator/test_pydantic_click_options.py b/tests/configurator/test_pydantic_click_options.py index 2e4832255..10b8162e8 100644 --- a/tests/configurator/test_pydantic_click_options.py +++ b/tests/configurator/test_pydantic_click_options.py @@ -186,6 +186,20 @@ def test_click_type_literal_non_auto_int_returns_string(): assert _click_type(Literal["disabled"] | int) == click.STRING +def test_click_type_int_literal_returns_int(): + """Literal[4, 8] must parse CLI values as ints before Pydantic validation.""" + from typing import Literal + + assert _click_type(Literal[4, 8]) == click.INT + + +def test_click_type_string_literal_plus_int_literal_returns_string(): + """String literals still need STRING so Pydantic can validate sentinel values.""" + from typing import Literal + + assert _click_type(Literal["disabled", 4]) == click.STRING + + def test_click_type_literal_auto_plus_other_int_returns_string(): """Literal['auto', 'manual'] | int must NOT use AutoParamType (multiple sentinels).""" from typing import Literal @@ -438,6 +452,20 @@ def cmd(**kwargs): assert captured["training"]["batch_size"] == 4 +def test_literal_int_override_end_to_end_via_click_runner(): + """A numeric Literal option must reach Pydantic as an int, not a string.""" + captured: dict = {} + + @pydantic_options(SafeSynthesizerParameters, field_separator="__") + @click.command() + def cmd(**kwargs): + captured.update(parse_overrides(kwargs)) + + result = CliRunner().invoke(cmd, ["--training__quantization_bits", "4"]) + assert result.exit_code == 0, result.output + assert captured["training"]["quantization_bits"] == 4 + + def test_deep_nested_override_end_to_end_via_click_runner(): """A deeply nested option (3+ segments) flows through decorator + parse_overrides.""" captured: dict = {} diff --git a/tests/llm/test_utils.py b/tests/llm/test_utils.py index b8edf88d4..05c6b3b90 100644 --- a/tests/llm/test_utils.py +++ b/tests/llm/test_utils.py @@ -10,6 +10,8 @@ from nemo_safe_synthesizer.llm.utils import ( ModelRef, + get_max_memory_map, + get_max_vram, get_quantization_config, load_fast_tokenizer, trust_remote_code_for_model, @@ -69,6 +71,17 @@ def test_trust_remote_code_for_model_requires_configured_cache_root(tmp_path: Pa assert trust_remote_code_for_model(spoofed_snapshot, cache_root=cache_root) is False +def test_vram_helpers_return_fraction_and_hf_memory_map() -> None: + gib = 1024**3 + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.device_count", return_value=1), + patch("torch.cuda.mem_get_info", return_value=(10 * gib, 16 * gib)), + ): + assert get_max_vram(max_vram_fraction=0.8) == {0: 0.5} + assert get_max_memory_map(max_vram_fraction=0.8) == {0: 8 * gib} + + def test_model_ref_trusts_snapshot_under_configured_cache_root(tmp_path: Path, hf_cached_snapshot_factory) -> None: """HF cache path recognition intentionally tracks Hub snapshot metadata.""" cache_root = tmp_path / "custom-cache" diff --git a/tests/preflight/test_preflight.py b/tests/preflight/test_preflight.py index ef1497912..143157c7e 100644 --- a/tests/preflight/test_preflight.py +++ b/tests/preflight/test_preflight.py @@ -7,18 +7,19 @@ from io import StringIO from pathlib import Path from types import SimpleNamespace +from typing import cast from unittest.mock import MagicMock, patch import pandas as pd import pytest from rich.console import Console -from transformers import PreTrainedTokenizerBase +from transformers import PretrainedConfig, PreTrainedTokenizerBase from nemo_safe_synthesizer.config.data import DataParameters from nemo_safe_synthesizer.config.parameters import SafeSynthesizerParameters from nemo_safe_synthesizer.config.time_series import TimeSeriesParameters from nemo_safe_synthesizer.config.training import TrainingHyperparams -from nemo_safe_synthesizer.defaults import PSEUDO_GROUP_COLUMN +from nemo_safe_synthesizer.defaults import DEFAULT_MAX_SEQ_LENGTH, PSEUDO_GROUP_COLUMN from nemo_safe_synthesizer.llm.metadata import ModelMetadata from nemo_safe_synthesizer.llm.utils import ModelRef from nemo_safe_synthesizer.preflight import ( @@ -104,9 +105,7 @@ def test_missing_torch_emits_torch_missing_not_no_gpu(self, default_config): @pytest.mark.unit class TestVRAMHeadroomCheck: @staticmethod - def _autoconfig(**overrides): - from types import SimpleNamespace - + def _autoconfig(**overrides) -> PretrainedConfig: defaults = dict( hidden_size=4096, num_hidden_layers=32, @@ -118,12 +117,38 @@ def _autoconfig(**overrides): tie_word_embeddings=False, ) defaults.update(overrides) - return SimpleNamespace(**defaults) + return cast(PretrainedConfig, SimpleNamespace(**defaults)) + + @staticmethod + def _metadata(autoconfig=None): + md = MagicMock(spec=ModelMetadata, autoconfig=autoconfig) + md.max_seq_length = DEFAULT_MAX_SEQ_LENGTH + return md + + @staticmethod + def _estimated_vram_gib(config: SafeSynthesizerParameters, autoconfig: PretrainedConfig) -> float: + from nemo_safe_synthesizer.preflight.checks.environment import ( + estimate_base_model_params, + estimate_training_vram_components, + ) + + result = estimate_base_model_params(autoconfig) + assert result is not None + n_params, _method = result + comp = estimate_training_vram_components( + n_params=n_params, + training_cfg=config.training, + batch_size=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), + ) + return comp.total_gib def test_low_vram_warns(self, default_config): - """Llama-3-8B-shaped config on a 1 GiB GPU must warn.""" - metadata = MagicMock(spec=ModelMetadata, autoconfig=self._autoconfig()) - fake_props = MagicMock(total_memory=1 * 1024**3) + """A marginally oversized model emits a warning, not a hard error.""" + metadata = self._metadata(autoconfig=self._autoconfig()) + fake_props = MagicMock(total_memory=16 * 1024**3) with ( patch("torch.cuda.is_available", return_value=True), patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: 1.0}), @@ -133,8 +158,13 @@ def test_low_vram_warns(self, default_config): assert any(i.code == "low_vram" and i.severity == "warning" for i in issues) def test_ample_vram_is_silent(self, default_config): - """Same config on an 80 GiB GPU must not warn (QLoRA ~5 GiB base + overhead).""" - metadata = MagicMock(spec=ModelMetadata, autoconfig=self._autoconfig()) + """Same config on an 80 GiB GPU must not warn (bf16 ~15 GiB base + overhead). + + The default config has quantize_model=False, so bytes_per_base_weight is + 2.0 (bf16) and the ~8B base weights estimate ~15 GiB; 80 GiB leaves ample + headroom for both the low_vram warning and the hard-fail threshold. + """ + metadata = self._metadata(autoconfig=self._autoconfig()) fake_props = MagicMock(total_memory=80 * 1024**3) with ( patch("torch.cuda.is_available", return_value=True), @@ -143,6 +173,89 @@ def test_ample_vram_is_silent(self, default_config): ): issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) 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_absurd_batch_errors(self, default_config): + """Per-device batch_size far too large must fail preflight.""" + default_config.training.batch_size = 100_000 + metadata = self._metadata(autoconfig=self._autoconfig()) + fake_props = MagicMock(total_memory=80 * 1024**3) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: 1.0}), + patch("torch.cuda.get_device_properties", return_value=fake_props), + ): + issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) + assert any(i.code == "vram_exceeds_capacity" and i.severity == "error" for i in issues) + + @pytest.mark.parametrize( + ("ratio", "expected_code", "expected_severity"), + [ + pytest.param(1.49, "low_vram", "warning", id="below-hard-fail-ratio"), + pytest.param(1.50, "vram_exceeds_capacity", "error", id="at-hard-fail-ratio"), + ], + ) + def test_vram_hard_fail_ratio_boundary(self, default_config, ratio, expected_code, expected_severity): + """The 1.5x hard-fail threshold is inclusive.""" + autoconfig = self._autoconfig() + metadata = self._metadata(autoconfig=autoconfig) + estimated_gib = self._estimated_vram_gib(default_config, autoconfig) + fake_props = MagicMock(total_memory=100 * 1024**3) + available_fraction = estimated_gib / (ratio * 100) + opposite_code = "vram_exceeds_capacity" if expected_code == "low_vram" else "low_vram" + with ( + patch("torch.cuda.is_available", return_value=True), + patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: available_fraction}), + patch("torch.cuda.get_device_properties", return_value=fake_props), + ): + issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) + assert any(i.code == expected_code and i.severity == expected_severity for i in issues) + assert not any(i.code == opposite_code for i in issues) + + @pytest.mark.parametrize("quantization_bits", (4, 8)) + def test_quantized_load_can_reduce_full_vram_check_below_warning(self, default_config, quantization_bits): + """Runtime quantization settings affect the complete VRAM preflight path.""" + metadata = self._metadata(autoconfig=self._autoconfig()) + fake_props = MagicMock(total_memory=14 * 1024**3) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: 1.0}), + patch("torch.cuda.get_device_properties", return_value=fake_props), + ): + default_config.training.quantize_model = True + default_config.training.quantization_bits = quantization_bits + quantized_issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) + default_config.training.quantize_model = False + unquantized_issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) + + assert not any(i.code in {"low_vram", "vram_exceeds_capacity"} for i in quantized_issues) + assert any(i.code in {"low_vram", "vram_exceeds_capacity"} for i in unquantized_issues) + + def test_gradient_accumulation_steps_does_not_fake_batch(self, default_config): + """Larger gradient_accumulation_steps alone should not match absurd batch activation.""" + metadata = self._metadata(autoconfig=self._autoconfig()) + default_config.training.gradient_accumulation_steps = 10_000 + fake_props = MagicMock(total_memory=80 * 1024**3) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: 1.0}), + patch("torch.cuda.get_device_properties", return_value=fake_props), + ): + issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) + assert not any(i.code == "low_vram" for i in issues) + + @pytest.mark.parametrize("fraction", (0.80, 0.50)) + def test_get_max_vram_receives_max_vram_fraction(self, default_config, fraction): + default_config.training.max_vram_fraction = fraction + metadata = self._metadata(autoconfig=self._autoconfig()) + fake_props = MagicMock(total_memory=80 * 1024**3) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: 1.0}) as mock_gmv, + patch("torch.cuda.get_device_properties", return_value=fake_props), + ): + VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) + mock_gmv.assert_called_once_with(max_vram_fraction=fraction) def test_stub_metadata_skips_silently(self, default_config): """Stubbed metadata (autoconfig=None) must skip without raising.""" @@ -154,8 +267,8 @@ def test_stub_metadata_skips_silently(self, default_config): issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) assert issues == [] - def test_lora_uses_more_bytes_per_param_than_qlora(self, default_config): - """LoRA (bf16 base) must estimate > QLoRA (4-bit base) on the same shape.""" + def test_unquantized_load_uses_more_bytes_per_param_than_quantized_load(self, default_config): + """bf16 base weights must estimate higher than a quantized base load.""" from nemo_safe_synthesizer.preflight.checks.environment import ( bytes_per_base_weight, estimate_base_model_params, @@ -168,12 +281,37 @@ def test_lora_uses_more_bytes_per_param_than_qlora(self, default_config): n, method = result assert method == "approximate" assert n > 7e9 # Llama-3-8B shape ~8B params - default_config.training.peft_implementation = "lora" - bpw_lora = bytes_per_base_weight(default_config.training) - default_config.training.peft_implementation = "QLORA" + default_config.training.quantize_model = False + bpw_unquantized = bytes_per_base_weight(default_config.training) + default_config.training.quantize_model = True default_config.training.quantization_bits = 4 - bpw_qlora = bytes_per_base_weight(default_config.training) - assert bpw_lora > bpw_qlora + bpw_quantized = bytes_per_base_weight(default_config.training) + assert bpw_unquantized > bpw_quantized + + def test_qlora_peft_label_without_quantize_model_uses_unquantized_base_weights(self, default_config): + """Only quantize_model controls k-bit base-weight memory.""" + from nemo_safe_synthesizer.preflight.checks.environment import bytes_per_base_weight + + default_config.training.peft_implementation = "QLORA" + default_config.training.quantize_model = False + + assert bytes_per_base_weight(default_config.training) == 2.0 + + def test_vram_component_estimate_falls_back_without_activation_shape(self, default_config): + from nemo_safe_synthesizer.preflight.checks.environment import estimate_training_vram_components + + comp = estimate_training_vram_components( + n_params=1_000_000_000, + training_cfg=default_config.training, + batch_size=default_config.training.batch_size, + seq_len=None, + hidden_size=4096, + num_hidden_layers=32, + ) + + assert comp.activation_gib is None + assert comp.overhead_gib == pytest.approx(2.0) + assert comp.total_gib == pytest.approx(comp.base_weights_gib + 2.0) @pytest.mark.parametrize( ("model_type", "fields", "expected_params_millions"), diff --git a/tests/training/test_huggingface_backend.py b/tests/training/test_huggingface_backend.py index 19a9a0335..6705221cd 100644 --- a/tests/training/test_huggingface_backend.py +++ b/tests/training/test_huggingface_backend.py @@ -779,19 +779,21 @@ def test_handles_single_loss(self): @patch("nemo_safe_synthesizer.training.huggingface_backend.get_device_map") -@patch("nemo_safe_synthesizer.training.huggingface_backend.get_max_vram") +@patch("nemo_safe_synthesizer.training.huggingface_backend.get_max_memory_map") @patch("nemo_safe_synthesizer.training.huggingface_backend.AutoConfig") class TestPrepareConfigIntegration: - def test_early_return_when_already_prepared(self, mock_autoconfig, mock_get_max_vram, mock_get_device_map, backend): + def test_early_return_when_already_prepared( + self, mock_autoconfig, mock_get_max_memory_map, mock_get_device_map, backend + ): """Test that prepare_config returns early when already prepared.""" backend.framework_load_params = {"already": "prepared"} backend.prepare_config() mock_autoconfig.from_pretrained.assert_not_called() - def test_prepares_config_correctly(self, mock_autoconfig, mock_get_max_vram, mock_get_device_map, backend): + def test_prepares_config_correctly(self, mock_autoconfig, mock_get_max_memory_map, mock_get_device_map, backend): """Test that prepare_config sets framework_load_params correctly.""" - mock_get_max_vram.return_value = {"cuda:0": "16GB"} + mock_get_max_memory_map.return_value = {0: 16 * 1024**3} mock_get_device_map.return_value = "auto" mock_config = MagicMock() mock_autoconfig.from_pretrained.return_value = mock_config @@ -800,8 +802,26 @@ def test_prepares_config_correctly(self, mock_autoconfig, mock_get_max_vram, moc backend.prepare_config() + mock_get_max_memory_map.assert_called_once_with(max_vram_fraction=backend.params.training.max_vram_fraction) assert backend.framework_load_params is not None assert backend.framework_load_params["pretrained_model_name_or_path"] == "test-model" + assert backend.framework_load_params["max_memory"] == {0: 16 * 1024**3} + + def test_prepare_config_max_vram_fraction_kwarg_overrides_config( + self, mock_autoconfig, mock_get_max_memory_map, mock_get_device_map, backend + ): + """Explicit model-load VRAM fraction overrides the configured default.""" + mock_get_max_memory_map.return_value = {0: 8 * 1024**3} + mock_get_device_map.return_value = "auto" + mock_autoconfig.from_pretrained.return_value = MagicMock() + backend.model_metadata.rope_scaling = None + + backend.prepare_config(max_vram_fraction=0.5) + + mock_get_max_memory_map.assert_called_once_with(max_vram_fraction=0.5) + assert backend.framework_load_params is not None + assert backend.framework_load_params["max_memory"] == {0: 8 * 1024**3} + assert "max_vram_fraction" not in backend.framework_load_params class TestPropagateMaxTokensPerExample: