⭐ [Feature] Support deepspeed training for DiT

[clownrat6]

Changed

• Code Style:
  • Rewrite dit modeling, split dit.py to modeling_dit.py and configuration_dit.py
• Accelerate w/ Deepspeed training: Support training "dit" on accelerate w/ deepspeed.

How to Train

1. Accelerate w/ zero2 training:

bash scripts/sky/train_256_dsz2_dit.sh

Specifically, we attempt to use v100 for training latte:

1. use zero2 training (bs = 5):

Memory cost

Loss curve

Tests

Training latte on accelerate w/ deepspeed zero 2 (bs 5 * 8, num_frames 16, sample_rate 3, 19500 step):

Click to expand

https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/58427300/d0ad508c-729a-47b6-9571-c823aff7badb

[kabachuha]

@sennnnn will you mind mirroring this to Open(MM)DiT? https://github.com/NUS-HPC-AI-Lab/OpenDiT/

[clownrat6]

@sennnnn will you mind mirroring this to Open(MM)DiT? https://github.com/NUS-HPC-AI-Lab/OpenDiT/

Thanks for your advice. I will complete it later.