Flux training: Implement scanned blocks, dynamic gradient checkpointing, and weight loading improvements