Add AuraFlow segmented checkpointing support

[bghira]

Summary

Splits the AuraFlow segmented-checkpointing integration out of #2925.

• adds AuraFlow segment-stride checkpoint selection
• updates AuraFlow ControlNet for unsloth backend matching
• adds AuraFlow to the stride safety-check allow-list

Stack

Base branch: agent/segmented-checkpointing-ace-step

Validation

• .venv/bin/python -m unittest tests.test_segmented_checkpointing_model_support tests.test_transformers.test_auraflow_transformer -v
• commit hooks: Black, isort, flake8, whitespace checks

[Copilot]

Pull request overview

Adds AuraFlow support for segmented (stride-window) gradient checkpointing, aligning AuraFlow’s transformer checkpointing behavior with the shared segmented-checkpointing helpers and expanding safety-check allow-list coverage.

Changes:

• Add gradient_checkpointing_segment_stride plumbing to AuraFlowTransformer2DModel and use should_checkpoint_block(...) to decide per-block checkpointing.
• Treat unsloth-* checkpoint backends consistently (e.g., unsloth-ffn) across AuraFlow transformer and ControlNet.
• Extend tests and the safety-check allow-list to cover AuraFlow segmented checkpointing controls.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_transformers/test_auraflow_transformer.py	Adds a backward regression test covering joint + single block checkpointing.
tests/test_segmented_checkpointing_model_support.py	Asserts AuraFlow exposes segmented checkpointing control setters.
simpletuner/helpers/training/default_settings/safety_check.py	Adds AuraFlow to the segment-stride allow-list.
simpletuner/helpers/models/auraflow/transformer.py	Implements segment-stride checkpoint selection and forwards attention_kwargs through checkpointed blocks.
simpletuner/helpers/models/auraflow/controlnet.py	Updates unsloth backend matching, but checkpointed paths still don’t forward attention_kwargs (see comments).

Suppressed comments (1)

simpletuner/helpers/models/auraflow/controlnet.py:417

• Same issue in the single-block checkpointed path: create_custom_forward(block) calls the block positionally, so attention_kwargs is omitted while the non-checkpoint path passes it. This can make outputs differ depending on whether checkpointing is enabled. Wrap the block with a custom_forward that forwards attention_kwargs.

                    def create_custom_forward(module):
                        def custom_forward(*inputs):
                            return module(*inputs)

                        return custom_forward

---

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