changing params to go fast on H200s

[daanelson]

In order to speed up training (losslessly!) in default situations on Replicate, I'm making a few tradeoffs for increased speed at the cost of increased memory usage. specifically:

• turning off gradient checkpointing
• not quantizing the flux DiT

---

Important

Optimize training speed on H200s by disabling gradient checkpointing and model quantization under specific conditions in train.py.

• Behavior:
  • Disable gradient checkpointing by default in train() to increase training speed.
  • Automatically enable gradient checkpointing if GPU memory < 100GB, batch size > 1, or resolution > 1024.
  • Disable model quantization by default; enable if GPU memory < 100GB.
• Parameters:
  • Add gradient_checkpointing parameter to train() with default False.
  • Set quantize to False by default in train().
• Misc:
  • Refactor resolution parsing in train() to use resolutions list.

^{This description was created by for db81504. It will automatically update as commits are pushed.}