Bump pytorch 2.8.0 (#3611)

AlexanderDokuchaev · daniil-lyakhov · anzr299 · web-flow · commit cc3487a15450 · 2025-08-12T09:00:16.000+03:00
### Changes Bump pytorch 2.8.0 ### Tests manual/job/post_training_quantization/69 https://github.com/openvinotoolkit/nncf/actions/runs/16882768007 https://github.com/openvinotoolkit/nncf/actions/runs/16885356994 --------- Co-authored-by: dlyakhov <daniil.lyakhov@intel.com> Co-authored-by: Aamir Nazir <aamir.nazir@intel.com>
diff --git a/constraints.txt b/constraints.txt
@@ -2,8 +2,8 @@
 openvino==2025.2.0
 
 # Pytorch
-torch==2.7.1
-torchvision==0.22.1
+torch==2.8.0
+torchvision==0.23.0
 
 # ONNX
 onnx==1.17.0; python_version < '3.13'
diff --git a/docs/Installation.md b/docs/Installation.md
@@ -49,7 +49,7 @@ as well as the supported versions of Python:
 
 | NNCF      | OpenVINO   | PyTorch  | ONNX     | TensorFlow | Python |
 |-----------|------------|----------|----------|------------|--------|
-| `develop` | `2025.2.0` | `2.7.1`  | `1.17.0` | `2.15.1`   | `3.10` |
+| `develop` | `2025.2.0` | `2.8.0`  | `1.17.0` | `2.15.1`   | `3.10` |
 | `2.17.0`  | `2025.2.0` | `2.7.1`  | `1.17.0` | `2.15.1`   | `3.10` |
 | `2.16.0`  | `2025.1.0` | `2.6.0`  | `1.17.0` | `2.15.1`   | `3.10` |
 | `2.15.0`  | `2025.0.0` | `2.5.1`  | `1.17.0` | `2.15.1`   | `3.10` |
@@ -69,4 +69,4 @@ as well as the supported versions of Python:
 
 > (*) Python 3.9 or higher is required for TensorFlow 2.15.1
 
-This repository is tested on Python* 3.10.14, PyTorch* 2.7.1 (NVidia CUDA\* Toolkit 12.6) and TensorFlow* 2.15.1 (NVidia CUDA\* Toolkit 11.8).
+This repository is tested on Python* 3.10.14, PyTorch* 2.8.0 (NVidia CUDA\* Toolkit 12.6) and TensorFlow* 2.15.1 (NVidia CUDA\* Toolkit 11.8).
diff --git a/examples/llm_compression/openvino/tiny_llama_synthetic_data/requirements.txt b/examples/llm_compression/openvino/tiny_llama_synthetic_data/requirements.txt
@@ -1,4 +1,4 @@
-torch==2.7.1
+torch==2.8.0
 datasets==3.0.1
 numpy>=1.23.5,<2
 openvino==2025.2.0
diff --git a/examples/llm_compression/torch/distillation_qat_with_lora/requirements.txt b/examples/llm_compression/torch/distillation_qat_with_lora/requirements.txt
@@ -1,5 +1,5 @@
 tensorboard==2.13.0
-torch==2.7.1
+torch==2.8.0
 numpy>=1.23.5,<2
 openvino==2025.2.0
 optimum-intel>=1.22.0
diff --git a/examples/llm_compression/torch/downstream_qat_with_nls/requirements.txt b/examples/llm_compression/torch/downstream_qat_with_nls/requirements.txt
@@ -1,5 +1,5 @@
 tensorboard==2.13.0
-torch==2.7.1
+torch==2.8.0
 numpy>=1.23.5,<2
 openvino==2025.2.0
 optimum-intel>=1.22.0
diff --git a/examples/llm_compression/torch_fx/tiny_llama/modelling.py b/examples/llm_compression/torch_fx/tiny_llama/modelling.py
@@ -112,6 +112,7 @@ def convert_and_export_with_cache(model: PreTrainedModel) -> tuple[ExportedProgr
             example_input_ids,
             example_cache_position,
         ),
+        strict=True,
         dynamic_shapes=dynamic_shapes,
     ).run_decompositions(decomp_table={})
     return exported_program, model_config, gen_config
diff --git a/examples/llm_compression/torch_fx/tiny_llama/requirements.txt b/examples/llm_compression/torch_fx/tiny_llama/requirements.txt
@@ -2,5 +2,5 @@ transformers==4.52.1
 datasets==2.14.7
 openvino==2025.2.0
 optimum==1.24.0
