Add segmented checkpointing base controls

documentation/OPTIONS.es.md

@@ -282,8 +282,14 @@ Donde `foo` es tu entorno de configuración; o simplemente usa `config/config.js
 
 ### `--gradient_checkpointing_interval`
 
-- **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.
+- **Qué**: Intervalo dependiente del modelo para checkpointing de bloques transformer. Un valor de 1 equivale básicamente a dejar `--gradient_checkpointing` habilitado.
+- **Nota**: Flux, Flux.2, Krea 2, LTXVideo2, MageFlow, Z-Image y Wan usan chunks contiguos de *n* bloques en rutas whole-block. Otras familias que exponen esta opción pueden seguir usando el comportamiento anterior de "checkpoint cada *n* bloques". Valores más altos pueden reducir recompute, pero normalmente dejan más activaciones en VRAM.
+
+### `--gradient_checkpointing_segment_stride`
+
+- **Qué**: Inicia un segmento con checkpoint cada *n* bloques en rutas segmented whole-block compatibles.
+- **Ejemplo**: Con `--gradient_checkpointing_interval=2` y `--gradient_checkpointing_segment_stride=4`, SimpleTuner checkpointa dos bloques, ejecuta los dos siguientes normalmente y repite.
+- **Nota**: Solo tiene efecto en familias de modelos que exponen soporte segmented whole-block en la version instalada de SimpleTuner. Las familias no soportadas registran una advertencia e ignoran el valor. El stride debe ser al menos igual al interval. Consulta [Segmented Checkpointing](experimental/SEGMENTED_CHECKPOINTING.md).
 
 ### `--gradient_checkpointing_backend`
 
@@ -293,7 +299,27 @@ Donde `foo` es tu entorno de configuración; o simplemente usa `config/config.js
   - `torch-ffn`: checkpoint solo del lado feed-forward en modelos con un límite FFN limpio.
   - `unsloth`: checkpoint del bloque completo compatible y offload de tensores guardados a CPU.
   - `unsloth-ffn`: checkpoint solo del lado feed-forward y offload de sus tensores guardados a CPU.
-- **Nota**: Solo efectivo cuando `--gradient_checkpointing` está habilitado. Las variantes `unsloth` requieren CUDA. Las variantes FFN-only soportan actualmente bloques estilo Flux.1 y MageFlow, y fallan de forma explícita si no existe ese scope. Consulta [Unsloth-style checkpointing](experimental/UNSLOTH_CHECKPOINTING.md) para tradeoffs medidos.
+- **Nota**: Solo efectivo cuando `--gradient_checkpointing` está habilitado. Las variantes `unsloth` requieren CUDA. Las variantes FFN-only soportan actualmente Chroma, Flux, Krea 2, LTXVideo2, MageFlow, Wan y Z-Image, y fallan de forma explícita si no existe ese scope. Consulta [Unsloth-style checkpointing](experimental/UNSLOTH_CHECKPOINTING.md) para tradeoffs medidos.
+
+### `--gradient_checkpointing_offload_attention`
+
+- **Qué**: Hace offload a CPU de las activaciones guardadas del lado attention en modelos con un límite attention/FFN limpio.
+- **Por qué**: Cuando transferir es más barato que recomputar attention, reduce VRAM sin pagar todo el coste de rematerializar attention.
+- **Nota**: Puede activarse por si solo. Tambien puede combinarse con cualquier checkpoint backend que soporte el modelo. Solo tiene efecto en familias de modelos que exponen una frontera attention/FFN limpia en la version instalada de SimpleTuner; las familias no soportadas fallan de forma explicita.
+
+### `--gradient_checkpointing_offload_pin_memory_max_buckets`
+
+- **Predeterminado**: `12`
+- **Qué**: Número máximo de buckets distintos de tensores CPU pinned usados por activation offload.
+- **Por qué**: Pinned memory mejora las transferencias CPU/GPU, pero resoluciones y longitudes de texto variables pueden crear formas raras. Al alcanzar este límite, las nuevas formas de bucket usan memoria CPU normal.
+- **Nota**: Usa `0` para desactivar el pooling de pinned memory para activation offload.
+
+### `--gradient_checkpointing_offload_prefetch`
+
+- **Predeterminado**: `false`
+- **Qué**: Aprende el orden de restore en backward para activations offloaded y precarga en GPU el tensor que probablemente venga después.
+- **Por qué**: El restore H2D justo a tiempo casi no se puede solapar. Con un orden estable, prefetch puede ocultar parte de la transferencia detrás del backward compute.
+- **Nota**: Experimental y solo activo con `--gradient_checkpointing_offload_attention`.
 
 ### `--refiner_training`
 
