Segmented Checkpointing¶

Segmented checkpointing हर block को checkpoint करने और कोई checkpointing न करने के बीच का mode है।

यह PyTorch का activation checkpointing उपयोग करता है। SimpleTuner लगातार Transformer blocks के समूह को एक checkpoint कॉल में चलाता है और लौटे hidden state को अगले समूह तक पहुँचाता है। बड़े समूह कम boundary states सहेजते हैं, लेकिन backward में checkpoint किए गए blocks की गणना फिर भी दोहराते हैं; गति और अधिकतम मेमोरी पर प्रभाव मापें।

CPU offload और FFN-only checkpointing के लिए Unsloth-style checkpointing देखें। छोटी rule of thumb Decision Rule में है।

Controls¶

{
  "gradient_checkpointing": true,
  "gradient_checkpointing_backend": "torch",
  "gradient_checkpointing_interval": 2
}

Supported whole-block paths पर gradient_checkpointing_interval segment width है। 2 का मतलब blocks 0-1, 2-3, 4-5 वगैरह checkpoint होंगे।

ज्यादा fine VRAM control के लिए stride जोड़ें:

{
  "gradient_checkpointing": true,
  "gradient_checkpointing_backend": "torch",
  "gradient_checkpointing_interval": 2,
  "gradient_checkpointing_segment_stride": 4
}

यह blocks 0-1 checkpoint करता है, 2-3 normal चलाता है, 4-5 checkpoint करता है, 6-7 normal चलाता है, और repeat करता है। stride कम से कम interval जितना होना चाहिए; overlapping schedules valid नहीं हैं।

Supported segmented whole-block paths: Flux.1, Flux.2, Flux.3 Image Klein, HunyuanVideo, Krea 2, LongCat Image, LongCat Video, LTXVideo 0.9, LTXVideo2, Lumina2, MageFlow, MiniMax H3, PixArt, Qwen Image 2.1, SD3, SanaVideo, Z-Image, ZLab I1, और Wan।

Stable Cascade stage C भी interval और stride support करता है, लेकिन schedule transformer whole-block groups की जगह UNet के Res/Timestep/Attention micro-block sequence पर apply होता है।

कुछ model families पुराने semantics use करते हैं:

जहां stride ignore होता है, वहां stride rows compare न करें। अगर numbers identical दिखें, तो usually option apply नहीं हुआ।

कब इस्तेमाल करें¶

जब normal per-block checkpointing fit हो जाए लेकिन step time बहुत महंगा हो, तब इसे इस्तेमाल करें। 2 से शुरू करें। अगर VRAM allow करे, बहुत deep models पर 2 with stride 4 try करें।

जब peak मुख्य रूप से trainable weights, optimizer state, validation, VAE caching, block swapping, या routing से हो, तो इससे मदद की उम्मीद न करें। जब model feature को per-block control चाहिए, SimpleTuner safer per-block path पर वापस जाता है।

dynamo_use_regional_compilation universal win नहीं है। कुछ image models पर यह helpful या neutral था, लेकिन नीचे Wan/RamTorch और LTXVideo2 profiles में खराब fit था।

Benchmarks¶

Real SimpleTuner examples से single-GPU H100, H200 और L40S pods पर measure किया गया। validation और checkpoint saves disabled थे, cache preparation exclude था, और post-warmup timing available होने पर train loop के first-step compile/setup को भी exclude किया गया।

हर measured cell post-warmup sec/step / peak VRAM GiB है। status-only cells का मतलब है: OOM GPU memory खत्म होना, failed measured training steps तक न पहुंचना, unsupported उस family में option wired न होना, और not run उस combination का sweep में न होना।

पहले एक ही family के अंदर modes compare करें। Cross-family comparison rough है क्योंकि resolution, frame count, attention backend, model depth, trainable adapter type, और dataset shape अलग हो सकते हैं।

नीचे की matrix इस sweep की source of truth है। Model-specific notes caveats बताते हैं जब कोई row recommendation के बजाय coverage data हो।

Family Sweep Results¶

ACE Step 1.5¶

Example: ace_step/v15-base/default/v1-5.peft-lora. Resolution: 512.

Note: This sweep did not produce a usable ACE Step throughput row. The status-only entries below should be treated as coverage gaps, not as a recommendation.

Anima पर AnyFlow Distillation¶

Example: anima/base-v1.0/default/anyflow-stage1.peft-lora। Resolution: 1024x1024। यह row Anima के साथ AnyFlow distillation मापती है, plain Anima LoRA example नहीं। plain 1024x1024 Anima image training के लिए anima/base-v1.0/default/peft-lora use करें।

AuraFlow¶

