parallel

class ParallelEmbedding(num_embeddings, embedding_dim, rank, world_size, weight_attr=None, name=None)[source]

Parallel Embedding.

Parameters

  • num_embeddings (int) – The size of embedding dictionary which dictates the maximum value of the input id.

  • embedding_dim (int) – The dimensions of each embedding vector.

  • rank (int) – The rank of the current part, which determines the start index of the vocab.

  • world_size (int) – The number of trainers.

  • weight_attr (Tensor, optional) – Specify the weight parameter property, including the initialization method. Defaults to None which means the default weight parameter property will be used.

  • name (str, optional) – Normally there is no need for user to set this property. Defaults to None.

forward(x)[source]
Parameters

x (Tensor) – A Tensor contains the id information. Its data type should be int32 or int64, and the value of the input id should be in [0, weight.shape[0]] .

Returns

Returns the embedding Tensor mapped by x.

Return type

Tensor

class ColumnParallelLiner(size, num_partitions=1, gather_out=True, param_attr=None, bias_attr=None, name=None)[source]

Parallel Linear, axis=1.

Parameters

  • size (int) – The size of embedding vector.

  • num_partitions (int, optional) – The number of parts within a model parallel group. Defaults to 1.

  • gather_out (bool, optional) – Whether to gather the output tensor. Defaults to True.

  • param_attr (Tensor, optional) – Specify the parameter property, including the initialization method. Defaults to None which means the default parameter property will be used.

  • bias_attr (Tensor, optional) – Specify the bias property. Defaults to None which means the default parameter property will be used.

  • name (str, optional) – Normally there is no need for user to set this property. Defaults to None.

forward(x)[source]
Parameters

x (Tensor) – The input tensor. Its data type can be int or float.

Returns

Returns the embedding Tensor mapped by x.

Return type

Tensor

class RowParallelLiner(size, num_partitions=1, input_is_parallel=False, param_attr=None, bias_attr=None, name=None)[source]

Parallel Linear, axis=0.

Parameters

  • size (int) – The size of embedding vector.

  • num_partitions (int, optional) – The number of parts within a model parallel group. Defaults to 1.

  • input_is_parallel (bool, optional) – Whether the input is parallel. Defaults to False.

  • param_attr (Tensor, optional) – Specify the parameter property, including the initialization method. Defaults to None which means the default parameter property will be used.

  • bias_attr (Tensor, optional) – Specify the bias property. Defaults to None which means the default parameter property will be used.

  • name (str, optional) – Normally there is no need for user to set this property. Defaults to None.

forward(x)[source]
Parameters

x (Tensor) – The input tensor. Its data type can be int or float.

Returns

Returns the embedding Tensor mapped by x.

Return type

Tensor