@@ -1741,6 +1767,7 @@ usage: train.py [-h] --model_family
                 [--text_encoder_3_precision {no_change,int8-quanto,int4-quanto,int2-quanto,int8-torchao,int8dq-torchao,int8dq-int4-torchao,nf4-bnb,int4-torchao,fp8-quanto,fp8uz-quanto,fp8-native,fp8-torchao,fp8wo-torchao,fp8-int4-torchao,fp8-transformerengine}]
                 [--text_encoder_4_precision {no_change,int8-quanto,int4-quanto,int2-quanto,int8-torchao,int8dq-torchao,int8dq-int4-torchao,nf4-bnb,int4-torchao,fp8-quanto,fp8uz-quanto,fp8-native,fp8-torchao,fp8wo-torchao,fp8-int4-torchao,fp8-transformerengine}]
                 [--gradient_checkpointing_interval GRADIENT_CHECKPOINTING_INTERVAL]
+                [--gradient_checkpointing_segment_stride GRADIENT_CHECKPOINTING_SEGMENT_STRIDE]
                 [--offload_during_startup [OFFLOAD_DURING_STARTUP]]
                 [--quantize_via {cpu,accelerator,pipeline}]
                 [--quantization_config QUANTIZATION_CONFIG]
@@ -2058,6 +2085,8 @@ options:
                         memory.
   --gradient_checkpointing_interval GRADIENT_CHECKPOINTING_INTERVAL
                         Checkpoint every N transformer blocks
+  --gradient_checkpointing_segment_stride GRADIENT_CHECKPOINTING_SEGMENT_STRIDE
+                        Start a checkpointed segment every N transformer blocks
   --offload_during_startup [OFFLOAD_DURING_STARTUP]
                         Offload text encoders to CPU during VAE caching
   --quantize_via {cpu,accelerator,pipeline}

documentation/OPTIONS.hi.md

@@ -282,8 +282,14 @@ simpletuner configure config/foo/config.json
 
 ### `--gradient_checkpointing_interval`
 
-- **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 रखती हैं।
+- **What**: Transformer block checkpointing के लिए model-dependent interval। 1 का मान लगभग `--gradient_checkpointing` enabled जैसा है।
+- **Note**: Flux, Flux.2, Krea 2, LTXVideo2, MageFlow, Z-Image, और Wan whole-block paths पर *n* contiguous block chunks use करते हैं। इस option को expose करने वाली दूसरी families अभी भी पुराने "हर *n*-th block checkpoint" behavior का उपयोग कर सकती हैं। Higher values recompute overhead घटा सकती हैं, लेकिन आम तौर पर VRAM में ज्यादा activations रखती हैं।
+
+### `--gradient_checkpointing_segment_stride`
+
+- **What**: Supported segmented whole-block paths पर हर *n* blocks में checkpointed segment शुरू करें।
+- **Example**: `--gradient_checkpointing_interval=2` और `--gradient_checkpointing_segment_stride=4` के साथ SimpleTuner दो blocks checkpoint करता है, अगले दो blocks normally चलाता है, और repeat करता है।
+- **Note**: यह केवल उन model families पर प्रभावी होता है जो installed SimpleTuner version में segmented whole-block support expose करती हैं। Unsupported families warning log करती हैं और value ignore करती हैं। stride interval से कम नहीं हो सकता। [Segmented Checkpointing](experimental/SEGMENTED_CHECKPOINTING.md) देखें।
 
 ### `--gradient_checkpointing_backend`
 
