Skip to content

Add Qwen Image segmented checkpointing support - #2954

Merged
bghira merged 1 commit into
agent/segmented-checkpointing-pixartfrom
agent/segmented-checkpointing-qwen-image
Aug 3, 2026
Merged

Add Qwen Image segmented checkpointing support#2954
bghira merged 1 commit into
agent/segmented-checkpointing-pixartfrom
agent/segmented-checkpointing-qwen-image

Conversation

@bghira

@bghira bghira commented Aug 2, 2026

Copy link
Copy Markdown
Owner

Summary

Splits the Qwen 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-pixart

Validation

  • .venv/bin/python -m unittest tests.test_segmented_checkpointing_model_support -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 Qwen Image’s transformer to the segmented gradient checkpointing support surface so it can participate in the shared checkpoint-interval/segment-stride runtime introduced in the base PR stack.

Changes:

  • Adds a Qwen Image coverage test ensuring the model exposes the expected segmented checkpointing control hooks.
  • Extends safety-check allow-lists so qwen_image can use --gradient_checkpointing_interval and --gradient_checkpointing_segment_stride.
  • Updates QwenImageTransformer2DModel to store/apply interval + segment stride and to use should_checkpoint_block(...) during the block loop (including recognizing unsloth* backends).

Reviewed changes

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

File Description
tests/test_segmented_checkpointing_model_support.py Adds a Qwen Image capability coverage test for interval + segment stride hooks.
simpletuner/helpers/training/default_settings/safety_check.py Allows qwen_image to use interval and segment-stride checkpointing flags.
simpletuner/helpers/models/qwen_image/transformer.py Implements interval/segment-stride checkpoint decision logic and backend handling for Qwen Image blocks.

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

@bghira
bghira force-pushed the agent/segmented-checkpointing-qwen-image branch from 036fec2 to a0115dc Compare August 3, 2026 10:45
@bghira
bghira force-pushed the agent/segmented-checkpointing-qwen-image branch from a0115dc to 595d2d6 Compare August 3, 2026 15:30
@bghira
bghira force-pushed the agent/segmented-checkpointing-qwen-image branch from 595d2d6 to c4c55db Compare August 3, 2026 15:32
@bghira
bghira force-pushed the agent/segmented-checkpointing-qwen-image branch from c4c55db to 5a32f39 Compare August 3, 2026 15:33
@bghira
bghira force-pushed the agent/segmented-checkpointing-qwen-image branch from 5a32f39 to 7cd94e4 Compare August 3, 2026 15:35
@bghira
bghira merged commit 3504e55 into main Aug 3, 2026
@bghira
bghira deleted the agent/segmented-checkpointing-qwen-image 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