Add segmented checkpointing controls

[bghira]

Summary:

• add gradient_checkpointing_segment_stride, attention activation offload, pinned offload bucket, and prefetch controls across CLI/env/WebUI/docs
• extend transformer checkpointing/offload support across supported image, video, and audio transformer families
• add CPU saved-tensor offload pooling/prefetch runtime and RamTorch/Musubi memory helper updates
• add segmented checkpointing documentation and translations, including the benchmark matrix

Validation:

• /Users/kash/src/SimpleTuner/.venv/bin/python -m unittest tests.test_gradient_checkpointing_backend tests.test_ramtorch tests.test_musubi_block_swap tests.test_stable_cascade_modules tests.test_transformers.test_auraflow_transformer tests.test_transformers.test_pixart_transformer tests.test_transformers.test_sanavideo_transformer tests.test_transformers.test_sd3_transformer tests.test_zlab_i1_model -v
• /Users/kash/src/SimpleTuner/.venv/bin/python -m black on changed Python files
• /Users/kash/src/SimpleTuner/.venv/bin/python -m isort on changed Python files
• /Users/kash/src/SimpleTuner/.venv/bin/python -m flake8 with repository hook args on changed Python files
• docs scan for local paths and session/LLM markers
• git diff --cached --check
• pre-commit hooks from git commit

[Copilot]

Pull request overview

This PR adds new memory-optimization controls (segmented gradient checkpointing stride plus attention-activation offload with pinned-bucket and prefetch knobs) and propagates support across many transformer model families, the WebUI field registry, configuration parsing/safety checks, and documentation/translations. It also extends the RamTorch/Musubi memory-management utilities and adds tests to lock in the new behaviors.

Changes:

• Introduces gradient_checkpointing_segment_stride and updates per-model checkpointing logic to support segmented “checkpoint N blocks, skip M blocks” patterns.
• Adds attention-side activation offload controls (gradient_checkpointing_offload_attention, pinned bucket pooling, and prefetch/autotune) and wires them through config, runtime, and several model blocks.
• Improves RamTorch forward-prefetch and DDP-ignore handling, expands Musubi swap support for SDNQ payloads, and adds targeted unit tests.

Reviewed changes

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

Show a summary per file

