595 lines
21 KiB
Python
595 lines
21 KiB
Python
import multiprocessing
|
|
import os
|
|
import threading
|
|
from multiprocessing.reduction import ForkingPickler
|
|
from multiprocessing.util import register_after_fork
|
|
from typing import Union
|
|
|
|
import torch
|
|
import torch.utils.hooks
|
|
from torch._namedtensor_internals import check_serializing_named_tensor
|
|
|
|
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:
|
|
r"""A weak reference to a Storage.
|
|
|
|
The cdata member is a Python number containing the integer representation of
|
|
the Storage pointer.
|
|
"""
|
|
|
|
__slots__ = ["cdata", "_free_weak_ref"]
|
|
|
|
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]
|
|
|
|
@classmethod
|
|
def from_weakref(cls, cdata):
|
|
instance = cls.__new__(cls)
|
|
instance.cdata = cdata
|
|
instance._free_weak_ref = torch.Storage._free_weak_ref # type: ignore[attr-defined]
|
|
return instance
|
|
|
|
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):
|
|
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.
|
|
#
|
|
|
|
# TODO: Handle distinguishing between subclass and non-subclass versions of NT better
|
|
# https://github.com/pytorch/pytorch/issues/110543
|
|
from torch.nested._internal.nested_tensor import NestedTensor
|
|
|
|
if tensor.is_nested and not isinstance(tensor, NestedTensor):
|
|
return reduce_nested_tensor(tensor)
|
|
|
|
if tensor.layout in {
|
|
torch.sparse_coo,
|
|
torch.sparse_csr,
|
|
torch.sparse_bsr,
|
|
torch.sparse_csc,
|
|
torch.sparse_bsc,
|
|
}:
|
|
return reduce_sparse_tensor(tensor)
|
|
|
|
storage = tensor._typed_storage()
|
|
|
|
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 rebuild_nested_tensor(
|
|
rebuild_buffer_func,
|
|
rebuild_buffer_args,
|
|
rebuild_sizes_func,
|
|
rebuild_sizes_args,
|
|
rebuild_strides_func,
|
|
rebuild_strides_args,
|
|
rebuild_offsets_func,
|
|
rebuild_offsets_args,
|
|
):
|
|
buffer = rebuild_buffer_func(*rebuild_buffer_args)
|
|
sizes = rebuild_sizes_func(*rebuild_sizes_args)
|
|
strides = rebuild_strides_func(*rebuild_strides_args)
|
|
offsets = rebuild_offsets_func(*rebuild_offsets_args)
|
|
return torch._nested_view_from_buffer_copy(buffer, sizes, strides, offsets)
|
|
|
|
|
|
def reduce_nested_tensor(nt):
|
|
rebuild_buffer_func, rebuild_buffer_args = reduce_tensor(nt.values())
|
|
rebuild_sizes_func, rebuild_sizes_args = reduce_tensor(nt._nested_tensor_size())
|
|
rebuild_strides_func, rebuild_strides_args = reduce_tensor(
|
|
nt._nested_tensor_strides()
|
|
)
|
|
rebuild_offsets_func, rebuild_offsets_args = reduce_tensor(
|
|
nt._nested_tensor_storage_offsets()
|
|
)
|
|
|
|
return (
|
|
rebuild_nested_tensor,
|
|
(
|
|
rebuild_buffer_func,
|
|
rebuild_buffer_args,
|
|
rebuild_sizes_func,
|
|
rebuild_sizes_args,
|
|
rebuild_strides_func,
|
|
rebuild_strides_args,
|
|
rebuild_offsets_func,
|
|
rebuild_offsets_args,
|
|
),
|
|
)
|
|
|
|
|
|
def rebuild_sparse_coo_tensor(
|
|
rebuild_indices_func,
|
|
rebuild_indices_args,
|
|
rebuild_values_func,
|
|
rebuild_values_args,
|
|
shape,
|
|
is_coalesced,
|
|
):
|
|
indices = rebuild_indices_func(*rebuild_indices_args)
|
|
values = rebuild_values_func(*rebuild_values_args)
|
|
return torch.sparse_coo_tensor(indices, values, shape, is_coalesced=is_coalesced)
|
|
|
|
|
|
def rebuild_sparse_compressed_tensor(
|
|
rebuild_compressed_indices_func,
|
|
rebuild_compressed_indices_args,
|
|
rebuild_plain_indices_func,
|
|
rebuild_plain_indices_args,
|
|
rebuild_values_func,
|
|
rebuild_values_args,
|
|
shape,
|
|
layout,
|
|
):
|
|
compressed_indices = rebuild_compressed_indices_func(
|
|
*rebuild_compressed_indices_args
|
|
)
|
|
plain_indices = rebuild_plain_indices_func(*rebuild_plain_indices_args)
|
|
values = rebuild_values_func(*rebuild_values_args)
|
|
return torch.sparse_compressed_tensor(
|
|
compressed_indices, plain_indices, values, shape, layout=layout
|
|
)
|
|
|
|
|
|
def reduce_sparse_tensor(sparse):
|
|
if sparse.layout is torch.sparse_coo:
|
|
rebuild_indices_func, rebuild_indices_args = reduce_tensor(sparse._indices())
|
|
rebuild_values_func, rebuild_values_args = reduce_tensor(sparse._values())
|
|
return (
|
|
rebuild_sparse_coo_tensor,
|
|
(
|
|
rebuild_indices_func,
|
|
rebuild_indices_args,
|
|
rebuild_values_func,
|
|
rebuild_values_args,
|
|
sparse.shape,
|
|
sparse.is_coalesced(),
|
|
),
|
|
)
|
|
else:
|
|
if sparse.layout in {torch.sparse_csr, torch.sparse_bsr}:
|
|
compressed_indices = sparse.crow_indices()
|
|
plain_indices = sparse.col_indices()
|
|
elif sparse.layout in {torch.sparse_csc, torch.sparse_bsc}:
|
|
compressed_indices = sparse.ccol_indices()
|
|
plain_indices = sparse.row_indices()
|
|
else:
|
|
raise NotImplementedError(sparse.layout)
|
|
(
|
|
rebuild_compressed_indices_func,
|
|
rebuild_compressed_indices_args,
|
|
) = reduce_tensor(compressed_indices)
|
|
rebuild_plain_indices_func, rebuild_plain_indices_args = reduce_tensor(
|
|
plain_indices
|
|
)
|
|
rebuild_values_func, rebuild_values_args = reduce_tensor(sparse.values())
|
|
return (
|
|
rebuild_sparse_compressed_tensor,
|
|
(
|
|
rebuild_compressed_indices_func,
|
|
rebuild_compressed_indices_args,
|
|
rebuild_plain_indices_func,
|
|
rebuild_plain_indices_args,
|
|
rebuild_values_func,
|
|
rebuild_values_args,
|
|
sparse.shape,
|
|
sparse.layout,
|
|
),
|
|
)
|
|
|
|
|
|
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)
|