| 540 | # all_encoder_layers.append(hidden_states) |
| 541 | # return all_encoder_layers |
| 542 | def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True, checkpoint_activations=False): |
| 543 | all_encoder_layers = [] |
| 544 | |
| 545 | def custom(start, end): |
| 546 | def custom_forward(*inputs): |
| 547 | layers = self.layer[start:end] |
| 548 | x_ = inputs[0] |
| 549 | for layer in layers: |
| 550 | x_ = layer(x_, inputs[1]) |
| 551 | return x_ |
| 552 | |
| 553 | return custom_forward |
| 554 | |
| 555 | if checkpoint_activations: |
| 556 | l = 0 |
| 557 | num_layers = len(self.layer) |
| 558 | chunk_length = 1 # math.ceil(math.sqrt(num_layers)) |
| 559 | while l < num_layers: |
| 560 | hidden_states = mpu.checkpoint(custom(l, l + chunk_length), hidden_states, attention_mask * 1) |
| 561 | l += chunk_length |
| 562 | # decoder layers |
| 563 | else: |
| 564 | for i, layer_module in enumerate(self.layer): |
| 565 | hidden_states = layer_module(hidden_states, attention_mask) |
| 566 | |
| 567 | if output_all_encoded_layers: |
| 568 | all_encoder_layers.append(hidden_states) |
| 569 | |
| 570 | if not output_all_encoded_layers or checkpoint_activations: |
| 571 | all_encoder_layers.append(hidden_states) |
| 572 | return all_encoder_layers |
| 573 | |
| 574 | |
| 575 | class BertPooler(nn.Module): |