@@ -293,7 +299,27 @@ simpletuner configure config/foo/config.json
   - `torch-ffn`: साफ FFN boundary वाले models पर केवल feed-forward side checkpoint करता है।
   - `unsloth`: पूरे supported block को checkpoint करता है और saved tensors CPU पर offload करता है।
   - `unsloth-ffn`: केवल feed-forward side checkpoint करता है और उसके saved tensors CPU पर offload करता है।
-- **Note**: केवल `--gradient_checkpointing` enabled होने पर प्रभावी। `unsloth` variants के लिए CUDA आवश्यक है। FFN-only variants अभी Flux.1-style blocks और MageFlow support करते हैं, और unsupported scope पर साफ error मिलता है। Measured tradeoffs के लिए [Unsloth-style checkpointing](experimental/UNSLOTH_CHECKPOINTING.md) देखें।
+- **Note**: केवल `--gradient_checkpointing` enabled होने पर प्रभावी। `unsloth` variants के लिए CUDA आवश्यक है। FFN-only variants अभी Chroma, Flux, Krea 2, LTXVideo2, MageFlow, Wan, और Z-Image support करते हैं, और unsupported scope पर साफ error मिलता है। Measured tradeoffs के लिए [Unsloth-style checkpointing](experimental/UNSLOTH_CHECKPOINTING.md) देखें।
+
+### `--gradient_checkpointing_offload_attention`
+
+- **What**: साफ attention/FFN boundary वाले models पर attention-side saved activations CPU पर offload करें।
+- **Why**: जब transfer attention recompute से सस्ता हो, तो full attention rematerialization cost दिए बिना VRAM घट सकती है।
+- **Note**: इसे अकेले enable किया जा सकता है। Model जिस भी checkpoint backend को support करे, उसके साथ भी combine किया जा सकता है। यह केवल उन model families पर प्रभावी होता है जो installed SimpleTuner version में साफ attention/FFN boundary expose करती हैं; unsupported model families साफ error देती हैं।
+
+### `--gradient_checkpointing_offload_pin_memory_max_buckets`
+
+- **Default**: `12`
+- **What**: activation offload द्वारा इस्तेमाल किए जाने वाले अलग-अलग pinned CPU tensor buckets की maximum संख्या।
+- **Why**: pinned memory CPU/GPU transfer में मदद करती है, लेकिन variable resolution और text length rare tensor shapes बना सकते हैं। limit पूरी होने के बाद नए bucket shapes normal CPU memory use करेंगे।
+- **Note**: activation offload के pinned-memory pooling को बंद करने के लिए `0` set करें।
+
+### `--gradient_checkpointing_offload_prefetch`
+
+- **Default**: `false`
+- **What**: offloaded activations के backward restore order को learn करके likely next tensor को GPU पर prefetch करता है।
+- **Why**: JIT H2D restore अक्सर overlap नहीं कर पाता। order stable होने के बाद prefetch transfer के कुछ हिस्से को backward compute के पीछे छिपा सकता है।
+- **Note**: Experimental है और केवल `--gradient_checkpointing_offload_attention` के साथ active होता है।
 
 ### `--refiner_training`
 
@@ -1739,6 +1765,7 @@ usage: train.py [-h] --model_family
                 [--text_encoder_3_precision {no_change,int8-quanto,int4-quanto,int2-quanto,int8-torchao,int8dq-torchao,int8dq-int4-torchao,nf4-bnb,int4-torchao,fp8-quanto,fp8uz-quanto,fp8-native,fp8-torchao,fp8wo-torchao,fp8-int4-torchao,fp8-transformerengine}]
                 [--text_encoder_4_precision {no_change,int8-quanto,int4-quanto,int2-quanto,int8-torchao,int8dq-torchao,int8dq-int4-torchao,nf4-bnb,int4-torchao,fp8-quanto,fp8uz-quanto,fp8-native,fp8-torchao,fp8wo-torchao,fp8-int4-torchao,fp8-transformerengine}]
                 [--gradient_checkpointing_interval GRADIENT_CHECKPOINTING_INTERVAL]
