Add Z-Image segmented checkpointing support

[bghira]

Summary

Splits the Z-Image 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-wan

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 Z-Image model-family integration for segmented gradient checkpointing controls introduced in the stacked base PR, including runtime hooks/flags, safety-check allow-listing, and a targeted support test to ensure the model exposes the required checkpointing control surface.

Changes:

• Add Z-Image checkpointing hooks/flags for interval + segment stride, FFN-only checkpointing scope, and attention activation offload in the Z-Image transformer.
• Allow-list z_image for segmented checkpointing and attention activation offload in safety checks.
• Extend the segmented-checkpointing model-support test suite to cover ZImageTransformer2DModel.

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 model-support test asserting Z-Image exposes the expected checkpointing control surface.
simpletuner/helpers/training/default_settings/safety_check.py	Allow-lists z_image for segmented checkpointing and attention activation offload safety validation.
simpletuner/helpers/models/z_image/transformer.py	Implements interval/stride segmented checkpointing path and attention/FFN boundary hooks for Z-Image.

---

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