| 83 | |
| 84 | |
| 85 | class CLIPEncoderLayer(Module): |
| 86 | |
| 87 | def __init__(self, hidden_size, num_attention_heads, |
| 88 | max_position_embeddings, norm_epsilon, intermediate_size, |
| 89 | hidden_act, mapping: Mapping, dtype): |
| 90 | super().__init__() |
| 91 | self.hidden_size = hidden_size |
| 92 | self.dtype = dtype |
| 93 | self.mapping = mapping |
| 94 | |
| 95 | self.input_layernorm = LayerNorm(normalized_shape=self.hidden_size, |
| 96 | eps=norm_epsilon, |
| 97 | dtype=self.dtype) |
| 98 | |
| 99 | self.attention = BertAttention( |
| 100 | hidden_size=self.hidden_size, |
| 101 | num_attention_heads=num_attention_heads, |
| 102 | max_position_embeddings=max_position_embeddings, |
| 103 | attention_head_size=self.hidden_size // num_attention_heads, |
| 104 | num_kv_heads=num_attention_heads, |
| 105 | dtype=self.dtype, |
| 106 | tp_group=self.mapping.tp_group, |
| 107 | tp_size=self.mapping.tp_size, |
| 108 | tp_rank=self.mapping.tp_rank, |
| 109 | cp_group=self.mapping.cp_group, |
| 110 | cp_size=self.mapping.cp_size) |
| 111 | |
| 112 | self.post_layernorm = LayerNorm(normalized_shape=self.hidden_size, |
| 113 | eps=norm_epsilon, |
| 114 | dtype=self.dtype) |
| 115 | |
| 116 | self.mlp = MLP(hidden_size=self.hidden_size, |
| 117 | ffn_hidden_size=intermediate_size, |
| 118 | hidden_act=hidden_act, |
| 119 | dtype=self.dtype, |
| 120 | tp_group=self.mapping.tp_group, |
| 121 | tp_size=self.mapping.tp_size) |
| 122 | |
| 123 | def forward(self, hidden_states): |
| 124 | |
| 125 | residual = hidden_states |
| 126 | hidden_states = self.input_layernorm(hidden_states) |
| 127 | hidden_states = self.attention(hidden_states) |
| 128 | hidden_states = residual + hidden_states |
| 129 | |
| 130 | residual = hidden_states |
| 131 | hidden_states = self.post_layernorm(hidden_states) |
| 132 | hidden_states = self.mlp(hidden_states) |
| 133 | hidden_states = residual + hidden_states |
| 134 | return hidden_states |
| 135 | |
| 136 | |
| 137 | class CLIPEncoder(Module): |