Refactor EMA to improve memory efficiency

[coryMosaicML]

What does this PR do?

This PR refactors composer's EMA algorithm. Changes:

• Avoids an extra copy of the model parameters, reducing extra device memory from 2x -> 1x the size of the model parameters
• Avoids duplicating non-trainable parameters which reduces memory required for models with many non trainable parameters.
• Avoids a deepcopy of state.model after it has been DDP-wrapped

A downside: it is now a bit more annoying to access the training weights or ema weights directly from the algorithm object.

What issue(s) does this change relate to?

CO-1525

[mvpatel2000]

@coryMosaicML lgtm. Before approval though, would like to see plots / some kind of actual run verifying it works and gives same results as before. Can you please run something like that?

[coryMosaicML]

@coryMosaicML lgtm. Before approval though, would like to see plots / some kind of actual run verifying it works and gives same results as before. Can you please run something like that?

Here is a wandb report showing the memory reduction for stable diffusion, and a comparison with/without the EMA changes on our ResNet50 mild recipe.

[mvpatel2000]

LGTM. Once lint is fixed, feel free to merge