Add segmented checkpointing base controls

[bghira]

Summary

Adds the reusable segmented-checkpointing and activation-offload infrastructure that model integrations can build on:

• registers the new checkpointing controls in CLI/env/WebUI field metadata
• adds the shared checkpoint scheduling, CPU activation offload, pinned-memory, prefetch, RamTorch, and Musubi helper updates
• wires model-foundation and trainer-level handling for the new controls without enabling new model families in safety checks
• adds translated option docs and the segmented-checkpointing benchmark reference

Split from #2925

This is the base PR for the #2925 split. Model-specific support is intentionally kept out of this branch so follow-up PRs can show each model family separately in the changelog.

Validation

• .venv/bin/python -m unittest tests.test_gradient_checkpointing_backend tests.test_ramtorch tests.test_musubi_block_swap -v
• pre-commit run --files <changed files>
• .venv/bin/python -m flake8 --max-line-length=125 --extend-ignore=E501,E266,E203,W503,F841,F401,E402,F811,E722,E721,F541 <changed python files>

[Copilot]

Pull request overview

This PR introduces baseline infrastructure for segmented gradient checkpointing and CPU activation offload (including pinned-memory pooling and learned/prefetched restores), wires the new controls into config/CLI/env/WebUI field metadata, and adds documentation + benchmarks to describe the new modes.

Changes:

• Add activation offload runtime primitives (pinned bucket pool, copy-stream pools, prefetch learning/autotune) and trainer-level autotune probing.
• Extend RamTorch/Musubi helpers to better handle CPU-resident params/buffers, multi-stream forward prefetch, and SDNQ payload streaming.
• Register new config fields and document them (OPTIONS + translations + new segmented-checkpointing doc and mkdocs nav entries).

Reviewed changes

Copilot reviewed 41 out of 41 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
tests/test_ramtorch.py	Adds coverage for RamTorch profiling reset, DDP ignore behavior, and prefetch-hook timing semantics.
tests/test_musubi_block_swap.py	Adds a regression test for SDNQ modules streaming without relying on _apply swaps.
tests/test_gradient_checkpointing_backend.py	Adds extensive tests for activation offload/prefetch internals, pin-memory pool behavior, and segmented checkpoint scheduling helpers.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/training.py	Registers activation offload + pinned-bucket + prefetch controls for the UI/CLI metadata.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py	Registers gradient_checkpointing_segment_stride and updates interval field help/tooltip.
simpletuner/helpers/utils/ramtorch.py	Adjusts DDP ignore-marking to include CPU-resident non-trainable params/buffers and RamTorch markers.
simpletuner/helpers/training/trainer.py	Adds activation-offload prefetch autotune probes and refactors loss computation into a helper.
simpletuner/helpers/training/offloaded_gradient_checkpointer.py	Replaces the previous minimal offload helper with full pinned-pool + prefetch runtime + copy-stream support.
simpletuner/helpers/training/gradient_checkpointing_interval.py	Adds segment_stride support to sequential checkpoint scheduling and a helper for per-block checkpoint decisions.
simpletuner/helpers/training/default_settings/safety_check.py	Adds validation / gating for new offload + stride controls.
simpletuner/helpers/ramtorch/profiling.py	Resets CUDA peak memory stats per device and snapshots activation-offload stats in profiling output.
simpletuner/helpers/ramtorch/modules/linear.py	Improves forward prefetch via multi-stream scheduling and safer stream-recording behavior.
simpletuner/helpers/ramtorch_extensions.py	Switches RamTorch prefetch to a “prefetch successor in pre-hook” model and adds discard/preserve logic.
simpletuner/helpers/musubi_block_swap.py	Extends Musubi streaming to handle SDNQ tensors/modules and quantized payload moves.
simpletuner/helpers/models/common.py	Wires new config values into model load/setup, including activation-offload settings and segment-stride setter calls.
simpletuner/helpers/configuration/env_file.py	Exposes new controls via env-var → CLI flag mappings.
mkdocs.yml	Adds Segmented Checkpointing doc to mkdocs navigation.
documentation/OPTIONS.md	Documents segment stride and activation-offload controls.
documentation/OPTIONS.zh.md	Updates Chinese OPTIONS translation for new controls and revised checkpointing text.
documentation/OPTIONS.pt-BR.md	Updates pt-BR OPTIONS translation for new controls and revised checkpointing text.
documentation/OPTIONS.ja.md	Updates Japanese OPTIONS translation for new controls and revised checkpointing text.
documentation/OPTIONS.hi.md	Updates Hindi OPTIONS translation for new controls and revised checkpointing text.
documentation/OPTIONS.es.md	Updates Spanish OPTIONS translation for new controls and revised checkpointing text.
documentation/index.md	Adds segmented checkpointing to the experimental features index section.
documentation/index.zh.md	Adds segmented checkpointing to the experimental features index section (zh).
documentation/index.pt-BR.md	Adds segmented checkpointing to the experimental features index section (pt-BR).
documentation/index.ja.md	Adds segmented checkpointing to the experimental features index section (ja).
documentation/index.hi.md	Adds segmented checkpointing to the experimental features index section (hi).
documentation/index.es.md	Adds segmented checkpointing to the experimental features index section (es).
documentation/experimental/UNSLOTH_CHECKPOINTING.md	Updates Unsloth doc to reference segmented stride and attention activation offload controls.
documentation/experimental/UNSLOTH_CHECKPOINTING.zh.md	Updates zh Unsloth doc with new control references.
documentation/experimental/UNSLOTH_CHECKPOINTING.pt-BR.md	Updates pt-BR Unsloth doc with new control references.
documentation/experimental/UNSLOTH_CHECKPOINTING.ja.md	Updates ja Unsloth doc with new control references.
documentation/experimental/UNSLOTH_CHECKPOINTING.hi.md	Updates hi Unsloth doc with new control references.
documentation/experimental/UNSLOTH_CHECKPOINTING.es.md	Updates es Unsloth doc with new control references.
documentation/experimental/SEGMENTED_CHECKPOINTING.md	Adds a new, detailed segmented checkpointing reference and benchmark matrix.

Suppressed comments (3)

simpletuner/helpers/training/offloaded_gradient_checkpointer.py:779

• Restoring a dense view after CPU offload uses torch.as_strided(..., storage_offset=0). This breaks views that originally had a non-zero storage offset; it should use the offset captured in _RestoreView.
  simpletuner/helpers/training/offloaded_gradient_checkpointer.py:697
• When creating _RestoreView for a flattened dense view, the original tensor storage_offset is not preserved, so the later restore cannot reconstruct offset views correctly.
  simpletuner/helpers/training/offloaded_gradient_checkpointer.py:792
• _restore_view_if_needed also hardcodes storage_offset=0 when rebuilding a view, which is incorrect for offset views. Use the stored offset instead.

---

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