Example: auraflow/v0.3/default/peft-lora. Resolution: 1024x1024.

Note: AuraFlow SDNQ और TorchAO quantization support करता है। Quantized none rows नीचे हैं; quantized checkpoint rows में numbers डालने से पहले fresh full-length benchmark coverage चाहिए।

Boogu Image¶

Example: boogu_image/v0.1-base/default/peft-lora. Resolution: 1024x1024.

Chroma¶

Example: chroma/default/default/peft-lora. Resolution: 1024x1024.

Note: Checkpointed Chroma rows use attention_mechanism=native-efficient, which was the stable attention path for this sweep.

Cosmos 2 Image¶

Example: cosmos/default/default/lycoris-lokr. Resolution: 512x512.

Cosmos 3¶

Example: cosmos3/edge/24g/image.lycoris-lokr. Resolution: 1024 px.

ERNIE 4.5 Image¶

Example: ernie/base/default/peft-lora. Resolution: 512x512.

Flux.3 Image Klein¶

उदाहरण: flux3_image/klein-base/default/peft-lora, Klein Base। प्रशिक्षण: 512x512 center crops, batch 1, rank/alpha 32 attention LoRA, BF16 AdamW, native SDPA और regional Inductor compilation (default, C++ wrapper)। CREPA और assistant adapter बंद हैं। हर मापी गई पंक्ति में 30 updates हैं; step time में पहले पाँच updates और cache preparation शामिल नहीं हैं। VRAM प्रशिक्षण updates के दौरान PyTorch की peak allocated memory है।

Dataset: RareConcepts/Domokun, 27 images। स्थिर caption: a photograph of a toy character (80 cached text tokens)।

PyTorch 2.14.1+cu130; Diffusers 0.41.0; Transformers 5.20.1; PEFT 0.21.2; Accelerate 1.15.0; SDNQ 0.2.8; TorchAO 0.17.0+cu130. dynamo_dynamic=true.

Image, text, reference-image और joint stream की block sequences अलग-अलग segment होती हैं। interval2 लगातार दो blocks को एक साथ checkpoint करता है; seg2-stride4 दो blocks checkpoint करता है और अगले दो की activations रखता है। Hidden-state capture per-block path इस्तेमाल करता है। Attention activation offload समर्थित नहीं है।

HeartMula¶

Example: heartmula/3b/default/peft-lora. Audio-token training.

HiDream¶

Example: hidream/full/default/peft-lora. Resolution: 512x512.

SDNQ Hadamard numbers use sdnq_compile_mode=eager. The compiled SDNQ path quantizes HiDream, but this sweep spent the first training step in Inductor dequantizer compilation, so it is not listed as a throughput row.

HunyuanVideo¶

Example: hunyuanvideo/t2v-480p/default/1.5-t2v.peft-lora. Training shape: 480 pixel-area video buckets, 48 frames, batch 2.

HunyuanVideo इस training shape पर activation-heavy है। Per-block और interval-2 checkpointing साफ चलते हैं; checkpointing बंद करने पर 80 GB H100 भी पर्याप्त नहीं है। seg2-stride4 इस sweep में attention activation offload के साथ ही चला, इसलिए वह speed recommendation नहीं बल्कि memory fallback है। SDNQ Hadamard काम करता है, लेकिन variable conditioning shapes measurement window में भी dynamic kernel compilation trigger करते हैं।

Ideogram 4.0¶

Example: ideogram/fp8/default/peft-lora. Resolution: 1024x1024. fp8 flavour Ideogram 4 के native weight-only fp8 checkpoint का उपयोग करता है (base_model_precision=no_change); bf16-upcast ideogram_fp8_base_upcast=true सेट करता है ताकि load पर base weights को bf16 में dequantize किया जाए। कोई भी अन्य base_model_precision (जैसे int8-sdnq) मांगने पर भी पहले dequantize होता है, ताकि quantizer असली weights पर काम करे।

इस table के पहले के संस्करण तब मापे गए थे जब Ideogram gradient_checkpointing_interval, gradient_checkpointing_segment_stride या gradient_checkpointing_backend को लागू नहीं करता था (उसका custom loader साझा wiring छोड़ देता था), और जब int8-sdnq वास्तव में fp8 checkpoint को quantize नहीं करता था — सभी checkpointing rows चुपचाप fp8-native weights पर full-layer torch checkpointing थीं। नीचे के H100 आंकड़े fix के बाद के हैं; L40S rows की re-measurement बाकी है।

torch-ffn/unsloth-ffn backends unsupported हैं: Ideogram 4 attention/FFN checkpointing boundary expose नहीं करता।

