Open
Description
🐛 Describe the bug
To get started:
$ pip install torch dillAfter you use torch.compile, Optimizer subclasses stop being dill picklable:
import dill
import torch
torch.save(torch.optim.AdamW, "a.pth", pickle_module=dill) # OK
@torch.compile
def bias_gelu_fused_(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
"""Bias-GeLU fused"""
x = inp + bias
return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
torch.save(torch.optim.AdamW, "fake.txt", pickle_module=dill) # PicklingError: Can't pickle <built-in function reset_code>: it's not found as torch._C._dynamo.eval_frame.reset_codeThis seems to be partly the reason:
pytorch/torch/_dynamo/eval_frame.py
Line 1506 in 8bf9e99
I think it boils down to torch._C._dynamo.eval_frame not being importable as a package, but as an attribute of torch._C._dynamo (like an imported module):
import dill
import torch._C._dynamo
print(torch._C._dynamo.eval_frame) # <module 'torch._C._dynamo.eval_frame'>
print(torch._C._dynamo.eval_frame._CacheEntry) # <class 'torch._C._dynamo.eval_frame._CacheEntry'>
torch.save(torch._C._dynamo.eval_frame._CacheEntry, "a.pth", pickle_module=dill) # PicklingError: Can't pickle <class 'torch._C._dynamo.eval_frame._CacheEntry'>: it's not found as torch._C._dynamo.eval_frame._CacheEntry
import torch._C._dynamo.eval_frame # ModuleNotFoundError: No module named 'torch._C._dynamo.eval_frame'; 'torch._C._dynamo' is not a packageeval_frame is defined here:
pytorch/torch/csrc/dynamo/init.cpp
Line 58 in 556e4ec
Is there some change to dill or C++-generated modules that solve this?
Versions
[conda] torch 2.3.0 pypi_0 pypi
cc @mruberry @mikaylagawarecki @msaroufim @ezyang @bdhirsh @anijain2305 @chauhang