merge

setup.py

@@ -336,7 +336,7 @@ def _collect_package_files(*directories: str):
     "vector-quantize-pytorch>=1.27.15",
     "cryptography>=41.0.0",
     "torchcodec>=0.10.0",
-    "sdnq>=0.2.2",
+    "sdnq>=0.2.3",
     "aiosqlite>=0.19.0",
     "httpx>=0.28.0",
     "psutil>=5.9.0",

simpletuner/helpers/models/krea2/quantized_loading.py

@@ -129,9 +129,6 @@ def _materialize_krea2_meta_buffers(model: nn.Module) -> None:
 
 def _load_sdnq_training_symbols():
     try:
-        from simpletuner.helpers.training.sdnq_compat import apply_sdnq_checkpointed_backward_fix
-
-        apply_sdnq_checkpointed_backward_fix()
         from sdnq.dequantizer import SDNQDequantizer
         from sdnq.layers import get_sdnq_wrapper_class
         from sdnq.training.forward import get_forward_func

simpletuner/helpers/models/z_image/quantized_loading.py

@@ -118,9 +118,6 @@ def _materialize_zimage_meta_buffers(model: nn.Module) -> None:
 
 def _load_sdnq_training_symbols():
     try:
-        from simpletuner.helpers.training.sdnq_compat import apply_sdnq_checkpointed_backward_fix
-
-        apply_sdnq_checkpointed_backward_fix()
         from sdnq.dequantizer import SDNQDequantizer
         from sdnq.layers import get_sdnq_wrapper_class
         from sdnq.training.forward import get_forward_func

simpletuner/helpers/training/quantisation/__init__.py

@@ -1058,10 +1058,6 @@ def _sdnq_model(
         # Silence sdnq startup logs
         logging.getLogger("sdnq").setLevel(logging.WARNING)
         import sdnq.common as sdnq_common
-
-        from simpletuner.helpers.training.sdnq_compat import apply_sdnq_checkpointed_backward_fix
-
-        apply_sdnq_checkpointed_backward_fix(logger)
         from sdnq.training import sdnq_training_post_load_quant
     except ImportError as e:
         raise ImportError(f"To use SDNQ, please install the sdnq library: `pip install sdnq`: {e}")