Update gp utils

[misko]

Refactor and update gputil functions

• Exclude as much possible to outside of our custom functions, this allows compiler to optimize
• Optimize padding operation
• Use reduce_scatter which is faster than all_reduce in all_gather backward
• Implement special all_gather backwards for GLOO only , for CPU GP + GP tests
• Update test cases to use these functions
• Add asynchronous version of all_gather in preparation of any overlapped compute + comms implementation
• Combine two sequential all_reduce calls in EFS head into one all_reduce call

Change layers to always output atom embeddings for the full system. Currently in GP mode we output a tensor the size of only the local atoms. This makes it hard to implement anything to overlap communication and computation.

• If we keep the current implementation you cannot chunk the communication and do compute in between.
• In this new implementation there is an additional all_reduce required to synchronize energy across systems (systems x 1) float