pytorch

1058b475 - Weak-ref-ify MetaConverter and FakeTensorConverter (#80544)

Weak-ref-ify MetaConverter and FakeTensorConverter (#80544) Make `MetaConverter` and `FakeTensorConverter` hold weak references to their memoized tensors, and also have `MetaConverter` hold weak reference to Tensor storage. Otherwise it can be tricky for users to make sure all existing FakeTensors or FakeTensorModes are deleted which otherwise will leak memory. I ran into https://github.com/pytorch/pytorch/issues/7733 which I was able to get around with the following (see comment from code): ``` # torch.Tensors cannot be used as a key in a dictionary # because they define a custom __eq__ function which when used # to resolve hash collisions will throw when comparing tensors: # "RuntimeError: bool value of Tensor with more than one value is ambiguous." # To avoid that, we use an object which will hold a Tensor and use # its id for both hashing and equality. # In order to use this as a weak key reference, we cannot # simply use weakref.WeakKeyDictionary because the newly constructed # WeakTensorRefKey only use would be a dictionary so it would have no strong # references. # To get around this issue, we can use it as a normal key, and then set # `weakref.finalize` to delete the key when its contained tensor dies. ``` While for the tensor memo we can set a `weakref.finalize` callback that will remove the corresponding `WeakTensorRefKey` from the tensor memo, at the point that this callback is invoked the tensor storage is not yet deallocated.. See comment from code: ``` # [expired-storages] # NB: even though the tensor has died, # the deallocation of its storage can take longer, # even when the storage has no other uses/views. # In this case, the StorageWeakRef object will be kept alive # longer than it needs to be, however the storage itself # will be deallocated. We retain the possibly dead storages # and periodically check if any of them are expired and # can be freed. ``` partial fix for https://github.com/pytorch/torchdynamo/issues/468 Pull Request resolved: https://github.com/pytorch/pytorch/pull/80544 Approved by: https://github.com/ezyang