-torch==2.7.1
-torchvision==0.22.1
+torch==2.8.0
+torchvision==0.23.0
diff --git a/examples/post_training_quantization/torch/mobilenet_v2/requirements.txt b/examples/post_training_quantization/torch/mobilenet_v2/requirements.txt
@@ -1,5 +1,5 @@
 fastdownload==0.0.7
 openvino==2025.2.0
 scikit-learn
-torch==2.7.1
-torchvision==0.22.1
+torch==2.8.0
+torchvision==0.23.0
diff --git a/examples/post_training_quantization/torch/ssd300_vgg16/requirements.txt b/examples/post_training_quantization/torch/ssd300_vgg16/requirements.txt
@@ -2,7 +2,7 @@ fastdownload==0.0.7
 onnx==1.17.0
 openvino==2025.2.0
 pycocotools==2.0.7
-torch==2.7.1
+torch==2.8.0
 torchmetrics==1.0.1
-torchvision==0.22.1
+torchvision==0.23.0
 numpy<2
diff --git a/examples/post_training_quantization/torch_fx/resnet18/requirements.txt b/examples/post_training_quantization/torch_fx/resnet18/requirements.txt
@@ -1,4 +1,4 @@
 fastdownload==0.0.7
 openvino==2025.2.0
-torch==2.7.1
-torchvision==0.22.1
+torch==2.8.0
+torchvision==0.23.0
diff --git a/examples/quantization_aware_training/torch/resnet18/requirements.txt b/examples/quantization_aware_training/torch/resnet18/requirements.txt
@@ -1,4 +1,4 @@
 fastdownload==0.0.7
 openvino==2025.2.0
-torch==2.7.1
-torchvision==0.22.1
+torch==2.8.0
+torchvision==0.23.0
diff --git a/src/nncf/experimental/torch/fx/nncf_graph_builder.py b/src/nncf/experimental/torch/fx/nncf_graph_builder.py
@@ -107,7 +107,7 @@ def get_node_type_and_metatype(node: torch.fx.Node, model: torch.fx.GraphModule)
                 node_type = "__getitem__"
             else:
                 # TODO(dlyakhov): get correct nodes types from this nodes as well
-                node_type = str(node.target)
+                node_type = node.target.__name__
             node_metatype = PT_OPERATOR_METATYPES.get_operator_metatype_by_op_name(node_type)
         else:
             node_type = node.op
