Embeddings and optimizers#

pylibwholegraph.torch.WholeMemoryOptimizer(...)

Sparse Optimizer for WholeMemoryEmbedding.

pylibwholegraph.torch.create_wholememory_optimizer(...)

Create WholeMemoryOptimizer.

pylibwholegraph.torch.destroy_wholememory_optimizer(...)

Destroy WholeMemoryOptimizer :param optimizer: WholeMemoryOptimizer to destroy :return: None

pylibwholegraph.torch.WholeMemoryCachePolicy(...)

Cache policy to create WholeMemoryEmbedding.

pylibwholegraph.torch.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.

pylibwholegraph.torch.create_builtin_cache_policy(...)

Create builtin cache policy

pylibwholegraph.torch.destroy_wholememory_cache_policy(...)

Destroy WholeMemoryCachePolicy :param cache_policy: WholeMemoryCachePolicy to destroy :return: None

pylibwholegraph.torch.WholeMemoryEmbedding(...)

WholeMemory Embedding

pylibwholegraph.torch.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

pylibwholegraph.torch.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:

pylibwholegraph.torch.destroy_embedding(...)

Destroy WholeMemoryEmbedding :param wm_embedding: WholeMemoryEmbedding to destroy :return: None

pylibwholegraph.torch.WholeMemoryEmbeddingModule(...)

torch.nn.Module wrapper of WholeMemoryEmbedding