Merge pull request #31 from FujitsuResearch/develop/v1-2-1

Fix unsafe .pt deserialization in model loaders (CWE-502)

Author: FKKimura

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Change log
 
+## [v1.2.1] 2026-07-03
+
+### Security
+
+- **Unsafe deserialization hardening (CWE-502)**: `QuantizedModelLoader.load_quantized_model_pt()` (alias `onecomp.load_quantized_model_pt()`) previously called `torch.load(model.pt, weights_only=False)` unconditionally, allowing arbitrary code execution when loading a malicious `.pt` checkpoint. It now refuses to load unless the caller explicitly opts in via `allow_unsafe_deserialization=True`, and emits a strong warning when it does load. For untrusted models, use the safetensors-based `load_quantized_model()`, which does not execute code.
+  - **Breaking change**: existing callers of `load_quantized_model_pt()` must pass `allow_unsafe_deserialization=True` for trusted `.pt` files.
+- **`Quantizer.load_results()` / `ResultLoader`**: same hardening applied. Loading with `weights_only=False` now requires `allow_unsafe_deserialization=True` (added as a `ResultLoader` field), and logs a warning. The safe `weights_only=True` path is unchanged.
+- Updated docstrings, docs, and the LoRA SFT example to document the risk and the required opt-in.
+
 ## [v1.2.0] 2026-06-08
 
 ### Save/Load Support for JointQ, RTN, and OneBit Quantizers

docs/api/quantized_model_loader.md

@@ -22,5 +22,20 @@ from onecomp import load_quantized_model, load_quantized_model_pt
 model, tokenizer = load_quantized_model("./saved_model")
 
 # Load a PyTorch .pt model (post-processed, e.g. LoRA-applied)
-model, tokenizer = load_quantized_model_pt("./saved_model_lora")
+# Requires explicit opt-in: the .pt loader uses torch.load(weights_only=False),
+# which can execute code from a malicious file (CWE-502). Only enable this for
+# model.pt files from a fully trusted source.
+model, tokenizer = load_quantized_model_pt(
+    "./saved_model_lora", allow_unsafe_deserialization=True
+)
 ```
+
+!!! warning "Unsafe deserialization (.pt loader)"
+    `load_quantized_model_pt()` loads `model.pt` with
+    `torch.load(..., weights_only=False)`. Because PyTorch `.pt` checkpoints use
+    Python `pickle`, a maliciously crafted `model.pt` can execute arbitrary code
+    during loading (CWE-502). The method refuses to load unless you pass
+    `allow_unsafe_deserialization=True`. Only opt in for models you produced
+    yourself or obtained from a fully trusted source. For untrusted or
+    third-party models, prefer the safetensors-based `load_quantized_model()`,
+    which does not execute code.

docs/user-guide/examples.md

@@ -579,11 +579,20 @@ LoRA-applied models use a dedicated save/load API:
 # Save after LoRA SFT
 runner.save_quantized_model_pt("./my_model_lora")
 
-# Load
+# Load (opt-in required: the .pt loader uses torch.load(weights_only=False),
+# which can execute code from a malicious file. Trusted sources only.)
 from onecomp import load_quantized_model_pt
-model, tokenizer = load_quantized_model_pt("./my_model_lora")
+model, tokenizer = load_quantized_model_pt(
+    "./my_model_lora", allow_unsafe_deserialization=True
+)
 ```
 
+!!! warning "Unsafe deserialization (.pt loader)"
+    `load_quantized_model_pt()` loads `model.pt` via `torch.load(weights_only=False)`
+    (Python `pickle`), so a malicious file can execute arbitrary code (CWE-502).
+    It refuses to load unless you pass `allow_unsafe_deserialization=True`; only
+    opt in for fully trusted models.
+
 !!! note
     For standard quantized models (without LoRA), use `save_quantized_model()` / `load_quantized_model()` instead.
 

docs/user-guide/post-process.md

@@ -335,9 +335,21 @@ runner.save_quantized_model_pt("./my_model_lora")
 ```python
 from onecomp import load_quantized_model_pt
 
-model, tokenizer = load_quantized_model_pt("./my_model_lora")
+# The .pt loader uses torch.load(weights_only=False) and requires an explicit
+# opt-in. Only enable it for models from a fully trusted source (see warning).
+model, tokenizer = load_quantized_model_pt(
+    "./my_model_lora", allow_unsafe_deserialization=True
+)
 ```
 
