DistributedOptimizer.load_state_dict copies the allocated optimizer state onto itself, which can OOM during checkpoint load

What Badgr ran

Read megatron/core/optimizer/distrib_optimizer.py on current main, and ran a plain PyTorch Adam check on CPU. The memory peak itself was not run: it needs Transformer Engine's FusedAdam and a multi-GPU checkpoint.

What came back

• The "Retrieve existing optimizer state" branch sets state_dict_state = inner_state_dict["state"], and self.optimizer.load_state_dict is then called with that state, as the issue describes.
• Plain torch.optim.Adam did not reallocate its state when loaded its own state_dict (data_ptr unchanged), so the extra copy depends on FusedAdam's behavior, which was not tested.