mage-flow: add FFN checkpointing logic

[bghira]

This pull request introduces several improvements and new tests for the MageFlow Transformer model, focusing on enhanced gradient checkpointing flexibility, support for feed-forward network (FFN) checkpointing, and increased test coverage for attention and packing logic.

MageFlow Transformer Model Enhancements

• Added support for FFN-specific gradient checkpointing by introducing the _supports_ffn_gradient_checkpointing flag and handling a new gradient_checkpointing_scope parameter. This allows selective checkpointing of feed-forward layers, improving memory efficiency during training. [1] [2] [3] [4]
• Updated imports and initialization to include get_checkpoint_backend_scope, ensuring the correct checkpointing scope is used throughout the model. [1] [2] [3]

Test Coverage Improvements

• Added a test (test_mageflow_attention_joint_pack_uses_single_flash_dtype) to verify that the MageFlow attention processor correctly enforces a single dtype (bfloat16) for all projections and outputs, ensuring consistency and preventing dtype mismatches in flash attention. [1] [2]
• Added a test (test_pack_latents_tracks_rectangular_target_shape) to validate that the model's latent packing logic correctly tracks and returns the expected shapes, indices, and lengths for rectangular targets, increasing reliability of input handling.