SDNQ: use v0.2.2 and compat-checked patch for smart checkpointing

[bghira]

This pull request updates the sdnq dependency and introduces a compatibility fix for SDNQ checkpointed backward support across several SDNQ integration points. The fix is applied during SDNQ initialization to ensure correct behavior when using checkpointed backward passes.

Dependency update:

• Updated the sdnq package requirement in setup.py from version >=0.1.2 to >=0.2.2 to ensure compatibility with the latest features and fixes.

SDNQ compatibility fix integration:

• Added import and invocation of apply_sdnq_checkpointed_backward_fix from simpletuner.helpers.training.sdnq_compat in the SDNQ loading routines for Krea2 and Z-Image models, ensuring the compatibility fix is applied before SDNQ symbols are loaded. [1] [2]
• Applied the SDNQ checkpointed backward compatibility fix during SDNQ initialization in the quantisation helper, passing the logger for improved diagnostics.

[Copilot]

Pull request overview

This pull request updates the SDNQ dependency floor and introduces a runtime compatibility patch (sdnq_compat.py) that monkeypatches SDNQ checkpointed-backward paths to support “smart checkpointing” behavior across multiple SDNQ integration points.

Changes:

• Bump sdnq minimum version from >=0.1.2 to >=0.2.2.
• Add simpletuner.helpers.training.sdnq_compat.apply_sdnq_checkpointed_backward_fix() implementing the checkpointed-backward compatibility patch.
• Invoke the compatibility patch during SDNQ initialization/loading in the quantisation helper and in Krea2 / Z-Image SDNQ symbol loaders.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
simpletuner/helpers/training/sdnq_compat.py	Adds SDNQ checkpointed-backward compatibility patch via monkeypatching multiple SDNQ linear ckpt modules.
simpletuner/helpers/training/quantisation/init.py	Applies the SDNQ compatibility patch during SDNQ initialization in _sdnq_model.
simpletuner/helpers/models/z_image/quantized_loading.py	Applies the SDNQ compatibility patch before importing SDNQ training symbols.
simpletuner/helpers/models/krea2/quantized_loading.py	Applies the SDNQ compatibility patch before importing SDNQ training symbols.
setup.py	Updates the sdnq dependency requirement to >=0.2.2.

Comments suppressed due to low confidence (5)

simpletuner/helpers/training/sdnq_compat.py:247

• ctx.save_for_backward(...) is passed optional/non-tensor values (new_input/input_scale may be None, bias is optional, and svd_up/svd_down are None when weight is not an SDNQTensor). This will raise at runtime. Store these as attributes on ctx (or otherwise ensure only tensors are saved) and update backward accordingly.

                weight,
                hadamard,
                do_grad_weight=ctx.needs_input_grad[1],
            )
            ctx.save_for_backward(new_input, new_weight, input_scale, weight_scale, bias, svd_up, svd_down)

simpletuner/helpers/training/sdnq_compat.py:359

• Same issue as above: ctx.save_for_backward(...) is passed None values (new_input/input_scale can be None, and bias is optional), which will raise at runtime. Store these on ctx and read from ctx in backward instead of unpacking ctx.saved_tensors.

                    matmul_dtype=matmul_dtype,
                )
            else:
                new_input = input_scale = None
            ctx.save_for_backward(new_input, weight, input_scale, bias)

simpletuner/helpers/training/sdnq_compat.py:495

• ctx.save_for_backward(...) is passed optional/non-tensor values (new_input/input_scale may be None, bias is optional, and svd_up/svd_down may be None). This will raise at runtime. Store these on ctx and unpack from ctx in backward.

                weight,
                hadamard,
                do_grad_weight=ctx.needs_input_grad[1],
            )
            ctx.save_for_backward(new_input, new_weight, input_scale, weight_scale, bias, svd_up, svd_down)

simpletuner/helpers/training/sdnq_compat.py:608

• ctx.save_for_backward(...) is passed None values (new_input/input_scale/input_zero_point can be None, and bias is optional). This will raise at runtime. Store these on ctx and read from ctx in backward rather than unpacking ctx.saved_tensors.

            if ctx.needs_input_grad[1]:
                new_input, input_scale, input_zero_point = module.get_uint8_matmul_backward_inputs(input, hadamard)
            else:
                new_input = input_scale = input_zero_point = None
            ctx.save_for_backward(new_input, weight, input_scale, input_zero_point, bias)

simpletuner/helpers/training/sdnq_compat.py:745

• ctx.save_for_backward(...) is passed optional/non-tensor values (new_input/input_scale/input_zero_point may be None, bias is optional, and svd_up/svd_down may be None). This will raise at runtime. Store these values on ctx and read them from ctx in backward instead of unpacking ctx.saved_tensors.

            ctx.save_for_backward(
                new_input,
                new_weight,
                input_scale,
                weight_scale,

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.