+!!! warning "Unsafe deserialization (.pt loader)"
+    `load_quantized_model_pt()` deserializes `model.pt` with
+    `torch.load(..., weights_only=False)`, which uses Python `pickle` and can
+    execute arbitrary code from a malicious file (CWE-502). It refuses to load
+    unless you pass `allow_unsafe_deserialization=True`. Only opt in for models
+    you produced yourself or trust completely; otherwise use the safetensors
+    `load_quantized_model()`, which does not execute code.
+
 !!! info "save_quantized_model vs save_quantized_model_pt"
     | Method | Format | Use Case |
     |--------|--------|----------|

example/post_process/example_lora_sft.py

@@ -133,7 +133,11 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=64):
 print(f"Step 4: Loading model from {SAVE_DIR}")
 print("=" * 70)
 
-loaded_model, loaded_tokenizer = load_quantized_model_pt(SAVE_DIR)
+# SAVE_DIR was produced by this script (trusted), so we opt in to the
+# pickle-based .pt loader. Never enable this for untrusted model.pt files.
+loaded_model, loaded_tokenizer = load_quantized_model_pt(
+    SAVE_DIR, allow_unsafe_deserialization=True
+)
 print(f"Loaded model type : {type(loaded_model).__name__}")
 print(f"Loaded model device: {next(loaded_model.parameters()).device}")
 

onecomp/__version__.py

@@ -6,4 +6,4 @@
 
 """
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"

onecomp/post_process/post_process_lora_sft.py

@@ -298,6 +298,8 @@ class PostProcessLoraSFT(PostQuantizationProcess):
             >>> import torch
             >>> from onecomp import ModelConfig, PostProcessLoraSFT
             >>> model_config = ModelConfig(model_id="meta-llama/Llama-2-7b-hf")
+            >>> # weights_only=False uses pickle and can execute code from a
+            >>> # malicious file (CWE-502); only load trusted checkpoints.
             >>> quantized_model = torch.load(
             ...     "quantized_model.pt",
             ...     map_location="cpu",

onecomp/quantized_model_loader.py

@@ -187,6 +187,7 @@ def load_quantized_model_pt(
         *,
         device_map: str = "auto",
         local_files_only: bool = True,
+        allow_unsafe_deserialization: bool = False,
     ) -> Tuple[Any, Any]:
         """Load a quantized model and tokenizer saved as a PyTorch .pt file.
 
