Skip to content

Conversation

@destefy
Copy link
Contributor

@destefy destefy commented Jul 9, 2024

When I try to compile with pytorch 2.0.0, I am unable to torch.save the model's inputs. I get the following error

Traceback (most recent call last):
  File "/home/ubuntu/local_compile/centml-python-client/venv-torch/lib/python3.10/site-packages/centml/compiler/backend.py", line 63, in _serialize_model_and_inputs
    torch.save(self.inputs, self.serialized_input_path, pickle_protocol=config_instance.PICKLE_PROTOCOL)
  File "/home/ubuntu/local_compile/centml-python-client/venv-torch/lib/python3.10/site-packages/torch/serialization.py", line 441, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/home/ubuntu/local_compile/centml-python-client/venv-torch/lib/python3.10/site-packages/torch/serialization.py", line 653, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object 'WeakValueDictionary.__init__.<locals>.remove'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/home/ubuntu/local_compile/centml-python-client/venv-torch/lib/python3.10/site-packages/centml/compiler/backend.py", line 147, in remote_compilation
    self._serialize_model_and_inputs()
  File "/home/ubuntu/local_compile/centml-python-client/venv-torch/lib/python3.10/site-packages/centml/compiler/backend.py", line 65, in _serialize_model_and_inputs
    raise Exception(f"Failed to save module or inputs with torch.save: {e}") from e
Exception: Failed to save module or inputs with torch.save: Can't pickle local object 'WeakValueDictionary.__init__.<locals>.remove'

Without having looked into it too deeply, this seems to be because in 2.0.0, the tensor gets turned into an object of type torch._subclasses.fake_tensor.FakeTensor.FakeTensor when passed to the backend, while 2.1.0 keeps it as just a torch.Tensor

@destefy destefy requested a review from anandj91 July 9, 2024 21:29
@destefy destefy self-assigned this Jul 9, 2024
@anandj91
Copy link
Contributor

anandj91 commented Jul 9, 2024

@yaxan was getting similar issue in the nightly version as well.

@yaxan
Copy link
Contributor

yaxan commented Jul 9, 2024

This issue exists in the latest nightly release as well when running torch.save on inputs containing SymInts

@destefy destefy merged commit 42e78a0 into main Jul 15, 2024
@destefy destefy deleted the stefan/pytorch-version-fix branch July 15, 2024 14:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants