| import torch | |
| import torch.utils.hooks | |
| from torch._namedtensor_internals import check_serializing_named_tensor | |
| import os | |
| import threading | |
| import multiprocessing | |
| from multiprocessing.util import register_after_fork | |
| from multiprocessing.reduction import ForkingPickler | |
| from typing import Union | |
| try: | |
| # Early load resource_sharer to prevent a partially initialized instance | |
| # from being inherited in a forked child process. The reduce_storage method | |
| # requires this module indirectly through DupFd(). The built-in mp.Queue | |
| # class pickles arguments in a background thread which may overlap with the | |
| # fork. | |
| import multiprocessing.resource_sharer | |
| except ImportError: | |
| pass | |
| class StorageWeakRef(object): | |
| r"""A weak reference to a Storage. | |
| The cdata member is a Python number containing the integer representation of | |
| the Storage pointer.""" | |
| def __init__(self, storage): | |
| self.cdata = storage._weak_ref() | |
| # Save a direct reference to _free_weak_ref because the `torch` module | |
| # might be cleared during Python shutdown before this module is cleared. | |
| self._free_weak_ref = torch.Storage._free_weak_ref # type: ignore[attr-defined] | |
| def expired(self): | |
| return torch.Storage._expired(self.cdata) # type: ignore[attr-defined] | |
| def __del__(self): | |
| self._free_weak_ref(self.cdata) | |
| def __hash__(self): | |
| return self.cdata | |
| def __eq__(self, other): | |
| if id(self) == id(other): | |
| return True | |
| return self.cdata == other.cdata | |
| class SharedCache(dict): | |
| """dictionary from multiprocessing handles to StorageWeakRef""" | |
| def __init__(self): | |
| # free_dead_references() is called if the len exceeds the current | |
| # limit. The limit scales with the number of remaining live objects. | |
| self.limit = 128 | |
| # `fork` inherits lock state, so in case we fork when the lock is held, | |
| # we register a function to reset the lock to a new object to avoid | |
| # possible deadlocks, following python multiprocessing library design. | |
| self._after_fork() | |
| register_after_fork(self, SharedCache._after_fork) | |
| def _after_fork(self): | |
| self.lock = threading.Lock() | |
| def get(self, key): | |
| with self.lock: | |
| return dict.get(self, key) | |
| def __setitem__(self, key, storage_ref): | |
| with self.lock: | |
| dict.__setitem__(self, key, storage_ref) | |
| if len(self) > self.limit: | |
| self.free_dead_references() | |
| def free_dead_references(self): | |
| live = 0 | |
| for key, storage_ref in list(self.items()): | |
| if storage_ref.expired(): | |
| del self[key] | |
| else: | |
| live += 1 | |
| self.limit = max(128, live * 2) | |
| # mapping from handles to StorageWeakRef objects | |
| shared_cache = SharedCache() | |
| def rebuild_event(device, handle): | |
| return torch.cuda.Event.from_ipc_handle(device, handle) | |
| def reduce_event(event): | |
| handle = event.ipc_handle() | |
| return (rebuild_event, (event.device, handle)) | |
| def rebuild_tensor(cls, storage, metadata): | |
| storage_offset, size, stride, requires_grad = metadata | |
| t = torch._utils._rebuild_tensor(storage, storage_offset, size, stride) | |
| if cls == torch.nn.parameter.Parameter: | |
| # we have to pass requires_grad into constructor, rather than set it as an | |
| # attribute later, because it's an important check for Integer Tensors to | |
| # have requires_grad=False (or else they raise an error) | |
| t = torch.nn.parameter.Parameter(t, requires_grad=requires_grad) | |
| else: | |
| t.requires_grad = requires_grad | |
| return t | |
| def rebuild_cuda_tensor(tensor_cls, tensor_size, tensor_stride, tensor_offset, | |
| storage_cls, dtype, storage_device, storage_handle, storage_size_bytes, storage_offset_bytes, | |
| requires_grad, ref_counter_handle, ref_counter_offset, event_handle, event_sync_required): | |
| # If storage_handle is None, storage points to nullptr. | |
| if storage_handle is None or storage_size_bytes == 0: | |
| storage = storage_cls(0, dtype=dtype, device=storage_device, _internal=True) | |
| else: | |
| storage = storage_from_cache(storage_cls, (storage_handle, storage_offset_bytes)) | |
| if storage is None: | |
| torch.cuda._lazy_init() | |
| storage = storage_cls._new_shared_cuda( | |
| storage_device, | |
| storage_handle, | |
| storage_size_bytes, | |
| storage_offset_bytes, | |
| ref_counter_handle, | |
| ref_counter_offset, | |
| event_handle, | |
| event_sync_required) | |
| shared_cache[(storage_handle, storage_offset_bytes)] = StorageWeakRef(storage) | |
| else: | |
| # We already ref counting this Storage, but producer needs new ref-counters to be released. | |
| storage_cls._release_ipc_counter(ref_counter_handle, ref_counter_offset, device=storage_device) | |
| _storage = storage if isinstance(storage, torch.UntypedStorage) else storage._untyped_storage | |
| t = torch._utils._rebuild_tensor( | |
| torch.storage.TypedStorage(wrap_storage=_storage, dtype=dtype, _internal=True), | |
| tensor_offset, tensor_size, tensor_stride) | |
| if tensor_cls == torch.nn.parameter.Parameter: | |
| # It is crucial for integer tensors to receive | |
| # the requires_grad=False as an argument in the constructor | |
| t = torch.nn.parameter.Parameter(t, requires_grad=requires_grad) | |
| else: | |
| t.requires_grad = requires_grad | |
| return t | |
| def reduce_tensor(tensor): | |
| storage = tensor._typed_storage() | |
| if tensor.requires_grad and not tensor.is_leaf: | |
| raise RuntimeError("Cowardly refusing to serialize non-leaf tensor which requires_grad, " | |
| "since autograd does not support crossing process boundaries. " | |
| "If you just want to transfer the data, call detach() on the tensor " | |
| "before serializing (e.g., putting it on the queue).") | |
| check_serializing_named_tensor(tensor) | |
| torch.utils.hooks.warn_if_has_hooks(tensor) | |
| # Note [CUDA IPC and the caching allocator] | |
| # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ | |
| # When you send a CUDA tensor over IPC, you might expect that you will | |
| # get out the same storage from the other end. However, the CUDA caching | |
| # allocator makes it difficult to preserve this invariant. Consider | |
| # the following situation: a tensor of size 0x100 points to offset 0x20 of | |
| # a storage at 0xA100 of size 0x100. (For simplicity, all of these | |
| # sizes are given in bytes). HOWEVER, with the caching allocator, this storage | |
| # might be part of a larger cudaMalloc allocation 0xA000 of size 0x4000. | |
| # | |
| # When we want to send this CUDA tensor over IPC, we must send the | |
| # *entire* cudaMalloc allocation, i.e., the 0xA000 region, not just | |
| # the storage 0xA100 (because that is what CUDA supports). So, on the | |
| # other end, there simply isn't any way to say, "Wait, you gave me | |
| # a bigger region (0xA000) than the one I wanted (0xA100)". | |
| # | |
| # OK, so if you sent the cudaMalloc allocation, can you just wrap that up as | |
| # one storage itself? No, because this cudaMalloc allocation might contain | |
| # storages of mixed types: float, bytes, double... If you make the entire | |
| # allocation a single storage of a type A, we'll hit an error when constructing | |
| # a tensor of type B on the storage. | |
| # | |
| # cudaIpcMemHandle is an identifier to access the sender cudaMalloc allocation on the | |
| # receiver side. However, cudaIpcMemHandles from each device in a given process may | |
| # only be opened by one context per device per other process. | |
| # If we open and close a memory handle multiples times in a process, CUDA is allowed | |
| # to give it a different address; similarly, once we close the memory, we're not | |
| # allowed to access it(and the storage/tensor built on top of it), even if it is | |
| # still live in the original process. As we cannot make a cudaMalloc allocation | |
| # to a single storage in one go, this requires us to cache the device pointer for | |
| # each cudaIpcMemHandle on C++ side to reconstruct types of storages, while keep | |
| # the old ones alives. | |
| # See [https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html] | |
| # | |
| # This is fine, because all we need to do is to save our position in the allocation, | |
| # and reconstruct storage and tensor from it. | |
| # 0xA000 -> -------CUDA Allocation------ | |
| # | | | |
| # | | | |
| # | | | |
| # | | | |
| # 0xA100 -> --------storage1 begin------ | |
| # | | | |
| # 0xA120 -> --------tensor1 begin ------ | |
| # | | | |
| # | | | |
| # | | | |
| # | | | |
| # | | | |
| # 0xA160 -> --------tensor1 end--------- | |
| # | | | |
| # | | | |
| # | | | |
| # 0xA200 -> --------storage1 end-------- | |
| # | | | |
| # 0xE000 -> --------CUDA allocation----- | |
| # | |
| # To send tensor1, the following info are required from sender to receiver for | |
| # storage recontruction. | |
| # 1. cudaIpcMemHandle of 0xA000(which can be mapped to a basePtr in receiver process). | |
| # basePtr may not be exactly 0xA000 since it's a different process. | |
| # 2. offset(0xA100) of storage1 in the CUDA allocation. | |
| # 3. size of storage1(0x100). | |
| # | |
| # On receiver side: | |
| # 1. Get the devPtr of the MemHandle to access the memory, reconstruct a storage | |
| # of the same type using (basePtr, offset, size). | |
| # 2. we can reconstruct the tensor on top of the reconstructed storage | |
| # Tensor(size=0x040, offset=0x020, storage=Storage(data=basePtr+0xA100, size=0x0100)) | |
| # | |
| # This strategy has a few implications: | |
| # | |
| # 1. When we serialize a CUDA tensor for IPC, we cannot do it all in one | |
| # go (non-compositionally), and this requires to have a global map | |
| # memHandle -> devPtr for each process. | |
| # | |
| # 2. We MUST NOT let the new IPC tensor be resizable. Originally, a resize | |
| # of the storage beyond 0x100 would merely have caused us to do a | |
| # reallocation. You don't really want to do this, but if you did, | |
| # all that would happen is that you would lose IPC sharing. But if | |
| # you do this in the new world, we will happily let you write out of | |
| # bounds of your "allocation", clobbering unrelated data in the cached | |
| # allocator block. BAD! | |
| # | |
| # By the way, in old versions of PyTorch, we supported this situation | |
| # natively using a "storage view", which permitted multiple storages to be | |
| # views on each other. But this was the *only* use of storage views, so we | |
| # eliminated it so that we could just use tensor views to implement the same | |
| # thing. | |
| # | |
| if storage._untyped_storage.device.type == 'cuda': | |
| (device, | |
| handle, | |
| storage_size_bytes, | |
| storage_offset_bytes, | |
| ref_counter_handle, | |
| ref_counter_offset, | |
| event_handle, | |
| event_sync_required) = storage._share_cuda_() | |
| tensor_offset = tensor.storage_offset() | |
| shared_cache[handle] = StorageWeakRef(storage) | |
| # _backward_hooks purposely omitted here, see | |
| # Note [Don't serialize hooks] | |
| return (rebuild_cuda_tensor, | |
| (type(tensor), | |
| tensor.size(), | |
| tensor.stride(), | |
| tensor_offset, # tensor offset in its storage | |
| type(storage), | |
| tensor.dtype, | |
| device, | |
| handle, # identifier which CUDA allocation is the storage in. | |
| storage_size_bytes, # size(in bytes) of the storage | |
| storage_offset_bytes, # offset(in bytes) of the storage in the CUDA allocation | |
| tensor.requires_grad, | |
| ref_counter_handle, | |
| ref_counter_offset, | |
| event_handle, | |
| event_sync_required)) | |
| # _backward_hooks purposely omitted here, see Note [Don't serialize hooks] | |
| metadata = (tensor.storage_offset(), tensor.size(), tensor.stride(), tensor.requires_grad) | |
| return (rebuild_tensor, ( | |
| type(tensor), | |
| storage, | |
| metadata)) | |
| def fd_id(fd): | |
| # Returns a tuple which uniquely identifies a file descriptor. In Mac OS, | |
| # this doesn't work with shared memory handles, which is why we don't | |
| # support the "file_descriptor" sharing method on that platform. | |
| stat = os.fstat(fd) | |
| return (stat.st_ino, stat.st_dev) | |
| def storage_from_cache(cls, key): | |
| storage_ref = shared_cache.get(key) | |
| if storage_ref is None: | |
| return None | |
| return torch.UntypedStorage._new_with_weak_ptr(storage_ref.cdata) | |
| def rebuild_storage_fd(cls, df, size): | |
| fd = df.detach() | |
| try: | |
| storage = storage_from_cache(cls, fd_id(fd)) | |
| if storage is not None: | |
| return storage | |
| storage = cls._new_shared_fd_cpu(fd, size) | |
| shared_cache[fd_id(fd)] = StorageWeakRef(storage) | |
| return storage | |
| finally: | |
| os.close(fd) | |
| def rebuild_storage_filename(cls, manager, handle, size, dtype=None): | |
| storage: Union[torch.TypedStorage, torch.UntypedStorage] = storage_from_cache(cls, handle) | |
| if storage is not None: | |
| return storage._shared_decref() | |
| if dtype is None: | |
| storage = torch.UntypedStorage._new_shared_filename_cpu(manager, handle, size) | |
| else: | |
| byte_size = size * torch._utils._element_size(dtype) | |
| untyped_storage: torch.UntypedStorage = torch.UntypedStorage._new_shared_filename_cpu(manager, handle, byte_size) | |
| storage = torch.TypedStorage( | |
| wrap_storage=untyped_storage, | |
| dtype=dtype, | |
| _internal=True) | |
| shared_cache[handle] = StorageWeakRef(storage) | |
| return storage._shared_decref() | |
| def rebuild_storage_empty(cls): | |
| return cls() | |
| def rebuild_typed_storage(storage, dtype): | |
| return torch.storage.TypedStorage(wrap_storage=storage, dtype=dtype, _internal=True) | |
| # Use for torch.storage.TypedStorage | |
| def reduce_typed_storage(storage): | |
| return (rebuild_typed_storage, (storage._untyped_storage, storage.dtype)) | |
| def rebuild_typed_storage_child(storage, storage_type): | |
| return storage_type(wrap_storage=storage, _internal=True) | |
| # Use for child classes of torch.storage.TypedStorage, like torch.FloatStorage | |
| def reduce_typed_storage_child(storage): | |
| return (rebuild_typed_storage_child, (storage._untyped_storage, type(storage))) | |
| def reduce_storage(storage): | |
| from . import get_sharing_strategy | |
| if storage.is_cuda: | |
| raise RuntimeError("Cannot pickle CUDA storage; try pickling a CUDA tensor instead") | |
| elif get_sharing_strategy() == 'file_system': | |
| metadata = storage._share_filename_cpu_() | |
| cache_key = metadata[1] | |
| rebuild = rebuild_storage_filename | |
| if isinstance(storage, torch.TypedStorage): | |
| metadata += (storage.dtype,) | |
| storage._shared_incref() | |
| elif storage.size() == 0: | |
| # This is special cased because Empty tensors | |
| # (with size 0) cannot be mmapped. | |
| return (rebuild_storage_empty, (type(storage),)) | |
| else: | |
| fd, size = storage._share_fd_cpu_() | |
| df = multiprocessing.reduction.DupFd(fd) | |
| cache_key = fd_id(fd) | |
| metadata = (df, size) | |
| rebuild = rebuild_storage_fd # type: ignore[assignment] | |
| shared_cache[cache_key] = StorageWeakRef(storage) | |
| return (rebuild, (type(storage),) + metadata) | |
| def init_reductions(): | |
| ForkingPickler.register(torch.cuda.Event, reduce_event) | |
| for t in torch._storage_classes: | |
| if t.__name__ == 'UntypedStorage': | |
| ForkingPickler.register(t, reduce_storage) | |
| else: | |
| ForkingPickler.register(t, reduce_typed_storage_child) | |
| ForkingPickler.register(torch.storage.TypedStorage, reduce_typed_storage) | |
| for t in torch._tensor_classes: | |
| ForkingPickler.register(t, reduce_tensor) | |
| # TODO: Maybe this should be in tensor_classes? :) | |
| ForkingPickler.register(torch.Tensor, reduce_tensor) | |
| ForkingPickler.register(torch.nn.parameter.Parameter, reduce_tensor) |