Segmented Checkpointing¶

Segmented checkpointing is the middle gear between checkpointing every block and checkpointing nothing.

It uses PyTorch's activation checkpointing backend. SimpleTuner runs a contiguous group of transformer blocks under one checkpoint call, then carries the returned hidden state into the next group. Wider groups save fewer boundary states but still recompute the checkpointed blocks during backward; measure the resulting speed and peak-memory tradeoff.

For CPU offload and FFN-only checkpointing, use Unsloth-style checkpointing. The short rule of thumb still lives in Decision Rule.

Controls¶

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

On supported whole-block paths, gradient_checkpointing_interval is the segment width. 2 means checkpoint blocks 0-1, 2-3, 4-5, and so on.

For finer VRAM control, add a stride:

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

That checkpoints blocks 0-1, runs 2-3 normally, checkpoints 4-5, runs 6-7 normally, and repeats. The stride must be at least the interval; overlapping schedules are not valid.

Supported segmented whole-block paths: Flux.1, Flux.2, 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, and Wan.

Stable Cascade stage C also supports interval and stride control, but it applies the schedule to the UNet Res/Timestep/Attention micro-block sequence instead of transformer whole-block groups.

Some model families use model-specific interval semantics:

Do not compare stride rows for families where stride is ignored. If the benchmark numbers look identical there, that is usually the option being ignored, not a useful performance result.

When To Use It¶

Use it after normal per-block checkpointing fits but costs too much step time. Start with 2. If VRAM allows, try 2 with stride 4 on very deep models.

Do not expect it to help when the peak is mostly trainable weights, optimizer state, validation, VAE caching, block swapping, or routing. SimpleTuner falls back to the safer per-block path when a model feature needs per-block control.

dynamo_use_regional_compilation is not a universal win. It helped or stayed neutral on several image-model runs, but it was a bad fit for the Wan/RamTorch and LTXVideo2 profiles below. Treat compile settings as part of the benchmark, not as background noise.

Benchmarks¶

Measured with real SimpleTuner examples on single-GPU H100, H200 and L40S pods. Validation and checkpoint saves were disabled, cache preparation was excluded, and first-step compile/setup inside the train loop is excluded when post-warmup timing is available.

Each measured cell is post-warmup sec/step / peak VRAM GiB. Status-only cells mean: OOM ran out of GPU memory, failed did not reach measured training steps, unsupported means that option was not wired for that family, and not run means the sweep did not include that combination.

Compare modes within a family first. Cross-family comparisons are rough because resolution, frame count, attention backend, model depth, trainable adapter type, and dataset shape differ.

The matrix below is the source of truth for this sweep. Model-specific notes call out caveats when a row is coverage data rather than a recommendation.

Family Sweep Results¶

ACE Step 1.5¶

Example: ace_step-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.

AnyFlow Distillation on Anima¶

Example: anima-anyflow-stage1.peft-lora. Resolution: 1024x1024. This row measures AnyFlow distillation using Anima, not the plain Anima LoRA example. Use anima.peft-lora for plain 1024x1024 Anima image training.

AuraFlow¶

Example: auraflow.peft-lora. Resolution: 1024x1024.

Note: AuraFlow supports SDNQ and TorchAO quantization. Quantized none rows are included below; quantized checkpoint rows need fresh full-length benchmark coverage before they get numbers here.

Boogu Image¶

Example: boogu-image-v0.1.peft-lora. Resolution: 1024x1024.

Chroma¶

Example: chroma.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: cosmos2image.lycoris-lokr. Resolution: 512x512.

Cosmos 3¶

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

ERNIE 4.5 Image¶

Example: ernie.peft-lora. Resolution: 512x512.

HeartMula¶

Example: heartmula.peft-lora. Audio-token training.

HiDream¶

Example: hidream.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-1.5-t2v.peft-lora. Training shape: 480 pixel-area video buckets, 48 frames, batch 2.

HunyuanVideo is activation-heavy at this training shape. Per-block and interval-2 checkpointing both fit cleanly; leaving checkpointing off does not fit on an 80 GB H100. seg2-stride4 only fit in this sweep when attention activation offload was enabled, and that row is a fallback rather than a speed recommendation. SDNQ Hadamard works, but variable conditioning shapes still trigger dynamic-kernel compilation in the measured window.

Ideogram 4.0¶