निष्कर्ष: seg2-stride4 full-layer से ~14% तेज है, बदले में ~24 GiB अधिक retained activations लगती हैं; unsloth offload ~1.2 GiB बचाता है और step-time में 1-3% की लागत आती है (int8 full-layer में यह torch से तेज है, क्योंकि offload quantized matmul के साथ overlap होता है); VRAM अनुमति दे तो bf16-upcast throughput का विजेता है।

Kandinsky 5 Image¶

Example: kandinsky5_image/default/default/6b-t2i.lycoris-lokr. Resolution: 1024x1024.

Note: batch 3 at 1024x1024 needs full checkpointing on both cards. SDNQ with Hadamard is the best low-VRAM row; H100 can also use partial checkpointing with SDNQ, but only near the top of the card.

Kandinsky 5 Video¶

Example: kandinsky5_video/default/default/2b-t2v.peft-lora. Resolution: 768x512, 81f.

Kandinsky 5 video is activation-heavy at this frame count. Full block checkpointing is the practical baseline on both cards. On H100, interval2 and seg2-stride4 are faster when they fit; on L40S, SDNQ interval2 is the only partial-checkpoint row here that fits without attention activation offload.

Kolors¶

Example: kolors/1.0/default/peft-lora. Resolution: 1024x1024.

Krea 2¶

Example: krea2/raw/default/peft-lora. Training resolution: 512 px square crop. Example validation setting: 1024x1024; benchmark में validation disabled था.

Main table regional compilation use करती है। यह Krea2 step speed के लिए अच्छा है, लेकिन VRAM comparison साफ नहीं रहता: compiled graph/workspace कई checkpoint modes में peak को uncheckpointed peak के करीब रखता है। regional compilation बंद करके bf16 control run ने दिखाया कि checkpointing wired है और memory/speed shape expected है:

यहां activation-offload row का मतलब full-block checkpointing plus attention activation offload है। सिर्फ full-block layer checkpointing के मुकाबले, इस Krea2 shape में attention offload ने peak VRAM घटाया नहीं; इसने mostly CPU transfer overhead जोड़ा।

LongCat Image¶

Example: longcat_image/default/default/peft-lora. Training resolution: 512 px square; validation resolution: 1024x1024. Rows use attention_mechanism=native-flash.

LongCat Video¶

Example: longcat_video/final/default/peft-lora+ramtorch. Resolution: 832x480, 81f. Rows attention_mechanism=native-flash use करते हैं।

LongCat Video इस shape पर activation-heavy है। Full per-block checkpointing practical row है। Partial checkpoint rows (interval2, seg2-stride4) यहाँ fit नहीं होते, strided row पर attention activation offload enabled होने पर भी नहीं। Plain attention activation offload bf16 और SDNQ में fit होता है, लेकिन full checkpointing से काफी धीमा है।

LTXVideo 0.9.5¶

Example: ltxvideo/0.9.5/default/t2v.peft-lora. Resolution: 768x512, 49f.

Numbers warm seconds/step / peak GiB हैं। full-run average setup और compile overhead include करता है, इसलिए वह sweep artifacts में रखा गया है।

इस sweep में LTXVideo 0.9 के लिए attention activation offload rows supported नहीं हैं।

LTXVideo2 2.3¶

Example: ltxvideo2/2.3-dev/default/720p-single-gpu.peft-lora+sdnq-hadamard. Resolution: 1280x704, 49f.

Note: LTXVideo2 2.3 should be read from the no-regional-compile rows in this sweep; regional compile raised memory pressure for this model.

Lumina2¶

Example: lumina2/default/default/peft-lora. Resolution: 512x512.

Note: Lumina2 अब segmented whole-block path इस्तेमाल करता है। interval2 हर two-block segment को checkpoint करता है; seg2-stride4 दो blocks checkpoint करता है, अगले दो blocks को activations रखने देता है, फिर repeat करता है। Attention activation offload इस Lumina2 run में शामिल नहीं था।

MageFlow¶

Example: mageflow/turbo/24g/image.peft-lora. Resolution: 1024x1024.

Note: MageFlow के 1024px variable-shape image path में सबसे ज्यादा फायदा attention activation offload और weight-only FP8 से मिला। Block checkpointing modes valid हैं, लेकिन इस sweep में measured peak residency कम नहीं हुई।

OmniGen¶

Example: omnigen/default/default/lycoris-lokr. Resolution: 1024x1024.

Note: OmniGen cached text embeddings की जगह token-ID prompts use करता है। ये rows supported no-checkpointing और full-block torch checkpointing paths को measure करती हैं; interval, segmented-stride, और attention-offload controls इस family के लिए इस sweep में implemented नहीं हैं।

