Add Cosmos segmented checkpointing support

[bghira]

Summary

Splits the Cosmos 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-chroma

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 Cosmos (T2I) model-family integration for the segmented gradient checkpointing controls introduced in the stacked base PR, ensuring Cosmos is recognized by safety checks and exposes the expected model hooks.

Changes:

• Adds a segmented-checkpointing support test for CosmosTransformer3DModel.
• Allows cosmos in safety-check allow-lists for --gradient_checkpointing_interval and --gradient_checkpointing_segment_stride.
• Wires Cosmos’ transformer forward pass to use should_checkpoint_block(...) gated by interval/segment-stride settings.

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 a capability/controls presence test for Cosmos segmented checkpointing support.
simpletuner/helpers/training/default_settings/safety_check.py	Extends safety-check allow-lists to permit interval + segment stride for the cosmos model family.
simpletuner/helpers/models/cosmos/transformer.py	Adds interval/stride setters and uses the shared should_checkpoint_block helper to gate checkpointing per block.

---

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