[BUG] OOM issue when loading `torch_dist` format distributed checkpoint

**Describe the bug**

For some scenarios that have GPU memory pressure, it is easy to have OOM error when loading dist-ckpt with `torch_dist` format.

1. The most severe one is loading optimizer states. The current implementation has a two-step workflow: it first generates a `sharded_state_dict` from the model, then loads the ckpt according to the dict. There are param allocations in both step, which makes it easy to OOM. A related issue is https://github.com/NVIDIA/Megatron-LM/issues/1746
2. Another one is that when `swiglu` is used, there is a merge operation here https://github.com/NVIDIA/Megatron-LM/blob/e4b7259801d221f668b689e3b33a293901b42819/megatron/core/transformer/mlp.py#L385. This operation needs extra memory usage. Even if there is a `try` block to alleviate the issue, it is still easy to get OOM.

**Expected behavior**

No obvious memory overhead when loading dist-ckpt.

**Additional context**

Add any other context about the problem here.