Example: ideogram-fp8.peft-lora. Resolution: 1024x1024. The fp8 flavour uses Ideogram 4's native weight-only fp8 checkpoint (base_model_precision=no_change); bf16-upcast sets ideogram_fp8_base_upcast=true to dequantize the base weights to bf16 at load. Requesting any other base_model_precision (e.g. int8-sdnq) also dequantizes first so the quantizer operates on real weights.

Earlier revisions of this table were measured before Ideogram honored gradient_checkpointing_interval, gradient_checkpointing_segment_stride, or gradient_checkpointing_backend (its custom loader skipped the shared wiring), and before int8-sdnq actually quantized the fp8 checkpoint — all checkpointing rows were silently full-layer torch checkpointing over fp8-native weights. The H100 numbers below are post-fix; L40S rows await re-measurement.

The torch-ffn/unsloth-ffn backends are unsupported: Ideogram 4 does not expose an attention/FFN checkpointing boundary.

Takeaways: seg2-stride4 is ~14% faster than full-layer at ~24 GiB more retained activations; unsloth offload saves ~1.2 GiB for a 1-3% step-time cost (and is faster than torch under int8 full-layer, where offload overlaps the quantized matmul); bf16-upcast is the throughput winner when VRAM allows.

Kandinsky 5 Image¶

Example: kandinsky5-image-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-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.peft-lora. Resolution: 1024x1024.

Krea 2¶

Example: krea2.peft-lora. Training resolution: 512 px square crop. Example validation setting: 1024x1024; validation was disabled for the benchmark.

The main table uses regional compilation, which is good for Krea2 step speed but not a clean VRAM comparison: the compiled graph/workspace keeps the peak close to the uncheckpointed peak for several checkpoint modes. A bf16 control run with regional compilation disabled showed checkpointing is wired and has the expected memory/speed shape:

The activation-offload row here means full-block checkpointing plus attention activation offload. Against full-block layer checkpointing alone, attention offload did not reduce Krea2 peak VRAM in this shape; it mostly added CPU transfer overhead.

LongCat Image¶

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

LongCat Video¶

Example: longcat-video.peft-lora+ramtorch. Resolution: 832x480, 81f. Rows use attention_mechanism=native-flash.

LongCat Video is activation-heavy at this shape. Full per-block checkpointing is the practical row. The partial checkpoint rows (interval2, seg2-stride4) do not fit here, even when attention activation offload is enabled for the strided row. Plain attention activation offload fits for bf16 and SDNQ, but it is much slower than full checkpointing.

LTXVideo 0.9.5¶

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

Numbers are warm seconds per step / peak GiB. The full-run average includes setup and compile overhead and is recorded in the sweep artifacts.

Attention activation offload rows are unsupported for LTXVideo 0.9 in this sweep.

LTXVideo2 2.3¶

Example: ltxvideo2-2.3-dev-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.peft-lora. Resolution: 512x512.

Note: Lumina2 now uses the segmented whole-block path. interval2 checkpoints every two-block segment; seg2-stride4 checkpoints two blocks, lets the next two keep activations, then repeats. Attention activation offload was not part of this Lumina2 run.

MageFlow¶

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

Note: MageFlow's 1024px variable-shape image path is mostly helped by attention activation offload and weight-only FP8. Block checkpointing modes are valid, but did not lower measured peak residency in this sweep.

OmniGen¶

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

Note: OmniGen uses token-ID prompts instead of cached text embeddings. These rows measure the supported no-checkpoint and full-block torch checkpointing paths; interval, segmented-stride, and attention-offload controls are not implemented for this family in this sweep.

PixArt¶

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

Qwen Image 2.1¶

Measured on verified H200 and L40S GPUs with qwen_image.peft-lora, Qwen-Image 2.1, Domokun at 512x512, BF16, rank-32 LoRA, Optimi Lion and regional Inductor compilation (default mode). Each run has 20 steps; timings exclude the first five. Cells are seconds/step / peak VRAM GiB; peaks include component setup.

Runtime: torch==2.11.0+cu128, diffusers==0.40.0; 10 text tokens and 1024 image tokens. Component setup dominates some batch-1 memory peaks.

BF16 checkpoint schedules¶

none disables checkpointing. layer checkpoints each block. interval2 checkpoints contiguous pairs 0–1, 2–3, …; seg2-stride4 checkpoints 0–1, retains activations for 2–3, and repeats. Grouped execution requires no block swapping, hidden-state capture or KV cache; those paths retain per-block interval/stride scheduling. Attention activation offload (activation-offload, seg2-stride4-offload) is unsupported. BF16 fits, so quantization was not included in this 2.1 sweep.

Attention backends¶

