Time needed for fine-tuning Bagel

Hi, thank you for your great work. I fine-tuned Bagel on 2 A100 40GB, training editing tasks only, with 512*512 images. For now, each sep takes 80-85 seconds, I wonder if there's any way to accelerate training, by the way I have cpu_offload=True.

No file, test, or entry point is named. Reproduce fine-tuning on two A100 40GB GPUs with editing tasks, 512×512 images, and cpu_offload=True, then measure the reported 80–85 seconds per step. Done means identifying and validating a way to reduce training time.

Written by the indexing model from the issue text.