cugraph_pyg.tensor.DistTensor#
- class cugraph_pyg.tensor.DistTensor(
- src: Tensor | str | List[str] | None = None,
- shape: list | tuple | None = None,
- dtype: dtype | None = None,
- device: Literal['cpu', 'cuda'] | None = 'cpu',
- partition_book: List[int] | None = None,
- backend: str | None = 'nccl',
- *args,
- **kwargs,
WholeGraph-backed Distributed Tensor Interface for PyTorch. Parameters ———- src: Optional[Union[torch.Tensor, str, List[str]]]
The source of the tensor. It can be a torch.Tensor on host, a file path, or a list of file paths. When the source is omitted, the tensor will be load later.
- shapeOptional[list, tuple]
The shape of the tensor. It has to be a one- or two-dimensional tensor for now. When the shape is omitted, the src has to be specified and must be pt or npy file paths.
- dtypeOptional[torch.dtype]
The dtype of the tensor. When the dtype is omitted, the src has to be specified and must be pt or npy file paths.
- deviceOptional[Literal[“cpu”, “cuda”]] = “cpu”
The desired location to store the embedding [ “cpu” | “cuda” ]. Default is “cpu”, i.e., host-pinned memory (UVA).
- partition_bookUnion[List[int], None] = None
1-D Range partition based on entry (dim-0). partition_book[i] determines the entry count of rank i and shoud be a positive integer; the sum of partition_book should equal to shape[0]. Entries will be equally partitioned if None.
- backendOptional[Literal[“vmm”, “nccl”, “nvshmem”, “chunked”]] = “nccl”
The backend used for communication. Default is “nccl”.
- Attributes:
- device
- dim
- dtype
- partition_book
- shape
Methods
from_file(file_path[, device, ...])Create a WholeGraph-backed Distributed Tensor from a file. Parameters ---------- file_path : str The file path to the tensor. The file can be in the format of PyTorch tensor or NumPy array. device : str, optional The desired location to store the embedding [ "cpu" | "cuda" ]. Default is "cpu". backend : str, optional The backend used for communication. Default is "nccl". Returns: ------- DistTensor The WholeGraph-backed Distributed Tensor.
from_tensor(tensor[, device, ...])Create a WholeGraph-backed Distributed Tensor from a PyTorch tensor. Parameters ---------- tensor : torch.Tensor The PyTorch tensor to be copied to the WholeGraph tensor. device : str, optional The desired location to store the embedding [ "cpu" | "cuda" ]. Default is "cpu". backend : str, optional The backend used for communication. Default is "nccl". Returns: ------- DistTensor The WholeGraph-backed Distributed Tensor.
get_comm()Get the communicator of the WholeGraph embedding. Returns: ------- WholeMemoryCommunicator The WholeGraph global communicator of the WholeGraph embedding.
get_local_offset()Get the local embedding tensor and its element offset at current rank. Returns: ------- (torch.Tensor, int) Tuple of local torch Tensor (converted from DLPack) and its offset.
get_local_tensor([host_view])Get the local embedding tensor and its element offset at current rank. Returns: ------- (torch.Tensor, int) Tuple of local torch Tensor (converted from DLPack) and its offset.
load_from_global_tensor
load_from_local_tensor
- __init__(
- src: Tensor | str | List[str] | None = None,
- shape: list | tuple | None = None,
- dtype: dtype | None = None,
- device: Literal['cpu', 'cuda'] | None = 'cpu',
- partition_book: List[int] | None = None,
- backend: str | None = 'nccl',
- *args,
- **kwargs,
Methods
__init__([src, shape, dtype, device, ...])from_file(file_path[, device, ...])Create a WholeGraph-backed Distributed Tensor from a file. Parameters ---------- file_path : str The file path to the tensor. The file can be in the format of PyTorch tensor or NumPy array. device : str, optional The desired location to store the embedding [ "cpu" | "cuda" ]. Default is "cpu". backend : str, optional The backend used for communication. Default is "nccl". Returns: ------- DistTensor The WholeGraph-backed Distributed Tensor.
from_tensor(tensor[, device, ...])Create a WholeGraph-backed Distributed Tensor from a PyTorch tensor. Parameters ---------- tensor : torch.Tensor The PyTorch tensor to be copied to the WholeGraph tensor. device : str, optional The desired location to store the embedding [ "cpu" | "cuda" ]. Default is "cpu". backend : str, optional The backend used for communication. Default is "nccl". Returns: ------- DistTensor The WholeGraph-backed Distributed Tensor.
get_comm()Get the communicator of the WholeGraph embedding. Returns: ------- WholeMemoryCommunicator The WholeGraph global communicator of the WholeGraph embedding.
get_local_offset()Get the local embedding tensor and its element offset at current rank. Returns: ------- (torch.Tensor, int) Tuple of local torch Tensor (converted from DLPack) and its offset.
get_local_tensor([host_view])Get the local embedding tensor and its element offset at current rank. Returns: ------- (torch.Tensor, int) Tuple of local torch Tensor (converted from DLPack) and its offset.
load_from_global_tensor(tensor)load_from_local_tensor(tensor)Attributes
devicedimdtypepartition_bookshape