[Feature] Support DeepSpeed/CorossalAI

[okotaku]

Motivation

MMEngine TRAINING BIG MODELS

Results (Optional)

Environment:

• A100 x 4 GPUs
• nvcr.io/nvidia/pytorch:23.11-py3

Settings:

• 1epoch training.

Model	total time
stable_diffusion_xl_pokemon_blip_fast (BS=4)	1 m 6 s
stable_diffusion_xl_pokemon_blip_deepspeed_stage3 (BS=8)	1 m 5 s
stable_diffusion_xl_pokemon_blip_deepspeed_stage2 (BS=8)	58 s
stable_diffusion_xl_pokemon_blip_colossal (stage=2, BS=8)	58s

Checklist

Before PR:

• Pre-commit or other linting tools are used to fix the potential lint issues.
• Bug fixes are fully covered by unit tests, the case that causes the bug should be added in the unit tests.
• The modification is covered by complete unit tests. If not, please add more unit test to ensure the correctness.
• The documentation has been modified accordingly, like docstring or example tutorials.

---

📚 Documentation preview 📚: https://DiffEngine--109.org.readthedocs.build/en/109/