Add PixArt segmented checkpointing support

[bghira]

Summary

Splits the PixArt 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-mageflow

Validation

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

[Copilot]

Pull request overview

Adds PixArt integration for segmented gradient checkpointing, while also extending the segmented-checkpointing/offload control surface across multiple transformer families, the WebUI field registry, safety checks, RamTorch utilities, tests, and documentation.

Changes:

• Adds PixArt model hooks for gradient_checkpointing_interval + gradient_checkpointing_segment_stride and verifies sequential-state checkpointing behavior via unit tests.
• Introduces/extends segmented checkpointing and attention-activation offload plumbing across several model families, plus trainer-side activation-offload prefetch autotuning.
• Wires new CLI/env/WebUI options and documents them (including translated OPTIONS + experimental docs).

Reviewed changes

Copilot reviewed 72 out of 73 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_transformers/test_pixart_transformer.py	Adds unit coverage ensuring PixArt segmented checkpointing uses checkpoint_sequential_state.
tests/test_transformers/test_auraflow_transformer.py	Adds regression coverage for gradient checkpointing backward across joint/single blocks.
tests/test_segmented_checkpointing_model_support.py	New coverage asserting model-level capability flags/setters for segmented checkpointing controls.
tests/test_ramtorch.py	Adds tests for CUDA peak-mem reset, DDP ignore marking behavior, and prefetch hook timing.
tests/test_musubi_block_swap.py	Adds SDNQ module streaming regression test for Musubi block swap manager.
tests/test_cosmos3_model.py	Adds dtype-preservation tests for Cosmos3 unpatchify/unpack path.
tests/test_boogu_image_model.py	Adds test ensuring patch embedding refine clones custom autograd outputs.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py	Adds WebUI/CLI fields for activation offload + pin-bucket + prefetch controls.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py	Updates interval field help/tooltip and adds segment-stride field.
simpletuner/helpers/utils/ramtorch.py	Updates DDP ignore marking behavior for RamTorch/CPU-resident params/buffers.
simpletuner/helpers/training/trainer.py	Adds activation-offload prefetch autotune probe + refactors prediction/loss computation.
simpletuner/helpers/training/gradient_checkpointing_interval.py	Adds should_checkpoint_block and extends sequential checkpointing with segment_stride.
simpletuner/helpers/training/default_settings/safety_check.py	Expands model allow-lists and validates segment stride + activation offload settings.
simpletuner/helpers/ramtorch/profiling.py	Resets CUDA peak memory stats and reports activation-offload stats in profiling snapshot.
simpletuner/helpers/ramtorch/modules/linear.py	Adds multi-stream forward prefetch, record-stream handling, and forward-residency discard.
simpletuner/helpers/ramtorch_extensions.py	Adds multi-stream transfer selection + record-stream support + discard/preserve hook behavior.
simpletuner/helpers/musubi_block_swap.py	Adds SDNQ tensor/module handling for device moves without swapping.
simpletuner/helpers/models/pixart/transformer.py	Implements segmented checkpointing for PixArt blocks and stride-aware per-block checkpointing decisions.
simpletuner/helpers/models/mageflow/vendor/models/modules/mage_layers.py	Wraps attention with activation-offload context; adds offload_attention arg to block forward.
simpletuner/helpers/models/mageflow/transformer.py	Adds segment stride + attention-offload plumbing and sequential checkpointing for MageFlow.
simpletuner/helpers/models/lumina2/transformer.py	Adds backend/interval/stride hooks and sequential segmented checkpointing support.
simpletuner/helpers/models/ltxvideo2/transformer.py	Adds FFN-only checkpoint + activation-offload context wiring; adds sequential segmented checkpointing.
simpletuner/helpers/models/ltxvideo2/model.py	Plumbs gradient_checkpointing_segment_stride config into model setup.
simpletuner/helpers/models/ltxvideo/transformer.py	Adds backend/interval/stride hooks; adds segmented sequential checkpointing and stride-aware block checkpointing.
simpletuner/helpers/models/longcat_video/transformer.py	Adds stride-aware checkpointing decisions and activation-offload contexts for attention.
simpletuner/helpers/models/longcat_image/transformer.py	Adds stride-aware checkpointing decisions and activation-offload contexts for attention.
simpletuner/helpers/models/krea2/transformer.py	Adds attention activation offload + FFN-only checkpointing; adds sequential segmented checkpointing path.
simpletuner/helpers/models/kolors/controlnet.py	Updates _set_gradient_checkpointing signature for diffusers compatibility and recursive propagation.
simpletuner/helpers/models/kandinsky5_video/transformer_kandinsky5.py	Adds activation-offload contexts, stride-aware checkpointing, and model setters for checkpointing controls.
simpletuner/helpers/models/ideogram/transformer.py	Adds interval/stride setters and stride-aware checkpointing decisions.
simpletuner/helpers/models/hunyuanvideo/transformer.py	Adds activation-offload + FFN checkpointing support, segmented sequential checkpointing, and stride-aware logic.
simpletuner/helpers/models/hunyuanvideo/autoencoder.py	Updates _set_gradient_checkpointing signature and recursive propagation.
simpletuner/helpers/models/hidream/transformer.py	Adds interval/stride setters and stride-aware checkpointing decisions; broadens unsloth backend detection.
simpletuner/helpers/models/flux2/transformer.py	Adds activation-offload contexts + model flags/setters; adds sequential segmented checkpointing.
simpletuner/helpers/models/flux/transformer.py	Adds activation-offload contexts + model flags/setters; extends segmented checkpointing handling.
simpletuner/helpers/models/ernie/transformer.py	Adds stride setter and uses should_checkpoint_block; broadens unsloth backend detection.
simpletuner/helpers/models/ernie/transformer_diffusers.py	Broadens unsloth backend detection.
simpletuner/helpers/models/cosmos3/transformer.py	Adds stride setter and uses should_checkpoint_block; adjusts unpatchify dtype handling.
simpletuner/helpers/models/cosmos/transformer.py	Adds interval/stride setters and stride-aware checkpointing decisions.
simpletuner/helpers/models/common.py	Plumbs activation-offload config + pin buckets + prefetch + segment stride into model load/setup.
simpletuner/helpers/models/chroma/transformer.py	Adds activation-offload contexts + FFN-only checkpointing support and stride-aware checkpointing.
simpletuner/helpers/models/boogu_image/transformer.py	Adds interval/stride setters and stride-aware checkpointing decisions; clones embed outputs before layout writes.
simpletuner/helpers/models/auraflow/transformer.py	Adds stride setter and uses should_checkpoint_block; broadens unsloth backend detection.
simpletuner/helpers/models/auraflow/controlnet.py	Broadens unsloth backend detection.
simpletuner/helpers/models/ace_step/transformer.py	Adds interval/stride setters and stride-aware checkpointing decisions; broadens unsloth backend detection.
simpletuner/helpers/configuration/env_file.py	Adds env var mappings for new checkpointing/offload controls.
mkdocs.yml	Adds nav entry for segmented checkpointing experimental doc.
documentation/OPTIONS.md	Documents interval semantics, segment stride, attention offload, pin buckets, prefetch.
documentation/OPTIONS.zh.md	Same as above (zh translation).
documentation/OPTIONS.ja.md	Same as above (ja translation).
documentation/OPTIONS.pt-BR.md	Same as above (pt-BR translation).
documentation/OPTIONS.es.md	Same as above (es translation).
documentation/OPTIONS.hi.md	Same as above (hi translation).
documentation/index.md	Adds segmented checkpointing link under experimental methods.
documentation/index.zh.md	Same as above (zh translation).
documentation/index.ja.md	Same as above (ja translation).
documentation/index.pt-BR.md	Same as above (pt-BR translation).
documentation/index.es.md	Same as above (es translation).
documentation/index.hi.md	Same as above (hi translation).
documentation/experimental/UNSLOTH_CHECKPOINTING.md	Documents stride interplay with unsloth and attention-activation offload controls.
documentation/experimental/UNSLOTH_CHECKPOINTING.zh.md	Same as above (zh translation).
documentation/experimental/UNSLOTH_CHECKPOINTING.ja.md	Same as above (ja translation).
documentation/experimental/UNSLOTH_CHECKPOINTING.pt-BR.md	Same as above (pt-BR translation).
documentation/experimental/UNSLOTH_CHECKPOINTING.es.md	Same as above (es translation).
documentation/experimental/UNSLOTH_CHECKPOINTING.hi.md	Same as above (hi translation).

Suppressed comments (1)

simpletuner/helpers/models/chroma/transformer.py:593

• Same as above: this FFN-only checkpoint path doesn’t pass use_reentrant=False to the checkpoint function. That can unintentionally enable re-entrant checkpointing when checkpoint_fn is torch.utils.checkpoint.checkpoint, diverging from the non-reentrant behavior used elsewhere.

---

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