Understanding stages of deepspeed zero integration

[mshubhankar]

I am using distributed training with FastDP and have questions about its integration with Deepspeed. This is my first time using Deepspeed, and I apologize if some of these questions are trivial:

1. Are all the three stages of Deepspeed zero necessary for distributed DPSGD training?
2. The image classification examples have two python files for stage 1 and stage 2and3. Are both of them private?
3. For stage 1, privacyengine.attach() is not recommended. Then, how is the dp_step() called?
4. The requirements file mentions older versions of torch and deepspeed. Have you tested them on any newer versions?

Thank you!

[woodyx218]

Thank you for using FastDP! To answer you questions:

1. You only need 1 of three stages for distributed DPSGD. As you move up the stage, you trade off the communication time for memory efficiency (i.e. slower but can train larger models).
2. Yes they are all private and should give you similar results.
3. For all stages, we don't need .attach() or dp_step. The modification is on the grad, not on the optimizer, so that regular step() would work directly.
4. I haven't tested on newer deepspeed version. It should work on torch>=2.2.0.