[models] allow specifying block numbers in enable_gradient_checkpointing() to speed up training

[sayakpaul]

As brought up in #9982 (comment) by @bghira, I think we could support this directly in the enable_gradient_checkpointing() method we expose for the models.

Users could specify the block interval they want gradient checkpointing to be applied in and we take care of the rest. The code for this is simple and doesn't require any hacks.

Gradient checkpointing is a crucial component for training/fine-tuning larger models and this technique allows for nice speed/memory trade-off.

Cc: @a-r-r-o-w @hlky

[a-r-r-o-w]

Taking this up soon. Will prio for next release

[bghira]

i've implemented this in simpletuner's modeling code for SDXL, Flux, SD3, and Sana, and even the smaller models benefit (maybe moreso) from having this option. this means that basically Sana goes from 4 seconds per step to 1.25 steps per second on a 24G card (int8) when training. I actually haven't checked the speedup for Flux but I assume on 80G card it'll be quite substantial.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.