PixArt¶

Example: pixart/default/default/lycoris-lokr. Resolution: 1024x1024.

Qwen Image 2.1¶

GPU मॉडल की पुष्टि के बाद H200 और L40S पर माप लिया गया: qwen_image/v2.1/default/peft-lora, Qwen-Image 2.1, 512x512 Domokun, BF16, rank-32 LoRA, Optimi Lion और regional Inductor compilation (default मोड)। हर रन में 20 steps हैं; पहले पाँच timing से बाहर हैं। हर cell सेकंड/step / अधिकतम VRAM GiB दिखाता है; peak में components की तैयारी शामिल है।

Runtime: torch==2.11.0+cu128, diffusers==0.40.0; 10 text tokens और 1024 image tokens। Batch 1 में कुछ memory peaks components की तैयारी से तय होते हैं।

BF16 checkpoint schedules¶

none checkpointing बंद करता है; layer हर block पर अलग checkpoint लगाता है। interval2 लगातार 0–1, 2–3, … blocks को समूह बनाता है; seg2-stride4 में 0–1 checkpoint होते हैं, 2–3 की activations रखी जाती हैं और क्रम दोहरता है। समूह execution के लिए block swapping, hidden-state capture और KV cache बंद होने चाहिए; उन paths पर प्रति-block interval/stride schedule बना रहता है। Attention activation offload (activation-offload, seg2-stride4-offload) समर्थित नहीं है। BF16 फिट होता है, इसलिए इस 2.1 sweep में quantization शामिल नहीं है।

Attention backends¶

Attention तुलना में checkpointing बंद है। Default native SDPA ने H200 पर cuDNN और L40S पर PyTorch Flash SDPA चुना। Qwen 2.1 text के लिए causal attention और target image के लिए सभी valid keys पर noncausal attention उपयोग करता है। समान लंबाई वाले बिना padding के prompts को varlen की ज़रूरत नहीं है। वर्तमान FlashAttention paths को बिना padding के text-to-image prompts चाहिए; padded prompts या interleaved conditioning के लिए native SDPA या FlexAttention उपयोग करें। केवल padding बताने वाला varlen mask उस causal संरचना को सुरक्षित नहीं रख सकता।

समान लंबाई वाले varlen metadata का helper upstream Diffusers से लिया गया है: lengths shapes से मिलती हैं और cumulative offsets के लिए arange उपयोग होता है। इससे प्रकाशित helper के GPU .item() synchronization और graph breaks हटते हैं। सभी positions valid होने पर अनावश्यक mask collate में हटा दिया जाता है। Forward/gradient समानता और graph capture की जाँच timing से अलग होती है।

attention_mechanism: "cudnn" को force करने पर PyTorch 2.11 Inductor में CantSplit त्रुटि आई; automatic native SDPA ने cuDNN सफलतापूर्वक उपयोग किया। FA4 Hub में nvidia-cutlass-dsl==4.7.1 के साथ cutlass.cute.core.ThrMma उपलब्ध नहीं था, इसलिए load विफल हुआ और timing उपलब्ध नहीं है। Hub backends में trust_remote_code: true है। FA3 Hopper के लिए है, L40S के लिए नहीं (upstream support)।

H200 batch scaling¶

उसी H200 पर native SDPA ने batch 10 में 0.612 सेकंड और batch 20 में 1.223 सेकंड लिए, यानी दोनों में लगभग 16.35 images/s। Throughput के स्थिर हो जाने के बाद batch दोगुना करने से काम भी लगभग दोगुना हुआ। 144 GB preset में checkpointing बंद ही है; बड़ा batch अतिरिक्त क्षमता उपयोग करता है, लेकिन अधिक images/s की गारंटी नहीं देता।

Cached Domokun inputs पर पूरे transformer की जाँच में redundant mask के साथ और बिना FP32 outputs समान थे; gradient relative L2 error 1.97e-6 था। BF16 backends संख्यात्मक रूप से समान नहीं हैं: इस एक batch में FP32 के मुकाबले पूरे LoRA gradient की cosine similarity native SDPA के लिए लगभग 0.92–0.93 और FA3 के लिए 0.80 थी। ये numerical diagnostics हैं, convergence की तुलना नहीं।

H200 batch-20 trace में captured interval के 93.1% समय kernels चल रहे थे। Kernel time का 70.5% GEMM और 9.7% attention में था; Inductor ने normalization/RoPE और pointwise epilogues पहले ही fuse कर दिए थे। सबसे बड़े खाली अंतराल batch की तैयारी के आसपास थे। Kernel active time, SM occupancy नहीं है और profiler timings throughput tables में शामिल नहीं हैं।

