Embeddings#

Typedefs

typedef struct wholememory_embedding_cache_policy_ *wholememory_embedding_cache_policy_t#

Opaque handle to WholeMemory Embedding Cache Policy.

An Opaque handle to WholeMemory Embedding Cache Policy

typedef struct wholememory_embedding_optimizer_ *wholememory_embedding_optimizer_t#

Opaque handle to WholeMemory Embedding Optimizer.

An Opaque handle to WholeMemory Embedding Optimizer

typedef struct wholememory_embedding_ *wholememory_embedding_t#

Opaque handle to WholeMemory Embedding.

An Opaque handle to WholeMemory Embedding

Enums

enum wholememory_access_type_t#

defines access type of WholeMemory Embedding

Values:

enumerator WHOLEMEMORY_AT_NONE#

Not defined

enumerator WHOLEMEMORY_AT_READONLY#

Only have readonly access to the WholeMemory

enumerator WHOLEMEMORY_AT_READWRITE#

May have write access to the WholeMemory

enum wholememory_optimizer_type_t#

defines optimizer type for WholeMemory Embedding

Values:

enumerator WHOLEMEMORY_OPT_NONE#

No optimizer needed

enumerator WHOLEMEMORY_OPT_SGD#

Use SGD optimizer

enumerator WHOLEMEMORY_OPT_LAZY_ADAM#

Use Lazy Adam optimizer

enumerator WHOLEMEMORY_OPT_RMSPROP#

Use RMSProp optimizer

enumerator WHOLEMEMORY_OPT_ADAGRAD#

Use AdaGrad optimizer

Functions

wholememory_error_code_t wholememory_create_embedding_optimizer(
wholememory_embedding_optimizer_t *optimizer,
wholememory_optimizer_type_t optimizer_type
)#

Create Optimizer

Parameters:
Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_optimizer_set_parameter(
wholememory_embedding_optimizer_t optimizer,
const char *parameter_name,
void *value
)#

Set parameter for optimizer.

Parameters:

  • optimizer – : Optimizer to set parameter

  • parameter_name – : parameter name

  • value – : parameter value

Returns:

: wholememory_error_code_t

void wholememory_destroy_embedding_optimizer(
wholememory_embedding_optimizer_t optimizer
)#

Destroy optimizer

Parameters:

optimizer – : optimizer to destroy.

wholememory_error_code_t wholememory_create_embedding_cache_policy(
wholememory_embedding_cache_policy_t *cache_policy,
wholememory_comm_t cache_level_comm,
wholememory_memory_type_t memory_type,
wholememory_memory_location_t memory_location,
wholememory_access_type_t access_type,
float cache_ratio
)#

Create WholeMemory Embedding Cache Policy

Parameters:

  • cache_policy – : Returned wholememory_embedding_cache_policy_t

  • cache_level_comm – : At which level to cache the full embedding. In most cases it should be same as wholememory_embedding_t’s comm. If access_type is WHOLEMEMORY_AT_READONLY, it can be different for multiple readonly caches. E.g. for a multi-node WHOLEMEMORY_MT_DISTRIBUTED WHOLEMEMORY_AT_READONLY embedding, it can have a intra-node WHOLEMEMORY_MT_CHUNKED cache. or a multi-node WHOLEMEMORY_MT_DISTRIBUTED cache.

  • memory_type – : Memory Type of the underlying WholeMemory for cache

  • memory_location – : Memory Location of the underlying WholeMemory for cache

  • access_type – : ReadOnly or ReadWrite

  • cache_ratio – : suggested cache ratio, values should be in range [1.0 / 512, 1.0]

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_embedding_cache_policy(
wholememory_embedding_cache_policy_t cache_policy
)#

Destroy WholeMemory Embedding Cache Policy

Parameters:

cache_policy – : WholeMemory Embedding Cache Policy to destroy.

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_create_embedding(
wholememory_embedding_t *wholememory_embedding,
wholememory_tensor_description_t *embedding_tensor_description,
wholememory_comm_t comm,
wholememory_memory_type_t memory_type,
wholememory_memory_location_t memory_location,
wholememory_embedding_cache_policy_t cache_policy,
size_t *embedding_entry_partition = nullptr,
int user_defined_sms = -1,
int round_robin_size = 0
)#

Create WholeMemory Embedding

Parameters:

  • wholememory_embedding – : Returned wholememory_embedding_t

  • embedding_tensor_description – : Description of the embedding, sizes and dtype used, stride and storage_offset ignored. Must be matrix

  • comm – : WholeMemory Communicator

  • memory_type – : Memory Type of the underlying WholeMemory

  • memory_location – : Memory Location of the underlying WholeMemory

  • cache_policy – : Cache policy for this embedding, if don’t use cache, use nullptr

  • embedding_entry_partition – Embedding entry count of each rank, the length must be world_size

  • user_defined_sms – : User-defined sms number for raw embedding gather/scatter

  • round_robin_size – : continuous embedding size in each rank under round-robin shard mode

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_embedding(
wholememory_embedding_t wholememory_embedding
)#

Destroy WholeMemory Embedding

Parameters:

wholememory_embedding – : WholeMemory Embedding to destroy

Returns:

: wholememory_error_code_t

wholememory_tensor_t wholememory_embedding_get_embedding_tensor(
wholememory_embedding_t wholememory_embedding
)#

Get WholeMemory Tensor from WholeMemory Embedding.

Parameters:

wholememory_embedding – : WholeMemory Embedding

Returns:

: WholeMemory Tensor

wholememory_error_code_t wholememory_embedding_set_optimizer(
wholememory_embedding_t wholememory_embedding,
wholememory_embedding_optimizer_t optimizer
)#

Set Optimizer for WholeMemory Embedding

Parameters:

  • wholememory_embedding – : WholeMemory Embedding

  • optimizer – : Optimizer to be set

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_embedding_gather(
wholememory_embedding_t wholememory_embedding,
wholememory_tensor_t indices,
wholememory_tensor_t output,
bool adjust_cache,
wholememory_env_func_t *p_env_fns,
int64_t stream_int
)#

Gather from WholeMemory Embedding

Parameters:

  • wholememory_embedding – : WholeMemory Embedding

  • indices – : indices to gather

  • output – : output tensor

  • adjust_cache – : if we should adjust cache in this gather

  • p_env_fns – : env fns

  • stream_int – : CUDA stream to use

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_embedding_gather_gradient_apply(
wholememory_embedding_t wholememory_embedding,
wholememory_tensor_t indices,
wholememory_tensor_t grads,
bool adjust_cache,
float lr,
wholememory_env_func_t *p_env_fns,
int64_t stream_int
)#

Gather backward for WholeMemory Embedding

Parameters:

  • wholememory_embedding – : WholeMemory Embedding

  • indices – : indices to gather

  • grads – : gradient of output tensor

  • adjust_cache – : if we should adjust cache in this gather

  • lr – : learning rate of current step.

  • p_env_fns – : env fns

  • stream_int – : CUDA stream to use

Returns:

: wholememory_error_code_t

const char *const *wholememory_embedding_get_optimizer_state_names(
wholememory_embedding_t wholememory_embedding
)#

Get optimizer internal state names

Parameters:

wholememory_embedding – : WholeMemory Embedding

Returns:

: nullptr terminated names.

wholememory_tensor_t wholememory_embedding_get_optimizer_state(
wholememory_embedding_t wholememory_embedding,
const char *name
)#

Get optimizer internal state

Parameters:

  • wholememory_embedding – : WholeMemory Embedding

  • name – : state name

Returns:

: internal state, nullptr for not exist.

wholememory_error_code_t wholememory_embedding_writeback_cache(
wholememory_embedding_t wholememory_embedding,
int64_t stream_int
)#

Writeback all cache WholeMemory Embedding

Parameters:

  • wholememory_embedding – : WholeMemory Embedding

  • stream_int – : CUDA stream to use.

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_embedding_drop_all_cache(
wholememory_embedding_t wholememory_embedding,
int64_t stream_int
)#

Drop all cache in WholeMemory Embedding

Parameters:

  • wholememory_embedding – : WholeMemory Embedding

  • stream_int – : CUDA stream to use.

Returns:

: wholememory_error_code_t