cugraph_pyg.tensor.DistTensor#

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]: Shape used when creating an empty tensor. It has to be one- or two-dimensional and is required when src is omitted.
expected_shapeOptional[list, tuple]: Optional expected shape when loading files. Parquet shape is inferred when omitted.
last_dim_sizeOptional[int]: Required for binary files and inferred for Parquet files when omitted. Zero requests a 1-D tensor; a positive value requests a 2-D tensor.
fail_on_dtype_mismatchbool: Raise an error instead of warning and converting when Parquet column dtypes differ from dtype.
dtypeOptional[torch.dtype]: The dtype of the tensor. Required for file input and inferred from an in-memory torch.Tensor source.
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”.
file_formatOptional[Literal[“auto”, “binary”, “parquet”]] = “auto”: Format used when src is a list of files or a Parquet file. auto detects Parquet files by extension and treats other files as raw binary. Parquet files must contain only scalar numeric columns in physical column order. Dtype is required for file input; Parquet shape is inferred unless expected_shape is supplied.

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,

file_format: Literal['auto', 'binary', 'parquet'] | None = 'auto',

expected_shape: list | tuple | None = None,

last_dim_size: int | None = None,

fail_on_dtype_mismatch: bool = False,

**kwargs,

)[source]#

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

`device`
`dim`
`dtype`
`partition_book`
`shape`