High memory usage on Llama3-70B full finetune during checkpoint save

[ebsmothers]

See the discussion on #993. @andyl98 has reported that when running on 8X A100 with 1 TB DRAM they hit CPU OOM during checkpoint save. They also point out that they do not see the OOM without the usage of FullOptimStateDictConfig.

In an ideal world I think we should only need (model params) + (optimizer params) = (70B * 2) + (70B * 2 * 2) = 420 GB in bf16, so seems like the unsharding is being done inefficiently (at least wrt CPU RAM)?

[RdoubleA]

Why not use DCP checkpointing which doesn't require unsharding and prevents CPU OOMs? cc @LucasLLC