27 images के इस dataset में batch 20 पर prefetch धीमा था: queue length 2 के साथ host prefetch का औसत 2.030 सेकंड/step और device prefetch (1 MiB threshold) का 2.036 सेकंड था, जबकि बिना prefetch 1.223 सेकंड लगा। Median लगभग 1.17 सेकंड होने पर भी दोनों में बीच-बीच में लगभग 3 सेकंड के steps आए; presets में prefetch बंद है।

250-step Domokun recipe throughput मापने का उदाहरण है, भरोसेमंद convergence recipe नहीं। पहले के एक checkpoint ने दोबारा लोड करने पर पहचानने योग्य Domokun बनाया, लेकिन नए 250-step runs वह परिणाम दोहरा नहीं पाए। Padding masks बनाए रखने, compilation बंद करने और पुराने RoPE expression का उपयोग करने वाले controls भी विफल रहे। Cached latents सही विषय में decode होते हैं। Training deterioration का कारण अभी स्पष्ट नहीं है; timing tables अलग attention backends की समान image quality साबित नहीं करतीं।

Qwen Image¶

ये Qwen-Image 1.0 (model_flavour: "v1.0") के पुराने माप हैं, जो पुराने qwen_image/v2.1/default/peft-lora उदाहरण पर 1024x1024 पर लिए गए थे। उदाहरण अब 2.1 चुनता है। नीचे interval/stride के लगभग समान परिणाम अलग-अलग schedules सक्रिय होने की पुष्टि नहीं करते और इन्हें 2.1 की सेटिंग का आधार नहीं बनाना चाहिए।

Sana¶

Example: sana/default/default/lycoris-lokr. Resolution: 1024x1024.

Note: Sana has interval checkpointing; stride is not a separate segmented schedule for this family in the measured rows.

SanaVideo¶

Example: sanavideo/2b-480p/default/peft-lora. Resolution: 832x480, 49f.

Note: SanaVideo linear attention use करता है, इसलिए attention activation offload unsupported रहता है। Standard path में segmented whole-block checkpointing supported है।

SD 1.x¶

Example: sd1x/dreamshaper/default/peft-lora. Resolution: 512x512.

Note: SD1x diffusers UNet path use करता है. Regular layer checkpointing supported है, लेकिन interval और segmented stride controls इस family में wired नहीं हैं.

SD3¶

Example: sd3/default/default/peft-lora. Resolution: 1024x1024.

Note: SD3 plain transformer path पर real contiguous segmented checkpointing use करता है. Attention activation offload supported है; यह VRAM काफी घटाता है, लेकिन throughput खर्च करता है.

SDXL¶

Example: sdxl/default/default/lycoris-lokr. Resolution: 1024x1024.

Note: SDXL has real layer checkpointing. Interval and stride rows are included as coverage data, not as segmented-support recommendations.

Stable Cascade¶

Example: stable_cascade/default/default/stage-c.lycoris-lokr. Resolution: 1024x1024.

Note: stage C full-precision prior path use करता है। ये rows mixed_precision=no और base_model_precision=no_change के साथ run हुईं; quantized base precision rows इस model के लिए meaningful नहीं हैं। interval और stride modes UNet के Res/Timestep/Attention micro-block sequence पर काम करते हैं।

Wan 2.1 T2V 1.3B¶

Example: wan/t2v-480p-1.3b-2.1/default/t2v-1.3b-480p-single-gpu.peft-lora+ramtorch. Resolution: 832x480, 81f.

Note: Wan 1.3B should be read from the no-regional-compile/RamTorch rows; regional compile was not a useful throughput setting in this sweep.

Wan 2.1 T2V 14B¶

Example: wan/t2v-480p-14b-2.1/default/t2v-14b-480p-single-gpu.peft-lora+ramtorch. Resolution: 832x480, 81f.

Note: Wan 14B is mainly a fit test for activation savings. Status-only cells are still useful because they show which combinations reached the memory limit.

Wan S2V¶

Example: wan_s2v/s2v-14b-2.2/default/14b-480p.peft-lora+ramtorch. Resolution: 832x480, 81f.

Note: Wan S2V is included as coverage data for the video/audio path. Treat failed cells as implementation coverage gaps.

Z-Image Turbo¶

Example: z_image/turbo/default/peft-lora. Resolution: 1024x1024.

ZLab I1¶

Example: zlab_i1/3b/default/peft-lora. Resolution: 1024x1024.

Note: ZLab I1 अपने U-Net-style skip tensors को segmented checkpoint state में carry करता है। इस family में attention activation offload wired नहीं है।