PyTorch Distributed Load Updates or Returns state_dict

[mvpatel2000]

🚀 The feature, motivation and pitch

Torch distributed checkpoint load_state_dict (https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/state_dict_loader.py#L20)
updates the passed in state_dict (and returns it). This function is deprecated in torch 2.3 in favor of load (https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/state_dict_loader.py#L48) which neither returns nor updates the passed in state_dict. Instead, it only calls load_state_dict on Stateful elements in the specified state_dict.

Unfortunately, this new API is greatly limiting. For example, in Composer's state_dict passed for checkpointing, we also store various RNG tensors in a dict for determinism. In order to use the new API, we have to rewrap everything in a Stateful class, which is a somewhat pointless abstraction. Instead, we prefer to receive a loaded state_dict and then manually call load_state_dict on appropriate subitems.

Can we modify load to update the passed in state_dict? This would entail adding:

state_dict[key] = statetful_sd[key]

after https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/state_dict_loader.py#L172-L177

Alternatives

No response

Additional context

No response

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @LucasLLC

[mvpatel2000]

@pytorchbot label "oncall: distributed"

[fegin]

@LucasLLC Any thought about this request?

[LucasLLC]

Sorry for the delay! Taking a look

[LucasLLC]

@mvpatel2000 this makes a lot of sense to me. Would you like to submit a PR or should I?

[mvpatel2000]

@mvpatel2000 this makes a lot of sense to me. Would you like to submit a PR or should I?

@LucasLLC if you can submit a PR that would be awesome :)