diff --git a/src/nncf/torch/function_hook/hook_storage.py b/src/nncf/torch/function_hook/hook_storage.py
@@ -43,10 +43,22 @@ def remove(self) -> None:
         - If no hooks left under the key, the key is also removed from the storage.
         """
         storage: Optional[nn.ModuleDict] = self.storage_ref()
-        if storage is not None and self.key in storage and self.id in storage[self.key]:
-            del storage[self.key][self.id]
-            if not storage[self.key]:
-                del storage[self.key]
+        if storage is None or self.key not in storage:
+            # hook storage has been garbage collected or key is not present
+            return
+
+        hooks_dict = storage[self.key]
+        if not isinstance(hooks_dict, nn.ModuleDict):
+            msg = f"Expected nn.ModuleDict for key={self.key}, got {type(hooks_dict)}"
+            raise TypeError(msg)
+
+        if self.id not in hooks_dict:
+            # hook with the specified id was already removed
+            return
+
+        del hooks_dict[self.id]
+        if not storage[self.key]:
+            del storage[self.key]
 
 
 class HookStorage(nn.Module):
@@ -106,8 +118,13 @@ def _insert_hook(
         if hook_key not in storage_dict:
             storage_dict[hook_key] = nn.ModuleDict()
 
-        hook_id = cls._get_next_hook_id(storage_dict[hook_key])
-        storage_dict[hook_key][hook_id] = hook
+        hooks_dict = storage_dict[hook_key]
+        if not isinstance(hooks_dict, nn.ModuleDict):
+            msg = f"Expected nn.ModuleDict for key={hook_key}, got {type(hooks_dict)}"
+            raise TypeError(msg)
+
+        hook_id = cls._get_next_hook_id(hooks_dict)
+        hooks_dict[hook_id] = hook
 
         return RemovableHookHandle(storage_dict, hook_key, hook_id)
 
@@ -148,7 +165,8 @@ def _execute_hooks(cls, storage_dict: nn.ModuleDict, op_name: str, port_id: int,
         hook_key = cls._generate_key(op_name, port_id)
         if hook_key not in storage_dict:
             return value
-        for hook in storage_dict[hook_key].values():
+        hooks_dict = cast(nn.ModuleDict, storage_dict[hook_key])
+        for hook in hooks_dict.values():
             value = hook(value)
         return value
 
diff --git a/src/nncf/version.py b/src/nncf/version.py
@@ -12,6 +12,6 @@
 __version__ = "2.18.0"
 
 
-BKC_TORCH_SPEC = "==2.7.*"
+BKC_TORCH_SPEC = "==2.8.*"
 BKC_TF_SPEC = "==2.15.*"
 STRICT_TF_SPEC = ">=2.9.3,<2.16.0"
diff --git a/tests/post_training/pipelines/base.py b/tests/post_training/pipelines/base.py
@@ -502,7 +502,7 @@ def save_compressed_model(self) -> None:
             )
             ov.serialize(ov_model, self.path_compressed_ir)
         elif self.backend in FX_BACKENDS:
-            exported_model = torch.export.export(self.compressed_model.cpu(), (self.dummy_tensor.cpu(),))
+            exported_model = torch.export.export(self.compressed_model.cpu(), (self.dummy_tensor.cpu(),), strict=True)
             # Torch export is used to save the model because ov.convert_model does not fully claim support for
             # Converting ExportedProgram
             torch.export.save(exported_model, self.output_model_dir / "model.pt2")
diff --git a/tests/post_training/pipelines/fx_modelling.py b/tests/post_training/pipelines/fx_modelling.py
@@ -100,5 +100,6 @@ def convert_and_export_with_cache(model: PreTrainedModel):
             example_cache_position,
         ),
         dynamic_shapes=dynamic_shapes,
+        strict=True,
     ).run_decompositions(decomp_table={})
     return exported_program, model_config, gen_config
diff --git a/tests/post_training/pipelines/image_classification_torchvision.py b/tests/post_training/pipelines/image_classification_torchvision.py
@@ -25,11 +25,11 @@
 
 
 def _torch_export_for_training(model: torch.nn.Module, args: tuple[Any, ...]) -> torch.fx.GraphModule:
-    return torch.export.export_for_training(model, args).module()
+    return torch.export.export_for_training(model, args, strict=True).module()
 
 
 def _torch_export(model: torch.nn.Module, args: tuple[Any, ...]) -> torch.fx.GraphModule:
-    return torch.export.export(model, args).module()
+    return torch.export.export(model, args, strict=True).module()
 
 
 @dataclass
@@ -103,7 +103,7 @@ def prepare_model(self) -> None:
         elif self.backend in [BackendType.OV, BackendType.FP32]:
             with torch.no_grad():
                 if self.model_params.export_torch_before_ov_convert:
-                    model = torch.export.export(model, (self.dummy_tensor,))
+                    model = torch.export.export(model, (self.dummy_tensor,), strict=True)
                 self.model = ov.convert_model(model, example_input=self.dummy_tensor, input=self.input_size)
             self.input_name = list(inp.get_any_name() for inp in self.model.inputs)[0]
 
@@ -125,7 +125,7 @@ def _dump_model_fp32(self) -> None:
             ov.serialize(ov_model, self.fp32_model_dir / "model_fp32.xml")
 
         if self.backend in FX_BACKENDS:
-            exported_model = torch.export.export(self.model.cpu(), (self.dummy_tensor.cpu(),))
+            exported_model = torch.export.export(self.model.cpu(), (self.dummy_tensor.cpu(),), strict=True)
             torch.export.save(exported_model, self.fp32_model_dir / "fx_model_fp32.pt2")
 
             if self.backend is BackendType.CUDA_FX_TORCH:
diff --git a/tests/torch2/data/fx/dynamic_shapes/post_quantization_compressed/swin_v2_t.dot b/tests/torch2/data/fx/dynamic_shapes/post_quantization_compressed/swin_v2_t.dot
diff --git a/tests/torch2/data/fx/dynamic_shapes/quantized/swin_v2_t.dot b/tests/torch2/data/fx/dynamic_shapes/quantized/swin_v2_t.dot
diff --git a/tests/torch2/fx/helpers.py b/tests/torch2/fx/helpers.py
diff --git a/tests/torch2/fx/test_model_transformer.py b/tests/torch2/fx/test_model_transformer.py
diff --git a/tests/torch2/fx/test_models.py b/tests/torch2/fx/test_models.py
diff --git a/tests/torch2/fx/test_sanity.py b/tests/torch2/fx/test_sanity.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-torch==2.7.1`
	`1`	`+torch==2.8.0`
`2`	`2`	`datasets==3.0.1`
`3`	`3`	`numpy>=1.23.5,<2`
`4`	`4`	`openvino==2025.2.0`