Skip to content

Add SanaVideo segmented checkpointing support - #2956

Merged
bghira merged 2 commits into
agent/segmented-checkpointing-sanafrom
agent/segmented-checkpointing-sanavideo
Aug 3, 2026
Merged

Add SanaVideo segmented checkpointing support#2956
bghira merged 2 commits into
agent/segmented-checkpointing-sanafrom
agent/segmented-checkpointing-sanavideo

Conversation

@bghira

@bghira bghira commented Aug 2, 2026

Copy link
Copy Markdown
Owner

Summary

Splits the SanaVideo 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-sana

Validation

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

@bghira
bghira requested a review from Copilot August 3, 2026 09:32
@bghira
bghira marked this pull request as ready for review August 3, 2026 09:32

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds SanaVideo model-family integration for the segmented gradient-checkpointing feature set introduced in the stacked base PR, including model hooks, safety-check allow-listing, and targeted tests to validate the wiring.

Changes:

  • Implemented segmented checkpointing support in SanaVideoTransformer3DModel (interval/stride/backend setters + segmented execution path).
  • Updated safety-check allow-lists so SanaVideo can use --gradient_checkpointing_interval and --gradient_checkpointing_segment_stride.
  • Added/extended unit tests to cover SanaVideo segmented-checkpointing capability flags and the sequential-state helper integration.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File Description
simpletuner/helpers/models/sanavideo/transformer.py Adds SanaVideo segmented checkpointing controls and forward-pass integration via checkpoint_sequential_state.
simpletuner/helpers/training/default_settings/safety_check.py Allow-lists sanavideo for interval + segment-stride safety checks.
tests/test_transformers/test_sanavideo_transformer.py Adds a unit test asserting the segmented-checkpointing path calls checkpoint_sequential_state with the expected args.
tests/test_segmented_checkpointing_model_support.py Adds SanaVideo to the model-family capability coverage test matrix.

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

Comment on lines +1057 to +1068
hidden_states = self._gradient_checkpointing_func(
block,
hidden_states,
attention_mask,
encoder_hidden_states,
encoder_attention_mask,
timestep,
post_patch_num_frames,
post_patch_height,
post_patch_width,
rotary_emb,
)
@bghira
bghira force-pushed the agent/segmented-checkpointing-sanavideo branch from 6b30d41 to 2ab3ec6 Compare August 3, 2026 10:45
@bghira
bghira force-pushed the agent/segmented-checkpointing-sanavideo branch from 2ab3ec6 to aa1a114 Compare August 3, 2026 15:30
@bghira
bghira force-pushed the agent/segmented-checkpointing-sanavideo branch from aa1a114 to 3792a6f Compare August 3, 2026 15:32
@bghira
bghira force-pushed the agent/segmented-checkpointing-sanavideo branch from 3792a6f to e066e25 Compare August 3, 2026 15:33
@bghira
bghira force-pushed the agent/segmented-checkpointing-sanavideo branch from e066e25 to fa6597b Compare August 3, 2026 15:35
@bghira
bghira merged commit 3504e55 into main Aug 3, 2026
@bghira
bghira deleted the agent/segmented-checkpointing-sanavideo branch August 3, 2026 15:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants