Implement pass-through `state_dict` and `load_state_dict` for dynamo `OptimizedModule`

[awaelchli]

🚀 The feature, motivation and pitch

I propose that we implement .state_dict() and .load_state_dict() in OptimizedModule in such a way that they bypass the wrapper and return/load the state from the inner _orig_mod module.

Motivation:

• User might accidentally save the state dict of a compiled model into their checkpoint, only to find out later they can't load because the keys don't match
• The torch-compiled module may be wrapped by FSDP, DDP, etc. This would make it cumbersome for the user to save a checkpoint while also removing these keys or unwrapping the model.

The sketch of the implementation:

class OptimizedModule(torch.nn.Module):
    ...
    
    def state_dict(self, *args, **kwargs):
        return self._orig_mod.state_dict(*args, **kwargs)

    def load_state_dict(self, *args, **kwargs):
        # TODO: add some backward-compatibility code here
        return self._orig_mod.load_state_dict(*args, **kwargs)

With this implementation, we can effortlessly load the state-dict of a regular nn.Module into a torch-compiled module and vice versa!

Alternatives

Keep as is.

Additional context

This would probably also close/fix #94575
I am happy to send a PR for this.

cc @mruberry @mikaylagawarecki @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519

[awgu]

Just some passing thoughts:

Since state_dict does not recursively call the state_dict of submodules, this approach does not work if a submodule is compiled into an OptimizedModule. This creates a divergence between the state dicts constructed when the root is OptimizedModule vs not. We previously handled this situation using state dict hooks, but we have learned the hard way that they are difficult to compose and get right.

cc: @fegin @wz337 for more context on the challenges around this kind of approach

[awaelchli]

@awgu I see state_dict gets called on submodules here:

pytorch/torch/nn/modules/module.py

Lines 1892 to 1894 in cb48f78

	for name, module in self._modules.items():
	if module is not None:
	module.state_dict(destination=destination, prefix=prefix + name + '.', keep_vars=keep_vars)

So I think even if there are compiled submodules, it would still retrieve the state dict correctly. Perhaps you meant that load_state_dict() is not called recursively on submodules? In this case, I think the approach would still work, because what it is trying to fix is really just the access of keys at the root module.

[awgu]

@awaelchli Good catch! That is my bad -- I misremembered and meant load_state_dict, and your point makes sense.

[awaelchli]

@awgu Great, thanks! I've opened a PR and so far I haven't discovered any bad edge cases.

[awaelchli]

What do you think about the approach in my PR #113423? I've tested different ways the module can be saved and loaded, and the rest of the CI also seems to be ok with the change. I appreciate the feedback :)