distributed#

class ParallelEmbedding(num_embeddings, embedding_dim, rank, world_size, weight_attr=None, name=None)[源代码]#

Parallel Embedding.

参数:
  • 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)[源代码]#
参数:

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 the embedding Tensor mapped by x.

返回类型:

Tensor

class ColumnParallelLiner(size, num_partitions=1, gather_out=True, param_attr=None, bias_attr=None, name=None)[源代码]#

Parallel Linear, axis=1.

参数:
  • 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)[源代码]#
参数:

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

返回:

Returns the embedding Tensor mapped by x.

返回类型:

Tensor

class RowParallelLiner(size, num_partitions=1, input_is_parallel=False, param_attr=None, bias_attr=None, name=None)[源代码]#

Parallel Linear, axis=0.

参数:
  • 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)[源代码]#
参数:

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

返回:

Returns the embedding Tensor mapped by x.

返回类型:

Tensor