| 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: |
| 89 | padding = None |
| 90 | |
| 91 | return embedding(x, |
| 92 | self.weight.value, |
| 93 | tp_size=self.tp_size, |
| 94 | tp_group=self.tp_group, |
| 95 | sharding_dim=self.sharding_dim, |
| 96 | tp_rank=self.tp_rank, |
| 97 | padding=padding) |
| 98 | |
| 99 | def weight_loader(self, mapping: Mapping, param: Parameter, |
| 100 | loaded_weight: torch.Tensor): |