segmented checkpointing for better cache hits / latency - #2920
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds segmented (chunked) gradient checkpointing for Flux and MageFlow transformers by checkpointing contiguous block groups, aiming to reduce checkpoint overhead/latency while preserving the existing memory–compute tradeoff options. It also aligns MageFlow’s public interface with Flux and updates tests + documentation accordingly.
Changes:
- Introduces
checkpoint_sequential_stateto checkpoint contiguous segments while carrying multi-tensor state through a block sequence. - Integrates segmented checkpointing into Flux and MageFlow transformer forward passes, and adds MageFlow’s
gradient_checkpointing_interval/setter for parity. - Expands tests and updates WebUI + multi-language docs describing
--gradient_checkpointing_interval.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_gradient_checkpointing_backend.py | Adds unit tests for segmented checkpointing correctness and contiguous chunk boundaries. |
| simpletuner/simpletuner_sdk/server/services/field_registry/sections/model.py | Updates WebUI help/tooltip text for gradient_checkpointing_interval. |
| simpletuner/helpers/training/gradient_checkpointing_interval.py | Adds checkpoint_sequential_state helper and refreshes module framing. |
| simpletuner/helpers/training/default_settings/safety_check.py | Adds mageflow to the allowlist for gradient_checkpointing_interval. |
| simpletuner/helpers/models/mageflow/transformer.py | Adds interval attribute/setter and segmented checkpoint execution for contiguous block chunks. |
| simpletuner/helpers/models/flux/transformer.py | Adds segmented checkpointing paths for both block sequences, plus a precedence fix in the checkpoint condition. |
| documentation/OPTIONS.md | Updates CLI option docs for interval checkpointing behavior. |
| documentation/OPTIONS.zh.md | Updates translated CLI option docs for interval checkpointing behavior. |
| documentation/OPTIONS.pt-BR.md | Updates translated CLI option docs for interval checkpointing behavior. |
| documentation/OPTIONS.ja.md | Updates translated CLI option docs for interval checkpointing behavior. |
| documentation/OPTIONS.hi.md | Updates translated CLI option docs for interval checkpointing behavior. |
| documentation/OPTIONS.es.md | Updates translated CLI option docs for interval checkpointing behavior. |
| documentation/experimental/UNSLOTH_CHECKPOINTING*.md | Updates experimental docs (all languages) to describe chunked behavior on supported whole-block paths. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+493
to
+494
| help_text="Checkpoint contiguous chunks of N transformer blocks (leave blank for per-block checkpointing)", | ||
| tooltip="Higher values reduce recompute overhead but keep more activations in VRAM. Flux and MageFlow use contiguous chunked checkpointing on whole-block paths.", |
Comment on lines
+285
to
+286
| - **What**: Checkpoint contiguous chunks of *n* blocks, where *n* is a value greater than zero. A value of 1 is effectively the same as just leaving `--gradient_checkpointing` enabled, and a value of 2 checkpoints two-block chunks. | ||
| - **Note**: Flux and MageFlow use contiguous chunked checkpointing on whole-block paths. Higher values reduce recompute overhead but keep more activations in VRAM. |
Comment on lines
+279
to
+280
| - **内容**：checkpoint 连续的 *n* 个 block chunk，*n* 必须大于 0。1 等同于启用 `--gradient_checkpointing`，2 会 checkpoint 两个 block 的 chunk。 | ||
| - **说明**：Flux 和 MageFlow 会在 whole-block 路径上使用连续 chunk checkpointing。值越大，重算开销越低，但 VRAM 里保留的 activation 越多。 |
Comment on lines
+279
to
+280
| - **内容**: 連続した *n* block chunk を checkpoint します。値は 0 より大きい必要があります。1 は `--gradient_checkpointing` と同等で、2 は 2-block chunk を checkpoint します。 | ||
| - **注記**: Flux と MageFlow は whole-block path で連続 chunk checkpointing を使います。値を大きくすると再計算 overhead は減りますが、VRAM に残る activation は増えます。 |
Comment on lines
+278
to
+279
| - **Qué**: Hace checkpoint de chunks contiguos de *n* bloques, donde *n* es un valor mayor que cero. Un valor de 1 equivale a dejar `--gradient_checkpointing` habilitado, y 2 checkpointa chunks de dos bloques. | ||
| - **Nota**: Flux y MageFlow usan checkpointing en chunks contiguos en rutas whole-block. Valores más altos reducen recompute pero dejan más activaciones en VRAM. |
Comment on lines
+278
to
+279
| - **O que**: Faz checkpoint de chunks contiguos de *n* blocos, onde *n* e um valor maior que zero. Um valor 1 e efetivamente o mesmo que deixar `--gradient_checkpointing` habilitado, e 2 faz checkpoint de chunks de dois blocos. | ||
| - **Nota**: Flux e MageFlow usam checkpointing em chunks contiguos nos caminhos whole-block. Valores maiores reduzem recompute, mas mantêm mais activations na VRAM. |
Comment on lines
+278
to
+279
| - **What**: *n* contiguous blocks के chunks checkpoint करें, जहाँ *n* शून्य से बड़ा मान है। 1 का मान `--gradient_checkpointing` enabled जैसा है, और 2 दो-block chunks checkpoint करता है। | ||
| - **Note**: Flux और MageFlow whole-block paths पर contiguous chunked checkpointing use करते हैं। Higher values recompute overhead घटाती हैं लेकिन VRAM में ज्यादा activations रखती हैं। |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces segmented (chunked) gradient checkpointing for the Flux and MageFlow transformer models, enabling more efficient memory and computation trade-offs during training. It updates both the implementation and documentation across multiple languages to clarify how the new chunked checkpointing works and which models support it.
Segmented Gradient Checkpointing Implementation:
flux/transformer.pyandmageflow/transformer.py, using the newcheckpoint_sequential_stateutility. This allows for more flexible and efficient memory usage during training. [1] [2] [3] [4]set_gradient_checkpointing_intervalmethod and thegradient_checkpointing_intervalattribute to MageFlow, aligning its interface with Flux for chunked checkpointing. [1] [2]gradient_checkpointing_interval.pyutility to support the new chunked checkpointing logic.Documentation Updates:
--gradient_checkpointing_intervalnow checkpoints contiguous chunks of blocks, and clarified the impact of higher values on memory and recompute. [1] [2] [3] [4] [5] [6]These changes make gradient checkpointing more efficient and flexible, and ensure that both the code and documentation accurately reflect the new capabilities.