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:
- srcOptional[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 loaded 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 should be a positive integer; the sum of partition_book should equal 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.
from_tensor(tensor[, device, ...])Create a WholeGraph-backed Distributed Tensor from a PyTorch tensor.
get_comm()Get the communicator of the WholeGraph embedding.
get_local_offset()Get the local embedding tensor offset at the current rank.
get_local_tensor([host_view])Get the local embedding tensor at the current rank.
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.
from_tensor(tensor[, device, ...])Create a WholeGraph-backed Distributed Tensor from a PyTorch tensor.
get_comm()Get the communicator of the WholeGraph embedding.
get_local_offset()Get the local embedding tensor offset at the current rank.
get_local_tensor([host_view])Get the local embedding tensor at the current rank.
load_from_global_tensor(tensor)load_from_local_tensor(tensor)Attributes
devicedimdtypepartition_bookshape