| 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. |