Update zero_to_fp32.py - to support deepspeed_stage_1

[PicoCreator]

Since deepspeed 1 checkpoint structure is identical to deepspeed 2 (AFAIK), we should just change the stage check and add support accordingly

However I am not 100% sure if this is intentional by design, or some coincidence in my use case - so might need someone with more knowledge on this topic to weigh in 🤔

[tjruwase]

@stas00, do you remember why stage 1 was excluded?

[stas00]

I think at the time I developed it I didn't think anybody used it, or at least I didn't, so I didn't have a use case for it.

[tjruwase]

@PicoCreator, thanks for this PR.

[PicoCreator]

Glad to see my wild guess, work and be of use (found a few issues on dependent projects that encounted this)

I currently use deepspeed 1 to train small toy models (<=3B), as fast as possible, and to test param/model architecture changes =) on consumer hardware (where the gpu-to-gpu communication of deepspeed 2+ is noticeable)

[tjruwase]

@PicoCreator, thanks for sharing your context and experience. We always appreciate hearing customer stories. I want to share a minor naming clarification in terms of DeepSpeed versus ZeRO.

1. DeepSpeed is the overall DL library that provides training, inference, compression, etc. optimizations.
2. ZeRO is a memory optimization of DeepSpeed with stages 1 to 3, in addition to CPU and NVMe offloading.