Does Accelerate support PyTorch Weight Averaging (SWA and EMA)?

[yeruoforever]

I'm working on a PyTorch project using the Accelerate library for distributed training. I'm interested in implementing Weight Averaging techniques such as Stochastic Weight Averaging (SWA) and Exponential Moving Average (EMA) to improve model stability. However, I'm unsure if Accelerate directly supports these features.

Questions:

1. Does Accelerate currently support PyTorch's Weight Averaging techniques like SWA and EMA?
2. If so, how can I correctly implement these techniques and ensure they work successfully in a multi-GPU environment on a single machine?
3. Are there any example codes or best practices available for implementing SWA and EMA with Accelerate?

I would appreciate any example code snippets or guidance on how to integrate SWA and EMA with Accelerate. Any guidance or example code on how to properly implement SWA and EMA with Accelerate would be extremely helpful. Thank you!

[SunMarc]

Hey @yeruoforever, we didn't implement those techniques in accelerate, maybe there are existing codebase that implement it using accelerate but i'm not aware of them

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[Om2005Prakash]

I don't think accelerate has anything to do with EMA, since accelerate syncs the gradient among all devices, so you can use torch EMA or any other library providing EMA without worrying.