@@ -198,18 +199,39 @@ def load_quantized_model_pt(
         - ``model.pt`` (serialized with ``torch.save``)
         - Tokenizer files
 
+        .. warning::
+            This method deserializes ``model.pt`` with
+            ``torch.load(..., weights_only=False)``. Because PyTorch ``.pt``
+            checkpoints use Python's ``pickle``, a maliciously crafted
+            ``model.pt`` can execute arbitrary code during deserialization
+            (CWE-502). ``weights_only=False`` is required here because the
+            ``.pt`` format preserves full custom module objects (e.g.
+            ``LoRAGPTQLinear``) that cannot be reconstructed from tensors
+            alone. Only load ``model.pt`` files that you produced yourself
+            or obtained from a fully trusted source. For untrusted or
+            third-party models, prefer the safetensors-based
+            :meth:`load_quantized_model`, which does not execute code.
+
         Args:
             save_directory: Path to the saved model directory.
             device_map: Device placement (default: ``"auto"``).
                 Set to ``""`` or ``None`` to skip device placement.
             local_files_only: Passed to ``AutoTokenizer.from_pretrained``.
+            allow_unsafe_deserialization: Must be explicitly set to ``True``
+                to acknowledge the unsafe-deserialization risk described
+                above and permit loading. Defaults to ``False``, in which
+                case this method raises before any code can be executed.
 
         Returns:
             (model, tokenizer)
 
+        Raises:
+            ValueError: If ``allow_unsafe_deserialization`` is not ``True``.
+
         Example:
             >>> model, tokenizer = QuantizedModelLoader.load_quantized_model_pt(
-            ...     "./quantized_model_lora"
+            ...     "./quantized_model_lora",
+            ...     allow_unsafe_deserialization=True,  # trusted source only
             ... )
         """
         save_directory = os.path.abspath(save_directory)
@@ -224,6 +246,24 @@ def load_quantized_model_pt(
                 "(safetensors format); use load_quantized_model() instead."
             )
 
+        if not allow_unsafe_deserialization:
+            raise ValueError(
+                f"Refusing to load '{model_path}': loading a .pt model uses "
+                "torch.load(weights_only=False), which deserializes arbitrary "
+                "Python objects via pickle and can execute code embedded in a "
+                "malicious file (CWE-502). Only load model.pt files you produced "
+                "yourself or obtained from a fully trusted source, then pass "
+                "allow_unsafe_deserialization=True to acknowledge this risk. "
+                "For untrusted or third-party models, use the safetensors-based "
+                "load_quantized_model() instead."
+            )
+
+        logger.warning(
+            "Loading '%s' with torch.load(weights_only=False); arbitrary code in "
+            "a malicious checkpoint can execute during deserialization. Ensure "
+            "this model.pt comes from a trusted source.",
+            model_path,
+        )
         model = torch.load(model_path, map_location="cpu", weights_only=False)
 
         if device_map:

onecomp/quantizer/_quantizer.py

@@ -641,22 +641,40 @@ def save_results(self, filepath):
         torch.save(self.results, filepath)
         self.logger.info("Saved quantization results to %s", filepath)
 
-    def load_results(self, filepath, weights_only=False):
+    def load_results(self, filepath, *, weights_only=False, allow_unsafe_deserialization=False):
         """Load the quantization results from a file into self.results.
 
         Loads saved quantization results and stores them in self.results.
 
+        .. warning::
+            With ``weights_only=False`` this method deserializes arbitrary
+            Python objects via ``torch.load`` / ``pickle``, so a maliciously
+            crafted file can execute code during deserialization (CWE-502).
+            Only load result files you produced yourself or obtained from a
+            fully trusted source, and set ``allow_unsafe_deserialization=True``
+            to acknowledge the risk.
+
         Args:
             filepath (str): The path to load the results from.
             weights_only (bool): If True, only load tensor weights (safer but limited).
                 Default is False to support loading QuantizationResult objects.
+            allow_unsafe_deserialization (bool): Required to be True when
+                ``weights_only`` is False, to acknowledge the
+                unsafe-deserialization risk. Ignored when ``weights_only`` is
+                True (that path is safe). Defaults to False.
 
         Returns:
             dict: A dict mapping layer name -> QuantizationResult (same reference as self.results).
 
+        Raises:
+            ValueError: If ``weights_only`` is False and
+                ``allow_unsafe_deserialization`` is not True.
+
         Example:
             >>> quantizer = JointQ()
-            >>> quantizer.load_results("quantization_results.pt")
+            >>> quantizer.load_results(
+            ...     "quantization_results.pt", allow_unsafe_deserialization=True
+            ... )
             >>> for layer_name, result in quantizer.results.items():
             ...     print(f"{layer_name}: {result.dequantized_weight.shape}")
 
@@ -667,6 +685,22 @@ def load_results(self, filepath, weights_only=False):
             - Loading files saved with older versions may fail if class
               definitions have changed.
         """
+        if not weights_only and not allow_unsafe_deserialization:
+            raise ValueError(
+                f"Refusing to load '{filepath}': loading with weights_only=False "
+                "deserializes arbitrary Python objects via pickle and can execute "
+                "code embedded in a malicious file (CWE-502). Only load result "
+                "files from a fully trusted source, then pass "
+                "allow_unsafe_deserialization=True to acknowledge this risk, or "
+                "use weights_only=True for the safe (tensors-only) path."
+            )
+        if not weights_only:
+            self.logger.warning(
+                "Loading '%s' with weights_only=False; arbitrary code in a "
+                "malicious file can execute during deserialization. Ensure this "
+                "file comes from a trusted source.",
+                filepath,
+            )
         self.results = torch.load(filepath, weights_only=weights_only)
         self.logger.info("Loaded quantization results from %s", filepath)
         return self.results
@@ -951,6 +985,9 @@ class ResultLoader(Quantizer):
     # Optional: load precomputed results at initialization time
     results_file: str = None
     weights_only: bool = False
+    # Required to be True when weights_only is False, to acknowledge the
+    # unsafe-deserialization risk of torch.load(weights_only=False) (CWE-502).
+    allow_unsafe_deserialization: bool = False
 
     # Ensure no layers are selected for quantization by default
     target_layer_types: tuple = field(default_factory=tuple)
@@ -962,7 +999,11 @@ class ResultLoader(Quantizer):
     def __post_init__(self):
         super().__post_init__()
         if self.results_file is not None:
-            self.load_results(self.results_file, weights_only=self.weights_only)
+            self.load_results(
+                self.results_file,
+                weights_only=self.weights_only,
+                allow_unsafe_deserialization=self.allow_unsafe_deserialization,
+            )
 
     def setup(self, model):  # pylint: disable=unused-argument
         """Select no layers (no-op).