Function embedding

tensorrt_llm/functional.py:2717–2851

Add an operation to perform embedding lookup. That operation performs the embedding lookup. The 'input' tensor contains the identifiers of the rows of 'weight' to gather. 1. Distribute the embedding lookup table over multiple GPU When 'tp_size' is greater than 1 and the 'tp_gr

(input: Tensor,
              weight: Tensor,
              tp_size=1,
              tp_group=None,
              sharding_dim=0,
              tp_rank=None,
              per_token_scale=None,
              padding=None)

def embedding(input: Tensor,
              weight: Tensor,
              tp_size=1,
              tp_group=None,
              sharding_dim=0,
              tp_rank=None,
              per_token_scale=None,
              padding=None) -> Tensor:
    '''
    Add an operation to perform embedding lookup.

    That operation performs the embedding lookup. The 'input' tensor contains
    the identifiers of the rows of 'weight' to gather.

    1. Distribute the embedding lookup table over multiple GPU
    When 'tp_size' is greater than 1 and the 'tp_group' is defined, this
    embedding lookup is distributed among multiple GPUs.

    When 'sharding_dim==0', each GPU stores a subset of the rows of the embedding
    table rows(that number of rows per GPU is given by weights.shape[0] and the offset to
    the 1st row stored on the GPU is given by rank * weights.shape[0]). Each
    parallel rank will query all the indices and set 0s for the weights that
    are not stored on the associated GPU. To compute the final result, a
    parallel all-reduce operation is added to the TensorRT graph. That lookup
    can be performed using either the plugin or the operators TensorRT support.

    When'sharding_dim==1', each GPU stores a subset of the embedding table's columns.
    Each rank can obtain a portion of the embedding results.
    Then the embedding is collected using the  all-gather operation.
    Related transposition operations are also used to obtain the final results.

    2. Store embedding lookup table as a whole
    When 'tp_size' is not greater than 1, the embedding lookup table will not
    be divided. In this case, when the default_net().plugin_config.lookup_plugin is set,
    the operation is implemented using a plugin (without the all-reduce operation).
    Otherwise, this operation is implemented using the standard IGatherLayer in TensorRT.

    Parameters:
        input : Tensor
            The input tensor the contains the indices to perform the lookup.

        weight : Tensor
            The table to gather from.

        tp_size : int
            The number of GPUs collaborating to perform that embedding.

        tg_group : Optional[List[int]]
            The group of world ranks participating in the all-reduce when
            tp_size > 1.

        sharding_dim : int
            sharding_dim = 0 means that we shard the embedding table in vocab dim;
            sharding_dim = 1 means that we shard the embedding table in embedding dim.

        tp_rank : int
            The tensor parallelism rank. Used to calculate offset in TP on vocab dim.