cugraph_pyg.tensor.DistEmbedding#
- class cugraph_pyg.tensor.DistEmbedding(
- src: torch.Tensor | str | List[str] | None = None,
- shape: list | tuple | None = None,
- dtype: torch.dtype | None = None,
- device: Literal['cpu', 'cuda'] | None = 'cpu',
- partition_book: List[int] | None = None,
- backend: str | None = 'nccl',
- cache_policy: pylibwholegraph.WholeMemoryCachePolicy | None = None,
- gather_sms: int | None = -1,
- round_robin_size: int = 0,
- name: str | None = None,
WholeGraph-backed Distributed Embedding 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. Whne 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”.
- cache_policyOptional[WholeMemoryCachePolicy] = None
The cache policy for the tensor if it is an embedding. Default is None.
- gather_smsOptional[int] = -1
Whether to gather the embeddings on all GPUs. Default is False.
- round_robin_size: int = 0
continuous embedding size of a rank using round robin shard strategy
- nameOptional[str]
The name of the tensor.
- Attributes:
- name
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". name : str, optional The name of the tensor. Returns: ------- DistTensor The WholeGraph-backed Distributed Tensor.
from_tensor(tensor[, device, ...])Create a WholeGraph-backed Distributed Embedding (hooked with PyT's grad tracing) 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". name : str, optional The name of the tensor. Returns: ------- DistEmbedding The WholeGraph-backed Distributed Tensor.
- __init__(
- src: torch.Tensor | str | List[str] | None = None,
- shape: list | tuple | None = None,
- dtype: torch.dtype | None = None,
- device: Literal['cpu', 'cuda'] | None = 'cpu',
- partition_book: List[int] | None = None,
- backend: str | None = 'nccl',
- cache_policy: pylibwholegraph.WholeMemoryCachePolicy | None = None,
- gather_sms: int | None = -1,
- round_robin_size: int = 0,
- name: str | None = None,
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". name : str, optional The name of the tensor. Returns: ------- DistTensor The WholeGraph-backed Distributed Tensor.
from_tensor(tensor[, device, ...])Create a WholeGraph-backed Distributed Embedding (hooked with PyT's grad tracing) 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". name : str, optional The name of the tensor. Returns: ------- DistEmbedding 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
devicedimdtypenamepartition_bookshape