(self, hidden_states, position_ids, attention_mask, memory_states=None, encoder_states=None,
return_memory=False, detach_memory=True)
| 740 | checkpoint = deepspeed.checkpointing.checkpoint |
| 741 | |
| 742 | def forward(self, hidden_states, position_ids, attention_mask, memory_states=None, encoder_states=None, |
| 743 | return_memory=False, detach_memory=True): |
| 744 | batch_size, query_length = hidden_states.size()[:2] |
| 745 | memory_length = memory_states[0].size(1) if memory_states else 0 |
| 746 | key_length = query_length + memory_length |
| 747 | # attention mask is the beginning postion of B region, \in [0, query_len) |
| 748 | is_scalar = torch.numel(attention_mask) == 1 |
| 749 | is_sep = is_scalar or torch.numel(attention_mask) == batch_size |
| 750 | if self.performer: |
| 751 | assert is_scalar, 'attention_mask should be a scalar to indicate the seperation position.' |
| 752 | assert memory_length == 0, 'Do not support transformer-xl.' |
| 753 | if is_sep: |
| 754 | sep = attention_mask.item() if is_scalar else attention_mask |
| 755 | |
| 756 | # conventional transformer |
| 757 | def build_mask_matrix(seq_length, sep, memory_length=0): |
| 758 | m = hidden_states.new_ones((1, seq_length, seq_length)) |
| 759 | m = torch.tril(m) |
| 760 | if is_scalar: |
| 761 | m[0, :, :sep] = 1 |
| 762 | else: |
| 763 | m = m.expand(batch_size, -1, -1) |
| 764 | ids = torch.arange(seq_length, device=sep.device, dtype=sep.dtype).view(1, -1) |
| 765 | mask = ids < sep.view(-1, 1) |
| 766 | m = m.masked_fill(mask.unsqueeze(1).expand_as(m), 1) |
| 767 | if memory_length > 0: |
| 768 | m = m.expand(batch_size, -1, -1) |
| 769 | m = torch.cat((hidden_states.new_ones((batch_size, seq_length, memory_length)), m), dim=2) |
| 770 | m = m.unsqueeze(1) |
| 771 | return m |
| 772 | |
| 773 | if not self.performer: |
| 774 | attention_mask = build_mask_matrix(query_length, sep, memory_length=memory_length) |
| 775 | else: |
| 776 | attention_mask = attention_mask.type_as(hidden_states) |
| 777 | attention_mask = attention_mask[:, :, :, -query_length - memory_length:] |
| 778 | |
| 779 | if self.relative_encoding: |
| 780 | position_sequence = torch.arange(key_length - 1, -1, -1.0, device=hidden_states.device, |
| 781 | dtype=hidden_states.dtype) |
| 782 | position_embeddings = self.position_embeddings(position_sequence) |
| 783 | # Apply dropout |
| 784 | position_embeddings = self.embedding_dropout(position_embeddings) |
| 785 | else: |
| 786 | if self.block_position_encoding: |
| 787 | position_ids, block_position_ids = position_ids[:, 0], position_ids[:, 1] |
| 788 | position_embeddings = self.position_embeddings(position_ids) |
| 789 | hidden_states = hidden_states + position_embeddings |
| 790 | if self.block_position_encoding: |
| 791 | block_position_embeddings = self.block_position_embeddings(block_position_ids) |
| 792 | hidden_states = hidden_states + block_position_embeddings |
| 793 | hidden_states = self.embedding_dropout(hidden_states) |
| 794 | |
| 795 | def check_detach(_hidden_states): |
| 796 | if detach_memory: |
| 797 | return _hidden_states.detach() |
| 798 | return _hidden_states |
| 799 |
nothing calls this directly
no test coverage detected