File	Description
tests/test_zlab_i1_model.py	Adds coverage ensuring ZLab I1 segmented checkpointing uses checkpoint_sequential_state.
tests/test_transformers/test_sd3_transformer.py	Adds tests for segmented checkpointing helper usage and attention offload context.
tests/test_transformers/test_sanavideo_transformer.py	Adds segmented checkpointing helper-call coverage for SanaVideo.
tests/test_transformers/test_pixart_transformer.py	Adds segmented checkpointing helper-call coverage for PixArt.
tests/test_transformers/test_auraflow_transformer.py	Adds a backward/grad correctness test under checkpointing.
tests/test_stable_cascade_modules.py	Adds interval/stride behavior validation for Stable Cascade UNet checkpointing.
tests/test_ramtorch.py	Adds tests for CUDA peak reset, DDP ignore filtering, and prefetch hook timing/order.
tests/test_musubi_block_swap.py	Adds test for SDNQ module streaming without _apply swap.
tests/test_cosmos3_model.py	Adds dtype-preservation tests for Cosmos3 unpatchify path.
tests/test_boogu_image_model.py	Adds test ensuring custom autograd outputs are cloned before layout writes.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py	Registers WebUI/CLI fields for attention offload, pinned buckets, and prefetch.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py	Adds gradient_checkpointing_segment_stride field and clarifies interval help/tooltip text.
simpletuner/helpers/utils/ramtorch.py	Refines mark_ddp_ignore_params to ignore RamTorch/CPU-resident frozen params/buffers more precisely.
simpletuner/helpers/training/trainer.py	Adds model-level enable/disable hooks and activation-offload prefetch autotune probe in training loop.
simpletuner/helpers/training/gradient_checkpointing_interval.py	Introduces should_checkpoint_block and extends checkpoint_sequential_state with segment_stride.
simpletuner/helpers/training/default_settings/safety_check.py	Validates/normalizes new stride and activation-offload CLI options; gates by model-family support.
simpletuner/helpers/ramtorch/profiling.py	Resets CUDA peak stats per device and includes activation-offload metrics in profiling snapshots.
simpletuner/helpers/ramtorch/modules/linear.py	Adds multi-stream forward prefetching, stream recording safeguards, and explicit forward residency discard.
simpletuner/helpers/ramtorch_extensions.py	Adds multi-stream transfer scheduling and stream recording for RamTorch extension modules and hooks.
simpletuner/helpers/musubi_block_swap.py	Adds SDNQ payload detection/movement support and refines “no swap” module move logic.
simpletuner/helpers/models/zlab_i1/transformer.py	Adds interval/stride fields and segmented whole-block checkpointing via checkpoint_sequential_state.
simpletuner/helpers/models/wan_s2v/transformer.py	Updates rotary embedding application and adds interval/stride checkpoint selection via helper.
simpletuner/helpers/models/stable_cascade/unet.py	Adds interval/stride state and centralized checkpoint decision via should_checkpoint_block.
simpletuner/helpers/models/sd3/transformer.py	Adds segment stride + segmented sequential checkpointing path and attention activation offload contexts.
simpletuner/helpers/models/sd3/expanded.py	Extends _set_gradient_checkpointing signature and broadens unsloth backend matching.
simpletuner/helpers/models/sanavideo/transformer.py	Adds interval/stride/backend fields; introduces segmented checkpointing path and helper-based per-block decisions.
simpletuner/helpers/models/sana/transformer.py	Adds segment stride support and migrates checkpoint decision to shared helper.
simpletuner/helpers/models/qwen_image/transformer.py	Adds interval/stride fields and helper-based checkpoint decision.
simpletuner/helpers/models/pixart/transformer.py	Adds interval/stride setters and segmented sequential checkpointing path.
simpletuner/helpers/models/mageflow/vendor/models/modules/mage_layers.py	Wraps attention compute in activation-offload context.
simpletuner/helpers/models/mageflow/transformer.py	Adds segment stride + attention offload toggles and uses sequential checkpointing for segments.
simpletuner/helpers/models/lumina2/transformer.py	Adds interval/stride/backend fields and segmented sequential checkpointing path.
simpletuner/helpers/models/ltxvideo2/model.py	Wires gradient_checkpointing_segment_stride into model load.
simpletuner/helpers/models/ltxvideo/transformer.py	Adds interval/stride/backend fields and segmented sequential checkpointing path.
simpletuner/helpers/models/longcat_video/transformer.py	Adds attention offload contexts, interval/stride controls, and propagates offload flag through blocks.
simpletuner/helpers/models/longcat_image/transformer.py	Adds attention offload contexts, interval/stride controls, and propagates offload flag through blocks.
simpletuner/helpers/models/krea2/transformer.py	Adds FFN-only checkpointing support, attention offload contexts, and segmented sequential checkpointing path.
simpletuner/helpers/models/kolors/controlnet.py	Extends _set_gradient_checkpointing signature and makes it recursive.
simpletuner/helpers/models/kandinsky5_video/transformer_kandinsky5.py	Adds attention offload plumbing and interval/stride-based checkpoint selection via helper.
simpletuner/helpers/models/ideogram/transformer.py	Adds interval/stride fields and applies helper-based checkpoint selection per layer.
simpletuner/helpers/models/hunyuanvideo/transformer.py	Adds segmented sequential checkpointing, attention offload, FFN-only checkpointing hooks, and trims unused conditioning tail.
simpletuner/helpers/models/hunyuanvideo/autoencoder.py	Extends _set_gradient_checkpointing signature and makes it recursive.
simpletuner/helpers/models/hidream/transformer.py	Adds interval/stride fields and helper-based checkpoint selection; updates unsloth backend matching.
simpletuner/helpers/models/flux2/transformer.py	Adds attention offload contexts, interval/stride/offload flags, and segmented sequential checkpointing for double/single streams.
simpletuner/helpers/models/flux/transformer.py	Adds attention offload contexts and segmented checkpointing gap logic (see review comments re: stride semantics in fallback path).
simpletuner/helpers/models/ernie/transformer.py	Adds segment stride and uses helper-based checkpoint decision; updates unsloth backend matching.
simpletuner/helpers/models/ernie/transformer_diffusers.py	Updates unsloth backend matching to accept prefixed variants.
simpletuner/helpers/models/cosmos3/transformer.py	Adds segment stride support via helper; fixes unpatchify to preserve materialized dtype when present.
simpletuner/helpers/models/cosmos/transformer.py	Adds interval/stride fields and helper-based checkpoint decision.
simpletuner/helpers/models/common.py	Wires new offload/prefetch/pinned bucket settings and segment stride into model setup; validates attention-offload support.
simpletuner/helpers/models/boogu_image/transformer.py	Adds interval/stride fields and helper-based checkpoint decisions; clones patch embed outputs to avoid in-place/layout hazards.
simpletuner/helpers/models/auraflow/transformer.py	Adds segment stride support and migrates per-block checkpoint decision to helper; broadens unsloth backend matching.
simpletuner/helpers/models/auraflow/controlnet.py	Updates unsloth backend matching to accept prefixed variants.
simpletuner/helpers/models/ace_step/transformer.py	Adds interval/stride fields and helper-based checkpoint decision; updates unsloth backend matching.
simpletuner/helpers/configuration/env_file.py	Adds env var mappings for new gradient-checkpointing offload and stride flags.
mkdocs.yml	Adds segmented checkpointing doc page to navigation.
documentation/OPTIONS.md	Documents new stride and activation-offload/prefetch/pinned-bucket options (English).
documentation/OPTIONS.zh.md	Documents new stride and activation-offload/prefetch/pinned-bucket options (Chinese).
documentation/OPTIONS.ja.md	Documents new stride and activation-offload/prefetch/pinned-bucket options (Japanese).
documentation/OPTIONS.pt-BR.md	Documents new stride and activation-offload/prefetch/pinned-bucket options (pt-BR).
documentation/OPTIONS.es.md	Documents new stride and activation-offload/prefetch/pinned-bucket options (Spanish).
documentation/OPTIONS.hi.md	Documents new stride and activation-offload/prefetch/pinned-bucket options (Hindi).
documentation/index.md	Adds segmented checkpointing link/mention to the documentation landing page (English).
documentation/index.zh.md	Adds segmented checkpointing link/mention to the documentation landing page (Chinese).
documentation/index.ja.md	Adds segmented checkpointing link/mention to the documentation landing page (Japanese).
documentation/index.pt-BR.md	Adds segmented checkpointing link/mention to the documentation landing page (pt-BR).
documentation/index.es.md	Adds segmented checkpointing link/mention to the documentation landing page (Spanish).
documentation/index.hi.md	Adds segmented checkpointing link/mention to the documentation landing page (Hindi).
documentation/experimental/UNSLOTH_CHECKPOINTING.md	Updates Unsloth checkpointing doc to reference segment stride and attention activation offload (English).
documentation/experimental/UNSLOTH_CHECKPOINTING.zh.md	Updates Unsloth checkpointing doc to reference segment stride and attention activation offload (Chinese).
documentation/experimental/UNSLOTH_CHECKPOINTING.ja.md	Updates Unsloth checkpointing doc to reference segment stride and attention activation offload (Japanese).
documentation/experimental/UNSLOTH_CHECKPOINTING.pt-BR.md	Updates Unsloth checkpointing doc to reference segment stride and attention activation offload (pt-BR).
documentation/experimental/UNSLOTH_CHECKPOINTING.es.md	Updates Unsloth checkpointing doc to reference segment stride and attention activation offload (Spanish).
documentation/experimental/UNSLOTH_CHECKPOINTING.hi.md	Updates Unsloth checkpointing doc to reference segment stride and attention activation offload (Hindi).

Suppressed comments (1)

simpletuner/helpers/models/flux/transformer.py:1413

• Same as the double-stream loop: the single-stream checkpoint decision uses index_block % gradient_checkpointing_interval == 0 and does not incorporate gradient_checkpointing_segment_stride unless the segmented path is taken. If segmented mode is disabled for any reason, stride semantics are lost. Aligning this conditional with should_checkpoint_block(...) (or equivalent inline logic) would keep --gradient_checkpointing_segment_stride behavior consistent.

---

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

[bghira]

Superseded by the split segmented checkpointing stack: base PR #2932, then model-specific PRs #2933 through #2961.