Integration of the deepspeed communication logging

[clasqui]

Morning,
I want to open a discussion on what's the best way to integrate the DeepSpeed Communitations Logging feature.
The provided guide asks to change the torch.dist calls into deepspeed.comm calls with the same name. I am not sure where in the code of axolotl these changes should be applied, and if there is a more graceful way to integrate this with axolotl.

[NanoCode012]

Hey! I'm unfamiliar with Deepspeed logging. Do you want to log the memory usage / latency across each epoch?

[clasqui]

Hey! No, I am talking about this feature: https://www.deepspeed.ai/docs/config-json/#communication-logging

[NanoCode012]

Thanks for getting back. It seems like, you can set in the deepspeed json file to see the logs? Am I understanding it right?

And to call, deepspeed.comm.log_summary(), you can add a simple callback to invoke at end of epoch/batch/step depending on your preference.

[clasqui]

There is an extra step/requirement for it to work: The application needs to use the deepspeed.comm package for the communication calls instead of the torch.dist ones. It implements the same interface, and it needs to be changed everywhere where it is used (this is what I understood from the tutorial). Thus, my question was relative to how to accomplish this in a respectful way in axolotl.

[NanoCode012]

Oh I see. We have our dist.* utility calls here: src/axolotl/utils/distributed.py. This should cover the majority of it.

Do you also need to replace calls such as dist.gather_object?