Add HiDream segmented checkpointing support

[bghira]

Summary

Splits the HiDream 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-flux2

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 HiDream model-family support for segmented gradient checkpointing controls, aligning it with the broader segmented-checkpointing feature set introduced in the stacked base branch.

Changes:

• Adds HiDream coverage to the segmented checkpointing model-support test suite.
• Allow-lists hidream for gradient_checkpointing_interval and gradient_checkpointing_segment_stride in safety checks.
• Updates HiDream transformer to expose interval/segment-stride setters, accept unsloth* backends, and gate per-block checkpointing via should_checkpoint_block.

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 HiDream model-family capability presence test for segmented checkpointing controls.
simpletuner/helpers/training/default_settings/safety_check.py	Adds hidream to the safety-check allow-lists for interval + segment stride options.
simpletuner/helpers/models/hidream/transformer.py	Implements HiDream-side checkpointing hooks (interval/stride + unsloth backend prefixes) and integrates should_checkpoint_block into the forward pass.

Suppressed comments (8)

simpletuner/helpers/models/hidream/transformer.py:1327

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

This issue also appears in the following locations of the same file:

• line 1442
• line 1501
• line 1546
• line 1603
• line 1717
• ...and 2 more

            if self.gradient_checkpointing_backend.startswith("unsloth"):
                from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                checkpoint_fn = offloaded_checkpoint
            else:

simpletuner/helpers/models/hidream/transformer.py:1446

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

            if self.gradient_checkpointing_backend.startswith("unsloth"):
                from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                checkpoint_fn = offloaded_checkpoint
            else:

simpletuner/helpers/models/hidream/transformer.py:1505

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

            if self.gradient_checkpointing_backend.startswith("unsloth"):
                from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                checkpoint_fn = offloaded_checkpoint
            else:

simpletuner/helpers/models/hidream/transformer.py:1607

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

            if self.gradient_checkpointing_backend.startswith("unsloth"):
                from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                checkpoint_fn = offloaded_checkpoint
            else:

simpletuner/helpers/models/hidream/transformer.py:1721

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

                if self.gradient_checkpointing_backend.startswith("unsloth"):
                    from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                    checkpoint_fn = offloaded_checkpoint
                else:

simpletuner/helpers/models/hidream/transformer.py:1858

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

                if self.gradient_checkpointing_backend.startswith("unsloth"):
                    from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                    checkpoint_fn = offloaded_checkpoint
                else:

simpletuner/helpers/models/hidream/transformer.py:1947

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

            if self.gradient_checkpointing_backend.startswith("unsloth"):
                from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                checkpoint_fn = offloaded_checkpoint
            else:

simpletuner/helpers/models/hidream/transformer.py:1550

• The non-"unsloth" branch always uses torch.utils.checkpoint.checkpoint, so any gradient_checkpointing_func passed via _set_gradient_checkpointing(...) (stored on self._gradient_checkpointing_func) will be ignored here. Use the stored function when available so the checkpointing hook actually works.

            if self.gradient_checkpointing_backend.startswith("unsloth"):
                from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

                checkpoint_fn = offloaded_checkpoint
            else:

---

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