+                [--gradient_checkpointing_segment_stride GRADIENT_CHECKPOINTING_SEGMENT_STRIDE]
                 [--offload_during_startup [OFFLOAD_DURING_STARTUP]]
                 [--quantize_via {cpu,accelerator,pipeline}]
                 [--quantization_config QUANTIZATION_CONFIG]
@@ -2056,6 +2083,8 @@ options:
                         memory.
   --gradient_checkpointing_interval GRADIENT_CHECKPOINTING_INTERVAL
                         Checkpoint every N transformer blocks
+  --gradient_checkpointing_segment_stride GRADIENT_CHECKPOINTING_SEGMENT_STRIDE
+                        Start a checkpointed segment every N transformer blocks
   --offload_during_startup [OFFLOAD_DURING_STARTUP]
                         Offload text encoders to CPU during VAE caching
   --quantize_via {cpu,accelerator,pipeline}

documentation/OPTIONS.ja.md

@@ -283,8 +283,14 @@ simpletuner configure config/foo/config.json
 
 ### `--gradient_checkpointing_interval`
 
-- **内容**: 連続した *n* block chunk を checkpoint します。値は 0 より大きい必要があります。1 は `--gradient_checkpointing` と同等で、2 は 2-block chunk を checkpoint します。
-- **注記**: Flux と MageFlow は whole-block path で連続 chunk checkpointing を使います。値を大きくすると再計算 overhead は減りますが、VRAM に残る activation は増えます。
+- **内容**: transformer block checkpointing のモデル依存 interval です。1 は `--gradient_checkpointing` を有効にした状態とほぼ同じです。
+- **注記**: Flux、Flux.2、Krea 2、LTXVideo2、MageFlow、Z-Image、Wan は whole-block path で連続した *n* block chunk を使います。このオプションを持つ他の family は、従来の「*n* block ごとに checkpoint」挙動のままの場合があります。値を大きくすると再計算 overhead は減ることがありますが、通常は VRAM に残る activation が増えます。
+
+### `--gradient_checkpointing_segment_stride`
+
+- **内容**: 対応する segmented whole-block path で、*n* block ごとに checkpointed segment を開始します。
+- **例**: `--gradient_checkpointing_interval=2` と `--gradient_checkpointing_segment_stride=4` では、SimpleTuner は 2 block を checkpoint し、次の 2 block を通常実行し、それを繰り返します。
+- **注記**: インストール済み SimpleTuner のバージョンで segmented whole-block support を公開している model family でのみ有効です。未対応 family では warning を記録し、値を無視します。stride は interval 以上である必要があります。[Segmented Checkpointing](experimental/SEGMENTED_CHECKPOINTING.md) を参照してください。
 
 ### `--gradient_checkpointing_backend`
 
@@ -294,7 +300,27 @@ simpletuner configure config/foo/config.json
   - `torch-ffn`: 明確な FFN 境界があるモデルで feed-forward 側だけを checkpoint します。
   - `unsloth`: 対応 block 全体を checkpoint し、保存 tensor を CPU に offload します。
   - `unsloth-ffn`: feed-forward 側だけを checkpoint し、保存 tensor を CPU に offload します。
