Save/Load OptimizedModule

[msaroufim]

This enables us to torch.save(opt_mod) by just unwrapping and saving the _orig_mod

torch.load("opt_model.pt") is semantically equivalent to loading an uncompiled model, we can make this better by saving a cache or use AOT Inductor

So torch.save(mod) == torch.save(opt_mod) and torch.load(opt_model) == torch.load(mod)

So users don't need to write code like this

if is_compiled_module(model):
    torch.save(model._orig_mod)

Idea 0 (rejected)

Don't bother saving an optimized model, just throw and error but make it human readable in that people need to torch.save(opt_model.state_dict()) #101997

Idea 1 (rejected)

Instantiate an nn module instead of an optimizedmodule, it works but it's confusing and breaks deepcopying

    def __getstate__(self):
        return {"_orig_mod": self._orig_mod}

    def __setstate__(self, state):
        orig_mod = state["_orig_mod"]
        self.__class__ = orig_mod.__class__
        self.__dict__.update(orig_mod.__dict__)

    def __new__(cls, *args, **kwargs):
        instance = super().__new__(cls)
        super(OptimizedModule, instance).__init__()
        return instance

Idea 2 (this PR now)

A lot more work but we can use wrapper classes around functions instead so they become picklable. Need to merge this fast though cause the merge conflicts make me sad

Test pass on modules test/dynamo/test_modules.py

Failing test tracker

  dynamo/test_logging
  dynamo/test_dynamic_shapes
  dynamo/test_misc

cc @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov @desertfire

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/101651

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 21 New Failures

As of commit 18cc651:

NEW FAILURES - The following jobs have failed:

• cuda11.8-py3.10-gcc7-sm86 / test (inductor_torchbench_dynamic, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• cuda11.8-py3.10-gcc7-sm86 / test (inductor_torchbench, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh)
• linux-bionic-cuda12.1-py3.10-gcc9 / test (default, 1, 5, linux.4xlarge.nvidia.gpu) (gh)
• linux-bionic-cuda12.1-py3.10-gcc9 / test (default, 4, 5, linux.4xlarge.nvidia.gpu) (gh)
• linux-bionic-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
• linux-bionic-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
• linux-bionic-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
• linux-bionic-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
• linux-bionic-py3.11-clang9 / test (default, 1, 3, linux.2xlarge) (gh)
• linux-bionic-py3.11-clang9 / test (default, 2, 3, linux.2xlarge) (gh)
• linux-bionic-py3.8-clang9 / test (default, 2, 3, linux.2xlarge) (gh)
• linux-bionic-py3.8-clang9 / test (default, 3, 3, linux.2xlarge) (gh)
• linux-focal-cpu-py3.8-gcc7-inductor / test (inductor_torchbench_cpu_accuracy, 1, 1, linux.4xlarge) (gh)
• linux-focal-py3.8-gcc7 / test (default, 1, 3, linux.2xlarge) (gh)
• linux-focal-py3.8-gcc7 / test (default, 3, 3, linux.2xlarge) (gh)
• linux-focal-py3.9-clang7-asan / test (default, 3, 6, linux.4xlarge) (gh)
• linux-focal-py3.9-clang7-asan / test (default, 4, 6, linux.4xlarge) (gh)
• linux-focal-py3.9-clang7-asan / test (default, 5, 6, linux.4xlarge) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[albanD]

You should definitely not do that haha

[jansel]

Silently converting one type to another in pickle load is confusing, let's not do that.

[msaroufim]

note from @albanD: check if it works with inductor as well. It's possible that this backend has a simple state

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

[msaroufim]

Closing this for now. I couldn't figure out why this changes introduces extra graph breaks. Might revisit