Add ACE-Step segmented checkpointing support

[bghira]

Summary

Splits the ACE-Step segmented-checkpointing integration out of #2925.

• adds ACE-Step interval and segment-stride controls
• adds ACE-Step to the relevant safety-check allow-lists
• keeps shared runtime/docs in the base PR so this diff stays model-specific

Stack

Base branch: agent/segmented-checkpointing-base

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 ACE-Step model support for segmented checkpointing controls by wiring interval/segment-stride into the ACE-Step transformer and enabling the corresponding safety-check allow-lists, plus a targeted unit test to ensure the model exposes the expected control surface.

Changes:

• Add gradient_checkpointing_interval and gradient_checkpointing_segment_stride setters/state to ACEStepTransformer2DModel, and gate per-block checkpointing via should_checkpoint_block(...).
• Allow ACE-Step to use --gradient_checkpointing_interval and --gradient_checkpointing_segment_stride via safety_check model-family allow-lists.
• Add a unit test covering ACE-Step’s checkpointing control availability.

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 small capability-coverage test for ACE-Step segmented checkpointing controls.
simpletuner/helpers/training/default_settings/safety_check.py	Adds ace_step to the interval and segment-stride allow-lists for validation/normalization.
simpletuner/helpers/models/ace_step/transformer.py	Implements interval/segment-stride fields + setters and applies should_checkpoint_block gating during training.

---

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