r""" The forward process of the block.
(
self,
x,
query_vector,
input_mask,
init_reset=True,
batch_valid_length=None,
)
| 138 | ) |
| 139 | |
| 140 | def construct( |
| 141 | self, |
| 142 | x, |
| 143 | query_vector, |
| 144 | input_mask, |
| 145 | init_reset=True, |
| 146 | batch_valid_length=None, |
| 147 | ): |
| 148 | r""" |
| 149 | The forward process of the block. |
| 150 | """ |
| 151 | # [bs * seq_length, embedding_size] |
| 152 | input_x = self.layernorm1(x) |
| 153 | input_x = F.cast(input_x, self.dtype) |
| 154 | |
| 155 | # indicate whether reset saved states |
| 156 | key_reset = None |
| 157 | value_reset = None |
| 158 | |
| 159 | if self.use_past: |
| 160 | # reset states, init_reset True for reuse and False for reset |
| 161 | key_reset = self.assign( |
| 162 | self.key_past, |
| 163 | self.mul(self.key_past, F.cast(init_reset, self.dtype)), |
| 164 | ) |
| 165 | value_reset = self.assign( |
| 166 | self.value_past, |
| 167 | self.mul(self.value_past, F.cast(init_reset, self.dtype)), |
| 168 | ) |
| 169 | # add dependency for desired execution order |
| 170 | input_x = F.depend(input_x, key_reset) |
| 171 | input_x = F.depend(input_x, value_reset) |
| 172 | |
| 173 | attention, layer_present = self.attention( |
| 174 | query_vector, |
| 175 | input_x, |
| 176 | input_x, |
| 177 | input_mask, |
| 178 | self.key_past, |
| 179 | self.value_past, |
| 180 | batch_valid_length, |
| 181 | ) |
| 182 | # For post-layernorm the inputs for residual path are output of self-attention and output of layernorm |
| 183 | if self.post_layernorm_residual: |
| 184 | x = self.add(input_x, attention) |
| 185 | # For pre-layernorm the inputs for residual path are output of self-attention and input of this layer |
| 186 | else: |
| 187 | x = self.add(x, attention) |
| 188 | |
| 189 | output_x = self.layernorm2(x) |
| 190 | output_x = F.cast(output_x, self.dtype) |
| 191 | mlp_logit = self.output(output_x) |
| 192 | |
| 193 | value_update = None |
| 194 | key_update = None |
| 195 | if self.use_past: |
| 196 | # current key and value |
| 197 | key_present, value_present = layer_present |