-- **注記**: `--gradient_checkpointing` が有効な場合のみ機能します。`unsloth` 系は CUDA が必要です。FFN-only 系は現在 Flux.1-style blocks と MageFlow に対応し、対応していない scope では明示的に失敗します。実測 tradeoff は [Unsloth-style checkpointing](experimental/UNSLOTH_CHECKPOINTING.md) を参照してください。
+- **注記**: `--gradient_checkpointing` が有効な場合のみ機能します。`unsloth` 系は CUDA が必要です。FFN-only 系は現在 Chroma、Flux、Krea 2、LTXVideo2、MageFlow、Wan、Z-Image に対応し、対応していない scope では明示的に失敗します。実測 tradeoff は [Unsloth-style checkpointing](experimental/UNSLOTH_CHECKPOINTING.md) を参照してください。
+
+### `--gradient_checkpointing_offload_attention`
+
+- **内容**: 明確な attention/FFN 境界があるモデルで、attention 側の保存 activations を CPU に offload します。
+- **理由**: attention を再計算するより転送が安い場合、完全な attention rematerialization cost を払わずに VRAM を減らせます。
+- **注記**: 単体でも有効化できます。そのモデルがサポートする任意の checkpoint backend とも組み合わせられます。インストール済み SimpleTuner のバージョンで明確な attention/FFN boundary を公開している model family でのみ有効です。未対応モデルでは明示的に失敗します。
+
+### `--gradient_checkpointing_offload_pin_memory_max_buckets`
+
+- **デフォルト**: `12`
+- **内容**: activation offload が使う pinned CPU tensor bucket の最大数です。
+- **理由**: pinned memory は CPU/GPU 転送に有利ですが、可変解像度や可変 text length では珍しい tensor shape が出ます。上限に達した後の新しい bucket shape は通常の CPU memory を使います。
+- **注記**: `0` にすると activation offload の pinned-memory pooling を無効化します。
+
+### `--gradient_checkpointing_offload_prefetch`
+
+- **デフォルト**: `false`
+- **内容**: offload された activations の backward restore order を学習し、次に必要になりそうな tensor を GPU に prefetch します。
+- **理由**: JIT H2D restore はほとんど overlap できません。順序が安定すると、prefetch は一部の転送を backward compute の裏に隠せます。
+- **注記**: 実験的機能で、`--gradient_checkpointing_offload_attention` が有効な場合のみ動作します。
 
 ### `--refiner_training`
 
@@ -1742,6 +1768,7 @@ usage: train.py [-h] --model_family
                 [--text_encoder_3_precision {no_change,int8-quanto,int4-quanto,int2-quanto,int8-torchao,int8dq-torchao,int8dq-int4-torchao,nf4-bnb,int4-torchao,fp8-quanto,fp8uz-quanto,fp8-native,fp8-torchao,fp8wo-torchao,fp8-int4-torchao,fp8-transformerengine}]
                 [--text_encoder_4_precision {no_change,int8-quanto,int4-quanto,int2-quanto,int8-torchao,int8dq-torchao,int8dq-int4-torchao,nf4-bnb,int4-torchao,fp8-quanto,fp8uz-quanto,fp8-native,fp8-torchao,fp8wo-torchao,fp8-int4-torchao,fp8-transformerengine}]
                 [--gradient_checkpointing_interval GRADIENT_CHECKPOINTING_INTERVAL]
+                [--gradient_checkpointing_segment_stride GRADIENT_CHECKPOINTING_SEGMENT_STRIDE]
                 [--offload_during_startup [OFFLOAD_DURING_STARTUP]]
                 [--quantize_via {cpu,accelerator,pipeline}]
                 [--quantization_config QUANTIZATION_CONFIG]
@@ -2058,6 +2085,8 @@ options:
                         memory.
   --gradient_checkpointing_interval GRADIENT_CHECKPOINTING_INTERVAL
                         Checkpoint every N transformer blocks
+  --gradient_checkpointing_segment_stride GRADIENT_CHECKPOINTING_SEGMENT_STRIDE
+                        Start a checkpointed segment every N transformer blocks
   --offload_during_startup [OFFLOAD_DURING_STARTUP]
                         Offload text encoders to CPU during VAE caching
   --quantize_via {cpu,accelerator,pipeline}