Add Cosmos 3 segmented checkpointing support

[bghira]

Summary

Splits the Cosmos 3 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-cosmos

Validation

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

[Copilot]

Pull request overview

Adds Cosmos 3 into the segmented gradient-checkpointing feature surface by wiring the model’s checkpointing decision logic into the shared should_checkpoint_block helper, exposing the segment-stride setter, and extending safety-check allow-lists. Also adds targeted unit coverage to ensure Cosmos 3 advertises the expected checkpointing controls and that unpatchify preserves/propagates prediction dtypes correctly.

Changes:

• Add gradient_checkpointing_segment_stride support to Cosmos3OmniTransformer and route layer checkpoint decisions through should_checkpoint_block(...).
• Extend safety-check allow-lists to recognize cosmos3 for interval/segment-stride checkpointing settings.
• Add unit tests for Cosmos 3 segmented-checkpointing capability flags and for _unpatchify_and_unpack_latents dtype behavior.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
tests/test_segmented_checkpointing_model_support.py	Adds a Cosmos 3 capability test ensuring the model exposes interval + segment-stride checkpointing hooks.
tests/test_cosmos3_model.py	Adds regression tests validating _unpatchify_and_unpack_latents preserves or adopts the materialized latent dtype.
simpletuner/helpers/training/default_settings/safety_check.py	Adds cosmos3 to the checkpointing segment-stride supported model allow-list (and confirms interval support in-file).
simpletuner/helpers/models/cosmos3/transformer.py	Implements segment-stride setter/state and uses shared should_checkpoint_block logic for per-layer checkpoint decisions; updates unpatchify output dtype selection.

---

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