support megatron_amp_O2 and tp
## Add the support 
* megatron_amp_O2 on salm-gpt, bestow-gpt, salm-t5. Follow this doc for example recipe https://docs.google.com/document/d/1VKP9j574B_BYwkUHcyZYsQUQDbwMmFEhu5p0uH-ns6o/edit#heading=h.f646we4te4u8
* tp>1 support on salm-gpt, bestow-gpt. Follow this doc for example recipe https://docs.google.com/document/d/1fr2UlqThlzqvqL_lpc_HIyrh0yCzL27A-QX1eCI_ojg/edit#heading=h.sccv9we30nts

## Container
/lustre/fsw/portfolios/llmservice/users/zhehuaic/containers/nemo-main-aug24.sqsh or gitlab-master.nvidia.com/zhehuaic/nemo_containers:nemo-main-aug24