Add Flux segmented checkpointing support

[bghira]

Summary

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

• adds Flux stride handling, attention activation offload, and FFN/offload block plumbing
• adds Flux to the stride and attention-offload safety-check allow-lists
• keeps shared runtime/docs in the base PR so this diff stays model-specific

Stack

Base branch: agent/segmented-checkpointing-ernie

Validation

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

[Copilot]

Pull request overview

Adds Flux-family support for the segmented-checkpointing feature set introduced in the base branch, wiring Flux’s transformer blocks into stride-based segmented checkpointing and attention activation offload, and extending the safety-check allow-lists and model-support tests accordingly.

Changes:

• Add Flux block plumbing for attention activation offload (activation_offload_context) and propagate offload flags through checkpointed and non-checkpointed execution paths.
• Add Flux segmented-checkpointing stride handling within the transformer block loops.
• Extend safety-check allow-lists and add unit tests validating Flux exposes the expected checkpointing controls and block call signatures.

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 coverage to the model-support assertions and verifies Flux blocks accept FFN-checkpoint and attention-offload flags.
simpletuner/helpers/training/default_settings/safety_check.py	Adds flux to the segmented-stride and attention-activation-offload allow-lists.
simpletuner/helpers/models/flux/transformer.py	Implements Flux attention activation offload plumbing and stride-aware segmented checkpointing behavior.

Suppressed comments (1)

simpletuner/helpers/models/flux/transformer.py:1419

• Same issue as above: the return_dict branch in create_custom_forward does not pass offload_attention, so attention activation offload is disabled when return_dict is set.

                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)

---

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