The embedding layer takes input indices (x) and the embedding lookup table (weight) as input. And output the corresponding embeddings according to input indices. The size of weight is [num_embeddings, embedding_dim] Four parameters (tp_size, tp_group, sharding_dim, tp_rank) are inv
| 29 | |
| 30 | |
| 31 | class Embedding(Module): |
| 32 | """ |
| 33 | The embedding layer takes input indices (x) and the embedding lookup table (weight) as input. |
| 34 | And output the corresponding embeddings according to input indices. |
| 35 | The size of weight is [num_embeddings, embedding_dim] |
| 36 | |
| 37 | Four parameters (tp_size, tp_group, sharding_dim, tp_rank) are involved in tensor parallelism. |
| 38 | Only when "tp_size > 1 and tp_group is not None", tensor parallelism is enabled. |
| 39 | When "sharding_dim == 0", the weight is shared in the vocabulary dimension. |
| 40 | tp_rank must be set when sharding_dim == 0. |
| 41 | When "sharding_dim == 1", the weight is shard in the hidden dimension. |
| 42 | """ |
| 43 | |
| 44 | def __init__(self, |
| 45 | num_embeddings: int, |
| 46 | embedding_dim: int, |
| 47 | dtype: Optional[str] = None, |
| 48 | tp_size: int = 1, |
| 49 | tp_group: Optional[list] = None, |
| 50 | sharding_dim: int = 0, |
| 51 | tp_rank: Optional[int] = None): |
| 52 | super().__init__() |
| 53 | # num_embeddings records the total vocab size no matter using TP or not |
| 54 | self.num_embeddings = num_embeddings |
| 55 | self.embedding_dim = embedding_dim |
| 56 | self.tp_size = tp_size |
| 57 | self.tp_group = tp_group |
| 58 | self.sharding_dim = sharding_dim |
| 59 | self.tp_rank = tp_rank |
| 60 | self.dtype = dtype |
| 61 | self.tp_dim = sharding_dim |
| 62 | |
| 63 | if sharding_dim == 1: |
| 64 | shape = (self.num_embeddings, self.embedding_dim // self.tp_size) |
| 65 | elif sharding_dim == 0: |
| 66 | shape = (math.ceil(self.num_embeddings / self.tp_size), |
| 67 | self.embedding_dim) |
| 68 | |
| 69 | self.weight = Parameter(shape=shape, dtype=dtype) |
| 70 | |
| 71 | self.weight_padding_size = ((8 - shape[0] % 8) % 8, shape[1]) |
| 72 | |
| 73 | set_obj_attrs(self.weight, { |
| 74 | "weight_loader": self.weight_loader, |
| 75 | }) |
| 76 | |
| 77 | def forward(self, x): |
| 78 | # The embedding weight is padded to the multiple of 8. |
| 79 | # The reason is that when lm_head and vocab_embedding are using the same embedding weight, |
| 80 | # previously weights can't be depulicated in the engine because gemm will pad the weight to the multiple of 8. |
| 81 | # If we also pad the embedding weight to the multiple of 8, the weights can be successfully deduplicated. |
| 82 | # This will not affect the input and output of the gather op and perf impact is negligible. |
| 83 | if self.weight_padding_size[0] != 0: |
| 84 | padding_values = np.zeros(self.weight_padding_size, |
| 85 | dtype=trt_dtype_to_np( |
| 86 | self.weight.value.dtype)) |
| 87 | padding = constant(padding_values) |
| 88 | else: |
no outgoing calls