segmented checkpointing for better cache hits / latency

[bghira]

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:

• Added support for checkpointing contiguous chunks of blocks (instead of every block or every n-th block) in flux/transformer.py and mageflow/transformer.py, using the new checkpoint_sequential_state utility. This allows for more flexible and efficient memory usage during training. [1] [2] [3] [4]
• Introduced the set_gradient_checkpointing_interval method and the gradient_checkpointing_interval attribute to MageFlow, aligning its interface with Flux for chunked checkpointing. [1] [2]
• Updated the safety check logic to recognize "mageflow" as a supported model for gradient checkpointing intervals.
• Refactored and clarified the gradient_checkpointing_interval.py utility to support the new chunked checkpointing logic.

Documentation Updates:

• Updated option descriptions in all supported languages to explain that --gradient_checkpointing_interval now checkpoints contiguous chunks of blocks, and clarified the impact of higher values on memory and recompute. [1] [2] [3] [4] [5] [6]
• Updated experimental documentation in all languages to describe the new chunked checkpointing behavior and its support in Flux and MageFlow. [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.