Attention comparisons disable checkpointing. The default native SDPA selected cuDNN on H200 and PyTorch Flash SDPA on L40S. Qwen 2.1 uses causal text attention and noncausal target-image attention over all valid keys. Equal, unpadded prompt lengths do not require varlen. FlashAttention paths currently require unpadded text-to-image prompts; use native SDPA or FlexAttention for padded prompts or interleaved conditioning. A padding-only varlen mask cannot preserve the latter causal structure.

The uniform-length varlen metadata helper is vendored from upstream Diffusers: lengths come from shapes and cumulative offsets use arange, removing the released helper’s GPU .item() synchronizations and graph breaks. Redundant all-valid masks are removed during collation. Forward/gradient equivalence and graph capture are checked separately from timing.

Forcing attention_mechanism: "cudnn" failed in PyTorch 2.11 Inductor with CantSplit; automatic native SDPA still used cuDNN successfully. FA4 Hub failed to load with nvidia-cutlass-dsl==4.7.1: cutlass.cute.core.ThrMma is missing, so no timing is available. Hub backends use trust_remote_code: true. FA3 is a Hopper backend; it is not an L40S option (upstream support).

H200 batch scaling¶

On the same H200, native SDPA took 0.612 s at batch 10 and 1.223 s at batch 20, or about 16.35 images/s in both cases. Doubling the batch doubled the work after throughput had plateaued. The 144 GB preset keeps checkpointing disabled; its larger batch uses the extra capacity without promising higher images/s.

A full-transformer check on cached Domokun inputs found identical FP32 outputs with and without the redundant mask; relative gradient L2 error was 1.97e-6. BF16 backends are not numerically interchangeable: on this single batch, aggregate LoRA-gradient cosine against FP32 was about 0.92–0.93 for native SDPA and 0.80 for FA3. These are numerical diagnostics, not convergence comparisons.

The H200 batch-20 trace had kernels active for 93.1% of the captured interval. GEMMs accounted for 70.5% of kernel time and attention for 9.7%; Inductor already fused normalization/RoPE and pointwise epilogues. The largest idle gaps were around batch preparation. Kernel active time is not SM occupancy, and profiler timings are excluded from the throughput tables.

Prefetching was slower for this 27-image dataset at batch 20: with queue length 2, host prefetch averaged 2.030 s/step and device prefetch (1 MiB threshold) 2.036 s/step, versus 1.223 s without prefetch. Both had periodic roughly 3-second steps despite medians near 1.17 s; the presets leave prefetch disabled.

The 250-step Domokun recipe is a throughput example, not a reliable convergence recipe. An earlier checkpoint produced recognizable Domokun images after reloading, but fresh 250-step runs did not reproduce that result. Controls retaining padding masks, disabling compilation, and restoring the earlier RoPE expression also failed. Cached latents decode to the correct subject. The cause of the training deterioration remains unresolved; the timing tables do not establish comparable image quality across attention backends.

Qwen Image¶

Historical Qwen-Image 1.0 measurements (model_flavour: "v1.0"), using the earlier qwen_image.peft-lora example at 1024x1024. The example now selects 2.1. The nearly identical interval/stride rows below do not establish that distinct checkpoint schedules were active and should not guide 2.1 configuration.

Sana¶

Example: sana.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.peft-lora. Resolution: 832x480, 49f.

Note: SanaVideo uses linear attention, so attention activation offload remains unsupported. Segmented whole-block checkpointing is supported for the standard path.

SD 1.x¶

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

Note: SD1x uses the diffusers UNet path. Regular layer checkpointing is supported, but interval and segmented stride controls are not wired for this family.

SD3¶

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

Note: SD3 uses true contiguous segmented checkpointing on the plain transformer path. Attention activation offload is supported; it cuts VRAM hard, but costs throughput.

SDXL¶

Example: sdxl.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: cascade-stage-c.lycoris-lokr. Resolution: 1024x1024.

Note: Stage C is a full-precision prior path. These rows ran with mixed_precision=no and base_model_precision=no_change; quantized base precision rows are not meaningful for this model. The interval and stride modes operate over the UNet's Res/Timestep/Attention micro-block sequence.

Wan 2.1 T2V 1.3B¶

Example: wan2.1-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: wan2.1-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-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.peft-lora. Resolution: 1024x1024.

ZLab I1¶

Example: zlab-i1.peft-lora. Resolution: 1024x1024.

Note: ZLab I1 carries its U-Net-style skip tensors through the segmented checkpoint state. Attention activation offload is not wired for this family.