feat: enable transformers v5 + add NVFP4/FP8/MXFP4 quantization schemes

docs/user-guide/configuration.md

@@ -75,7 +75,8 @@ for the full field list.
 | `training.lora_alpha_over_r` | `1.0` | LoRA scaling ratio (alpha / rank) | Leave at 1.0 |
 | `training.pretrained_model` | `"HuggingFaceTB/SmolLM3-3B"` | HuggingFace model ID or local path | See supported families below; `TinyLlama/TinyLlama-1.1B-Chat-v1.0` for fast CPU/low-VRAM iteration |
 | `training.quantize_model` | `false` | Enable quantization to reduce VRAM usage | Enable if VRAM is limited; 8-bit has lower quality impact than 4-bit |
-| `training.quantization_bits` | `8` | Bit width (4 or 8) when `training.quantize_model` is `true` | Prefer 8 over 4 for quality |
+| `training.quantization_scheme` | `null` | Quantization scheme: `bnb-4bit`, `bnb-8bit`, `fp8`, `nvfp4`, `mxfp4` | See [Quantization schemes](#quantization-schemes) below; leave unset to fall back to `quantization_bits` |
+| `training.quantization_bits` | `8` | Legacy bit width (4 or 8) used when `quantization_scheme` is unset | Prefer setting `quantization_scheme` explicitly for new configs |
 | `training.attn_implementation` | `"sdpa"` | Attention backend for model loading | Leave at default |
 | `training.rope_scaling_factor` | `"auto"` | Scale the base model's context window via RoPE (`"auto"` or int) | Leave at `"auto"` |
 | `training.validation_ratio` | `0.0` | Fraction of training data held out for validation loss monitoring | Leave at 0.0 unless you specifically want to monitor validation loss |
@@ -108,6 +109,38 @@ When `training.pretrained_model` is set to a Hugging Face Hub model ID, the mode
 !!! warning "Security Note: Pretrained models from Hugging Face Hub"
     Loading and using pretrained models from Hugging Face Hub (or any public source) can expose your environment to significant risks, including arbitrary code execution (ACE) or remote code execution (RCE) vulnerabilities. Only use models you have reviewed yourself or from organizations and authors you explicitly trust. Malicious or modified models may contain embedded code, backdoors, or privacy-leaking mechanisms.
 
+### Quantization schemes
+
+When `training.quantize_model: true`, Safe Synthesizer applies one of the
+following quantization schemes via transformers v5's
+`quantization_config=` interface:
+
+| Scheme | Backend | Bits/param | Hardware | Notes |
+|--------|---------|-----------:|----------|-------|
+| `bnb-4bit` | bitsandbytes NF4 + bf16 compute | 4 | Ampere+ (sm_80+) | Default for QLoRA fine-tuning. LoftQ-compatible. |
+| `bnb-8bit` | bitsandbytes int8 | 8 | Ampere+ (sm_80+) | Higher quality than 4-bit; uses ~2x the VRAM. LoftQ-compatible. |
+| `fp8` | `FineGrainedFP8Config` | 8 | Hopper (sm_90+) / Blackwell | Float8 with block-wise scaling. Inference-leaning. |
+| `nvfp4` | `TorchAoConfig(NVFP4WeightOnlyConfig())` | 4 | Blackwell (sm_100+) | Weight-only NVIDIA FP4. Requires recent torchao. |
+| `mxfp4` | `Mxfp4Config` | 4 | Varies by torch/torchao version | OCP Microscaling FP4. |
+
+Select a scheme with `training.quantization_scheme`:
+
+```yaml
+training:
+  quantize_model: true
+  quantization_scheme: nvfp4   # or bnb-4bit / bnb-8bit / fp8 / mxfp4
+```
+
+If `quantization_scheme` is unset, Safe Synthesizer falls back to
+`quantization_bits` for backward compatibility (`4` → `bnb-4bit`,
+`8` → `bnb-8bit`).
+
+!!! note "LoftQ + non-BNB schemes"
+    `peft_implementation: loftq` is incompatible with `fp8`, `nvfp4`, and
+    `mxfp4` — LoftQ requires the bitsandbytes runtime. Setting both will
+    raise a `ParameterError` at startup. Use `qlora` or `lora` with the
+    non-BNB schemes.
+
 ---
 
 ## Generation

docs/user-guide/troubleshooting.md

@@ -66,6 +66,49 @@ configuration reference.
 
 ## Installation
 
+### Transformers v5 + vLLM Version Selection
+
+`uv sync` fails with an error mentioning incompatible `transformers` and
+`vllm` requirements.
+
+Safe Synthesizer requires `transformers>=5.6,<6`. vLLM 0.20.0 accepts
+transformers v5, but excludes several early 5.x releases that are not
+compatible with its runtime. Keep vLLM's exclusions intact and resolve to a
+newer transformers v5 release.
+
+```toml
+[project]
+dependencies = [
+  "transformers>=5.6,<6",
+  "vllm==0.20.0",
+]
+```
+
+If you've vendored or copied parts of `pyproject.toml` into another project,
+avoid adding a broad `transformers>=5.0,<6` override for vLLM. That can erase
+vLLM's explicit exclusions and allow incompatible early v5 releases.
+
+### Slow Tokenizer Warning
+
+After upgrading to transformers v5 you may see a log line like:
+
+> Loaded slow (Python) tokenizer for `<model>` — no Rust backend available.
+
+This means the model's tokenizer has no Rust (`tokenizers` crate)
+implementation and v5 fell back to the SentencePiece/Python backend. Data
+prep continues to work but tokenization is ~5–10× slower than the fast
+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.
+- `trust_remote_code=True` model with a custom slow tokenizer class.
+
+The warning is informational. To suppress it, switch to a model with a
+fast tokenizer (most popular models do; check
+`AutoTokenizer.from_pretrained(model).is_fast`).
+
 ### Python 3.14 Is Not Supported
 
 Safe Synthesizer requires **Python 3.11, 3.12, or 3.13**. Python 3.14+ is not
@@ -97,11 +140,15 @@ Training OOM errors appear during the "Training" phase with HuggingFace Trainer
 stack traces. If you see `torch.cuda.OutOfMemoryError`:
 
 1. Enable 4-bit quantization -- the single largest memory saver. Set
-   `training.quantize_model: true` and `training.quantization_bits: 4`. QLoRA
-   stores the frozen base model in 4-bit NF4 while training LoRA adapters in
-   full precision, cutting model weight memory by ~4x. Quantization reduces
+   `training.quantize_model: true` and `training.quantization_scheme: bnb-4bit`
+   (or for back-compat, `training.quantization_bits: 4`). QLoRA stores the
+   frozen base model in 4-bit NF4 while training LoRA adapters in full
+   precision, cutting model weight memory by ~4x. Quantization reduces
    precision in the frozen weights; in practice QLoRA typically produces
-   results close to full-precision LoRA, but verify with your evaluation report
+   results close to full-precision LoRA, but verify with your evaluation
+   report. On Blackwell hardware, `nvfp4` and `mxfp4` schemes offer similar
+   4-bit footprint with hardware-accelerated matmul — see
+   [Quantization schemes](configuration.md#quantization-schemes)
 2. Reduce the context window -- see
    [Context Length and Record Fitting](#context-length-and-record-fitting) for
    how to lower `training.rope_scaling_factor`, truncate records, or simplify

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
 dependencies = [
   "faker>=20.0",
   "httpx>=0.27.0",
-  "huggingface-hub>=0.34.4,<1",
+  "huggingface-hub>=1.3.0,<2",
   "pandas>=2.1.3",
   "pydantic[email]>=2.12.5",
   "pydantic-settings>=2.6.1",
@@ -89,7 +89,7 @@ engine = [
     "dateparser",
     "faker",
     "datasets>=4.8.4",
-    "huggingface-hub>=0.34.4,<1",
+    "huggingface-hub>=1.3.0,<2",
     "json-repair",
     "matplotlib",
     "outlines>=1.0.0",
@@ -117,13 +117,13 @@ engine = [
 # and allow explicit installation of the cpu depset on linux.
 
 cpu = [
-  "accelerate",
-  "bitsandbytes==0.49.1",
+  "accelerate>=1.1.0",
+  "bitsandbytes>=0.46.1",
   "flashinfer-python==0.6.8.post1; sys_platform=='linux'",
   "flashinfer-cubin==0.6.8.post1; sys_platform=='linux'",
   "gliner",
   "kernels>=0.12.1",
-  "peft",
+  "peft>=0.18.0",
   "opacus",
   "sentence-transformers",
   "torch==2.11.0; sys_platform == 'darwin'",
@@ -133,15 +133,15 @@ cpu = [
   "torchvision==0.26.0; sys_platform == 'darwin'",
   "torchvision==0.26.0+cpu; sys_platform == 'linux'",
   "torchao==0.17.0",
-  "transformers==4.57.3",
+  "transformers>=5.6,<6",
   "triton>=2.0.0; sys_platform=='linux'",
   "trl>=0.23.0",
   "vllm==0.20.0; sys_platform=='linux'",
 ]
 
 cu129 = [
-  "accelerate",
-  "bitsandbytes==0.49.1",
+  "accelerate>=1.1.0",
+  "bitsandbytes>=0.46.1",
   "flashinfer-python==0.6.8.post1; sys_platform == 'linux'",
   "flashinfer-cubin==0.6.8.post1; sys_platform == 'linux'",
   "flashinfer-jit-cache==0.6.8.post1+cu129; sys_platform == 'linux'",
@@ -150,14 +150,14 @@ cu129 = [
   "nvidia-cublas-cu12; sys_platform == 'linux'",
   "nvidia-ml-py; sys_platform == 'linux'",
   "opacus",
-  "peft",
+  "peft>=0.18.0",
   "sentence-transformers",
   "torch==2.11.0+cu129; sys_platform == 'linux'",
   "torch-c-dlpack-ext",
   "torchaudio==2.11.0+cu129; sys_platform == 'linux'",
   "torchvision==0.26.0+cu129; sys_platform == 'linux'",
   "torchao==0.17.0+cu129; sys_platform == 'linux' and platform_machine == 'x86_64'",
-  "transformers==4.57.3",
+  "transformers>=5.6,<6",
   "triton>=2.0.0; sys_platform == 'linux'",
   "trl>=0.23.0",
   "vllm==0.20.0+cu129; sys_platform == 'linux'",

src/nemo_safe_synthesizer/config/training.py

@@ -3,7 +3,10 @@
 
 from __future__ import annotations
 
+import importlib
+from enum import StrEnum
 from typing import (
+    TYPE_CHECKING,
     Annotated,
     Literal,
 )
@@ -25,10 +28,126 @@
     OptionalAutoInt,
 )
 
+if TYPE_CHECKING:
+    from transformers.utils.quantization_config import QuantizationConfigMixin
+
 __all__ = [
+    "QuantizationScheme",
     "TrainingHyperparams",
 ]
 
+
+class QuantizationScheme(StrEnum):
+    """Quantization schemes supported when ``quantize_model=True``.
+
+    Members are string values so they serialize cleanly through pydantic
+    and JSON configs. The enum also owns construction of the corresponding
+    transformers ``quantization_config`` object; optional ML dependencies stay
+    locally imported in that construction path.
+
+    Selection guide:
+    - ``bnb-4bit`` / ``bnb-8bit``: bitsandbytes NF4 / int8. Widest hardware
+      support (Ampere+), works with QLoRA and LoftQ. Default for training.
+    - ``fp8``: transformers ``FineGrainedFP8Config``. Float8 with block-wise
+      scaling. Requires Hopper (sm_90+) or Blackwell. Inference-leaning.
+    - ``nvfp4``: NVIDIA FP4 via ``torchao.prototype.mx_formats.NVFP4WeightOnlyConfig``
+      wrapped in ``TorchAoConfig``. Requires Blackwell (sm_100+). Weight-only.
+    - ``mxfp4``: OCP Microscaling FP4 via transformers ``Mxfp4Config``.
+      Hardware support varies by torch/torchao version.
+    """
+
+    BNB_4BIT = "bnb-4bit"
+    BNB_8BIT = "bnb-8bit"
+    FP8 = "fp8"
+    NVFP4 = "nvfp4"
+    MXFP4 = "mxfp4"
+
+    @property
+    def effective_bits(self) -> int:
+        """Per-parameter bit width for memory estimation."""
+        return {
+            QuantizationScheme.BNB_4BIT: 4,
+            QuantizationScheme.BNB_8BIT: 8,
+            QuantizationScheme.FP8: 8,
+            QuantizationScheme.NVFP4: 4,
+            QuantizationScheme.MXFP4: 4,
+        }[self]
+
+    @property
+    def is_bitsandbytes(self) -> bool:
+        """Whether the scheme is implemented via bitsandbytes (QLoRA-compatible)."""
+        return self in (QuantizationScheme.BNB_4BIT, QuantizationScheme.BNB_8BIT)
+
+    @classmethod
+    def from_alias(cls, scheme: QuantizationScheme | str | Literal[4, 8]) -> QuantizationScheme:
+        """Normalize string and legacy bit-count aliases to a scheme."""
+        if isinstance(scheme, int):
+            legacy_aliases = {
+                4: cls.BNB_4BIT,
+                8: cls.BNB_8BIT,
+            }
+            try:
+                return legacy_aliases[scheme]
+            except KeyError as exc:
+                raise ValueError(f"Unknown quantization bit-count alias: {scheme!r}. Expected 4 or 8.") from exc
+        return cls(scheme)
+
+    def to_transformers_config(self) -> QuantizationConfigMixin:
+        """Build the transformers quantization config for this scheme."""
+        match self:
+            case QuantizationScheme.BNB_4BIT:
+                return self._bnb_4bit_config()
+            case QuantizationScheme.BNB_8BIT:
+                return self._bnb_8bit_config()
+            case QuantizationScheme.FP8:
+                return self._fp8_config()
+            case QuantizationScheme.NVFP4:
+                return self._nvfp4_config()
+            case QuantizationScheme.MXFP4:
+                return self._mxfp4_config()
+        raise ValueError(f"Unknown quantization scheme: {self!r}")
+
+    @staticmethod
+    def _bnb_4bit_config() -> QuantizationConfigMixin:
+        import torch
+        from transformers import BitsAndBytesConfig
+
+        return BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+
+    @staticmethod
+    def _bnb_8bit_config() -> QuantizationConfigMixin:
+        from transformers import BitsAndBytesConfig
+
+        return BitsAndBytesConfig(load_in_8bit=True)
+
+    @staticmethod
+    def _fp8_config() -> QuantizationConfigMixin:
+        from transformers import FineGrainedFP8Config
+
+        return FineGrainedFP8Config()
+
+    @staticmethod
+    def _nvfp4_config() -> QuantizationConfigMixin:
+        from transformers import TorchAoConfig
+
+        # Keep this dynamic because the experimental torchao module is runtime-only
+        # in some type-checker environments.
+        torchao_mx_formats = importlib.import_module("torchao.prototype.mx_formats")
+        nvfp4_weight_only_config = torchao_mx_formats.NVFP4WeightOnlyConfig
+        return TorchAoConfig(quant_type=nvfp4_weight_only_config())
+
+    @staticmethod
+    def _mxfp4_config() -> QuantizationConfigMixin:
+        from transformers.utils.quantization_config import Mxfp4Config
+
+        return Mxfp4Config()
+
+
 ValueGTZero = ValueValidator(lambda p: range_validator(p, lambda v: v >= 0))
 
 
@@ -218,10 +337,29 @@ class TrainingHyperparams(Parameters):
         Literal[4, 8],
         Field(
             title="quantization_bits",
-            description="The number of bits to use for quantization if ``quantize_model`` is ``True``. Accepts 8 or 4.",
+            deprecated=True,
+            description=(
+                "Deprecated: use ``quantization_scheme`` instead. Bit width for "
+                "bitsandbytes quantization when ``quantization_scheme`` is not set "
+                "(back-compat alias: 4 → bnb-4bit, 8 → bnb-8bit)."
+            ),
         ),
     ] = 8
 
+    quantization_scheme: Annotated[
+        QuantizationScheme | None,
+        Field(
+            title="quantization_scheme",
+            description=(
+                "Quantization scheme to use when ``quantize_model=True``. Accepts "
+                "``bnb-4bit``, ``bnb-8bit``, ``fp8``, ``nvfp4``, or ``mxfp4``. "
+                "If unset, falls back to ``quantization_bits`` for backward "
+                "compatibility. Non-bitsandbytes schemes are incompatible with "
+                "``peft_implementation='loftq'``."
+            ),
+        ),
+    ] = None
+
     peft_implementation: Annotated[
         str,
         Field(