pylibwholegraph API doc

APIs

torch.initialize.init_torch_env(world_rank, ...)	Init WholeGraph environment for PyTorch.
torch.initialize.init_torch_env_and_create_wm_comm(...)	Init WholeGraph environment for PyTorch and create
torch.initialize.finalize()	Finalize WholeGraph.
torch.comm.WholeMemoryCommunicator(wmb_comm)	WholeMemory Communicator.
torch.comm.set_world_info(world_rank, ...)	Set the global world's information.
torch.comm.create_group_communicator([...])	Create WholeMemory Communicator. For example: 24 ranks with group_size = 4 and comm_stride = 2 will create following groups: [0, 2, 4, 6], [1, 3, 5, 7], [8, 10, 12, 14], [9, 11, 13, 15], [16, 18, 20, 22], [17, 19, 21, 23] :param group_size: Size of each group, -1 means to use all ranks in just one single group. :param comm_stride: Stride of each rank in each group :return: WholeMemoryCommunicator.
torch.comm.destroy_communicator(wm_comm)	Destroy WholeMemoryCommunicator :param wm_comm: WholeMemoryCommunicator to destroy :return: None
torch.comm.get_global_communicator([...])	Get the global communicator of this job :return: WholeMemoryCommunicator that has all GPUs in it.
torch.comm.get_local_node_communicator()	Get the local node communicator of this job :return: WholeMemoryCommunicator that has GPUs in the same node.
torch.comm.get_local_device_communicator()	Get the local device communicator of this job :return: WholeMemoryCommunicator that has only the GPU belonging to current process.
torch.tensor.WholeMemoryTensor(wmb_tensor)	WholeMemory Tensor
torch.tensor.create_wholememory_tensor(comm, ...)	Create empty WholeMemory Tensor. Now only support dim = 1 or 2 :param comm: WholeMemoryCommunicator :param memory_type: WholeMemory type, should be continuous, chunked or distributed :param memory_location: WholeMemory location, should be cpu or cuda :param sizes: size of the tensor :param dtype: data type of the tensor :param strides: strides of the tensor :param tensor_entry_partition: rank partition based on entry; tensor_entry_partition[i] determines the entry count of rank i and shoud be a positive integer; the sum of tensor_entry_partition should equal to total entry count; entries will be equally partitioned if None :return: Allocated WholeMemoryTensor.
torch.tensor.create_wholememory_tensor_from_filelist(...)	Create WholeMemory Tensor from list of binary files. :param comm: WholeMemoryCommunicator :param memory_type: WholeMemory type, should be continuous, chunked or distributed :param memory_location: WholeMemory location, should be cpu or cuda :param filelist: list of binary files :param dtype: data type of the tensor :param last_dim_size: 0 for create 1-D array, positive value for create matrix column size :param last_dim_strides: stride of last_dim, -1 for same as size of last dim. :param tensor_entry_partition: rank partition based on entry; tensor_entry_partition[i] determines the entry count of rank i and shoud be a positive integer; the sum of tensor_entry_partition should equal to total entry count; entries will be equally partitioned if None :return: WholeMemoryTensor.
torch.tensor.destroy_wholememory_tensor(...)	Destroy allocated WholeMemory Tensor :param wm_tensor: WholeMemory Tensor :return: None
torch.embedding.WholeMemoryOptimizer(global_comm)	Sparse Optimizer for WholeMemoryEmbedding.
torch.embedding.create_wholememory_optimizer(...)	Create WholeMemoryOptimizer.
torch.embedding.destroy_wholememory_optimizer(...)	Destroy WholeMemoryOptimizer :param optimizer: WholeMemoryOptimizer to destroy :return: None
torch.embedding.WholeMemoryCachePolicy(...)	Cache policy to create WholeMemoryEmbedding.
torch.embedding.create_wholememory_cache_policy(...)	Create WholeMemoryCachePolicy NOTE: in most cases, create_builtin_cache_policy() can support. This function is a more flexible interface :param cache_comm: WholeMemory communicator of the cache :param memory_type: WholeMemory type of cache :param memory_location: WholeMemory location of cache :param access_type: Access type needed :param ratio: Ratio of cache :return: WholeMemoryCachePolicy.
torch.embedding.create_builtin_cache_policy(...)	Create builtin cache policy
torch.embedding.destroy_wholememory_cache_policy(...)	Destroy WholeMemoryCachePolicy :param cache_policy: WholeMemoryCachePolicy to destroy :return: None
torch.embedding.WholeMemoryEmbedding(...)	WholeMemory Embedding
torch.embedding.create_embedding(comm, ...)	Create embedding :param comm: WholeMemoryCommunicator :param memory_type: WholeMemory type, should be continuous, chunked or distributed :param memory_location: WholeMemory location, should be cpu or cuda :param dtype: data type :param sizes: size of the embedding, must be 2D :param cache_policy: cache policy :param embedding_entry_partition: rank partition based on entry; embedding_entry_partition[i] determines the entry count of rank i and shoud be a positive integer; the sum of embedding_entry_partition should equal to total entry count; entries will be equally partitioned if None :param gather_sms: the number of SMs used in gather process :param round_robin_size: continuous embedding size of a rank using round robin shard strategy :return: WholeMemoryEmbedding
torch.embedding.create_embedding_from_filelist(...)	Create embedding from file list :param comm: WholeMemoryCommunicator :param memory_type: WholeMemory type, should be continuous, chunked or distributed :param memory_location: WholeMemory location, should be cpu or cuda :param filelist: list of files :param dtype: data type :param last_dim_size: size of last dim :param cache_policy: cache policy :param embedding_entry_partition: rank partition based on entry; embedding_entry_partition[i] determines the entry count of rank i and shoud be a positive integer; the sum of embedding_entry_partition should equal to total entry count; entries will be equally partitioned if None :param gather_sms: the number of SMs used in gather process :param round_robin_size: continuous embedding size of a rank using round robin shard strategy :return:
torch.embedding.destroy_embedding(wm_embedding)	Destroy WholeMemoryEmbedding :param wm_embedding: WholeMemoryEmbedding to destroy :return: None
torch.embedding.WholeMemoryEmbeddingModule(...)	torch.nn.Module wrapper of WholeMemoryEmbedding
torch.graph_structure.GraphStructure()	Graph structure storage Actually, it is the graph structure of one relation, represented in CSR format.