Add Chroma segmented checkpointing support

[bghira]

Summary

Splits the Chroma 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-boogu-image

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 Chroma model-family integration for segmented gradient checkpointing, including FFN-only checkpointing and attention activation offload support, and wires Chroma into the training safety-check allow-lists. This continues the model-specific extraction from the broader segmented-checkpointing controls work in #2925.

Changes:

• Add Chroma model flags and setters for segmented checkpointing stride and attention activation offload.
• Add FFN-only checkpointing and attention activation offload hooks inside Chroma transformer blocks.
• Extend safety-check allow-lists and model-support tests to include Chroma.

Reviewed changes

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

File	Description
tests/test_segmented_checkpointing_model_support.py	Adds a Chroma model-family test to assert the new checkpointing/offload controls are exposed.
simpletuner/helpers/training/default_settings/safety_check.py	Allows Chroma for segment-stride checkpointing and attention activation offload gating.
simpletuner/helpers/models/chroma/transformer.py	Implements Chroma segmented checkpointing stride support, FFN-only checkpointing, and attention activation offload integration.

Suppressed comments (1)

simpletuner/helpers/models/chroma/transformer.py:1236

• Same as above: this call passes True into should_checkpoint_block while also checking self.gradient_checkpointing separately. Passing self.gradient_checkpointing directly matches the established calling pattern in other model families and keeps the checkpoint gating logic in one place.

            use_checkpoint = (
                torch.is_grad_enabled()
                and self.gradient_checkpointing
                and should_checkpoint_block(
                    index_block,
                    True,
                    self.gradient_checkpointing_interval,
                    self.gradient_checkpointing_segment_stride,
                )

---

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