Add Flux.2 segmented checkpointing support

[bghira]

Summary

Splits the Flux.2 segmented checkpointing support model integration out of #2925.

• applies the model-family checkpointing hooks and support flags
• adds this family to the relevant safety-check allow-lists
• keeps the shared runtime/docs in the base PR so this diff stays model-specific

Stack

Base branch: agent/segmented-checkpointing-flux

Validation

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

[Copilot]

Pull request overview

This PR adds Flux.2 model-family integration for the segmented gradient-checkpointing controls introduced in the stacked base branch, wiring Flux.2’s transformer forward pass into the shared checkpointing/offload infrastructure and ensuring the safety-check allow-lists recognize Flux.2.

Changes:

• Add a Flux.2 model-support test ensuring required checkpointing control hooks/flags are present.
• Extend safety-check allow-lists to permit Flux.2 usage of gradient-checkpointing interval/segment-stride and attention activation offload.
• Implement Flux.2 transformer support for segmented checkpointing (interval + segment stride) and attention activation offload plumbing.

Reviewed changes

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

File	Description
tests/test_segmented_checkpointing_model_support.py	Adds Flux.2 to the model-support coverage checks for segmented checkpointing controls.
simpletuner/helpers/training/default_settings/safety_check.py	Allows Flux.2 in the relevant safety-check lists for checkpointing interval/stride and attention activation offload.
simpletuner/helpers/models/flux2/transformer.py	Adds Flux.2 runtime support for segmented checkpointing and threads offload_attention through attention processors/blocks.

Suppressed comments (1)

simpletuner/helpers/models/flux2/transformer.py:439

• Same as above for Flux2ParallelSelfAttnProcessor: the maybe_metal_flash_rope_attention(...) path executes outside activation_offload_context, so offload_attention=True won’t offload saved tensors when this backend is used. Wrap that call in activation_offload_context too.

                getattr(attn, "to_qkv_mlp_proj", None) and attn.to_qkv_mlp_proj.weight,
            )

        if hidden_states is None:
            with activation_offload_context(offload_attention, label=f"{attn.__class__.__qualname__}:attention"):

---

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